Collapse the multi-line comment blocks across the image, video, sd.cpp and diffusion-training code to one or two lines each, and drop comments that only restate the statement below them. Comments only, no code or behaviour changes.
Five items from the latest review; four were real.
An unload or arbiter eviction only cancels the generation holding
_generate_lock. A second request queued behind it holds no cancel event yet,
and Python locks are not FIFO, so it could take the lock the instant the
active denoise released it, still see a loaded pipeline, and run a whole new
denoise after the model was told to go away: the eviction then waits minutes
for it and an image lands after the eject. Unload and a superseding load now
raise a fence under _lock before they queue, and a generation that wins the
lock while one is pending refuses instead.
The cached-model scan judged pipeline completeness across every revision, so
a repo holding an older complete snapshot plus a newer companion-only one
read as complete while the snapshot from_pretrained actually opens has no
transformer. Both scans now look at the revision the loader will open.
Deleting a dataset image deleted its caption sidecar unconditionally, which
for cat.jpg alongside cat.png removed the caption the survivor still resolves
to. The sidecar now goes only with the last image of that stem, matching what
the thumbnail cleanup beside it already did.
Importing an example into a folder that holds no images but does hold files
fell back to promoting the staging dir one file at a time, so an interruption
left a partial dataset that the image_count check accepts as complete on
retry. Those files are folded into the staging dir instead and the promotion
stays a single atomic rename.
The MPS generator report does not apply: torch.Generator(device="mps") has
worked since PyTorch 2.0 (pytorch/pytorch#91348) and the studio installer
pins torch>=2.4.
An fp8 text-encoder request loads a hosted PRE-CAST checkpoint, but the image
download plan never received text_encoder_quant, so the manager staged the base
repo's dense encoder (FLUX.2-dev's Mistral-24B is ~48 GB, Qwen-Image's
Qwen2.5-VL ~16.6 GB) and the load then pulled the pre-cast file inline, outside
the manager's progress and disk preflight. The plan now takes the field,
resolves the hosted artifact with the same resolver the injection uses, stages
that file, and drops only those components' dense weight shards. The load's own
prefetch takes the same treatment, since it paid the same cost. Only a
checkpoint that really resolves on the Hub earns the drop, so a gated or renamed
artifact still stages the dense encoder the load will fall back to.
The two trainers admitted each other with independent check-then-act guards:
the diffusion route checks the LLM backend several network-bound preflights
before it reserves, and the LLM route checks the diffusion service well before
it spawns, so two near-simultaneous starts could both pass and train on one GPU.
reserve() now re-tests the LLM backend under its own lock, and the LLM route
holds the diffusion service's gpu_load_admission across its spawn, so exactly
one of the two wins. Both halves fail open, so a chat-only install still
trains.
Comment-only pass over the Python this PR touches: drop what the code already
says, collapse multi-line explanations that still read on one line, and keep
the reasoning that is not recoverable from the code. No code, docstring
semantics or behaviour changes; verified with an AST comparison against the
previous revision, and the backend suite is unchanged (same 37 environment
failures as before: the API integration tests that need a live keyed server,
the flash-attn install hooks, and the GPU memory fields).
Reports the repos and exact files a pick needs so the download manager can stage
them with the loader's own file scope. A plain snapshot would add the packaged
root single, transformer shards and fp16 twins the loader never opens.
Four bugs in the batched inference path, all found by review:
- A mixed-prompt batch sent a scalar negative prompt against a prompt
list. Z-Image asserts on the length, and Qwen-Image, Krea 2 and FLUX
true-CFG encode a batch-1 negative against batch-N latents and fail in
the transformer's text/image concat. Broadcast it to match the batch.
- The FBCache step-cache reset sat above the chunk loop. diffusers only
resets that state at the end of a successful call, so a forward that
raised (the OOM the backoff is meant to recover) left its own residual
behind and the halved retry died on a shape mismatch. Reset before
every forward instead.
- The conditioning cache keyed on the checkpoint alone, but a GGUF or
single-file load takes its text encoders from the companion base, so
the same checkpoint against a different base reused the previous
base's embeddings. Key the base too.
- Gallery records stored the base seed and the requested batch size even
when a prompts/seeds list drove the run, so restoring the second image
of seeds=[5, 99] replayed seed 5. List-driven outputs now record as
single-image recipes on their own seed.
Also bound strength above 0: every img2img pipeline derives its step
count from it, so 0 leaves zero denoising steps and either raises or, on
SDXL, crashes on empty latents.
Batched generation: /images/generate takes a prompts list (one image per
prompt, txt2img only) or a seeds list (one prompt, one image per seed);
the legacy batch_size path derives per-image seeds base..base+n-1 like
the native engine. Every image gets its own torch.Generator so any batch
member replays alone from its gallery recipe; the whole list runs as one
forward by default with OOM backoff that halves a failed chunk, and an
explicit batch_size caps images per forward. Validated 10-22x over
serial engines on 32-image suites with LPIPS deltas within 0.002.
Conditioning cache on the inference path: UNSLOTH_DIFFUSION_COND_CACHE_DIR
(the inference sibling of the trainers' cond_cache_dir, same persistent
store) wraps encode_prompt so repeated prompts skip the text-encoder
forward entirely; verified bit-identical outputs. Bypassed while LoRA
adapters are attached; tensor-argument calls pass through uncached.
Compile cache: GGUF loads fingerprint their own bundles (quant=gguf, a
different compiled graph than the dense family) and batched calls
register every distinct (w, h, batch) chunk shape they ran, so the heavy
GGUF batched warmups (~159 s at batch 32 on 12B-class, ~655 s on 20B
CFG-batched) are paid once ever.
GGUF loader: strip the sd.cpp model.diffusion_model. container prefix in
the single-file converter; diffusers' FLUX.2 converter KeyErrors on it
and the Qwen-Image identity mapping strands the model on meta.
Round 2 of the hosted TE set, each bit-identical to dense-load-then-cast
and gated through the real backend (marker + status fp8 + same-seed LPIPS
vs dense TEs):
- FLUX.1 T5-XXL (text_encoder_2): 9.52 -> 5.90 GB, one artifact for
schnell/dev/Krea-dev (T5 shards byte-identical across all three,
verified sha256). 220 tensors, 144 fp8, LPIPS 0.109.
- Lumina Gemma2-2B: fp32 hub store 10.46 -> 3.20 GB (3.3x download cut).
288 tensors, 182 fp8, LPIPS 0.041.
- Z-Image Qwen3-4B: 8.04 -> 4.41 GB. 399 tensors, 252 fp8, LPIPS 0.112.
NOT shared with flux.2-klein-4B: klein retrained layer 35's MLP
(verified tensor diff, maxdiff 0.86), so klein hosts no entry.
- Krea-2 Qwen3-VL-4B: 8.88 -> 4.83 GB. 713 tensors, 460 fp8, LPIPS 0.082.
The constructor-assembled krea pipeline takes the encoder directly
(load_krea2_pipeline text_encoder kwarg); the loader remaps 5.x
rope_parameters and re-ties weights after assign so the rebuilt encoder
matches the builder's structure.
HunyuanImage 2.1 reuses the Qwen-Image artifact outright: its Qwen2.5-VL
text encoder is byte-identical (every shard sha256, 16,584,414,544 bytes),
recorded in the new component-level base-equivalence table the checkpoint
validator consults. The injection loop now covers text_encoder.._3 so a
family can host several components. Live check: LPIPS 0.123 vs dense.
Two live-test findings on the images load path:
- transformer_quant with baked LoRAs, when the dense quantized build is
declined for memory or fails: the load completed as a plain GGUF with the
adapters silently dropped (HTTP success, supports_lora=false after the
fact) -- wrong output with no signal. The load now fails with the recovery
options (drop the adapters, free VRAM, or pick a smaller model). Weight-0
adapters still count as no bake request, and the plain no-LoRA decline
keeps its silent GGUF fallback.
- A fresh GGUF load on a small GPU prefetched the base repo's full bf16
transformer shards (~47 GB on Qwen-Image) because the dense-quant prefetch
widening only checked scheme viability, not whether the device could ever
hold the candidate resident. Gate the widening on total device capacity
(reserve + 0.85 margin, the plan_fits_total_capacity bar) so a card that is
certain to decline the dense build never pays the download; capable devices
keep the prefetch.
Alpha-VLLM/Lumina-Image-2.0 is a 2.6B single-stream DiT with a Gemma2-2B
encoder and a standard 16-channel VAE, all transformers-4.x-compatible, so the
generic from_pretrained pipeline path loads it as a new lumina-2 family:
- Family entry (Lumina2Pipeline / Lumina2Transformer2DModel), aliased to
lumina-image-2.0 / lumina-image-2 / lumina2. No bare lumina alias: Lumina-Next
checkpoints are a different arch and must stay unknown rather than crash
mid-load. bf16-only upstream, so the fp16 fallback stays off like z-image.
- Trust the official repo for non-GGUF loads; bf16 component table entry
(ships fp32, ~5.2 GB transformer + 5.2 GB encoder bf16-resident).
- Generation defaults 50 steps / guidance 4.0 per the model card, and the
generate call passes the card's cfg_trunc_ratio=0.25 itself (family-gated,
signature-gated): the pipeline default (1.0) runs the CFG double-forward on
every step and oversaturates output.
- Catalog group with the single ungated bf16 pipeline artifact (11 GB resident)
plus routing assertions; images page defaults row.
- No GGUF artifact: none exists upstream (only finetune/LLM quants), so the
dense transformer_quant fast path (GGUF-kind-only) stays unreachable for now.
Offline probes of the future prequant campaign: int8 and fp8 both engage and
render cleanly (fp8 LPIPS 0.11 vs bf16, int8 0.33 from 50-step trajectory
drift with intact quality), so neither scheme is family-denied.
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.
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).
_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.
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).
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.
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.
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.
img2img and inpaint take their output size from the uploaded image and only
snap it to a multiple of 16, so an ordinary phone photo (up to the 4096/side
decode cap, 4x the txt2img 2048 ceiling and ~16x the area) drove an OOM-scale
latent and an opaque 500 on a normal card, while txt2img, upscale, edit, and
FLUX.2-klein inpaint are all already megapixel-bounded. Clamp the init longest
side to 2048 (the txt2img ceiling) before deriving width/height; edit is exempt
since its pipeline resizes to ~1MP internally.
_cast_nvfp4 quantized every nn.Linear with no filter, unlike the int8 and fp8
torchao text-encoder modes which exclude the VLM vision tower / lm_head / T5 wo.
On qwen-image / qwen-image-edit that 4-bit quantized the Qwen2.5-VL image tower,
degrading the edit/image conditioning the sibling schemes protect. Apply the
same make_filter_fn exclusion (require_bf16, mirroring _cast_fp8_dynamic).
Two evict/OOM fixes on the diffusion load paths:
- The video load moved a pipeline onto the GPU (apply_memory_plan) and
committed it while holding no lock, so an unload / GPU-arbiter eviction --
which bumps the load token and then barriers on _generate_lock before
freeing -- could hand VIDEO to chat/images and let the new owner allocate
concurrently with the in-flight placement, OOMing. Hold _generate_lock
across placement + the locked commit, mirroring the image backend, so an
evicting owner waits until this worker's placement is torn down or
committed. Lock order stays _generate_lock -> _lock (unload takes _lock
then releases it before the barrier), so there is no deadlock.
- resolve_local_single_file reinterpreted an On-Device folder as a base
single_file load whenever it held exactly one .safetensors, so a PEFT LoRA
adapter folder (adapter_config.json + adapter_model.safetensors) with a
family-token name was picked as a base checkpoint, evicting the resident
model before from_single_file failed on the adapter weights. Skip adapter
folders (adapter_config.json) and the adapter_model basename so the pick
stays a pipeline load and 400s in validation, before the GPU handoff.
Adds regression tests for both.
Fold PR #6872's image-generation fixes into the branch, deduped against the
round-12 dataset-upload and gallery integrity work already on image-generation.
Fixes carried forward from #6872:
- fp8 single-file transformer memory estimate: an fp8 checkpoint loads with no
quantization_config and diffusers upcasts it to bf16 (~2x resident), so budget
it accordingly in _plan_memory and estimate_safetensors_dense_mib.
- dense-quant OOM-evict preflight: when the GGUF fits resident but the dense bf16
transformer this path materializes does not, skip the fast path up front rather
than evict the current pipeline and OOM in finalization. Combined with the
existing offload->resident candidate re-plan so both the family-table estimate
and the on-disk shard measurement gate engagement (unified on the
transformer_resident_override_mib plan override).
- ControlNet: evict the previous module and its from_pipe wrapper before loading a
new one so swapping ControlNets within a base-model load cannot accumulate to OOM.
- ControlNet union_control_mode: raise on an unknown control type instead of
silently defaulting to canny.
- edit-family mask rejection: raise instead of silently dropping a mask on an
image-editing model that has no inpaint pipeline.
- companion cache: walk the snapshot dir and exclude transformer/ so the
dense-quant prefetch's cached shards do not inflate the companion total and
wrongly force offload.
- training: drop piecewise_constant from the LR scheduler enum and force bf16 for
fp16-incompatible families.
- dataset upload: batch-atomic staging with the same-stem duplicate guard.
- images page: guard negative-prompt restore on guidance>0, clear stale ControlNet
selection on restore, and revert an optimistic quant label when a pipeline load
never starts.
- uninstall (sh + ps1): keep the owner-marker guard on sd.cpp removal.
Conflicts resolved in favour of image-generation's evolved memory system,
loadSpecFor catalog, and stop-and-save (lora_path) run detection; #6872's fp8 and
dense-preflight fixes carried forward on top. All affected backend tests pass
(test_diffusion_backend, test_diffusion_training, test_diffusion_lora_trainer,
test_video_gallery, test_diffusion_controlnet).
Several image/video/training preflights ran before the route acquires the GPU or
frees resident models, but let a doomed local pick through and only failed deep in
the background load, after the user's chat/Images/Video model was already evicted.
- Local base_repo / base_model: _is_trusted_diffusion_repo accepts any existing
local path, but the base loads via from_pretrained (needs model_index.json). A
local dir that is not a diffusers pipeline passed the trust gate, evicted the
resident model, then failed. Add a shared _assert_local_base_is_pipeline check
and call it in the image, video, and training preflights.
- Dataset images: discover_image_caption_pairs only checked filenames, so a
corrupt or zero-byte upload passed the start-route preflight, freed the GPU, then
crashed the spawned trainer in PIL. Add an opt-in verify_images decode probe
(cheap PIL header check) that the start route enables; the trainers leave it off
since they decode every image anyway.
- Local single-file safetensors: the On-Device scanner advertises a bare
.safetensors directory (no model_index.json) as a text-to-image model, but the
picker starts it as a pipeline with no filename, so every click 400s. Reinterpret
such a pick as a single_file load of the sole checkpoint (resolve_local_single_file)
so the advertised model is actually loadable.
Regression tests for each: local non-pipeline base (image/video/training), the
verify_images decode gate, and resolve_local_single_file.
_reset_step_cache looked up reset_stateful_hooks on the transformer, but on a
diffusers CacheMixin transformer (Flux, QwenImage) that method lives only on the
HookRegistry; the transformer-level entry point is _reset_stateful_cache. So with
FBCache engaged on an image model the reset was a silent no-op, and the next
generation reused the previous request's first-block residual: a tensor-shape
mismatch (crash) when the resolution or batch changed, or stale cached output
otherwise. Prefer _reset_stateful_cache and fall back to reset_stateful_hooks,
matching the video backend. Update the tests to the real hook name.
- The companion base for a GGUF/single-file image load is resolved from the GGUF
repo's base_model card tag when no base_repo is passed, and that value loads via
from_pretrained. The explicit base_repo is already trust-gated, but the card tag is
attacker-controlled metadata on any remote repo, so it now clears the same
unsloth/allowlist/local trust bar; an untrusted tag is dropped in favour of the
curated family default and never reaches from_pretrained. This closes a pickle
deserialization vector on the normal GGUF load path (a user loading an attacker's
GGUF repo whose card points base_model at a malicious pipeline), matching the
trust discipline the ControlNet path already applies via evaluate_file_security.
The allowlist already contains every legitimate variant base, so variant
resolution for the supported unsloth GGUFs is unchanged.
- The images/unload route ran the slow VRAM-freeing unload on a thread and then
released the DIFFUSION arbiter owner unconditionally. release() is owner-guarded
and identity-less, so a concurrent /images/load that re-acquired DIFFUSION while
the unload ran would have its ownership cleared by the trailing release, and a
later chat load would then see no owner, skip eviction, and OOM against the newly
resident pipeline. The route now releases only when nothing is resident again.
- _apply_group_offload placed the resident companions before attaching the
transformer's group-offload hooks so a companion OOM returns with no hooks and the
whole-module fallback stays valid, but the streamed loop itself installs hooks on
each DiT in turn. On a dual-DiT pipeline where the second tower failed after the
first got its hooks, it returned False with hooks already installed, and the
caller's enable_model_cpu_offload fallback then crashed (diffusers rejects it on a
partially group-offloaded pipe). It now propagates the real failure once any hook
is installed, so the load fails with its actual cause instead of a misleading crash.
Adds regression tests: the untrusted card tag dropped to the family default (trusted
tag still honoured, explicit base still wins), unload keeping ownership when a model
is still resident (and releasing when not), and the partial dual-DiT hook set
propagating rather than falling through to a crashing whole-module offload.
- Image load now trust-gates a client-supplied base_repo. validate_load_request
rejects a base_repo that is not an unsloth/* repo, an allowlisted official base, or a
local path, mirroring the repo_id gate and the video loader. The route passes
base_repo into that pre-eviction validation, so an authenticated client can no longer
keep model_path on a trusted GGUF while pointing base_repo at an arbitrary remote repo
that the server would download and deserialize (a from_pretrained pickle/config path),
and no resident model is evicted for the rejected load.
- The keepwarm middleware now tracks the image and video generation routes
(/images/generate, /images/generations, /video/generate), so
other_inference_request_count() sees an in-flight generation and an API-key training
start is refused (409) before its unload would cancel that generation. endswith keeps
the GET *-progress and */cancel variants untracked.
- The OpenAI-compatible /v1 surface is now blanket body-capped like /api/inference,
instead of only /v1/chat/completions and /v1/completions. Every /v1 POST route
(images/generations, audio, embeddings, responses, messages, ...) buffers a JSON body
and none is a multipart-upload passthrough, so an unbounded ImageGenerationRequest
prompt on /v1/images/generations can no longer be buffered outside the request limit.
Adds regression tests: the base_repo trust gate at both the backend (untrusted remote
raises, local passes) and the route (untrusted base_repo returns 400 with no load), the
keepwarm tracking of the image/video generation paths (and not the progress/cancel
variants), and the /v1 surface being body-protected.
* Auto policies: deferred dense compile, video compile default, step cache and precision auto
Image dense loads with speed unset no longer sit at plain off: the load stays
bit-identical eager, and the 3rd generation in a session engages the default
compile profile plus the cuDNN attention upgrade mid-session (a one-off image
never pays the warmup, repeated use amortises it). Video dense loads resolve
straight to the default profile since a clip denoise amortises the compile
within a single run, and never to max.
Video also gains the image backend's tri-state auto policies: unset step cache
now decides from the default schedule and re-checks the actual step count per
generation, and unset precision (transformer_quant) hands the decision to the
hardware ladder instead of staying off. Memory badge reason now says plainly
that everything fits when no offload is planned.
* Rename Dtype to Precision, add the video Precision control, step cache Auto option
The images Advanced panel's Dtype row is now Precision (same control, clearer
name), and the video Advanced panel gains the matching Precision select wired
to the load route's existing transformer_quant field, gated to full-pipeline
loads the way the image control gates to GGUF. Step cache selects on both
pages gain an explicit Auto option as the default (the previous Off default
silently behaved as auto and never let anyone pin off), and the Speed and
Attention tooltips now state the deferred dense compile and the SageAttention
black-frame caveat.
* Model catalog: canonical diffusion model groups with device-aware routing
One canonical name per image/video model, its published artifacts (GGUF, FP8,
bnb-4bit, official BF16) as data, and pure routing helpers: suffix-stripped
canonical keys (owner-preserving; cross-owner merges only via explicit
aliases), group/artifact lookups, a flat back-compat options shim, load-spec
resolution replacing the pages' lookup tables, search matching over old ids
and format tokens, the GGUF fit ladder extracted from the variant expander,
and pickDefaultArtifact/pickDefaultQuant deciding what a bare group click
loads (downloaded first, then the best quality that fits 70 percent of VRAM,
GGUF as the safe fallback). Checked by npm run catalog:check, following the
i18n:check pattern.
* Picker: one canonical row per diffusion model with a format second level
The Images and Video pickers now render the curated catalog as one row per
model in Recommended: clicking loads the best artifact for the device (the
routed GGUF quant, a prequant FP8/bnb-4bit that fits, or the official BF16),
and a chevron opens the per-format list, with the GGUF row nesting the usual
quant expander. Live HF listing rows that belong to a group are deduplicated,
search collapses member repos into their group (old ids and format tokens
still match), and the On Device sections group cached member repos under the
same canonical name with the per-repo rows inside. Curated groups render from
the catalog rather than the HF listing, which finally surfaces LTX-2.3 in the
video Recommended list (its hub pipeline_tag is image-to-video, so the
text-to-video listing always missed it) and exposes the HunyuanVideo 720p
repack next to 480p.
Backend: /cached-models now tags trusted video-family repos text-to-video
instead of blanket text-to-image, and the pickers admit catalog-known
non-unsloth repos On Device, so cached Lightricks/Wan/Hunyuan pipelines
finally appear in the Video picker. Chat pickers pass no catalog and are
unchanged.
* Download formats, tab icons, plain-language train tips, 3-loop autoplay
The image Download button becomes a menu: PNG saves the original bytes with
the embedded recipe, JPEG and WebP re-encode client-side from the fetched
blob (JPEG flattened onto white). The video Download button gains MP4
(original, keeps audio), WebM and GIF; the latter two transcode server-side
from the stored MP4 via PyAV (VP9 realtime profile for WebM, ~12 fps adaptive
palette for GIF) behind a new gallery export route that 501s with a readable
message when a codec is missing.
Generated clips no longer loop forever: the player replays a clip three times
per selection, then pauses with controls up; a new generation or a refresh
gets its own three plays. The Create/Train tabs reuse the sidebar's New Chat
and Train icons (TestTubeOutlineIcon moved to a shared lib module), and every
Train tab helper text is now one plain sentence.
* [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
* Keep the Create/Train tab icon and label on one line
TabsTrigger renders its children inside a plain inline span and the
Tailwind preflight gives svg display:block, so the HugeiconsIcon forced
the label onto a second line. Wrap icon plus label in their own
inline flex row inside each trigger.
* Strip -int8 and -nvfp4 prequant suffixes in the model catalog key
canonicalKeyFor already lowercases before matching, so -GGUF/-FP8 in any
case were covered; -int8 and -nvfp4 were not in the suffix table, so
such repos rendered as standalone rows in Recommended and On Device
instead of standardizing into their base-name group and routing through
pickDefaultArtifact. Added both suffixes plus case-insensitivity and
routing assertions to the catalog check.
* Standardize non-catalog picker rows to their base model name
The curated catalog already collapses its own groups, but hub listing
rows and cached repos outside the catalog (ERNIE-Image, FLUX.2-klein,
Qwen-Image-Edit-2509, FLUX.2-dev) still rendered raw ids with -GGUF /
-FP8 style suffixes in Recommended and On Device.
- model-catalog.ts: new stripArtifactSuffixesForDisplay, a
case-preserving twin of canonicalKeyFor's stripping that keeps the
owner prefix and original casing for display.
- pickers.tsx: recommended hub rows and the downloaded GGUF/model rows
pass their labels through it when a catalog is present, so only the
diffusion pickers change; chat rows keep raw ids. Click targets keep
the full repo id, and the format badge still shows the artifact kind.
- Catalog check covers the new helper across GGUF/FP8/int8/nvfp4 in
both cases plus no-op and suffix-only names.
* Offer official BF16/FP8 artifacts per model group and fix gallery label clipping
Model picker changes so groups are not limited to unsloth quant repos:
- model-catalog.ts: each image group that has an official vendor pipeline
now carries its BF16 (official) artifact as the top (highest quality)
entry - Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image, Qwen/Qwen-Image-2512,
Qwen/Qwen-Image-Edit-2511, black-forest-labs/FLUX.1-dev, FLUX.1-schnell
and FLUX.1-Kontext-dev. The LTX-2.3 video group now lists Lightricks'
own bf16 and fp8 distilled single-file checkpoints alongside the GGUF.
Resident sizes are set from the actual weight totals (FLUX ships a
duplicate single-file that from_pretrained ignores, so FLUX bf16 is ~32
GB not 54). The repos that used to be aliases are now real artifacts.
- The router already prefers the highest-quality artifact that fits the
0.7 x GPU budget, so a datacenter GPU now defaults to official BF16
while consumer GPUs still route to the fitting quant or GGUF. That is
why bnb-4bit was the Z-Image-Turbo default before: it was the only
non-GGUF artifact and it was already downloaded.
- diffusion.py: allowlist the four official image repos not previously
trusted (qwen/qwen-image-2512, qwen/qwen-image-edit-2511,
black-forest-labs/flux.1-schnell, flux.1-kontext-dev). All verified as
safetensors-only diffusers model_index pipelines. The LTX-2.3
checkpoints are already on the video trust list.
- catalog check: BF16-wins-on-datacenter, quant-wins-on-consumer, and the
single-file load specs for the LTX-2.3 checkpoints.
Also fixes the video gallery thumbnail caption: the leading duration was
clipped by the rounded corner and selection border, so the strip now has
enough left/bottom padding to clear the curve.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* video gallery: guard export transcode against a stream-less clip
_transcode_webm and _transcode_gif indexed src.streams.video[0] before
checking the stream list, so a container with no video stream raised a bare
IndexError that the broad handlers then re-labeled as a missing libvpx or
decoder. Raise an explicit RuntimeError naming the real cause in both the
WebM and GIF paths.
* Studio: honor explicit attention/format choices, fix distilled-LTX defaults and On Device catalog routing
* Remove stray planning notes accidentally committed to the branch
* video: add transformerQuant to the load callback deps
handleLoad reads transformerQuant but omitted it from the useCallback dep array,
so after the user changes only Precision and then selects a model or clicks
Reapply, the memoized callback keeps the stale closure and loads the previous
precision. The image page's equivalent callback already lists it.
* model picker: honor the format filter when routing catalog clicks; add catalog rows to the roving list
- routedArtifactFor now scopes a group's artifacts to the active format filter
(the same matchesFormatFilter predicate the visibility check uses) before
pickDefaultArtifact, so a group shown only because it owns a GGUF no longer
routes a click to a large non-GGUF download. Covers both the Recommended and
On Device grouped paths.
- hubOptionKeys now includes the catalog-group, search-catalog-group, and grouped
On Device row keys in exact render order, so arrow/Home/End roving reaches the
catalog rows instead of giving them a duplicate missing id and skipping them.
* model picker: don't treat a partial base cache as downloaded
A partially-cached base repo (a cancelled download that left only some weights)
was counted as downloaded, so an On Device click routed to a fresh multi-GB
re-download instead of the complete GGUF. The picker's endpoint (/api/models/
cached-models) did not carry a partial flag at all, so a frontend-only guard
could not see it. Surface partial from that endpoint by reusing the hub inventory
scan's snapshot-partial detector, plumb it through CachedModelRepo (backend +
frontend types), and skip partial base repos when building the downloaded set.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* model picker + diffusion: drop partial/unloadable cached rows, skip defer-compile before a LoRA gen
- On Device (cached non-GGUF) rows filtered partial-download snapshots back in: sortedCachedModels
gated on passesTaskGate + a groupForRepoId key match but, unlike downloadedSet, never checked
c.partial, so an incomplete unsloth snapshot showed as a loadable On Device row (click errors or
silently re-fetches multi-GB). It also admitted repos that only match the catalog by group KEY
(a base / uncurated-quant sibling like Qwen/Qwen-Image-2512) which have no loadable artifact and
dead-end at the trust gate. Add !c.partial and gate on artifactForRepoId (what loadSpecFor
resolves) instead of groupForRepoId, so a cached row shows only when the backend can load it.
- Deferred speed-auto engaged the compile profile on the 3rd generation BEFORE _apply_loras. A
compiled transformer rejects LoRA (supports_lora is False) and _apply_loras raises before its
unchanged-selection no-op, so once compile engaged every LoRA generation on that load failed
permanently. Skip the deferral when a LoRA is requested (compile and LoRA are mutually exclusive)
and let it engage on a later LoRA-free generation.
* Scope the cached-model partial probe to the listed snapshot dir
list_cached_models builds each row from the largest/complete copy across HF cache
roots, but _cached_repo_partial probed is_snapshot_partial with no repo_cache_dir,
so the scan spanned every root: a stale .incomplete copy in one root would flag a
complete copy in another as partial and hide the usable model from the picker (the
click then routes to a re-download). Forward the winning snapshot's repo_path so all
three partial signals are scoped to that copy, matching the sibling inventory paths
(models/dataset cache_inventory, local_inventory).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Do not auto-route to gated repos, prefer complete cached copies, defer compile past attached LoRA, scope group expand keys
Four fixes:
- pickDefaultArtifact's not-downloaded ladder returned the gated BF16 FLUX.1-dev / Kontext-dev
before the open GGUF on a large GPU, so a bare group click routed to a repo the user may lack
license/token access to. Add a gated flag and skip gated artifacts in the not-downloaded ladder
(an already-downloaded gated artifact is still returned).
- list_cached_models picked the largest duplicate cache copy and computed partial only on it, so a
larger partial copy shadowed a smaller complete one; since partial rows are dropped from the
picker the usable model vanished. Prefer completeness, then size.
- the deferred-speed compile engaged on a no-LoRA generation while an adapter from a prior
generation was still attached, baking it into the compiled graph (the later unload is swallowed
on a compiled pipe); also defer while adapters remain attached.
- routeGroupClick's GGUF fallback toggled the context-free canonicalId while the chevron toggles
the context-scoped expandKey, leaving the format list un-collapsible in one context, dead in the
other, and risking cross-context expansion; thread expandKey through.
* Guard video pipeline repos from deletion, drop the always-failing LTX FP8 artifact, prefer 720p Hunyuan
Three round-6 fixes:
- cached non-GGUF video repos now surface in the Video On-Device picker with the normal delete
action, but /delete-cached only guarded chat + the Images engine, so a loaded/loading Wan / LTX /
Hunyuan pipeline could have its HF snapshot removed from under it. Add a VideoBackend
loading_repo_ids accessor and a video loaded/loading guard mirroring the Images one.
- the catalog advertised Lightricks/LTX-2.3-fp8 as loadable, but the LTX-2.3 loader refuses the
official scaled-FP8 single file (.weight_scale/.input_scale) and points to GGUF/BF16, so a pick
routed to a ~76 GB download that always fails on load. Remove the FP8 artifact.
- pickDefaultArtifact only sorts by format, so the HunyuanVideo group's 480p (listed first) beat
the 720p even on GPUs where 720p fits the budget. List 720p first so the fit loop prefers it and
falls back to 480p only on smaller cards.
* diffusion: add compute int8/fp8_dynamic text-encoder quant, wire into video
Add two torchao compute text-encoder quant modes to the diffusion precision
engine, alongside the existing layerwise fp8 and weight-only nvfp4:
- int8: per-token activation + per-channel weight (torch._int_mm), with per-layer
keep-bf16 selection. int8 degrades on large encoders unless the most
quant-sensitive decoder blocks stay bf16, so it engages only for families with
a measured keep-bf16 schedule (qwen-image / qwen-image-edit keep first+last 6,
flux.2-dev keeps first 3); a family without one falls back to fp8.
- fp8_dynamic: per-row fp8 compute (torch._scaled_mm), keeping the matmul in fp8
on the tensor cores instead of upcasting each forward like the layerwise fp8.
The selective int8 caster reuses the committed transformer-quant factory
(_make_quant_config / make_filter_fn / exclude_tokens_for_scheme) plus a small
structural first/last-N block skip, so it depends only on committed APIs.
Wire text-encoder quant into the video backend, which previously loaded the
companion encoder (Gemma3 / UMT5 / Qwen2.5-VL) dense bf16 while quantising only
the DiT. text_encoder_quant is plumbed through the load request, validation, the
load chain, the resolved record, and status, mirroring the image backend; it
applies for every load kind (the encoder is dense regardless of how the DiT was
sourced). Widen the image and video load request Literals and add the video
status field.
Tests: int8 family-schedule routing and fp8 fallback, fp8_dynamic routing,
hardware gates (int8 sm_80+, fp8_dynamic sm_89+), the structural block selection,
the real int8 filter closure (keeps the first blocks plus the vision tower /
lm_head / T5 wo dense), and the video route threading and 422 validation.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* text-encoder quant: skip the torchao modes under offload (both backends)
quantize_text_encoders applied int8-with-schedule / fp8_dynamic / nvfp4 (all torchao) to the
text encoder regardless of the offload policy. An offload placement then moves the quantized encoder
with Module.to(), which torchao tensor subclasses reject (aten._has_compatible_shallow_copy_type is
unimplemented) -- a hard crash, the same one the DiT path already skips torchao quant under offload to
avoid. Add offload_active to quantize_text_encoders and skip the torchao modes when set; layerwise fp8
is not torchao and still streams under offload. Both the video and image loaders pass
offload_active = (offload policy != none).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* diffusion: skip non-bf16 linears for scaled_mm quant schemes
The fp8 / mxfp8 / nvfp4 schemes run on torch._scaled_mm and the fp4 / mx GEMMs,
which assert a bfloat16 input weight. On a mixed-precision DiT that keeps some
linears in fp32 for numerical stability (the Wan and Hunyuan video transformers
do this), quantize_ hits the first fp32 linear, raises, and the best-effort
wrapper swallows it to None, so the whole transformer stays dense with no error
and no speedup or memory saving.
Add a require_bf16 gate to make_filter_fn and pass it for the scaled_mm schemes
in quantize_transformer (and the fp8_dynamic text-encoder caster). The gate
skips non-bf16 linears so the scheme engages on the bf16 ones. int8 uses
torch._int_mm, which quantizes fp32/fp16 weights fine, so it leaves the gate off
and keeps its current coverage.
Verified on Wan2.2-TI2V-5B: fp8 and mxfp8 now quantize 303 linears via the
committed quantize_transformer path where they previously engaged 0.
* prequant builder: mirror the scaled-mm bf16 gate offline
The runtime DiT quantizer skips non-bf16 Linears for the scaled_mm schemes (fp8,
nvfp4, mxfp8) so the scheme engages on a mixed-precision transformer instead of
aborting on the first fp32 Linear. The offline prequant builder reused make_filter_fn
without that gate, so building an fp8/nvfp4/mxfp8 checkpoint for a mixed-precision DiT
(Wan, Hunyuan keep _keep_in_fp32_modules in fp32 even under torch_dtype=bf16) would hit
the same fp32 Linear and abort, breaking the builder's stated offline == runtime,
LPIPS-0 invariant. Thread require_bf16 = scheme in _SCALED_MM_SCHEMES through the builder,
record it in the checkpoint metadata, and verify it on load (mirrors the existing
exclude_name_tokens guard) so a future _SCALED_MM_SCHEMES change cannot silently load a
checkpoint built under the old filter.
* Keep nvfp4 fp32 linears quantised (bf16 gate is fp8/mxfp8 only)
Verified on torchao 0.17 / B200: fp8 per-row asserts 'PerRow quantization only
works for bfloat16 precision input weight' and mxfp8 asserts 'Only supporting bf16
out dtype', but NVFP4's high-precision conversion quantises an fp32 weight fine
(forward included). So the bf16 skip-gate must be fp8/mxfp8 only, not all scaled_mm
schemes -- otherwise nvfp4 leaves large fp32 projections dense, losing the intended
memory/speed gain. Rename _SCALED_MM_SCHEMES -> _REQUIRE_BF16_SCHEMES = (fp8, mxfp8)
and thread it through the runtime filter, the offline builder, and the loader
require_bf16 verification (offline == runtime preserved).
* [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>
_dense_quant_prefetch_needed widened the prefetch to pull the base repo's transformer/
shards whenever a dense-quant candidate resolved, but balanced/low_vram (and the legacy
cpu_offload flag) force load_pipeline onto offload unconditionally in plan_diffusion_memory,
so its re-plan never flips to OFFLOAD_NONE and the dense build never runs. The offloaded GGUF
path then never touches those shards, so the widened prefetch only wastes a multi-GB download,
and a disk-full on that begin_load pull has no GGUF fallback (unlike the in-load_pipeline dense
failure). Mirror plan_diffusion_memory's definite-offload gates so the prefetch stays scoped.
- _dense_quant_prefetch_needed widened the transformer/ prefetch to pull the base
repo's full dense bf16 shards even when a prequant checkpoint is configured
(candidate.prequant), contradicting its own docstring. That both defeats the
prequant download savings and can hard-fail begin_load on a disk-full (no GGUF
fallback there). Only widen for a real dense build (candidate is not None and
not candidate.prequant).
- DiffusionStatusResponse declared no 'resolved' field, so Pydantic's default
extra='ignore' silently dropped the per-control auto-policy provenance the
backend records (build_resolved_record / state.resolved) -- the plumbing never
reached any client. Declare the field so it round-trips.
- _dense_quant_prefetch_needed widened the prefetch to pull the base repo's bf16 transformer/ shards
whenever a dense-quant scheme could resolve, with no disk check. On the offload path that can fill
the cache volume mid-download and hard-fail the load in a spot unload/cancel cannot preempt, instead
of the disk guard falling back to running the GGUF as-is (the Dtype hint's documented disk fallback).
Defer to resolve_dense_quant_candidate, the same disk-aware resolver load_pipeline re-plans against,
so the prefetch widens only when the dense build would really run.
- An explicit Speed=off (bit-exact) load with an unset dtype was promoted to auto-quant by the Dtype
default, silently engaging int8/fp8 + compile and breaking the bit-exact request (an auto DEFAULT
overriding an EXPLICIT control). Suppress the auto-dtype default when speed is explicitly off, in both
load_pipeline and the prefetch.
The dense-quant re-plan passes transformer_resident_override_mib (the bf16 build
peak) AND computes companions via _companion_cache_bytes(base), which sums every
flat blob in the HF cache. Because the dense path prefetches the base transformer/
shards into that same cache before load_pipeline runs, the transformer is counted
twice, inflating the footprint (~44 GB instead of ~20 GB in the reproduction) and
wrongly forcing offload for models that fit resident -- the case this path exists
to enable. Add companion_override_mib and pass the auto-policy's own text-encoder
plus VAE estimate on the re-plan so the cache (with its prefetched transformer) is
not read for this artifact.
Ideogram 4 assembles two DiTs per-component (a conditional transformer plus a
separate unconditional_transformer), so there is no transformer-only single-file
or GGUF artifact that could supply both. Add a pipeline_only family flag and
reject the gguf/single_file kinds in validate_load_request, before a load evicts
the current model, instead of assembling a pipeline missing its second DiT.
Extend the fp8 bf16-resident size override to a LOCAL directory mirror of the
ideogram-4-fp8 base: such a path never string-matches base_repo, so detect the
fp8 layout from the transformer shard headers (a *.weight_scale marker) and
reserve the bf16 footprint, matching the remote-base behaviour. A local nf4
mirror has no fp8 scales and correctly stays planned against its compressed bytes.
