Commit graph

6,260 commits

Author SHA1 Message Date
Daniel Han
e8c9568e6c Merge remote-tracking branch 'origin/diffusion-lora-training-api' into diffusion-lora-training-ui 2026-07-02 01:25:04 +00:00
Daniel Han
e9e9d82836 Merge remote-tracking branch 'origin/diffusion-lora-training' into diffusion-lora-training-api 2026-07-02 01:24:12 +00:00
Daniel Han
e0db8f4612 Merge remote-tracking branch 'origin/diffusion-lora-ux' into diffusion-lora-training 2026-07-02 01:23:23 +00:00
Daniel Han
76b5ee354d Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch
The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.
2026-07-02 01:23:15 +00:00
Daniel Han
2ac1e77dfc Merge remote-tracking branch 'origin/diffusion-sdxl' into diffusion-lora-ux 2026-07-02 01:21:52 +00:00
Daniel Han
74c3aa8a6d Merge remote-tracking branch 'origin/diffusion-controlnet' into diffusion-sdxl
# Conflicts:
#	studio/backend/core/inference/diffusion.py
2026-07-02 01:21:41 +00:00
Daniel Han
582e2dcf39 Merge remote-tracking branch 'origin/diffusion-lora' into diffusion-controlnet 2026-07-02 01:18:40 +00:00
Daniel Han
bb78742718 Merge remote-tracking branch 'origin/diffusion-image-workflows' into diffusion-lora 2026-07-02 01:17:02 +00:00
Daniel Han
96575cebf5 Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh
Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
  image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
  opens; the keep-alive dialog otherwise kept its mount-time default after a
  model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
  a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
  the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
  the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
  tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
  appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
  reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.

Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.
2026-07-02 01:14:28 +00:00
pre-commit-ci[bot]
dcf6cf5117 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-02 01:14:22 +00:00
Daniel Han
e05c9cc947 Images: preserve restored LoRAs through model load and never send hidden LoRAs
- The LoRA effect cleared the selection on every load->capable transition, which
  wiped adapters restored from a gallery recipe before the model finished loading.
  Track the previously-loaded family in a ref and clear only on a real family swap;
  keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
  hidden (loaded model does not support LoRA) is never sent to the backend.
