Commit graph

1,050 commits

Author SHA1 Message Date
Daniel Han
08b853df15 Add FLUX.1 Krea dev to the image model catalog
Krea's guidance-distilled FLUX.1-dev finetune keeps the exact dev layout, so it
runs under the existing flux.1 family unchanged. Wire it up end to end:

- Catalog group with the gated official bf16 pipeline and the open QuantStack
  GGUF quants; the gated artifact is skipped on auto-routing when undownloaded.
- Trust the official repo for non-GGUF from_pretrained loads, next to the other
  black-forest-labs bases.
- Generation defaults: 28 steps at guidance 4.5 per the model card. The generic
  "krea" defaults key (Krea-2-Turbo's 8-step no-CFG recipe) used to swallow the
  id, which would have produced garbage output; the new flux.1-krea key precedes
  it on both the backend table and the images page table.
- The flux.1 prequant checkpoints are schnell-based; the loader's baked
  base_model_id validation refuses them for the Krea-dev base, so int8/fp8
  requests dense-quantize instead (covered by existing prequant tests).
2026-07-17 10:49:04 +00:00
Daniel Han
362cacc448 Support LoRA adapters on torchao int8/fp8 quantized image pipelines
Adapters are baked at load time: they attach to the dense transformer,
then quantize_ converts only the frozen base linears (the lora_ side
path is excluded by name), then the loader compiles. Post-quant PEFT
injection is not possible on a manually quantized module, so the
prequant shortcut is skipped for a baked load and the memory plan is
sized for the dense build (force_dense on the quant candidate).

At generation time the baked topology is frozen: weight tweaks and
disabling (scale 0 reproduces the quantized base exactly) go through
set_adapters, while adding or removing adapters returns a clean 400
telling the client to reload with the new selection.

supports_lora now returns True for int8/fp8 diffusers loads (checked
before the gguf-kind early return, since the quant fast path keeps the
picker kind); nvfp4/mxfp8 and GGUF-via-diffusers stay blocked. The
load request model takes an optional loras list, threaded through
begin_load on both engines (native ignores it and keeps applying LoRA
at generation).

Verified end to end on GPU: Z-Image GGUF picker + int8 + trained
adapter loads through the API, bake marker logged, weight 1.0 vs 0
renders differ visibly, weight 0.5 accepted live, unknown adapter
rejected as 400. Affected suites: 296 passed.
2026-07-17 09:37:17 +00:00
Daniel Han
07076430b3 Add FLUX.2 Klein and FLUX.2-dev DiT LoRA training
Register flux.2-klein and flux.2-dev in the DiT trainer following the
upstream DreamBooth references: latents train patchified and batch-norm
normalized from the VAE posterior mode, the packed forward reuses
step-invariant position ids, and the guidance vector (3.5) is gated on
the variant's guidance_embeds config. Conditioning stacks load per
variant (Mistral via Flux2Pipeline for dev, Qwen3 via Flux2KleinPipeline
for Klein) and are encoded and freed before the transformer lands on the
device. The fused single-stream to_qkv_mlp_proj joins the attention
projections in the LoRA targets; the single-stream out projection stays
dense because its to_out suffix would also match the double-stream
ModuleList container.

Wire both families through the training registry (family set, labels,
VRAM notes, rank 16 / lr 1e-4 defaults, bf16-only preflight), mark them
trainable with train base repos in the family registry, add FLUX.2-dev
to the gated-repo token check, and trust both official bases for
training downloads.

Verified on B200: 30-step klein int8 (19.6s) and nf4 (20.9s) and dev
int8 (52.0s) runs train with finite decreasing loss and the saved
adapters apply on the bf16 base pipeline (weight 0 reproduces the base
image exactly, weight 1 visibly restyles it).
2026-07-17 09:34:01 +00:00
Daniel Han
b35985162f Harden the diffusion memory plan against transient free-VRAM undercounts
A cold FLUX.2-dev int8 load on an idle 183 GB B200 planned offload=model
(companions exceed budget) and silently served the GGUF as-is; the identical
retry went resident and engaged the hosted prequant. The plan arithmetic was
byte-identical across both loads (required 90,228 MiB, resident needs free of
about 124 GB); the only divergent input was torch.cuda.mem_get_info, which is
device-wide and instantaneous: a transient foreign CUDA context briefly held
about 100 GB at the first snapshot, and the planner trusted that single read.

Three changes:
- settled_snapshot_device_memory: on cuda, synchronize + empty_cache
  (best-effort) and take the MAX free over up to 3 spaced reads. A transient
  can only shrink free, so the max rejects transient undercounts while a
  persistent tenant still caps every read. _plan_memory now uses it.
- plan_fits_total_capacity + one replan retry: when the dense/prequant
  candidate fits TOTAL device capacity under the standard reserve and the 0.85
  resident margin, an offload verdict can only stem from the free reading, so
  the loader re-snapshots and replans once before declining the fast path.
  Explicit balanced/low_vram modes skip the retry (they offload by mode).
- diffusion.transformer_quant_declined log line with required/budget/free and
  the plan reasons, so the next decline is diagnosable from the server log
  (previously silent).

Verified: cold FLUX.2-dev int8 first load in a fresh server now engages the
hosted prequant resident (offload=none).
2026-07-17 08:46:12 +00:00
Daniel Han
5d8bf5d094 Keep Qwen-Image's text-stream linears bf16 on int8 (short prompts break torch._int_mm)
Qwen-Image's MMDiT runs every text-stream Linear at M = actual prompt tokens: the
Qwen2.5-VL embeds are not padded to a fixed length like FLUX's 512-token T5. A short
prompt (13 tokens) or the near-empty negative prompt drives torch._int_mm below its
M > 16 floor and the first denoise step raises 'self.size(0) needs to be greater than
16, but got 13' (measured on B200 through the Studio images tab).

Add per-family int8 exclusions (txt_in, add_q/k/v_proj, to_add_out, txt_mlp) for
qwen-image and qwen-image-edit, threaded through exclude_tokens_for_scheme(scheme,
family) and the prequant checkpoint validation, so a checkpoint baked under the old
token list is rejected and re-quantised instead of loaded crashing. The text stream
runs at M = tens vs the image stream's M ~ 4k, so the exclusion costs nothing; the
rebuilt hosted checkpoint gates 28/28 PASS with LPIPS mean 0.057 (was 0.069).
2026-07-17 07:16:23 +00:00
Daniel Han
91be795a9b Route krea-2 through its per-component loader on the transformer-quant fast path
_assemble_pipe used Pipeline.from_pretrained for every family, but the krea repo
ships transformers-5.x configs and no top-level tokenizer files, so the tokenizer
dies with vocab_file=None. The pre-quantized checkpoint loaded fine and then the
assembly crashed, dropping the load to the GGUF build, which krea-2 cannot take
(Krea2Transformer2DModel has no from_single_file). Assemble per-component via
load_krea2_pipeline like the pipeline-kind and single-file paths already do.

Verified live: Krea-2-Turbo int8 and fp8 hosted prequant loads now assemble and
render through the Studio images tab.
2026-07-17 06:36:24 +00:00
Daniel Han
5e1614259c Wire hosted pre-quantized DiT checkpoints into the image families
Point prequant_repos for flux.1, flux.2-klein, flux.2-dev, qwen-image
(int8 only there; fp8 is family-denied), z-image and krea-2 at the
unsloth/<Model>-FP8 Hub repos carrying gate-validated int8 and fp8
transformer checkpoints, so the fast quant path loads the small
pre-quantized file instead of materialising the dense bf16 transformer
and quantising on device. Measured on FLUX.2-dev int8: build peak drops
from 60.7 GB (dense + quantize) to 30.7 GB (hosted prequant), identical
30.7 GB resident after either path since loading a checkpoint is
bit-identical to on-the-fly quantisation.

The hosted repos name files <Model>-<SCHEME>.pt, so resolve_prequant_source
now derives that model-name filename from the repo id (scheme suffix
stripped case-insensitively) and carries the legacy transformer_<scheme>.pt
as a fallback the resolver tries when the primary 404s, keeping older
repos loadable.

Wiring a repo also exposed a fallback hazard: with a prequant source
present, the dense-fit preflight used to be skipped entirely, so a failed
prequant download would fall through to the dense bf16 load the memory
plan never budgeted, OOMing after eviction. The preflight now always runs
and gates an allow_dense_fallback flag through _load_dense_quant_pipeline:
a dense misfit still skips the fast path when no prequant exists, but with
one it proceeds and a prequant failure raises to the GGUF build instead of
loading dense. The same flag is set when the auto-policy replans an
offloaded GGUF against a prequant-sized transient.

Tests updated to the new filename convention plus new coverage for the
derivation and the legacy-name fallback; the prequant-skips-refit test now
asserts the re-check runs and forbids the dense fallback. Verified end to
end on GPU: z-image int8 resolves the hosted repo, downloads the
model-name file and renders (6.8s load, 5.9 GB peak).
2026-07-17 05:47:45 +00:00
Daniel Han
117abb320a Studio: restore Reapply target on async image/video load errors; recheck training state before dataset commit; reclaim partial sd.cpp installs on retry
Images/Video: a background model load that fails AFTER starting (error/eviction during download) leaves the previous pipeline resident, but handleLoad had already overwritten lastLoad.current with the failed pick, so "Reapply to loaded model" reloaded the failed model. Carry the prior Reapply target into the poller and restore it on the async error/null paths, mirroring the quant rollback.

Training: an in-flight diffusion dataset upload passed _require_diffusion_dataset_mutable() at entry but could still commit files after a concurrent /diffusion/start reserved the training slot, mutating the dataset underneath the trainer. Re-check the interlock immediately before the commit phase; a 409 there leaves the staged temps for the finally to clean.

sd.cpp install: an interrupted extraction (disk full, killed process, a raising post-extract cudart fetch) left the target non-empty with no owner marker, so the next lazy install tripped the "not a Studio-managed directory" refusal and wedged native install. Write the ownership marker before the partial writes when the target is reclaimable, so a retry recognises the debris as ours and re-extracts.
2026-07-13 15:27:46 +00:00
Daniel Han
90827bc05c Studio: gate gallery serve/export on ownership; keep image progress active until persisted; reserve diffusion training before the dataset scan 2026-07-13 14:40:18 +00:00
Daniel Han
53912b2f99 Studio: tighten image-generation fix comments and docstrings 2026-07-13 13:32:28 +00:00
Daniel Han
be0bd00064 Studio: gate local-pipeline image tagging on a real family; validate video sidecars before delete/clear 2026-07-13 12:28:52 +00:00
Daniel Han
06b543b880 Studio: track the loaded GGUF filename so native companion resolution reproduces the load identity 2026-07-13 11:21:34 +00:00
pre-commit-ci[bot]
3fa9fdbc0d [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 10:04:07 +00:00
Daniel Han
5eef2f4003 Studio: preserve foreign gallery files, force safetensors on remote ControlNets, and close dataset/seed/GPU gaps
Gallery clear/delete now scope to Studio-owned files: image_gallery and
video_gallery skip PNGs / MP4s without a readable recipe (a hand-dropped or
orphan file the listing already hides), so clear() and a guessed-id delete no
longer destroy files the gallery never surfaced.

Remote ControlNets now force use_safetensors: a bare owner/name reaches
from_pretrained without the base trust gate, and the Hub scan fails open when
unavailable, so requiring safetensors closes the pickle deserialization vector.

POSIX uninstall now stops resident sd-server / sd-cli under an owned sd.cpp root
before removing the tree (marker-gated), mirroring the Windows stop-before-delete
scan; a live native server no longer survives unlinking its binary.

Diffusion dataset containment: the training-start read path and the discovery
picker route bare names through the protected resolver, so a symlinked dataset
is rejected / not advertised like the caption/delete routes already do. Uploads
gain the inference decode guard (oversized real images 400 before OOMing the
trainer) and dataset upload/caption/delete/import are blocked with 409 while a
diffusion run is active.

JSONL readers (trainer + routes) tolerate non-object JSON and invalid UTF-8
instead of raising AttributeError / 500.

LoRA family compatibility is enforced in the shared resolver, not only the
picker, so a direct API client cannot apply a mismatched-family adapter.

GPU arbiter gains release_if so the image/video unload idle-check and release
are atomic against a concurrent same-owner load's registration. Native batch
recipes persist the base batch_seed and restore replays from it, so a native
batch_index>0 image no longer advances its seed twice.

FLUX.2-klein selects its sd.cpp text encoder by variant (4B -> Qwen3-4B,
9B -> Qwen3-8B) instead of the single family default.
2026-07-13 10:02:42 +00:00
Daniel Han
b127256eb4 Studio: surface local pipeline scan roots, tag single-file checkpoints by filename, mark companion-only pipelines partial
- _scan_models_dir: admit a scan folder that is itself a diffusers pipeline
  (root model_index.json, weights in transformer/ vae/ subdirs). _is_model_directory
  rejects such a root, so the child scan would list the component subdirs as bogus
  models and hide the real pipeline; treat the root as one model via _local_pipeline_index.

- _local_is_diffusers / _local_model_task: include the sole checkpoint filename in the
  family-detection needles (_local_family_needles, resolved via resolve_local_single_file).
  A generically named folder holding one loadable qwen-image-*.safetensors / ltx-*.safetensors
  identifies its family only from the filename; the load route already resolves that file, so
  tag it or the task-scoped picker (which rejects task=null) hides the on-device model.

- list_cached_models: mark a companion-only base snapshot partial. A GGUF image load prefetches
  the base repo's VAE / text-encoder / model_index.json but skips the transformer (the GGUF
  supplies it); the snapshot has a pipeline manifest yet is not a loadable BF16 pipeline, and
  _cached_repo_partial misses it. _repo_pipeline_missing_denoiser flags a pipeline snapshot whose
  transformer/ or unet/ component carries no weight, so the picker drops it instead of advertising
  it as fully on-device.
2026-07-13 08:26:43 +00:00
pre-commit-ci[bot]
11330b8de5 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 06:31:51 +00:00
Daniel Han
899465ed80 Studio: close arbiter load-registration race and surface native progress + local pipeline folders
Publish native sd.cpp generate progress (_gen) before LoRA resolution so a reload probe reads active during setup, matching the diffusers path.

Register the diffusion/video GPU load under the arbiter lock (acquire_for now takes a register callback) so a competing acquire cannot evict an owner before its load is marked in-flight and let two loaders allocate VRAM at once.

Admit local diffusers pipeline folders (root model_index.json, weights in component subdirs) in the local model scan so they reach task tagging and the On Device picker.
2026-07-13 06:31:00 +00:00
Daniel Han
e0ef488f47 Tighten comments and docstrings added by the image-generation fixes 2026-07-13 05:29:09 +00:00
Daniel Han
dabb12e198 Studio: gate dataset uploads on the symlink check and surface local video single-file checkpoints 2026-07-13 04:21:36 +00:00
Daniel Han
da1770bb44 Refuse sd.cpp install into unowned non-empty target dir
When the install target already exists, is non-empty and lacks the
.unsloth-studio-owned marker (a user's own stable-diffusion.cpp checkout,
or unrelated files beside a custom Studio root), install() previously still
extracted the release into it. Skipping the ownership marker only stopped the
uninstaller from deleting the directory; extraction still merged binaries into
the user's working tree and could overwrite same-named files.

Fail up front with a clear message pointing the user at a fresh/empty location
before any download or extraction, leaving their directory untouched. Update
the ownership test suite to assert the refusal.
2026-07-13 03:17:49 +00:00
Daniel Han
5a17614b51 Reject native batch seeds outside the JSON-safe range 2026-07-13 02:09:15 +00:00
pre-commit-ci[bot]
499bada598 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-13 01:23:46 +00:00
Daniel Han
7ec9e77a5d Studio: fix diffusion install ownership, dataset upload atomicity, gallery pagination, and teardown races
install_sd_cpp_prebuilt: only write the .unsloth-studio-owned marker when the
install created the target directory or it was empty. Adopting a pre-existing,
unowned, non-empty directory (a user's own stable-diffusion.cpp checkout) made
it eligible for the uninstaller's recursive delete.

routes/training upload: make the multi-file promotion transactional. Back up
each displaced original and roll every destination back on any failure, so a
mid-loop rename error can no longer partially overwrite the live dataset.

routes/training _resolve_dataset_folder: reject a symlinked dataset directory
and prove the resolved folder stays under the datasets root, so image
read/caption/delete cannot escape the root through a link.

routes/training delete: escape glob metacharacters in the thumbnail filename so
deleting an image named like [ab].png removes only its own thumbnails.

image_gallery / video_gallery listing: filter records against the response
schema inside the pager via a valid callback, so offset/limit/has_more all count
over accepted records. A leading schema-invalid record no longer returns an
empty page with has_more=true and stalls infinite scroll at offset 0.

image_gallery / video_gallery save: publish via a temp file plus atomic rename
(the sidecar is the video pair's commit marker) and clean up on failure, so a
partial write never surfaces a truncated PNG or strands an orphan MP4.

diffusion_train_common discovery: treat an empty caption sidecar as a metadata
tombstone that still falls through to the dreambooth instance prompt, so
clearing every metadata caption no longer fails with no captioned images found.

diffusion backend unload: wait for an in-flight denoise to exit before tearing
down process-wide patches and state, mirroring the load path.

diffusion_engine_router: serialize the whole check/unload/publish transition so
a concurrent selection cannot return the engine being unloaded.

uninstall.ps1: gate the default sd.cpp process stop on the owner marker so a
user's own sd-server is not terminated for a directory we then keep.
2026-07-13 01:22:54 +00:00
Daniel Han
21052db120 Publish image generation active state before pre-denoise setup
generate() assigned self._gen only at the pipe() call, after deferred
compile, LoRA resolution/application, and ControlNet download/build had
run. Across that setup window generate_progress() reported inactive even
though _generate_lock was held, so a reloaded page's mount probe showed
idle and let a second generate queue behind the first.

Publish an active step-0 _GenState the moment the generation lock is
acquired, before the setup work, and clear it in the outer finally so a
setup-time error cannot leave the UI stuck active. Mirrors the video
backend's queued phase and the training start guard.
2026-07-12 13:42:24 +00:00
pre-commit-ci[bot]
7776148463 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-12 12:44:19 +00:00
Daniel Han
23a71b1c2c Close video single-file, training reservation, and image mount-resume gaps
Route on-device single-checkpoint video folders through the single_file loader:
a bare local .safetensors directory (no model_index.json) is advertised as a
pipeline with no filename, so validation rejected it before it could load.
Reinterpret the pick as a single_file load of the sole checkpoint, mirroring the
image load route.

Treat a reserved-but-not-yet-spawned LLM training start as active in
is_training_active() so /images/load, /video/load, and /diffusion/start cannot
race the reserved run for VRAM during the pre-spawn free window. Mirrors the
diffusion training service reservation.

Resume an in-flight image generation on the Images page mount: probe
generate-progress, re-enter the poll loop, and refresh the gallery on completion
so a run started elsewhere is reflected and its saved image appears without a
manual refresh. Seed resident image defaults from the resolved base_repo rather
than a possibly path-shaped repo_id so the first resident generation uses the
right recipe.
2026-07-12 12:41:21 +00:00
Daniel Han
4cc35aa87a Tighten comments in the image stack tests and scripts 2026-07-12 12:21:14 +00:00
Daniel Han
e1fa4fec04 Studio diffusion: fix static compile shape registration and prequant path validation
Register the dims the forward actually compiled with: image-conditioned
workflows (img2img, inpaint, upscale, edit) run at the input image's size,
not the slider's, so recording the slider values marked never-compiled
shapes as covered and warm restarts kept paying compile for the real one.

Validate a request-supplied transformer_prequant_path (existence plus the
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist) before treating prequant as
available at the resident-fit re-check: an unusable path skipped the dense
fit check up front and then fell back to materializing dense bf16 after
the previous pipeline was evicted, recreating the post-eviction OOM path.
Shared as usable_prequant_source, also used by the auto-policy planner.
2026-07-11 16:50:19 +00:00
Daniel Han
57f08ebfb2 Merge remote-tracking branch 'origin/main' into ig_merge 2026-07-11 15:15:54 +00:00
Daniel Han
6412efd7d9
Studio: auto-detect completion masking markers, stop silent full-sequence training (#7054)
* Auto-detect completion masking markers with template table fallback

Studio's train_on_completions previously relied only on the hardcoded
MODEL_TO_TEMPLATE_MAPPER / TEMPLATE_TO_RESPONSES_MAPPER tables and
silently disabled masking when a model was not in the table, so unmapped
models (LFM2-8B-A1B, DeepSeek, and others) trained on full sequences
without telling the user. Several mapped templates (glm, mistral, llama,
starling, zephyr, qwen3-thinking) also carried markers that mask every
assistant token, which made every row drop in the post-masking filter.

Both training callsites (CUDA trainer.py and MLX worker.py) now share
utils.datasets.completion_masking.apply_completion_masking:

- Try unsloth_zoo chat template auto-detection first; it raises loudly
  when the template cannot be parsed and never masks the EOS token.
- gpt-oss models keep their manual markers so non-final assistant
  <|end|> tokens stay trained, matching current behavior.
- If auto-detection raises, fall back to the template table exactly as
  before.
- If the table also misses, emit an explicit user-visible warning that
  completion masking could not be applied and full-sequence training
  will occur, instead of a quiet log line.

The >30 percent dropped-rows safety net in trainer.py now guards the
auto path as well. Table consumers for inference and chat templates are
unchanged. Validated against one representative tokenizer for every
template in TEMPLATE_TO_RESPONSES_MAPPER plus the unmapped models:
no template regresses; unit tests cover the four decision paths.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Restrict masking fallback to marker detection failures

The auto branch wrapped the whole train_on_responses_only call, so a real
failure while applying the masking (dataset map, tokenization) was treated
as a detection miss and training silently proceeded on full sequences.
Detect markers separately via get_chat_template_parts (test seam via
detect_fn), then apply them with errors propagating, matching the manual
path. Tokenizers with preset unsloth marker attrs skip detection and call
bare so zoo reuses the stored parts.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fail the run when applying completion masking raises

The helper already falls back internally on detection failures and returns
applied=False on a double miss, so an exception reaching the callsites is a
real failure applying the masking. Remove the callsite catches that
downgraded it to full-sequence training; the run now fails visibly instead.

Also use the explicit re-export alias form in utils/datasets/__init__.py for
the two new names, satisfying the import-hoist source lint.

* Import completion masking from its submodule

The import-hoist source lint counts only real name loads, so package-level
re-exports of the two new names cannot satisfy it. Import
apply_completion_masking from utils.datasets.completion_masking directly at
both callsites and leave utils/datasets/__init__.py untouched.

* Completion masking: gpt-oss renames and MLX raw/alpaca parity

Renamed or private gpt-oss checkpoints are name-detected as gpt-oss but miss
the exact-name table; default them to the gpt-oss template markers instead of
falling through to full-sequence training.

Gate the MLX masking call on not raw_text_mode and format_type != alpaca,
mirroring the CUDA path: raw/CPT text has no chat turns to mask and
Alpaca-rendered text lacks the tokenizer's chat markers.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Define raw_text_mode outside the MLX feature-detect block

With an older zoo lacking the append_eos config field, the masking
gate referenced raw_text_mode before assignment. Hoist the assignment
above the feature detection so both consumers see it.

* Gate MLX masking on the formatter's resolved format

format_type auto can resolve to alpaca or raw text; the masking skip
checked only the requested value, so auto-detected Alpaca data got
chat-template markers applied to rendered prompt text. Track the
final_format returned by format_and_template_dataset and gate on it,
matching the CUDA path.

* Unwrap the mlx-lm TokenizerWrapper before marker checks

The wrapper delegates plain reads to the wrapped HF tokenizer but hides
underscore attrs, so preset unsloth markers were invisible and detection
relied on the loader's call patch. Unwrap to the real tokenizer first,
as the zoo MLX resolver does.

* Tighten masking comments

* gpt-oss: auto-detect markers first like every other template

The quantized and BF16 gpt-oss checkpoints ship a chat template without
the channel final header, so the pinned manual markers match nothing
there and masking trained zero tokens. Auto-detection derives markers
from whichever template the checkpoint ships and keeps the final
terminator trained; the manual gpt-oss markers remain the detection
failure fallback, including for renamed checkpoints.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-11 05:13:45 -07:00
Daniel Han
352fb40089 Warm-save the compile cache by default, compile U-Net denoisers whole-module
diffusion_compile_cache: auto mode now saves the Mega-cache bundle after the
first compiled generation (UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE=0 opts out), so
users get warm restarts without the distributor env; a bundle hit starts clean
(no pointless rewrite of the just-loaded artifacts) and explicit mode 1/on keeps
the distributor-style re-save. New register_shape + manifest shape coverage: a
STATIC compile produces new artifacts per (width, height, batch), so the
generate path registers each generation's shape and an uncovered shape
re-dirties the context, growing the bundle to cover every shape the session
used. Measured (B200, real backend): Qwen-Image deferred gen-3 hitch 29.1 ->
22.2 s warm with bit-identical output (7.9 MB bundle, ~0.5 s save); SDXL gen-3
115.7 -> 24.7 s and a mid-session 768px recompile 65.8 -> 12.6 s (bundle 63.6 ->
98.7 MB after the 768 re-save).

diffusion_speed: U-Net denoisers (UNet2DConditionModel; no _repeated_blocks, so
the regional compile never reached them) now get a whole-module STATIC
torch.compile on the default tier, plus fused QKV projections and a compiled VAE
decode. Measured on SDXL (30 steps / 7.0 / 1024px, 4 prompts, LPIPS vs the
bit-exact reference): 6.16 -> 3.14 s end to end (1.96x) at LPIPS 0.035, steady
state 0.70-0.88 s/image through the real backend. Rejected on measurement:
dynamic=True whole-module (366 s compile for 39.3 ms/step vs static's 73 s for
26.9), regional BasicTransformerBlock only (45.0 ms/step; ResNet convs stay
eager), max-autotune + inductor flags (25.9 ms/step for a 445 s warmup),
channels-last UNet alone (neutral). DiT tiers unchanged: fused QKV measured
exactly neutral under the regional compile (Qwen-Image 6.53 vs 6.52 s), so it
stays max-only there, and the DiT VAE decode stays eager (a few % of a DiT
generation). compiled_shapes_are_static tells the cache layer which loads are
per-shape (max tier, U-Net whole-module).

diffusion: register each generation's shape with the compile cache before the
save, pass pipe.unet to the cache fingerprint when the pipe has no transformer,
and correct the transformer_quant resolved reason on dense loads (it claimed a
GGUF transformer was loaded on every non-quantized pipeline load).

Tests: 333 passing across the related suites (speed 42, compile_cache 27, cache
40, precision 20, backend, base_precision, transformer_quant, memory); ruff
clean. Full measurement record: outputs/image_optim_round2_audit.md.
2026-07-11 06:18:29 +00:00
oobabooga
d105bd7b42
Studio: detect Windows Intel GPUs via the registry before WMI (#7064) 2026-07-10 17:59:04 -03:00
oobabooga
7bfa209623
Studio: hint at Model auto-switch in the OpenAI "No model loaded" 400 (#7006) 2026-07-10 17:48:27 -03:00
Apoze
fef37cb25b
Studio: queue local GGUF OpenAI-compatible requests before llama-server (#7047)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-07-10 17:05:48 -03:00
Daniel Han
6cb44270fc Reject extension-case sidecar collisions, gate untrainable families, hide dead Reapply
Reject an image whose name differs from an existing one only by extension
case (cat.PNG vs cat.png): the stems are exactly equal, so on a
case-sensitive filesystem both files land and both resolve to one cat.txt
caption sidecar, silently sharing and corrupting the caption. Stem case
variants (Pic.png vs pic.png) stay exempt: they are one file on
case-insensitive filesystems and write separate sidecars on Linux.

Treat an empty precision_modes list on a DiT family as the backend's
deliberate cannot-train signal (a non-bf16 CUDA GPU fails the trainer's
preflight for every mode) instead of falling back to the full mode list:
the precision selector disables and the start button reads not supported,
so the form no longer offers a run that always 400s. An absent field still
means an older backend and keeps the fallback.

Hide the Images page Reapply button when no reload target is known: a
resident GGUF or single_file model discovered by refresh carries no
checkpoint filename in status, so clicking was a silent no-op. A resident
full pipeline keeps the button (it reloads by repo id alone).
2026-07-10 19:03:17 +00:00
Daniel Han
39b256f8c9 Fix root-only pipeline index detection, seed range overflow, local safetensors loads
Require the model_index.json to sit at the snapshot ROOT before flagging a
cached repo as pipeline-loadable: CachedFileInfo.file_name is the basename,
so the previous name match also claimed nested copies (subdir/model_index.json)
and the picker then sent a from_pretrained load that fails only after the GPU
handoff. Scope by file_path against the revision's snapshot_path.

Validate the maximum derived seed before the multi-run image loop: an explicit
seed near 2**53-1 plus the per-run offset (base + i*batchSize) exceeded the
backend cap and 422'd a later run after earlier images had already generated.

Route local single-file .safetensors picks on the Images and Video pages
through the single_file load path (parent dir + basename), matching the local
GGUF branch: the pipeline route rejects a bare file with no model_index.json,
and only after evicting the resident model.
2026-07-10 17:01:24 +00:00
Daniel Han
de2f22df2b perf(image): compile numeric parity, cache-hook compile arming, FBCache toggle crash fix, TE fp8 zero-row guard
Applies the video round-2 accuracy findings to the image diffusion stack and fixes
two real image-path bugs found while measuring. All numbers B200, production
settings (family default steps/guidance, 1024px, seed 42, 4 fixed prompts), LPIPS
(AlexNet) via the new scripts/image_speedmem_bench.py, which drives the production
lever functions in the loader's own order.

- inductor precision parity: emulate_precision_casts=True on the regional-compile
  path (fused pointwise kernels keep fp32 intermediates where eager rounds to bf16
  between ops). Pairwise LPIPS of the compiled tier vs the same-stack eager tier:
  Qwen-Image 0.019 to 0.006 at identical speed (72.4 vs 72.5 ms/step), FLUX.1-dev
  0.046 to 0.029 at +2% step time (69.8 vs 68.3, reproduced), FLUX.2-klein-4B
  0.018 to 0.017 at identical speed. Snapshot/restored with the other process-wide
  backend flags so an off load never inherits it.
- cache x compile composition: re-point each cache hook's fn_ref.original_forward
  at a torch.compile'd wrapper of the same bound method (armed only where the
  speed layer compiled the block; restored before every disable_cache and before
  the partial-hook cleanup). Qwen-Image FBCache computed steps 91.8 to 71.2 ms
  (back at the uncached compiled rate), 1.21x end to end (7.36 to 6.06 s per 4
  images); FLUX.1-dev already traced through its FBCache hook and is measured
  neutral (same-process armed vs unarmed latents bit-identical). Skip counts
  within noise (13 vs 11 of 76; pairwise LPIPS 0.005).
- FBCache mid-session toggle crash: diffusers 0.39 caches the HookRegistry child
  list on first cache_context use, so an uncached generation followed by a
  20+-step generation (the auto toggle path) enabled hooks the context never
  reached and crashed with "No context is set" (reproduced live on FLUX.1-dev).
  Invalidate the stale child cache after every enable_cache.
- TE fp8_dynamic zero-row guard: torchao per-row fp8 derives a per-output-channel
  scale from the row amax, so an all-zero weight row is 0/0 = NaN. SDXL's
  text_encoder_2 (OpenCLIP bigG) ships exactly such a row, and every explicit
  fp8_dynamic SDXL render came out black; keep zero-row Linears dense (LPIPS
  0.976 black to 0.096 working). Other families' encoders have no such rows and
  are byte-identical.
- No AUTO TE quant exists on the image branch (text_encoder_quant defaults dense,
  explicit-only), so the video round's auto-dense retune has no image analogue;
  the explicit lever's cost is now measured (TE fp8_dynamic alone, LPIPS vs
  bit-exact: Qwen-Image 0.038, FLUX.1-dev 0.084, SDXL 0.096; no speed win, VRAM
  -6.5 GB on Qwen-Image) for the docs.

Tests: 96 passing across the cache/speed/precision suites (11 new arming, 2
child-registry, 2 zero-row, 4 inductor-flag); ruff clean.
2026-07-10 16:07:46 +00:00
pre-commit-ci[bot]
cee2bf6ed2 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-10 09:20:11 +00:00
Daniel Han
daaac9e10b Run video generation as a background job so secure mode's tunnel cap cannot 524 it
POST /video/generate previously held the response open for the whole
generation (multi-minute for 720p), so in --secure mode the Cloudflare
quick tunnel's ~100s origin-response cap returned a 524 while the server
kept generating, and the frontend treated the run as failed.

Generation now follows the same return-at-once pattern as /video/load:
begin_generate validates synchronously (409 on no model or on a second
concurrent generate via a new busy sentinel) and runs the existing
generate + gallery-persist pipeline, with the route's exact error
mapping, on a daemon thread. GET /video/generate-progress gains optional
terminal fields: phase completed carries the saved gallery record, phase
failed a client-safe error; active only drops together with a terminal
phase. The cancel event is registered before the worker starts so
/video/generate/cancel keeps working across the whole job.

VideoGenerateResponse becomes an accepted acknowledgement (status
started, video kept as an always-null compat field). The video page
fires the POST, then drives completion off the progress poll it already
runs (completed prepends the clip, failed surfaces the error, the
cancelled sentinel stays toast-free). The API-key training-start guards
now also probe the video backend for an in-flight background clip, since
it is no longer visible as an in-flight HTTP request to the keep-warm
counter.

Route tests keep the fake backend for load/generate/status but inherit
the real job machinery, covering immediate accept, concurrent 409, the
terminal completed record, sanitized/ValueError/cancelled failures, and
cancel of a running job.
2026-07-10 09:19:02 +00:00
pre-commit-ci[bot]
783c0c1af6 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-10 08:07:50 +00:00
Daniel Han
c249d0c50a Fix picker dead-end for single-file repos, stale LoRA state after deferred compile, slash upload cap
Tag cached diffusion repos that ship no model_index.json with single_file in the
cached-models listing, and keep them out of the task-scoped On Device pickers
unless the curated catalog carries their artifact: the selection fall-through
loads uncataloged rows as a full pipeline and from_pretrained fails on a
single-file checkpoint repo after the GPU handoff.

Refresh diffusion status after a successful generation run on the Images page.
Speed Auto compiles the transformer on the third LoRA-free generation and flips
supports_lora to false; without the refresh the LoRA picker stayed enabled and
the next LoRA generation failed on the backend.

Match upload passthrough exact paths with trailing slashes normalized: the
trailing-slash variant of /api/train/diffusion/dataset reaches MaxBodyMiddleware
before the router's redirect_slashes 307, so it fell through to the default
/api/train body cap and 413ed large uploads. JSON sub-routes keep extra path
components after normalization and stay on the small cap.
2026-07-10 08:07:01 +00:00
pre-commit-ci[bot]
07d78c61a8 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-10 07:01:57 +00:00
Daniel Han
5f65f01e2e Fix training start blocking, dataset resolution shadowing, duplicate uploads, partial-caption gate
Run backend.start_training off the event loop with asyncio.to_thread so the
synchronous diffusion/video unload calls (which wait on engine generation
locks) cannot freeze concurrent requests; guard against overlapping starts
with a _start_in_progress compare-and-set under the service lock.

Resolve bare diffusion dataset names directly under datasets_root() before
falling back to the generic resolver, so an unrelated LLM upload file or
recipe folder sharing the name cannot shadow the image dataset.

Reject exact duplicate filenames within one multipart upload batch: two
parts staged to the same destination would let the later tmp.replace
silently discard the earlier file. Case variants stay exempt per the
existing stem-guard contract.

Require an instance prompt in the train panel when only some images have
captions, since backend discovery silently skips uncaptioned images.
2026-07-10 06:56:32 +00:00
Daniel Han
b9ebfe089b Merge remote-tracking branch 'origin/main' into ig_merge
# Conflicts:
#	scripts/scan_packages_baseline.json
2026-07-10 06:28:04 +00:00
Apoze
6a9b77ee37
Studio: harden OpenAI-compatible GGUF streaming (#6950)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-07-09 12:09:08 -03:00
Daniel Han
b5aef63c03
Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename (#7031)
* Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename

The Gemma 4 QAT GGUF repos renamed the higher-precision MTP/ subdir
copies from gemma-4-...-<quant>-MTP.gguf to mtp-gemma-4-...-<quant>.gguf,
so their basenames now start with the same mtp- prefix as the small
repo-root drafter (mtp-gemma-4-E4B-it.gguf).

The drafter selectors filtered candidates by a mtp- basename prefix and
took the first in sort order. With the new names the MTP/ copies also
match, and because MTP/ (uppercase) sorts before the lowercase root file,
selection flipped to the large BF16 copy under MTP/ instead of the root
drafter both functions document they should pick.

Restrict both selectors, and the companion byte estimate, to root-level
mtp-*.gguf so the MTP/ copies stay explicit-selection only:
- core/inference/llama_cpp.py _pick_mtp (loader auto-download)
- hub/utils/gguf_plan.py preferred_mtp_sibling (Hub variant plans)
- routes/inference.py _remote_gguf_companion_bytes (VRAM headroom)

Also reuse a drafter already in the local cache before downloading, so a
device that already holds a copy on disk does not re-fetch it.

Old-scheme names keep working (they have no root-level mtp- sibling to
mis-select). Adds regression tests for the new naming, both selection
paths, and the on-disk reuse.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: gate MTP drafter cache reuse to offline mode

Reuse the cached drafter only when HF is offline. Online, route back
through _download_companion_gguf/hf_hub_download so the current revision
is checked (etag) and a changed drafter is refetched, matching the
offline-only cross-snapshot reuse already used for the main GGUF. This
avoids pairing freshly downloaded weights with a stale cached draft.
Make the reuse tests offline and add an online-skips-reuse test.

* Studio: prefer a root MTP drafter across all cached snapshots

Offline reuse scanned snapshots one at a time and returned the first
snapshot that held any drafter, only preferring root within it. A newer
partial snapshot with just the MTP/ copy could shadow the small root
drafter in an older snapshot. Collect drafters across all snapshots and
prefer any repo-root file before an MTP/ copy.

* Studio: keep newest-first snapshot order when reusing cached drafters

Collecting root candidates and sorting by absolute snapshot path could
pick a drafter from an older snapshot. _iter_hf_cache_snapshots yields
newest first and the main GGUF is resolved in that order, so preserve it
(root still preferred over MTP/ copies) to avoid pairing a fresh main
weight with a stale drafter revision.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-09 06:46:00 -07:00
oobabooga
3502335120
Studio: add Vulkan llama.cpp support (#5819)
* Studio: add Vulkan llama.cpp support

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address gemini's feedback

* Studio: move the Vulkan VRAM probe into a standalone script

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Improve Vulkan probe error reporting

* Resolve llama-server symlink so Vulkan build is detected

* Drop unreachable Vulkan fallback in GPU free-memory dispatcher

* Skip the Intel GPU probe when NVIDIA or ROCm is present

* Reserve host RAM headroom for Vulkan integrated GPUs

* Add a `UNSLOTH_FORCE_VULKAN` environment variable

* [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

* Honor GGML_VK_VISIBLE_DEVICES, reserve discrete Vulkan VRAM headroom, and clear Intel GPU on --cpu-fallback

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Route Intel and forced-Vulkan hosts to the upstream Vulkan prebuilt, add arm64 Vulkan, keep Vulkan out of RAG auto-detect

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Clear the fork release pin when routing a Vulkan host to the upstream repo

* Gate auto-Vulkan routing on no physical NVIDIA so hidden CUDA devices aren't used

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Pin Vulkan launches with --device Vulkan<i> instead of the raw GGML_VK_VISIBLE_DEVICES index space

* Let user --device override the Vulkan pin, and gate direct Vulkan asset picks on no physical NVIDIA

* Update RAG auto-backend test mocks for the _resolve_auto binary and Vulkan probes

* Keep the add_dll_directory handle alive through the Vulkan probe DLL loads

* Revert RAG auto Vulkan guard, guard multi-backend Vulkan detection, and preserve forced Vulkan across updates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use getattr for RTLD_GLOBAL in the Vulkan probe CDLL mode

* Skip CUDA/ROCm APU and datacenter GPU tuning on Vulkan builds

On a Vulkan llama.cpp build gpu_indices are ggml compact ordinals, not
CUDA/ROCm physical ids, so _amd_apu_wants_unified_memory and
_apply_datacenter_env were reading the wrong device. On a mixed AMD APU
plus discrete GPU host that could raise a spurious system-RAM shortfall
and block a valid discrete-GPU load. Gate all three call sites on
not is_vulkan_backend; the Vulkan path already reserves iGPU host
headroom and the backend ignores GGML_CUDA_* anyway.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten Vulkan-guard comment in load_model

* Reduce comments in Vulkan support to be more succinct

* Resolve shell-wrapper llama-server entrypoint to the real lib dir

create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install
root when it cannot symlink into build/bin. _find_llama_server_binary
returns that root entrypoint, but Path.resolve() does not follow a shell
wrapper, so _llama_lib_dir returned the install root and _is_vulkan_backend
missed libggml-vulkan.so -- silently skipping the Vulkan probe and --device
pin on an otherwise valid Vulkan install. Follow the wrapper's exec target
to build/bin. Regression test: test_shell_wrapper_entrypoint_resolves_to_real_lib_dir.

* [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: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
2026-07-09 03:39:48 -07:00
Daniel Han
eb775d3207
Studio /v1/messages: accept thinking and unknown content blocks (#7017)
* Studio /v1/messages: accept thinking and unknown content blocks

The Anthropic-compatible /v1/messages endpoint modeled a message's content as
Union[str, list[{text|image|tool_use|tool_result}]], so any other block type
made Pydantic reject the whole request with
`messages.N.content.str: Input should be a valid string`. Resuming a Claude
session commonly replays assistant turns that carry `thinking` (extended
thinking) blocks, and sometimes a null content for a tool-only turn, both of
which tripped this and returned a 400.

Accept them:
- Add a permissive AnthropicUnknownBlock fallback (any block whose type is not
  one of the four known ones), so thinking/redacted_thinking/provider-specific/
  future blocks validate. A validator keeps known types on their typed models,
  so a malformed known block (e.g. a tool_use without id) still fails cleanly.
- Coerce a null message (and tool_result) content to "" so the converter's
  `for block in content` stays safe.

The converter already drops block types it does not translate, so a thinking
block is not forwarded to the model.

* Studio /v1/messages: keep user content validation strict

Make the thinking/null leniency role-aware so it never silently drops real
user input. Assistant turns (replayed history) still accept unknown/thinking
blocks and coerce a null tool-only turn to empty. User turns keep the strict
boundary: a null user content is rejected, and a content block the converter
cannot translate is rejected instead of being dropped into an empty prompt.

Also remove an empty file committed by accident.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio /v1/messages: coalesce resumed user turns and tighten content checks

- The /v1/messages count and generation paths now coalesce the adjacent user
  turns that dropping an empty or null assistant turn can leave behind, so a
  strict GGUF chat template no longer 400s on non-alternating roles.
- A user content block with a non-string type (list / dict) is rejected as a
  clean 400 instead of raising TypeError and escaping as a 500.
- The assistant null-to-empty coercion only applies to an explicit null; an
  assistant turn that omits content entirely still fails required-field
  validation instead of being silently coerced to an empty string.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio /v1/messages: tighten comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-09 12:20:02 +02:00
Daniel Han
c00eb20958 diffusion: address review round (FBCache context guard, aiter/ROCm, video cleanup, prequant + ControlNet gating)
- diffusion_cache: do not engage FBCache when the selected pipeline opens no cache_context.
  A CacheMixin transformer is necessary but not sufficient -- Flux Kontext / img2img /
  inpaint / controlnet reuse the CacheMixin FluxTransformer2DModel yet their __call__ never
  opens a cache_context, so the First-Block-Cache hook raised 'No context is set' on the
  first forward, crashing every default FLUX.1-Kontext edit (28 steps, above the FBCache
  threshold). Detect it from the pipeline __call__ source, resolved off the instance so the
  per-expert proxy view delegates to the real pipe.
- diffusion_attention: honor an explicit aiter backend on ROCm/AMD targets instead of
  dropping it via the NVIDIA-only guard (aiter is the AMD ROCm kernel; it only works there).
- video: clear the CUDA cache on a failed load so a partially built pipeline's reserved VRAM
  does not OOM the next load (mirrors the image backend), and re-check cancellation after the
  export/mux so a clip cancelled during the blocking encode is discarded, not persisted.
- diffusion_auto_policy / diffusion_prequant: validate a request-supplied prequant path
  override (present AND allowlisted) before budgeting the small prequant plan, so the loader
  does not skip the dense shards and then rebuild dense after evicting the resident pipeline.
- diffusion_controlnet: family-gate a curated ControlNet addressed by its full repo id, not
  only its short catalog id, so a cross-family repo id 400s up front instead of downloading
  and loading through the wrong ControlNet class.
2026-07-09 08:52:16 +00:00
Daniel Han
3b659ba075 Merge branch 'main' into image-generation
# Conflicts:
#	studio/frontend/src/components/app-sidebar.tsx
#	studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
2026-07-09 08:04:50 +00:00