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.
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.
The SDXL trainer drew min(train_batch_size, len(pairs)) indices, so a dataset with
fewer images than the batch trained at a smaller effective batch than configured
while the scheduler and samples-per-second still assumed the full batch. The shared
PermutationBatchSampler already refills across permutation cycles to return exactly k
indices, and the DiT trainer calls it with the full batch size, so drop the clamp and
pass train_batch_size through for parity and to honor the configured batch.
Two evict/corruption fixes surfaced by review of the diffusion training path:
- krea-2 sets force_bf16 in the DiT trainer spec, but the route-level
_FORCE_BF16_FAMILIES preflight listed only qwen-image and z-image, so a
krea-2 start with mixed_precision=fp16 passed the route check, reserved
training and evicted resident GPU models, and only the child trainer then
raised. Add krea-2 to the set and a drift-guard test asserting it equals
the trainer specs whose force_bf16 is set.
- The dataset-upload same-stem duplicate check compared stems
case-sensitively, so on case-insensitive filesystems (Windows / default
macOS) sample.png and Sample.jpg both passed even though their caption
sidecars sample.txt / Sample.txt resolve to the same file, silently
sharing and corrupting one caption. Compare stems and the same-name guard
with casefold at both the on-disk and in-batch sites.
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).
Two overlapping /diffusion/start requests can interleave between the is_active()
check and the reservation, so reserve() itself must reject a second reservation
atomically. Otherwise both callers reserve, both free the GPU's resident chat or
image model, and the loser only 409s after the eviction -- the evict-then-fail the
reservation exists to prevent. reserve() now raises under the lock if a start is
already reserved or a job is already running.
start_diffusion_training freed resident GPU models and only then called
service.start(config), which is where is_active() first flips true. During that
free-then-spawn window a concurrent /images/load or /video/load saw training as
inactive, passed its training guard, acquired the GPU, and began a background load,
so the trainer and that pipeline both allocated VRAM. Add reserve()/unreserve() to
the training service (is_active() also reports the reservation) and reserve BEFORE
the free, in a try/finally so a failed start rolls the reservation back. An
overlapping load's guard now refuses during the window. Regression tests: the route
reserves before the free (and the free sees an active service), and the service
reservation marks active then rolls back.
The dataset upload took Path(filename).name, which on POSIX does not split on a
backslash, so a Windows client sending a backslash path in the multipart filename
stored the name verbatim. The caption/thumbnail/delete endpoints then run it
through _safe_dataset_image_path, which rejects backslashes and '..', so the
labeling grid could list an image it could never preview, caption, or delete (an
orphan). Fold backslashes to forward slashes before taking the basename so the
true name is stored, and reject a basename that still contains '..' at upload
rather than persisting an unmanageable entry. Regression test covers a Windows
backslash path (stored and served under the clean basename) and a '..' rejection.
start_diffusion_training is async but called _preflight_gated_base inline; it does
a blocking urllib urlopen HEAD to Hugging Face (up to a 5s timeout) to detect a
gated/unauthorized base repo. On a slow or unreachable network that stalled the
FastAPI event loop, freezing every concurrent status/progress/cancel request until
it returned or timed out. Wrap it in asyncio.to_thread, matching the dataset
preflight and GPU cleanup just below it. Regression test asserts it runs off the
coroutine thread.
start_diffusion_training is an async route, but it called the blocking
_free_gpu_for_diffusion_training() inline. That teardown waits on generation
locks and joins the export subprocess, so it can block for seconds and freeze the
FastAPI event loop, stalling every concurrent status/progress/cancel request
until it finishes. Wrap it in asyncio.to_thread, mirroring the dataset-preflight
call just above it and the inference routes' load/unload offloading. Add a
regression test asserting the cleanup runs off the coroutine thread.
- 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.
- Chat load defers the CHAT arbiter handoff until after identifier / gpu_ids /
training-memory validation, so a doomed chat load (bad id, unsupported gpu_ids on
GGUF, or a training 409) no longer evicts a resident Images/Video pipeline and then
errors. The already-loaded fast paths re-assert CHAT ownership themselves. Mirrors
the image and video loaders, which validate before acquire_for.
- Both training-start GPU cleanups now unload a resident Video pipeline and release the
VIDEO arbiter owner, not just Images/DIFFUSION, so starting LLM or diffusion training
after a video generation session no longer competes with the still-resident video
model and OOMs the run.
- Example dataset import materializes into a private staging dir and promotes into the
dataset folder only after the whole import succeeds. A materialize that fails partway
no longer leaves a partial dataset that a retry would treat as complete (imported=0),
stranding the user with a truncated dataset. Hidden dirs are skipped by the dataset
scan so the staging dir never surfaces as a dataset.
- Cached/local Wan GGUFs are now classified name-aware: arch "wan" alone is ambiguous
between the loadable single-DiT TI2V-5B and the dual-expert A14B MoE the loader
refuses, so _arch_to_task falls back to the repo/file name (as the loader's own
detect_video_family does) and tags only a non-MoE match text-to-video, surfacing
loadable Wan GGUFs in the Video picker without surfacing unloadable A14B files.
- The diffusion dataset upload route is added to the MaxBodyMiddleware upload
passthrough, so its own get_upload_limit_bytes() cap (plus multipart overhead, and a
raised max_upload_size_mb) applies instead of the default body limit rejecting
near-limit batches with 413 before the handler runs.
Adds regression tests for each: the chat handoff not evicting on a doomed load, the
video unload on diffusion-training start, the atomic import leaving no partial dataset,
the name-aware Wan GGUF classification (TI2V-5B video, A14B unsupported, bare arch
unsupported), and the diffusion upload passthrough cap.
- Diffusion dataset upload now streams each file into a sibling temp file and
atomically os.replace()s it into place only after the whole file is written and
within the size cap. A mid-batch 413 (or any abort) removes the temp, never an
example already stored under the same name, so re-uploading a too-large batch can
no longer truncate or delete a previously uploaded image.
- _diffusion_dataset_summary counts an image as captioned only when it resolves to a
non-empty caption via the same sidecar-over-metadata precedence the trainer uses. An
empty (tombstone) sidecar shadows a metadata row and makes the trainer skip the
image, so counting it over-reported caption_count and mislabeled an effectively
uncaptioned dataset as captioned.
- uninstall.sh/.ps1 now remove a custom/env-mode Studio's native diffusion build that
installs beside the root as a stable-diffusion.cpp sibling (find_sd_cpp_binary
resolves it from the Studio home's parent), guarded by the same unsafe-path check,
and stop processes locking the default-mode stable-diffusion.cpp before removing it.
Adds regression tests for the upload data-loss and caption-count paths and a hermetic
shell test for the custom-root stable-diffusion.cpp removal.
The start route preflight only rejected non-bf16 GPUs; an explicit int8 request on
a host with a missing or stub torchao passed the preflight, evicted resident GPU
workloads, then died in the trainer child (its int8 base quantizer has no fallback).
Fold both gates into training_precision_preflight_error so int8-without-torchao fails
fast before eviction. Also empty the advertised DiT precision_modes (and surface the
reason in vram_note, drop compile) whenever the bf16 preflight would reject the family,
so /info never offers an nf4 DiT option the route always 400s.
- normalized() + family_train_infos() mirror the inference fp8 deny for
Qwen-Image (activation outliers exceed fp8's range and corrupt the trained
result); int8 stays allowed and the UI no longer advertises fp8 for it.
- _resolve_base_precision() gates an explicit int8 on a FUNCTIONAL torchao, the
same gate auto and /info already apply, so a missing/stub torchao fails fast
instead of silently loading dense with compile disabled.
- train_precision_modes() gates the dense modes (bf16/int8/fp8/auto) on
torch.cuda.is_bf16_supported(), so a non-bf16 CUDA GPU (T4/V100/RTX 20xx) is
offered only nf4 instead of a start that evicts resident models and then fails.
- start_diffusion_training preflights bf16 support for the DiT families BEFORE
_free_gpu_for_diffusion_training(), so any DiT start (nf4 included, since the
trainer requires bf16 unconditionally on CUDA) fails fast without eviction.
Review follow-ups on the image-generation PR:
- ControlNet: resolve_controlnet accepts a bare owner/name repo without the
non-GGUF base trust gate, and _controlnet_pipe hands it straight to
from_pretrained. A malicious pickle .bin would deserialize on load, so run
the same Hugging Face malware preflight (evaluate_file_security) the chat and
export loaders use before any remote ControlNet load; local dirs are exempt.
- Dataset thumbnails: key the cache on the full filename instead of the stem so
sample.png and sample.jpg no longer collide on one .thumbs file (which could
serve or delete the wrong image); the delete cleanup globs the same key.
- Diffusion training start: mirror start_training's API-key guard so an API
client cannot start training (which frees VRAM by unloading chat) while an
inference request is streaming; it now returns 409 before any GPU is freed.
- Model picker: include the curated safetensors row keys in the recommended
roving key list so arrow-key navigation reaches those rows instead of hitting
the duplicate option-missing id.
Tests: ControlNet malware gate (remote blocked before from_pretrained, local
skipped), thumbnail same-stem cache separation, API-key diffusion-start 409
before GPU free. Full diffusion suites green.
Replace the with-replacement per-batch index draw in the SDXL and DiT LoRA
trainers with a shared PermutationBatchSampler that visits every image once per
cycle before repeating, so short runs cover the whole dataset. The sampler
reshuffles from the run's rng so the index stream stays seed-deterministic.
Guard the diffusion run detail route against a valid-JSON non-object record,
which previously raised TypeError and returned a 500; it now 404s like the list
path's shape check.
Add regression tests for both.
- Run the trainer's caption discovery in the start route BEFORE freeing GPU
residents, so a missing or uncaptionable dataset 400s without evicting the
loaded chat/Images model.
- sd.cpp unload now waits out a cancelled one-shot generation on the generate
lock before reporting the device free, matching the diffusers backend.
- Clearing a caption that came from metadata.jsonl writes an empty sidecar
tombstone instead of unlinking (both readers treat an existing sidecar as
authoritative), so the cleared label cannot resurface.
- The ControlNet wrapper pipe is only cached while its load is still current,
closing the unload race the model cache already handled.
- list_diffusion_runs skips wrong-shape records and the runs route tolerates
per-record ValidationError so one bad file never breaks the panel
- max_steps: 0 epoch-mode sentinel no longer trips train_steps validation
before epochs are resolved
- numberField keeps an explicit 0 (Seed, LR warmup) instead of falling back
- previous-runs list refetches once more shortly after a terminal status so
the just-finished run appears even if the record write races the fetch
- Unload the ACTIVE image engine (sd_cpp or diffusers) before diffusion training starts, not just the diffusers singleton
- Count metadata.jsonl captions in dataset summaries so metadata-captioned datasets are not reported as uncaptioned
- Sidecar captions now override metadata rows everywhere (grid edits win); trainer and dataset API agree
- Tag local diffusers image checkpoints with text-to-image so they appear in the Images picker
- Family LoRA targets (_FLUX_TARGETS etc) apply when the config carries the generic defaults; explicit overrides still win
- Latent caches (DiT + SDXL) now hold the posterior mean/std in fp32 and draw the
per-step sample in fp32, casting only the result to the training dtype. This
matches the in-loop path (encode fp32 -> sample fp32 -> cast) exactly instead of
sampling in bf16; the cache is tiny so the doubled RAM is negligible.
- The training service nulls non-finite floats (NaN/Inf loss, avg_loss,
learning_rate) at its single ingestion point so status snapshots and persisted
run records stay strict-JSON serializable; the metric history skips non-finite
loss points. Test covers NaN/Inf progress followed by a finite point.