2026-07-02 01:13:39 +00:00
Daniel Han
a4fb348f9f SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers
Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
  before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
  (config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
  longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
  but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
  degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
  access state.family.denoiser_attr directly.

Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.
2026-07-02 01:13:18 +00:00
Daniel Han
94d74e8bbe Diffusion: guard trust check against OSError and validate conditioning inputs
- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
  characters (or a bare owner/name id) can't raise OSError; treat any failure as
  not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
  still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
  input image, and reject reference_images on a family that does not support
  reference conditioning, instead of silently degrading to txt2img / img2img.
2026-07-02 01:13:04 +00:00
pre-commit-ci[bot]
1d4461b23b [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-02 01:12:15 +00:00
Daniel Han
50a313f93d Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown
Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
  and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
  so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
  foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
  <lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
  the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
  dropped wholesale (freeing adapters), so the previous call could race an in-flight
  denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
  when the catalog refresh fails.
2026-07-02 01:11:07 +00:00
pre-commit-ci[bot]
a914b96c88 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-02 01:10:34 +00:00
Daniel Han
74e7854e15 Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs
Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
  start while an LLM run was active put two trainers on the same GPU. Add a
  symmetric interlock: diffusion start returns 409 when LLM training is active,
  and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
  pre-spawn free of the export subprocess, the resident Images pipeline (with an
  arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
  failed and absolute paths bypassed containment. Resolve them with
  resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
  path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
  set them trained with defaults. Add both fields.

The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.
2026-07-02 01:09:56 +00:00
pre-commit-ci[bot]
c898473620 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-02 01:08:47 +00:00
Daniel Han
cdfc6f0f56 ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call
  cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
  to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
  download / VRAM cost; the control image is decoded and validated BEFORE the
  ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
  not a torch.dtype, so it silently fell back to float32) and honor the base
  offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.

Adds regression tests for family enforcement and the union control-mode map.
2026-07-02 01:08:02 +00:00
pre-commit-ci[bot]
f28be14639 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-02 01:07:40 +00:00
Daniel Han
af638f98c3 Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
  official bases, or a local path) before from_pretrained, so an untrusted remote repo
  is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
  cancel during download is honoured; a stop may carry save=False to cancel without
  leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
  the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
  of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
  anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
  lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
  lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.

Adds regression tests for the config coercion/validation and aliasing.
2026-07-02 01:06:42 +00:00
pre-commit-ci[bot]
67dfd3a2bb [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-02 00:00:05 +00:00
pre-commit-ci[bot]
aa6d4ff5ef [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-01 23:57:54 +00:00
Daniel Han
1ffcab52b8 Merge remote-tracking branch 'origin/diffusion-lora-training-api' into diffusion-lora-training-ui 2026-07-01 23:57:42 +00:00
Daniel Han
c681c976c3 Merge remote-tracking branch 'origin/diffusion-lora-training' into diffusion-lora-training-api 2026-07-01 23:57:41 +00:00
Daniel Han
e77ae6fecf Merge remote-tracking branch 'origin/diffusion-lora-ux' into diffusion-lora-training 2026-07-01 23:57:39 +00:00
Daniel Han
3c37299180 Merge remote-tracking branch 'origin/diffusion-sdxl' into diffusion-lora-ux 2026-07-01 23:57:26 +00:00
Daniel Han
16c7c7bccb Merge remote-tracking branch 'origin/diffusion-controlnet' into diffusion-sdxl 2026-07-01 23:57:25 +00:00
Daniel Han
85008e424c Diffusion training service: join the old pump outside the lock
start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
2026-07-01 23:57:10 +00:00
Daniel Han
b8da49118d Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16
The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.
2026-07-01 23:55:41 +00:00
Daniel Han
c4191569df Pipeline prefetch: fetch only the default torch weights
A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).
2026-07-01 23:55:12 +00:00
Daniel Han
060fac0a9d ControlNet: reject filesystem-like ids and do not cache a model past an unload race
Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
  path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
  Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
  from_pretrained: an unload that raced the download had already cleared the
  caches, so caching the late module would pin it past the unload.
2026-07-01 23:52:44 +00:00
Daniel Han
0028be413e Merge remote-tracking branch 'origin/diffusion-lora-training-api' into diffusion-lora-training-ui 2026-07-01 23:46:49 +00:00
Daniel Han
03fdebd389 Merge remote-tracking branch 'origin/diffusion-lora-training' into diffusion-lora-training-api 2026-07-01 23:46:16 +00:00
Daniel Han
b4b66907fe Merge remote-tracking branch 'origin/diffusion-lora-ux' into diffusion-lora-training 2026-07-01 23:45:29 +00:00
Daniel Han
eec9655630 Merge remote-tracking branch 'origin/diffusion-sdxl' into diffusion-lora-ux 2026-07-01 23:44:40 +00:00
Daniel Han
7262330f9a Merge remote-tracking branch 'origin/diffusion-controlnet' into diffusion-sdxl 2026-07-01 23:43:04 +00:00
Daniel Han
52ab033c00 Merge remote-tracking branch 'origin/diffusion-lora' into diffusion-controlnet 2026-07-01 23:42:09 +00:00
Daniel Han
c2b1c8a5a1 Merge remote-tracking branch 'origin/diffusion-image-workflows' into diffusion-lora 2026-07-01 23:40:21 +00:00
Daniel Han
68cad1bfbf Remove stray async task scratch outputs committed by mistake 2026-07-01 23:40:14 +00:00
Daniel Han
a4197d24c0 Merge remote-tracking branch 'origin/image-generation' into diffusion-image-workflows
# Conflicts:
#	studio/backend/core/inference/diffusion.py
#	studio/backend/core/inference/diffusion_families.py
#	studio/backend/core/inference/diffusion_speed.py
#	studio/backend/core/inference/sd_cpp_engine.py
#	studio/backend/tests/test_diffusion_speed.py
2026-07-01 23:39:29 +00:00
Daniel Han
c86f0cc2a1 Merge branch 'diffusion-lora-training-ui' of https://github.com/unslothai/unsloth into diffusion-lora-training-ui 2026-07-01 22:12:28 +00:00
Daniel Han
18f71583e6 Merge remote-tracking branch 'origin/diffusion-lora-training-api' into diffusion-lora-training-ui 2026-07-01 22:12:21 +00:00
Daniel Han
6ff245f3e1 Merge branch 'diffusion-lora-training-api' of https://github.com/unslothai/unsloth into diffusion-lora-training-api 2026-07-01 22:11:51 +00:00
Daniel Han
2e56ee7920 Import diffusion training schemas from models.training directly
The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.
2026-07-01 22:11:43 +00:00
pre-commit-ci[bot]
a34d48d1fd [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-01 21:05:56 +00:00
oobabooga
8d16ef977b Fix diffusion GGUF memory over-estimate and torch.compile crashes
Three chained bugs that made Z-Image (and other GGUF DiTs) crash at generation
on anything but a huge, fully-idle GPU. Verified end to end on an RTX 6000 Ada:
Q2_K now plans resident and generates a real 1024x1024 PNG on both the resident
and forced-group-offload paths.

- Memory planner over-estimated the GGUF transformer's resident size. diffusers
  keeps GGUF weights PACKED (uint8 GGUFParameter) and dequantises per-matmul
  transiently, so resident VRAM is ~= the on-disk size, not the unpacked bf16
  size (measured: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). The old per-quant
  expansion (x8 for Q2) over-estimated ~7.6x, so a 3.6 GB model on a 48 GB-free
  card was judged a "tight fit" and forced into group offload. Replace the
  multiplier table with estimate_gguf_resident_mib = storage * 1.05 (matches
  diffusers' own get_memory_footprint of a loaded GGUF model).

- torch.compile with fullgraph=True crashed under CPU offload: group/model/
  sequential offload installs a @torch.compiler.disable'd ModuleGroup.onload_
  hook, which graph-breaks. Drop fullgraph when offloading is planned, same as
  the existing step-cache case (fullgraph = not (cache_active or offload_active)).
  This mirrors diffusers' documented compile+offload guidance.

- compile_repeated_blocks compiles one graph per distinct block shape, but
  Z-Image's "repeated" blocks are heterogeneous (~11 variants), above dynamo's
  default recompile_limit of 8, so a resident load hard-errored under fullgraph.
  Raise the limit (diffusers' documented fix for regional-compile recompilation).
  Confirmed force_parameter_static_shapes=False is the wrong lever: same variant
  count, ~6x slower compile.

Also drops the now-dead infer_gguf_quant_label / gguf_filename plumbing and adds
regression tests for the estimate and the offload fullgraph drop.
2026-07-01 18:02:34 -03:00
pre-commit-ci[bot]
c34020bfec [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-01 20:02:38 +00:00
oobabooga
d0c5cf6e07 Fix diffusion flag leak, sd-cli orphan, and native family fallback
Six correctness fixes to the diffusion stack, found reviewing the merged
phase PRs on this branch:

- load_pipeline: restore the try/finally guard around the speed/quant/
  placement span. A failure after apply_speed_optims (e.g. OOM in quant or
  the memory plan) left TF32/cudnn flags flipped process-wide and the
  half-built pipe resident in VRAM. Now restores the flags and frees VRAM
  on a failed load.
- sd-cli Popen binds to the parent (PR_SET_PDEATHSIG via child_popen_kwargs,
  matching the llama.cpp sites), so a parent crash mid-generation can't
  orphan it holding VRAM/RAM.
- Native begin_load uses the filename-fallback family detector the route
  validated with, so a local .gguf whose family keyword lives only in the
  basename no longer dead-ends 400 on a no-GPU host.
- Generate error handler matches exact sentinel messages instead of the
  "cancelled" substring, fixing a 409 misroute and a raw sd-cli output leak.
- find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME/STUDIO_HOME like the
  installer, so a custom Studio home resolves.
- Drop the redundant _tf32_prev bookkeeping; snapshot/restore_backend_flags
  is now the single owner of the TF32/cudnn restore.

The two client-state messages are now shared constants so the 409-vs-500
contract can't drift. Adds a regression test for each behavioral fix.
2026-07-01 17:00:22 -03:00
oobabooga
a8b3e86020 Merge commit '692d3c1975' into HEAD
# Conflicts:
#	studio/backend/core/inference/diffusion.py
#	studio/backend/core/inference/diffusion_engine_router.py
#	studio/backend/core/inference/diffusion_families.py
#	studio/backend/core/inference/diffusion_memory.py
#	studio/backend/core/inference/diffusion_speed.py
#	studio/backend/core/inference/sd_cpp_args.py
#	studio/backend/core/inference/sd_cpp_backend.py
#	studio/backend/core/inference/sd_cpp_engine.py
#	studio/backend/models/inference.py
#	studio/backend/routes/inference.py
#	studio/backend/tests/test_diffusion_backend.py
#	studio/backend/tests/test_diffusion_speed.py
#	studio/backend/tests/test_sd_cpp_install.py
#	studio/install_sd_cpp_prebuilt.py
2026-07-01 16:00:27 -03:00