Images page Train tab: multi-family training UI, loss charts, labeling, deploy (#6822)

* Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler

Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
  model is not left resident-but-unreachable (the evictor only targets the active
  engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
  output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
  so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
  stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
  offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
  now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.

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

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

* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in

load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.

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

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

* Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa

- apply_attention_backend now restores the native default when no backend is requested or a
  kernel fails. diffusers keeps a process-wide active attention backend that
  set_attention_backend updates, and a fresh transformer's processors follow it, so a load
  that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
  load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
  Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
  request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
  valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.

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

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

* Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF

- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
  CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
  cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
  transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
  transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
  now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
  quantized FBCache threshold when the caller leaves it unset, instead of the dense default
  that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
  fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
  relative to the script, not a hardcoded path.

* Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder

_is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as
consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead
of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already
classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the
same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate.

Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0).

* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config

Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.

quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.

* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load

- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
  missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
  whole snapshot. restore_backend_flags restores each flag independently so one failure can't
  leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
  apply_speed_optims mutated the process-wide flags but before _state captured them for unload
  to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
  contaminated later off generations.

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

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

* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output

Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.

* Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn

The filter callback can be invoked without a module name, so fqn.lower() would raise
AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens,
so the linear is kept) instead of crashing the quantise pass.

* Studio diffusion (Phase 16) review fixes: native engine robustness

- sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64);
  a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds
  aliased to the same image. Random seeds stay 53-bit (JS-safe).
- sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub
  fall back to anonymous instead of failing auth on a blank token.
- sd_cpp_backend: a superseding load now cancels the in-flight generation, so the
  old sd-cli can no longer return/persist an image from the previous model.
- diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a
  slow 10+ GB free / CUDA sync does not block engine selection.
- diffusion_engine_router: probe sd-cli runnability (version()) before committing to
  native, so a present-but-unrunnable binary falls back to diffusers at selection.
- diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a
  CPU-only install can still reach the native sd.cpp engine instead of failing load.
- tests updated for the runnability probe + a not-runnable fallback case.

* Studio diffusion (Phase 9) review fixes: prequant safety + validation

- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
  resolves inside an operator-configured ALLOWLIST of directories
  (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
  once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
  any path a load request named (arbitrary code execution). realpath() blocks symlink
  escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
  checkpoint that quantised a different layer set is rejected instead of silently
  loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
  path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
  author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.

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

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

* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts

- diffusion_memory: when group offload is unavailable and the plan falls back to
  whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
  is the low-VRAM path where the decode spike can OOM). Covers both the group and
  sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
  balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
  the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.

* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path

- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
  resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
  it replaced), so when transformer_quant engaged and speed resolved to off, promote to
  default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
  CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
  otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.

* Studio diffusion (Phase 10) review fixes: attention gating + probe isolation

- diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere
  NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation,
  so auto now stays on native SDPA there.
- diffusion_attention: _active_attention_backend handles get_active_backend() returning an
  enum/None (not a tuple); the old  unpack always raised and was swallowed, so
  the native-restore short-circuit never fired.
- perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS
  on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul
  so the inductor_flags variant doesn't leak into later compiled rows.
- tests for the SM80 cuDNN gate.

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

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

* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening

- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
  timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
  extract through a per-member containment check (Zip-Slip guard); expanduser the
  --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
  the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
  installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
  sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
  crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.

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

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

* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs

Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.

Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.

* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats

Codex review on the native engine arg builder:

- build_sd_cpp_command emitted --width/--height unconditionally, so an
  img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
  the input. width/height are now Optional (None = unset): an image-conditioned
  run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
  the size from the input image (set_width_and_height_if_unset); a plain txt2img
  run with unset dims keeps the prior 1024x1024 default; explicit dims are always
  honored. width/height are read only by the builder, so the type change is local.

- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
  that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
  explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
  and emits the flag for any explicit value != 1.

Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)

* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc

Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.

Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).

* Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset

Codex review on attention-backend selection:

- Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on
  pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first
  generation with no fallback. select_attention_backend now applies
  _cudnn_attention_supported() to an explicit cuDNN request too.

- flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a
  Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3
  is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a
  (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+.

- apply_attention_backend's success path left diffusers' process-wide active
  backend pinned to the kernel it set; a later component whose processors are
  unconfigured (backend None) would inherit it. It now resets the global registry
  to native after a successful per-transformer set (the transformer keeps its own
  backend), best-effort. Also fixed _active_attention_backend: get_active_backend()
  returns a (name, fn) tuple, so the prior code stringified the tuple and never
  matched a name, defeating the native-restore short-circuit.

Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on
SM90; global registry reset after a successful set; _active_attention_backend
reads the tuple return.

* Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder

Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300
(Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So
_is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the
auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting
llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they
are treated as data-center class and keep the intended fp8-first behavior.

Test: extends the datacenter parametrize with 'NVIDIA B300' and
'NVIDIA GH200 480GB' (now _is_consumer_gpu False).

* Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder

Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.

Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).

Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.

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

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

* Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race

Codex review on the native-engine routing:

- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
  unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
  never touches the GPU, so that needlessly tore down the resident chat model. The
  handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
  CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.

- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
  _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
  the setting was unverifiable. status now derives them from state.offload_flags
  (still 'none' on CPU, where the flags are empty).

- _run_load committed the new state without cancelling/waiting on a generation that
  started during the (slow) asset download, so a stale sd-cli run against the OLD
  model could finish afterward and persist an image from the previous model once the
  new load reported ready. The commit now signals the in-flight cancel and waits on
  _generate_lock before swapping _state (taken only at commit, so the download never
  serialises against generation), mirroring the diffusers load path.

Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.

* Studio diffusion (Phase 14) review round 2: align helper name with the stack

Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.

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

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

* Studio diffusion: eager patches + torch.compile cache speed phase

Adds the opt-in speed path for the GGUF diffusion transformer behind a
selectable speed mode (default off, so output is unchanged until a profile
is chosen):

- diffusion_eager_patches.py: shared eager fast-paths (channels_last,
  attention/backend selection, fused norms and QKV) installed at load and
  rolled back on unload or failed load.
- diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent
  torch.compile cache and the GGUF-transformer compile wiring.
- diffusion_arch_patches.py: architecture-specific patches.
- diffusion_patch_backend.py: shared install/restore plumbing.
- diffusion_speed.py: speed-profile planning.

Tests for each module plus the benchmarking and probe scripts used to
measure speed, memory, and accuracy of the path.

* Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI

Backend:
- Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file
  fp8 transformers, gated to the unsloth org plus a curated allowlist.
- Image-conditioned workflows built with Pipeline.from_pipe so they reuse the
  loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint,
  outpaint, and a hires-fix upscale pass.
- Instruction editing as its own family kind (Qwen-Image-Edit-2511,
  FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and
  multi-reference) plus klein inpaint.
- Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched
  mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the
  decoded image size and cap upscale output to avoid OOM on large inputs.
- Fixes: from_pipe defaulting to a float32 recast that crashed torchao
  quantized transformers; image-conditioned calls forcing the slider size
  onto the input image. Native sd.cpp engine rejects image-conditioned and
  reference requests it cannot serve.

Frontend:
- Redesigned Images page with capability-gated workflow tabs (Create,
  Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor,
  client-side outpaint, and a multi-reference picker.
- Advanced options moved to a right-docked panel mirroring Chat: closed by
  default, toggled by a single fixed top-bar button that stays in place.

sd.cpp installer: pin the release, verify each download's sha256, add a
download timeout, and make the source repo configurable for a future mirror.

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

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

* Studio Images: correct the Advanced panel comment (closed by default, fixed toggle)

* Studio diffusion: LoRA adapters for the Images workflow

Add community LoRA support across both diffusion backends, the single
biggest step toward broad image-workflow coverage.

Backend
- New shared module core/inference/diffusion_lora.py: adapter discovery
  (local scan + curated catalog + owner/name[:file] Hub refs), download
  via hf_hub_download_with_xet_fallback, alias sanitization, native
  managed-dir materialization with collision-broken aliases, prompt-tag
  injection (deduped against user-typed tags), and a supports_lora gate.
- Native sd-cli: resolve + materialize selected LoRAs into a per-run
  managed dir, inject <lora:ALIAS:w> tags, pass --lora-model-dir with
  --lora-apply-mode auto. The arg builder already emitted these flags.
- Diffusers: non-fused load_lora_weights + set_adapters manager, tracked
  on the pipe so an unchanged selection is a no-op and a model swap
  resets; cleared on unload. Never fuses (breaks quantized transformers
  and blocks live weight tweaks).
- Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers,
  and native Qwen-Image (no LoRA name-conversion branch upstream).
- Request contract: optional loras on DiffusionGenerateRequest; empty or
  omitted is identical to today. supports_lora surfaced in status; chosen
  LoRAs persisted in gallery recipe metadata.
- New GET /api/models/diffusion-loras for the picker (family-filtered).

Frontend
- Repeatable multi-LoRA picker (adapter select + weight slider 0..2 +
  remove), gated by the loaded model's supports_lora and family, max 8.

Tests
- New test_diffusion_lora.py (14): helpers, request validation, native
  tag/dir wiring, diffusers set_adapters manager, supports_lora matrix.

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

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

* Studio diffusion: ControlNet for the Images workflow (diffusers)

Add ControlNet conditioning, the #2 most-used diffusion workflow after
LoRA, on the diffusers backend for the families with ControlNet pipelines
(FLUX.1 and Qwen-Image), with Union models as the default picks.

Backend
- New core/inference/diffusion_controlnet.py: family-gated discovery
  (curated Union models + local dirs + bare owner/name repos), resolution
  to a loadable repo/dir, control-image preprocessing (passthrough +
  a dependency-free canny edge map), and a supports_controlnet gate.
- diffusion.py: a ControlNet manager parallel to the LoRA one. Loads the
  (small) ControlNet model once via from_pretrained (cached by id) and
  builds the family's ControlNet pipeline via Pipeline.from_pipe(base,
  controlnet=model), reusing the resident base modules at their loaded
  dtype (no reload, no recast). Passes the control image + conditioning
  scale + guidance start/end at generate time; cleared on unload.
- Families: FLUX.1 -> FluxControlNetPipeline/Model, Qwen-Image ->
  QwenImageControlNetPipeline/Model. Others declare none (gated off).
- Gated off for the native engine, GGUF-via-diffusers, and torchao
  fp8/int8 dense (same rule as LoRA). v1 conditions txt2img only.
- Request contract: optional controlnet on DiffusionGenerateRequest;
  supports_controlnet in status; the choice persisted in gallery meta.
- New GET /api/models/diffusion-controlnets for the picker.

Frontend
- A ControlNet control in the Images rail (model select + control-image
  upload + control-type select + strength slider), gated by the loaded
  model's supports_controlnet + family, shown for text-to-image.

Tests
- New test_diffusion_controlnet.py (10): discovery/resolve/preprocess/gate
  helpers, request validation, family wiring, and the diffusers pipe
  manager (loads once, caches, from_pipe with controlnet, rejects
  unsupported families).

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

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

* Studio ControlNet: show the picker on the Create tab (workflow id is 'create', not 'txt2img')

The ControlNet control gated on workflow === 'txt2img', but the Images workflow tab ids are create/transform/inpaint/extend/upscale/reference/edit -- there is no 'txt2img'. So the picker never rendered even with a ControlNet-capable model loaded. Gate on 'create' (the text-to-image tab) for both the picker and the request wiring. Found via a live Playwright capture of the running Studio.

* Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander

Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a
stray "gguf" tag on the Hub but ship no .gguf files. The model search classified
them as GGUF from the bare tag, so the picker rendered the GGUF variant expander,
which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when
the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata
(populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected.

* Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op

handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other
non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo
surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline
load when the id is unsloth-hosted or on-device (the backend infers the family + base repo
and gates loads to unsloth/* or local paths), and show a clear message otherwise instead
of silently ignoring the click. Curated and GGUF paths are unchanged.

* Studio Images: keep curated safetensors models in Recommended after download

The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's
Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet),
so they vanished from the picker after the first load and could only be found by typing an
exact search. The row already renders a downloaded badge, matching how GGUF Recommended
rows stay visible when cached. Drop the exclusion so the curated safetensors always list.

* Studio diffusion LoRA: sanitize dots out of adapter aliases

The LoRA alias is used as the diffusers PEFT adapter name, and PEFT rejects names
containing "." (module name can't contain "."). sanitize_alias kept dots, so a LoRA whose
filename carries a version tag (e.g. Qwen-Image-2512-Lightning-8steps-V1.0-bf16) failed to
apply with a 400. Replace dots too; the alias stays a valid native <lora:NAME:w> filename
stem. Adds regression coverage for internal dots.

* Studio Images: clarify the GGUF transformer-quant Advanced control

Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode"
with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/
FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF,
falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only.

* Studio Images: list on-device unsloth diffusion models in the picker

The Images picker's On Device tab hid every non-GGUF cached repo whenever a
task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit
and FP8 safetensors) never showed up there. List cached repos that pass the
task gate, limited under a filter to unsloth-hosted ones so base repos (which
fail the diffusion load trust gate) don't appear only to dead-end on click.
Chat behavior is unchanged: the task gate still drops image repos there.

* Studio: hide single-file image checkpoints from the chat model picker

The chat picker treats a cached repo as an image model, and hides it, only
when it ships a diffusers model_index.json. Single-file, ComfyUI, and
ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a
Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models.
Fall back to resolving the repo id against the known diffusion families, the
same resolver the Images backend loads from, so these checkpoints are tagged
text-to-image and stay in the Images picker only.

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

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

* Studio Images: add the FLUX.2-dev model family

Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the
Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline)
resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline
+ Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated,
reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text
encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror.
text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for
dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from
klein's turbo defaults. Verified live: GGUF load resolves the family + gated base
repo and generates a real 1024x1024 image on GPU.

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

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

* Studio Images: clearer error for an unsupported diffusion model

When a repo id resolves to no diffusion family the load raised 'Could not infer a
diffusion family... Pass family_override (z-image)', which points at an unrelated
family and doesn't say what is supported. Replace it with a message that lists the
supported families (from a new supported_family_names helper) and notes that video
models and image models whose diffusers transformer has no single-file loader are
not supported. Applies to both the diffusers and native sd.cpp load paths. Also
refreshes two stale family-registry comments that still called FLUX.2-dev omitted.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* Add SDXL diffusion family (U-Net pipeline support)

SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.

DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.

The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.

The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.

Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).

Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.

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

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

* Images: LoRA free-text Hugging Face entry + recipe round-trip

The backend has always accepted a bare Hugging Face repo id (owner/name, or
owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the
picker only rendered when the curated catalog had entries, and the catalog is empty,
so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded
model supports LoRA, and replace the curated-only dropdown with a text input: type a
Hub repo id, or pick a discovered adapter from a datalist of suggestions when the
catalog is populated.

Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's
"id:weight" strings (splitting on the last colon, since the id itself may contain one
for a specific weight file) back into the selection, so replaying a saved image
reproduces its adapters. The generate payload trims hand-typed ids and drops empty /
zero-weight rows, and a model swap clears the selection (a LoRA is family-specific)
without discarding a free-text pick that is not in the curated list.

* Add diffusion LoRA training (SDXL text-to-image)

First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.

core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
  accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
  mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
  .txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
  projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
  timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
  AdamW + get_scheduler + grad accumulation + grad clipping, then export via
  save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
  polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
  so the training worker can spawn it; plus a CLI entry point.

Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.

Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).

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

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

* diffusion trainer: emit learning_rate in progress events (Studio pump compatibility)

The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.

* Wire diffusion LoRA training into the Studio API

Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).

core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.

routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.

Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.

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

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

* Images: add a Train LoRA (SDXL) dialog

Surface the diffusion training API in the Images page. A "Train LoRA" button in the top
bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick
the base model, dataset folder, output folder, an optional instance prompt, and the core
hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog
polls the training status while open and shows a progress bar, step count, live loss, and
the saved adapter path, with a Stop button for a clean stop.

The dialog is independent of the loaded generation model (training runs in its own
subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else
the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining /
getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes.

* Import diffusion training schemas from models.training directly

The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.

* Remove stray async task scratch outputs committed by mistake

* ControlNet: reject filesystem-like ids and do not cache a model past an unload race

Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
  path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
  Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
  from_pretrained: an unload that raced the download had already cleared the
  caches, so caching the late module would pin it past the unload.

* Pipeline prefetch: fetch only the default torch weights

A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).

* Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16

The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.

* Diffusion training service: join the old pump outside the lock

start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.

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

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

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

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

* Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety

Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
  official bases, or a local path) before from_pretrained, so an untrusted remote repo
  is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
  cancel during download is honoured; a stop may carry save=False to cancel without
  leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
  the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
  of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
  anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
  lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
  lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.

Adds regression tests for the config coercion/validation and aliasing.

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

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

* ControlNet: address review findings on the diffusers path

- resolve_controlnet enforces catalog family compatibility so a direct API call
  cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
  to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
  download / VRAM cost; the control image is decoded and validated BEFORE the
  ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
  not a torch.dtype, so it silently fell back to float32) and honor the base
  offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.

Adds regression tests for family enforcement and the union control-mode map.

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

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

* Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs

Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
  start while an LLM run was active put two trainers on the same GPU. Add a
  symmetric interlock: diffusion start returns 409 when LLM training is active,
  and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
  pre-spawn free of the export subprocess, the resident Images pipeline (with an
  arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
  failed and absolute paths bypassed containment. Resolve them with
  resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
  path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
  set them trained with defaults. Add both fields.

The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.

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

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

* Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown

Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
  and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
  so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
  foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
  <lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
  the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
  dropped wholesale (freeing adapters), so the previous call could race an in-flight
  denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
  when the catalog refresh fails.

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

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

* Diffusion: guard trust check against OSError and validate conditioning inputs

- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
  characters (or a bare owner/name id) can't raise OSError; treat any failure as
  not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
  still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
  input image, and reject reference_images on a family that does not support
  reference conditioning, instead of silently degrading to txt2img / img2img.

* SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers

Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
  before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
  (config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
  longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
  but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
  degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
  access state.family.denoiser_attr directly.

Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.

* Images: preserve restored LoRAs through model load and never send hidden LoRAs

- The LoRA effect cleared the selection on every load->capable transition, which
  wiped adapters restored from a gallery recipe before the model finished loading.
  Track the previously-loaded family in a ref and clear only on a real family swap;
  keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
  hidden (loaded model does not support LoRA) is never sent to the backend.

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

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

* Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh

Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
  image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
  opens; the keep-alive dialog otherwise kept its mount-time default after a
  model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
  a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
  the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
  the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
  tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
  appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
  reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.

Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.

* Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch

The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.

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

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

* Train LoRA dialog: stop suggesting absolute paths the backend rejects

The dataset and output placeholders showed /path/to/... examples, but the
training routes resolve those fields inside the Studio home and reject
absolute paths outside the approved roots, so following the placeholder
produced a 400. Use folder-name placeholders and say in the labels and the
dialog description where each folder resolves.

* Align the VAE to the denoiser's first FLOATING dtype, not its first parameter

A GGUF-quantized transformer's leading parameters are packed uint8 storage,
so reading next(parameters()).dtype handed nn.Module.to() an integer dtype
and every image-conditioned generation on a GGUF model (Qwen-Image-Edit)
failed with a 500. Probe the parameters for the first floating dtype, treat
an all-integer module as a no-op, and also catch TypeError so an unexpected
dtype can never break generation. Regression test included.

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

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

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

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

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

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

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

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

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

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

* Count LR scheduler warmup/decay in optimizer steps, not micro-steps

lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.

* Address Codex review findings on the image-workflows PR

Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.

Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.

Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.

Adds regression tests for segment matching and kind/extension validation.

* Guard inference loads and worker lifetime against diffusion training

Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.

Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.

Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.

* Harden diffusion LoRA handling on the diffusers and native paths

Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max):
diffusers requires the adapter loaded before compilation, so applying one to
the already-compiled module fails with adapter-key mismatches. The status
gate now hides the picker and generate raises a clear message instead.

Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the
diffusion cancellation sentinel in resolve_specs, so an unload/superseding
load during resolution maps to a 409 instead of a generic server error.

Drop weight-0 LoRA rows before the native support gate so a request carrying
only disabled adapters stays a no-op on families where native LoRA is
unsupported, matching the diffusers path.

Reject duplicate LoRA ids in the request model: both apply paths suffix
colliding names, so a repeated id would stack the same adapter past its
per-adapter weight bound.

Strip all user-typed <lora:...> prompt tags on the native path (only the
selected adapters are materialized in the managed lora-model-dir, so an
unselected tag can never resolve), and restore saved LoRA selections from a
gallery recipe so restore reproduces a LoRA image.

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

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

* Harden ControlNet resolve, gallery metadata, and the control-type picker

Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.

Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.

Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.

Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.

* Address further Codex findings on the image-workflows PR

- Persist the actual output image size in the gallery recipe instead of the
  request sliders: Transform/Inpaint/Edit derive the size from the uploaded
  image, Extend grows the canvas, and Upscale resizes it, so the sliders
  recorded (and later restored) the wrong dimensions for those workflows.
- Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file
  name) in validate_load_request, so the unloadable pick fails before chat is
  evicted rather than deep in from_pretrained.
- Only publish an image-conditioned from_pipe wrapper to the shared aux cache
  when the load is still current: from_pipe runs under the generate lock but
  not the state lock, so an unload racing its construction could otherwise
  cache a wrapper over torn-down modules that a later load would reuse.
- Verify the Windows CUDA runtime archive checksum before extracting it, like
  the main sd-cli archive, so a corrupt or tampered runtime is rejected rather
  than extracted next to the binary.

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

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

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

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

* Refuse non-SDXL base models at diffusion training start

The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.

* Add diffusion dataset upload and training info endpoints

Training an image LoRA required knowing the Studio home layout and copying
files onto the server by hand, which is the most confusing step of the whole
flow. Two small endpoints fix that:

- GET /api/train/diffusion/info reports the datasets and outputs roots plus
  every dataset folder that contains images (with image/caption counts), so
  the UI can offer a picker instead of a blind free-text path.
- POST /api/train/diffusion/dataset uploads images and optional caption
  .txt / metadata.jsonl files into a named folder under the datasets root,
  creating it on first use and accumulating on repeat uploads so large sets
  can arrive in batches. Names are validated to a single path component and
  files stream to disk under the same per-upload size cap as LLM dataset
  uploads. The returned name is a valid data_dir for /diffusion/start.

* Rework the Train LoRA dialog into a guided SDXL flow

The dialog assumed users knew the Studio home layout and that only SDXL is
trainable, and hid both facts behind free-text fields. Restructure it around
the three real decisions:

- Base model is a dropdown of the trainable SDXL picks (Base 1.0, Turbo, the
  loaded SDXL pipeline when there is one) with a custom repo/path escape
  hatch, instead of a bare text field defaulting to a repo id.
- Training images come from an in-browser upload (new dataset endpoints) or
  a picker over existing dataset folders with image/caption counts. No shell
  access or knowledge of the datasets root is needed any more, and the
  captioning rules are explained inline.
- The output field is now Adapter name and the instance prompt is labelled
  as the trigger prompt, with a no-captions warning wired to the selected
  dataset's actual caption count.

Hyperparameters collapse behind a training settings toggle since the
defaults suit a first run. A completed run says where the adapter went and
offers Done / Train another, and the top-bar button gets an icon and a
plainer description. The dialog title states the SDXL-only scope and that
other families load LoRAs but cannot train them yet.

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

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

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

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

* Validate diffusion training config before freeing the GPU

The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.

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

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

* Refactor diffusion LoRA training into a family-aware platform

Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:

- New core/training/diffusion_train_common.py holds the config + validation, dataset
  discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
  registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
  re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
  from the base model (or an explicit model_family) via the diffusion family registry,
  and a known-but-not-yet-trainable family is refused with a clear message. Unknown
  custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
  DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
  trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
  picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
  and Qwen-Image repos (safetensors-only, no remote code).

* Retain diffusion training loss history and expose it in status

The training service kept only the latest loss, so a live loss chart could show a
single point. Fold each progress event into bounded (step, loss, lr) history arrays
(capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM,
and record the family / base model / catalog path on completion. The status endpoint
returns these as a nested metric_history object the UI can chart directly, and the
start request accepts an optional model_family override.

* Tests for the diffusion training platform

Cover the trainer registry (get_trainer resolves SDXL, unknown family raises),
family resolution (explicit model_family validation, resolved_family on the config),
the metadata sidecar write + scan read with family gating, and the service loss-history
folding (append, bad-point skipping, decimation at cap, family/perf fields) plus the
status route nesting metric_history.

* Add diffusion dataset labeling and example-import endpoints

The Train tab needs to let users caption small datasets in the browser and
pull in a ready-made set to see training work end to end, neither of which
the upload-only endpoint supported.

Add, under /api/train/diffusion/dataset:
- GET {name}/images lists every image with its resolved caption (metadata
  beats a per-image sidecar, matching the trainer's discovery order) so
  uncaptioned images are visible and flaggable.
- GET {name}/image/{filename} serves an image, with ?thumb=<px> returning a
  cached downscaled JPEG kept in a hidden .thumbs subdir (regenerated when
  the source is newer) so the labeling grid stays light.
- PUT {name}/caption/{filename} writes, or when blank clears, the .txt
  sidecar; DELETE {name}/image/{filename} removes the image plus its
  sidecars and thumbnails.
- GET dataset-examples lists a curated, license-labelled registry, and
  POST dataset/import-example materializes one into a dataset folder as
  numbered images + .txt captions. Two loaders cover the shapes seen in the
  wild: streaming rows from datasets.load_dataset (dog-example, Tuxemon) and
  a snapshot + jsonl walk for imagefolder repos whose captions live in a
  non-standard *.jsonl (the public-domain tarot set). Imports are idempotent
  and cap the image count.

Filenames and dataset names are validated against path traversal and pinned
inside the datasets root.

* Test diffusion dataset labeling and example-import endpoints

Cover caption precedence, thumbnail generation and .thumbs exclusion,
caption write/clear, image delete cleanup, path-traversal rejection on
names and filenames, and example import with a mocked datasets.load_dataset
(files plus sidecars written, idempotent second call, cap respected, load
failure mapped to 502).

* Add flow-matching DiT LoRA trainers (FLUX.1-dev, Qwen-Image, Z-Image)

Extends diffusion LoRA training beyond SDXL to the three popular DiT families
via a single shared flow-matching loop parameterised by small per-family specs
(loading, prompt/latent encoding, transformer forward, save). Verified against
diffusers 0.38.0:

- FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly
  nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the
  user's HF token).
- Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std,
  img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16
  base).
- Z-Image: list I/O with the reversed timestep convention and a negated
  prediction, bf16 only.

The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now
route these families to the DiT trainer; the SDXL blocklist guard is replaced by
a positive family resolution that also rejects GGUF repos (inference-only) and
still-unsupported families. Per-family defaults + labels + VRAM notes are exposed
via family_train_infos for the Train UI.

Memory: caption embeddings are precomputed once and the text encoders freed
before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit)
and 8-bit AdamW are on by default.

* Speed up + shrink SDXL LoRA training (precompute text embeds, 8-bit AdamW)

SDXL re-encoded every caption with both CLIP text encoders on every step (pure
waste, since captions are constant) and kept the encoders resident. Precompute
each unique caption's embeddings once, then free the text encoders before the
loop: numerically identical (embeddings are deterministic and this consumes no
torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB
lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32
fallback, halving optimizer state with no meaningful LoRA quality cost. Env
toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard
A/B the paths.

* Expose trainable families in /diffusion/info and preflight gated bases

The training info endpoint now returns the trainable model families (name,
label, default + allowed base repos, recommended defaults, and a VRAM/access
note) so the Train UI can offer a base picker with realistic guidance. The start
route preflights a gated base repo (HEAD model_index.json with the user's token)
BEFORE freeing resident GPU workloads, so a missing FLUX.1-dev license/token
fails fast with an actionable 400 instead of evicting the loaded model and then
hitting a confusing mid-load 401.

* Tests for DiT trainers, family resolution, info families, gated preflight

Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only
guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image
are trainable (and GGUF repos are rejected as inference-only), the families list
in /diffusion/info, and the gated-base 400 preflight that leaves the GPU
untouched.

* Add diffusion training API client: metrics, families, dataset labeling, examples

Extends the Images training client for the Train tab: the status type now carries
metric_history (step/loss/lr) plus catalog_path/family/base_model/samples_per_second/
peak_memory_gb; the start request gains model_family; and info gains an optional
families list (per-family bases + defaults). Adds typed calls for the dataset
labeling and one-click example endpoints: list images with captions, thumbnail URL,
write/clear a caption, delete an image, list example datasets, and import an example.

* Add diffusion Train panel: config, dataset labeling, live charts, deploy

New full-page training workspace for the Images tab. Left column configures the run:
model family (FLUX.1-dev, Qwen-Image, Z-Image, SDXL in popularity order, with per-family
VRAM/license notes and defaults, backfilled from the backend families list when present),
base repo, dataset (existing folder, browser upload, or one-click example import), an
in-browser caption labeling grid (per-image thumbnail + caption saved on blur, delete,
uncaptioned highlight), adapter name, trigger prompt, and collapsed training settings.
Right column shows the live run: progress + loss/avg/speed/peak-VRAM readouts, the reused
training loss/LR charts fed from metric_history, and a completion card that deploys the
adapter into Create or starts another run.

* Wire Create/Train tab switch into the Images page and deploy flow

Replaces the Train LoRA dialog with a top-bar Create | Train segmented control next to
the model selector. Create renders the existing generation workspace unchanged; Train
renders the full-page training panel (unmounted in Create so its polling stops while the
backend run and its retained metric history survive a tab switch). Adds a deploy handler:
loading the trained adapter's base as a pipeline, queueing the adapter so the LoRA
discovery effect applies it once the base is loaded and LoRA-capable for the matching
family (with a mismatch warning), seeding the prompt with the trigger, and switching back
to Create. Removes the now-unused dialog.

* Wrap the DiT training forward in bf16 autocast

The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single
compute dtype during the forward, exactly like the diffusers dreambooth
scripts run under accelerator.autocast. Without it the 4-bit backward on
FLUX.1-dev fails with an illegal-address CUBLAS error partway into the
first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL
path (its own trainer) is untouched.

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

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

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

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

* Fix Train tab example cards and Create/Train tab layout

The example-dataset cards used a two-column grid in the ~340px config
column, which wrapped titles one word per line and let the long license
text overrun into the neighbouring card. Switch to one card per row with a
horizontal layout: title with a compact truncated license badge (full text
in the tooltip), a two-line clamped description, and the Import button on
the right.

The Create/Train switch had an icon inside the Train trigger that overhung
the pill corner. Drop the icon, make both triggers a fixed equal width so
the active pill sits flush in the top bar.

* Show only loss and learning-rate charts for diffusion training

The Train tab reused the LLM charts section, which also rendered an empty
Grad Norm card and an Eval Loss card showing an Evaluation not configured
placeholder with a red smear. Neither applies to diffusion LoRA training.
Add a diffusion-only two-card view that reuses the loss and learning-rate
cards directly with fixed presentation defaults, and note under the loss
chart that per-step loss is noisy by design so users read the smoothed
line for the trend rather than the raw jitter.

* Add a dataset preview strip to the Train tab

When a dataset with images is selected, show a strip of up to 8 sampled
thumbnails with a +N more tile, so users can see what is in the folder
before training. Clicking the strip opens the existing caption review
grid. Samples are drawn evenly across the folder and refresh on dataset
change or after an upload/import.

* Stop example cards from overflowing the Train config column

The example-dataset cards still overran the ~340px config column: the
license used the Badge component whose baked-in w-fit and whitespace-nowrap
ignored the max-width and truncate, and the grid children had the default
min-width auto so wide content pushed past the column edge and clipped the
Import buttons. Replace the badge with a plain truncating pill span, and
give the config column min-w-0 with overflow-x-hidden so nothing escapes
its width.

* Add Smithsonian Butterflies and Nouns example datasets

Two permissive ~100-image sets for the Train tab: huggan/smithsonian_butterflies_subset
(CC0, the classic diffusers-docs training set, imported as a subject set with a trigger
prompt since its metadata columns are species names not captions) and m1guelpf/nouns
(CC0, captioned pixel-art avatars via the text column). Both cap at 100 images.

* Paginate the Train tab caption grid with prev/next controls

Large example datasets (100+ images) rendered every tile at once, so the
caption review grid grew unbounded. Show 24 images per page with < >
chevrons and an x-y of N indicator; a new dataset or refresh resets to
the first page.

* Offer example datasets in the Train dropdown with previews

Add an Examples group to the training-images dropdown that imports a
curated dataset in one pick, alongside the existing cards. Cards now show
up to three preview thumbnails pulled from the public HF datasets-server
so the set is visible before download. Hide the trigger prompt when every
image already has a caption (a captioned style set needs no trigger), and
turn the training-settings toggle into a ghost button with a rotating
chevron.

* Clamp the training base repo to the selected family

The base-model select's state could briefly hold the previous family's
repo after a family switch (the reseed effect runs a beat later, and a
value with no matching option makes the browser display the first option
anyway). The request then carried the stale repo: picking Qwen or Z-Image
still sent black-forest-labs/FLUX.1-dev and surfaced FLUX's gated-repo
error under the wrong family. Derive an effectiveBase clamped to the
current family's repos and use it for the select value, the start
request, and the deploy fallback.

Also move the Trigger prompt above Adapter name: the trigger describes
the dataset, the name only labels the output.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-03 12:26:27 -07:00 committed by GitHub
commit 17e55db92b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 1886 additions and 72 deletions

View file

@ -257,10 +257,14 @@ export async function fetchGalleryObjectUrl(url: string): Promise<string> {
return URL.createObjectURL(await res.blob());
}
// ── Diffusion (SDXL) LoRA training ────────────────────────────────────────────
// ── Diffusion LoRA training ───────────────────────────────────────────────────
// Mirrors DiffusionTrainingStartRequest on the backend; only the paths are required.
export interface DiffusionTrainingStartRequest {
base_model: string;
// Explicit family (sdxl / flux.1 / qwen-image / z-image). Optional: the backend
// resolves it from base_model when omitted, but the Train tab always sends it so a
// custom base still trains under the intended family.
model_family?: string | null;
data_dir: string;
output_dir: string;
instance_prompt?: string | null;
@ -277,10 +281,18 @@ export interface DiffusionTrainingStartRequest {
mixed_precision?: "bf16" | "fp16" | "no";
gradient_checkpointing?: boolean;
lr_scheduler?: string;
// Forwarded to StableDiffusionXLPipeline.from_pretrained for a gated/private base repo.
// Forwarded to the pipeline's from_pretrained for a gated/private base repo (e.g. FLUX).
hf_token?: string | null;
}
// Paired step-indexed history arrays for the live loss + LR charts. `lr` entries may be
// null so a sparse learning-rate series still aligns with `steps` by index.
export interface DiffusionMetricHistory {
steps: number[];
loss: number[];
lr: Array<number | null>;
}
// A snapshot of the current diffusion training job (GET /api/train/diffusion/status).
export interface DiffusionTrainingStatus {
active: boolean;
@ -298,6 +310,16 @@ export interface DiffusionTrainingStatus {
lora_path: string | null;
started_at: number | null;
updated_at: number | null;
// Where the trained adapter was mirrored into the Studio LoRA catalog, and the family /
// base it was trained from -- lets the Train tab deploy the adapter onto the right base.
catalog_path?: string | null;
family?: string | null;
base_model?: string | null;
// Live throughput + peak VRAM (from the trainer's progress events).
samples_per_second?: number | null;
peak_memory_gb?: number | null;
// Bounded step/loss/lr history for the live charts.
metric_history?: DiffusionMetricHistory | null;
}
export async function startDiffusionTraining(
@ -328,11 +350,33 @@ export interface DiffusionDatasetSummary {
caption_count: number;
}
// Per-family training defaults (from GET /api/train/diffusion/info families[], added by
// the DiT-trainer backend). Absent on older backends; the Train tab falls back to a
// hardcoded family list when it is.
export interface DiffusionTrainableFamily {
name: string;
label: string;
default_base: string;
base_repos: string[];
defaults?: {
lora_rank?: number;
learning_rate?: number;
resolution?: number;
train_steps?: number;
train_batch_size?: number;
mixed_precision?: "bf16" | "fp16" | "no";
} | null;
vram_note?: string | null;
gated?: boolean | null;
}
// Where diffusion training reads/writes on this Studio, plus usable dataset folders.
export interface DiffusionTrainingInfo {
datasets_root: string;
outputs_root: string;
datasets: DiffusionDatasetSummary[];
// Added by the multi-family trainer backend; tolerate its absence.
families?: DiffusionTrainableFamily[];
}
export async function getDiffusionTrainingInfo(): Promise<DiffusionTrainingInfo> {
@ -357,3 +401,115 @@ export async function uploadDiffusionDataset(
await authFetch("/api/train/diffusion/dataset", { method: "POST", body: form }),
);
}
// ── Dataset labeling + example imports (GET/PUT/DELETE .../dataset/{name}/...) ──
// One image in a training dataset folder, with its resolved caption. `caption_source`
// records where the caption came from ("metadata" beats a per-image "sidecar"; "none"
// when uncaptioned) so the labeling grid can highlight images that still need one.
export interface DiffusionDatasetImageRecord {
filename: string;
caption: string | null;
caption_source: "sidecar" | "metadata" | "none";
width: number;
height: number;
size_bytes: number;
}
export interface DiffusionDatasetImages {
name: string;
path: string;
images: DiffusionDatasetImageRecord[];
}
/** List every image in a dataset folder (including uncaptioned ones) for the grid. */
export async function listDiffusionDatasetImages(
name: string,
): Promise<DiffusionDatasetImages> {
return parseJson(
await authFetch(`/api/train/diffusion/dataset/${encodeURIComponent(name)}/images`),
);
}
/** Build the auth-protected thumbnail URL for a dataset image. Fetch it via
* fetchGalleryObjectUrl (Bearer auth) into an object URL; it can't be a plain <img src>. */
export function diffusionDatasetImageUrl(
name: string,
filename: string,
thumb = 256,
): string {
const q = thumb > 0 ? `?thumb=${thumb}` : "";
return `/api/train/diffusion/dataset/${encodeURIComponent(name)}/image/${encodeURIComponent(filename)}${q}`;
}
/** Write (or, when blank, clear) a per-image caption sidecar. Returns the updated record. */
export async function setDiffusionDatasetCaption(
name: string,
filename: string,
caption: string,
): Promise<DiffusionDatasetImageRecord> {
return parseJson(
await authFetch(
`/api/train/diffusion/dataset/${encodeURIComponent(name)}/caption/${encodeURIComponent(filename)}`,
{
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ caption }),
},
),
);
}
/** Delete an image (and its caption + thumbnail) from a dataset folder. */
export async function deleteDiffusionDatasetImage(
name: string,
filename: string,
): Promise<void> {
const res = await authFetch(
`/api/train/diffusion/dataset/${encodeURIComponent(name)}/image/${encodeURIComponent(filename)}`,
{ method: "DELETE" },
);
if (!res.ok) throw new Error(await readFastApiError(res));
}
// A curated, one-click-importable example image dataset. `license` is shown verbatim so
// users see the terms before importing; `suggested_trigger` seeds the trigger prompt.
export interface DiffusionDatasetExample {
id: string;
label: string;
repo: string;
description: string;
license: string;
image_cap: number;
suggested_trigger?: string | null;
}
export async function listDiffusionDatasetExamples(): Promise<DiffusionDatasetExample[]> {
const data = await parseJson<{ examples: DiffusionDatasetExample[] }>(
await authFetch("/api/train/diffusion/dataset-examples"),
);
return data.examples;
}
export interface DiffusionDatasetImportResult {
name: string;
path: string;
image_count: number;
caption_count: number;
imported: number;
license: string;
source_repo: string;
}
/** Materialize a curated example dataset (by id) into a Studio dataset folder. */
export async function importDiffusionDatasetExample(
id: string,
name?: string,
): Promise<DiffusionDatasetImportResult> {
return parseJson(
await authFetch("/api/train/diffusion/dataset/import-example", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ id, name }),
}),
);
}

View file

@ -3,7 +3,6 @@
import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react";
import {
AiMagicIcon,
ArrowLeftRightIcon,
ArrowReloadHorizontalIcon,
Delete02Icon,
@ -32,6 +31,7 @@ import {
import { Slider } from "@/components/ui/slider";
import { Spinner } from "@/components/ui/spinner";
import { Switch } from "@/components/ui/switch";
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { Textarea } from "@/components/ui/textarea";
import { InfoHint } from "@/components/ui/info-hint";
import { ModelSelector } from "@/components/assistant-ui/model-selector";
@ -67,7 +67,7 @@ import {
loadDiffusionModel,
unloadDiffusionModel,
} from "./api";
import { DiffusionTrainDialog } from "./diffusion-train-dialog";
import { DiffusionTrainPanel } from "./train/diffusion-train-panel";
// Curated diffusion GGUFs the picker recommends. The backend resolves each one's
// pipeline + base diffusers repo from its repo id, so the rail just lists them;
@ -464,34 +464,41 @@ function Field({
function AdvancedSelect({
label,
hint,
desc,
value,
onValueChange,
options,
}: {
label: string;
hint?: ReactNode;
// A short always-visible description under the row (the hint tooltip carries the full
// detail). Used for controls whose label alone does not convey what they do.
desc?: string;
value: string;
onValueChange: (v: string) => void;
options: Array<[string, string]>;
}) {
return (
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
{label}
{hint && <InfoHint>{hint}</InfoHint>}
</span>
<Select value={value} onValueChange={onValueChange}>
<SelectTrigger className="h-8 w-[160px] text-xs">
<SelectValue />
</SelectTrigger>
<SelectContent>
{options.map(([v, l]) => (
<SelectItem key={v} value={v} className="text-xs">
{l}
</SelectItem>
))}
</SelectContent>
</Select>
<div className="flex flex-col gap-1">
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
{label}
{hint && <InfoHint>{hint}</InfoHint>}
</span>
<Select value={value} onValueChange={onValueChange}>
<SelectTrigger className="h-8 w-[160px] text-xs">
<SelectValue />
</SelectTrigger>
<SelectContent>
{options.map(([v, l]) => (
<SelectItem key={v} value={v} className="text-xs">
{l}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
{desc && <p className="text-[11px] leading-snug text-muted-foreground/70">{desc}</p>}
</div>
);
}
@ -948,8 +955,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
// offers. Applied at generate time; available adapters are refreshed per loaded family.
const [loras, setLoras] = useState<LoraSpecInput[]>([]);
const [availableLoras, setAvailableLoras] = useState<DiffusionLoraInfo[]>([]);
// "Train a LoRA" dialog (SDXL). Independent of the loaded generation model.
const [trainOpen, setTrainOpen] = useState(false);
// Page mode: "create" is the generation workspace; "train" is the full-page LoRA
// training workspace. Independent of the loaded generation model.
const [pageMode, setPageMode] = useState<"create" | "train">("create");
// Bumped when a training run completes, to force the LoRA discovery effect to rescan so
// a freshly-trained adapter appears in the picker without a model reload.
const [loraRefreshKey, setLoraRefreshKey] = useState(0);
@ -1025,6 +1033,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
// loaded, so the poll must roll the label back rather than advertise the failed
// quant. `{ prev }` distinguishes "revert to null" from "nothing pending".
const quantRevert = useRef<{ prev: string | null } | null>(null);
// A trained adapter awaiting deployment: after Deploy loads the base, the LoRA discovery
// effect applies this once the model is loaded + LoRA-capable for the matching family.
const pendingDeploy = useRef<{ loraId: string; family: string } | null>(null);
const dismissLoadToast = useCallback(() => {
if (loadToastId.current != null) toast.dismiss(loadToastId.current);
@ -1064,6 +1075,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
setLoras([]);
}
prevLoraFamilyRef.current = fam;
// A just-deployed adapter: now that the base is loaded + LoRA-capable, apply it (after
// the family-swap clear above so it isn't wiped). Only when the family matches what it
// was trained for; otherwise warn instead of silently applying an incompatible adapter.
const deploy = pendingDeploy.current;
if (deploy) {
pendingDeploy.current = null;
if (!deploy.family || deploy.family === fam) {
setLoras([{ id: deploy.loraId, weight: 1 }]);
} else {
toast.error(
`The trained adapter is for ${deploy.family}, but the loaded model is ` +
`${fam ?? "a different family"}, so it was not applied.`,
);
}
}
let cancelled = false;
listDiffusionLoras(status?.family ?? undefined)
.then((list) => {
@ -1565,6 +1591,36 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
[busy, handleLoad, quant],
);
// Deploy a freshly-trained adapter from the Train tab: switch to Create, load the base as
// a pipeline, and queue the adapter so the LoRA discovery effect applies it once the base
// is loaded + LoRA-capable. Seeds the prompt with the trigger phrase when provided.
const handleDeployAdapter = useCallback(
(args: { baseRepo: string; family: string; catalogPath: string; trigger: string }) => {
if (busy !== null) {
toast.error("Finish the current model load before deploying the adapter.");
return;
}
// The picker keys a local adapter by its filename stem (see diffusion_lora scan).
const base = args.catalogPath.replace(/\\/g, "/").split("/").pop() ?? "";
const stem = base.replace(/\.(safetensors|gguf)$/i, "");
if (!stem) {
toast.error("Could not resolve the trained adapter's name.");
return;
}
pendingDeploy.current = { loraId: stem, family: args.family };
if (args.trigger.trim()) setPrompt(args.trigger.trim());
setPageMode("create");
setQuant(null);
const d = defaultsFor(args.baseRepo);
setSteps(d.steps);
setGuidance(d.guidance);
void handleLoad(args.baseRepo, { kind: "pipeline" }).then((started) => {
if (!started) pendingDeploy.current = null;
});
},
[busy, handleLoad],
);
const handleUnload = useCallback(async () => {
// Ejecting cancels any in-flight replacement load on the backend, so tear
// down its client-side tracking too: the load poll reschedules on phase
@ -1815,7 +1871,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
to GGUF (or nothing loaded) and otherwise show why it is unavailable. */}
{!status?.loaded || status.model_kind === "gguf" ? (
<AdvancedSelect
label="GGUF speed mode"
label="GGUF compute"
desc="Off runs the GGUF as-is. INT8/FP8/FP4 dequantise the transformer onto low-precision tensor cores for a faster step, at the cost of a larger download and more VRAM."
hint="Optional speed-up for GGUF models. Off runs the GGUF as-is. FP8/INT8/FP4 instead load the FULL base model and quantise its transformer onto low-precision tensor cores: faster per step, but a larger download and more VRAM, and it falls back to the GGUF if it can't fit. Needs CUDA."
value={transformerQuant}
onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}
@ -1829,9 +1886,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
]}
/>
) : (
<div className="flex items-center justify-between gap-2 text-xs">
<span className="text-muted-foreground">GGUF speed mode</span>
<span className="text-muted-foreground/70">GGUF models only</span>
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
GGUF compute
</span>
<span className="text-xs text-muted-foreground/60">GGUF models only</span>
</div>
)}
<AdvancedSelect
@ -1910,54 +1969,57 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
onOpenChange={(o) => setSelectorOpen(active && o)}
/>
<div className="flex items-center gap-2">
{/* Train a LoRA (SDXL): opens a self-contained dialog; available regardless of
whether a generation model is loaded. */}
<Button
type="button"
variant="ghost"
size="sm"
className="h-[34px]"
onClick={() => setTrainOpen(true)}
title="Teach SDXL your own style or subject from a folder of images"
>
<HugeiconsIcon icon={AiMagicIcon} className="mr-1.5 size-3.5" />
Train LoRA
</Button>
{/* Create | Train page-mode switch, next to the model selector. Create is the
generation workspace; Train is the full-page LoRA training workspace. */}
<Tabs value={pageMode} onValueChange={(v) => setPageMode(v as "create" | "train")}>
<TabsList className="h-[34px]">
<TabsTrigger value="create" className="w-[64px]">
Create
</TabsTrigger>
<TabsTrigger value="train" className="w-[64px]">
Train
</TabsTrigger>
</TabsList>
</Tabs>
{/* Single fixed toggle for the right-docked Advanced panel (mirrors Chat's settings
toggle, same icon in both states so it never moves). Highlighted when open. */}
<button
type="button"
onClick={() => setAdvancedOpen((o) => !o)}
aria-label={advancedOpen ? "Hide advanced options" : "Show advanced options"}
aria-pressed={advancedOpen}
title="Advanced options"
className={cn(
"flex h-[34px] w-[34px] items-center justify-center rounded-full transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
advancedOpen
? "bg-muted text-foreground"
: "text-muted-foreground hover:bg-muted hover:text-foreground",
)}
>
<HugeiconsIcon icon={LayoutAlignRightIcon} className="size-4" />
</button>
toggle, same icon in both states so it never moves). Highlighted when open.
Only meaningful in Create mode (load-time tuning), so hidden while training. */}
{pageMode === "create" && (
<button
type="button"
onClick={() => setAdvancedOpen((o) => !o)}
aria-label={advancedOpen ? "Hide advanced options" : "Show advanced options"}
aria-pressed={advancedOpen}
title="Advanced options"
className={cn(
"flex h-[34px] w-[34px] items-center justify-center rounded-full transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
advancedOpen
? "bg-muted text-foreground"
: "text-muted-foreground hover:bg-muted hover:text-foreground",
)}
>
<HugeiconsIcon icon={LayoutAlignRightIcon} className="size-4" />
</button>
)}
</div>
</div>
<DiffusionTrainDialog
open={active && trainOpen}
onOpenChange={setTrainOpen}
defaultBaseModel={
status?.family === "sdxl"
? // Prefer base_repo (the full diffusers pipeline) over repo_id: for a GGUF or
// single-file SDXL load repo_id is the checkpoint path, which the trainer's
// from_pretrained cannot open. base_repo is the companion pipeline.
status?.base_repo ?? status?.repo_id ?? undefined
: undefined
}
onTrainingComplete={() => setLoraRefreshKey((k) => k + 1)}
/>
{/* Controls rail + preview canvas. Padding mirrors the other tabs
(Export, Data Recipes): px-5 / sm:px-9, with a roomy bottom. */}
{/* Train mode: the full-page training workspace. Kept unmounted in Create mode so its
polling stops; Create's own state (gallery, model, workflow) is untouched. */}
{pageMode === "train" ? (
<DiffusionTrainPanel
active={active && pageMode === "train"}
loadedFamily={status?.family ?? null}
loadedBaseRepo={
// Prefer base_repo (the full diffusers pipeline) over repo_id: for a GGUF or
// single-file load repo_id is the checkpoint path, not a trainable base.
status?.base_repo ?? status?.repo_id ?? null
}
onTrainingComplete={() => setLoraRefreshKey((k) => k + 1)}
onDeploy={handleDeployAdapter}
/>
) : (
/* Controls rail + preview canvas. Padding mirrors the other tabs
(Export, Data Recipes): px-5 / sm:px-9, with a roomy bottom. */
<div className="flex min-h-0 min-w-0 flex-1 gap-4 overflow-hidden px-5 pb-8 sm:px-9">
{/* The controls rail. Plain card (the gray surface) with no header
the prompt + Generate button make the panel self-explanatory. */}
@ -2640,6 +2702,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
</div>
)}
</div>
)}
</div>
);
}

View file

@ -0,0 +1,299 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useCallback, useEffect, useRef, useState } from "react";
import { ArrowLeft01Icon, ArrowRight01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Button } from "@/components/ui/button";
import { Spinner } from "@/components/ui/spinner";
import { Textarea } from "@/components/ui/textarea";
import { cn } from "@/lib/utils";
import { toast } from "@/lib/toast";
// One batch of images shown at a time in the labeling grid; larger sets page with < >.
const PAGE_SIZE = 24;
import {
type DiffusionDatasetImageRecord,
deleteDiffusionDatasetImage,
diffusionDatasetImageUrl,
fetchGalleryObjectUrl,
listDiffusionDatasetImages,
setDiffusionDatasetCaption,
} from "../api";
// One tile: an auth-fetched thumbnail (object URL, revoked on unmount) plus a caption
// Textarea saved on blur. Uncaptioned tiles get a highlighted ring so a user labeling a
// small set can see at a glance what still needs a caption.
function LabelTile({
dataset,
record,
onSaved,
onDeleted,
}: {
dataset: string;
record: DiffusionDatasetImageRecord;
onSaved: (rec: DiffusionDatasetImageRecord) => void;
onDeleted: (filename: string) => void;
}) {
const [thumb, setThumb] = useState<string | null>(null);
const [caption, setCaption] = useState(record.caption ?? "");
const [saving, setSaving] = useState(false);
const [savedTick, setSavedTick] = useState(false);
const [deleting, setDeleting] = useState(false);
// The last caption we persisted, so blur only writes when the text actually changed.
const persisted = useRef(record.caption ?? "");
// Load the thumbnail once; revoke the object URL on unmount to avoid a leak.
useEffect(() => {
let url: string | null = null;
let cancelled = false;
fetchGalleryObjectUrl(diffusionDatasetImageUrl(dataset, record.filename, 256))
.then((u) => {
if (cancelled) {
URL.revokeObjectURL(u);
return;
}
url = u;
setThumb(u);
})
.catch(() => {
/* a missing thumbnail just leaves the placeholder */
});
return () => {
cancelled = true;
if (url) URL.revokeObjectURL(url);
};
}, [dataset, record.filename]);
const save = useCallback(async () => {
const next = caption.trim();
if (next === persisted.current.trim()) return;
setSaving(true);
try {
const updated = await setDiffusionDatasetCaption(dataset, record.filename, next);
persisted.current = updated.caption ?? "";
setCaption(updated.caption ?? "");
onSaved(updated);
setSavedTick(true);
window.setTimeout(() => setSavedTick(false), 1500);
} catch (e) {
toast.error(e instanceof Error ? e.message : "Failed to save caption");
} finally {
setSaving(false);
}
}, [caption, dataset, record.filename, onSaved]);
const remove = useCallback(async () => {
setDeleting(true);
try {
await deleteDiffusionDatasetImage(dataset, record.filename);
onDeleted(record.filename);
} catch (e) {
toast.error(e instanceof Error ? e.message : "Failed to delete image");
setDeleting(false);
}
}, [dataset, record.filename, onDeleted]);
const uncaptioned = caption.trim().length === 0;
return (
<div
className={cn(
"flex flex-col gap-1.5 rounded-lg border p-2",
uncaptioned ? "border-amber-500/60 bg-amber-500/5" : "border-border",
)}
>
<div className="relative aspect-square w-full overflow-hidden rounded-md bg-muted">
{thumb ? (
<img src={thumb} alt={record.filename} className="size-full object-cover" />
) : (
<div className="flex size-full items-center justify-center">
<Spinner className="size-4 text-muted-foreground" />
</div>
)}
<button
type="button"
onClick={remove}
disabled={deleting}
title="Remove this image from the dataset"
className="absolute right-1 top-1 rounded-md bg-background/80 px-1.5 py-0.5 text-[11px] text-muted-foreground opacity-0 transition-opacity hover:bg-background hover:text-destructive group-hover:opacity-100 focus:opacity-100"
>
{deleting ? "..." : "Remove"}
</button>
</div>
<Textarea
value={caption}
onChange={(e) => setCaption(e.target.value)}
onBlur={() => void save()}
rows={2}
spellCheck={false}
placeholder="Describe this image..."
className="min-h-[3rem] resize-none text-[11px]"
aria-label={`Caption for ${record.filename}`}
/>
<div className="flex items-center justify-between text-[10px] text-muted-foreground">
<span className="truncate" title={record.filename}>
{record.filename}
</span>
{saving ? (
<span>Saving...</span>
) : savedTick ? (
<span className="text-emerald-500">Saved</span>
) : uncaptioned ? (
<span className="text-amber-600">No caption</span>
) : null}
</div>
</div>
);
}
// A responsive grid over a dataset folder's images with per-image caption editing. Fetches
// the image list on open (and whenever `refreshKey` changes, e.g. after an upload/import).
export function DatasetLabelingGrid({
dataset,
refreshKey = 0,
onCountsChanged,
}: {
dataset: string;
refreshKey?: number;
// Fired after a caption save or delete so the parent can refresh dataset counts.
onCountsChanged?: () => void;
}) {
const [records, setRecords] = useState<DiffusionDatasetImageRecord[] | null>(null);
const [error, setError] = useState<string | null>(null);
const [page, setPage] = useState(0);
useEffect(() => {
let cancelled = false;
setRecords(null);
setError(null);
setPage(0); // a new dataset (or refresh) always starts at the first batch
listDiffusionDatasetImages(dataset)
.then((r) => {
if (!cancelled) setRecords(r.images);
})
.catch((e) => {
if (!cancelled) setError(e instanceof Error ? e.message : "Failed to list images");
});
return () => {
cancelled = true;
};
}, [dataset, refreshKey]);
const onSaved = useCallback(
(rec: DiffusionDatasetImageRecord) => {
setRecords((cur) =>
cur ? cur.map((r) => (r.filename === rec.filename ? rec : r)) : cur,
);
onCountsChanged?.();
},
[onCountsChanged],
);
const onDeleted = useCallback(
(filename: string) => {
setRecords((cur) => (cur ? cur.filter((r) => r.filename !== filename) : cur));
onCountsChanged?.();
},
[onCountsChanged],
);
if (error) {
return <p className="text-[11px] text-destructive">{error}</p>;
}
if (records === null) {
return (
<div className="flex items-center gap-2 py-4 text-xs text-muted-foreground">
<Spinner className="size-4" /> Loading images...
</div>
);
}
if (records.length === 0) {
return <p className="text-[11px] text-muted-foreground">This dataset has no images yet.</p>;
}
const uncaptioned = records.filter((r) => !r.caption || r.caption.trim() === "").length;
const total = records.length;
const pageCount = Math.max(1, Math.ceil(total / PAGE_SIZE));
const clampedPage = Math.min(page, pageCount - 1);
const start = clampedPage * PAGE_SIZE;
const pageRecords = records.slice(start, start + PAGE_SIZE);
return (
<div className="flex flex-col gap-2">
<div className="flex items-center justify-between gap-2 text-[11px] text-muted-foreground">
<span className="min-w-0 truncate">
{total} image{total === 1 ? "" : "s"}
{uncaptioned > 0 ? ` · ${uncaptioned} without a caption` : " · all captioned"}
</span>
{pageCount > 1 && (
<div className="flex shrink-0 items-center gap-1">
<span className="tabular-nums">
{start + 1}-{Math.min(start + PAGE_SIZE, total)} of {total}
</span>
<Button
type="button"
variant="ghost"
size="icon"
className="size-6"
disabled={clampedPage === 0}
onClick={() => setPage((p) => Math.max(0, p - 1))}
aria-label="Previous images"
>
<HugeiconsIcon icon={ArrowLeft01Icon} className="size-3.5" />
</Button>
<Button
type="button"
variant="ghost"
size="icon"
className="size-6"
disabled={clampedPage >= pageCount - 1}
onClick={() => setPage((p) => Math.min(pageCount - 1, p + 1))}
aria-label="Next images"
>
<HugeiconsIcon icon={ArrowRight01Icon} className="size-3.5" />
</Button>
</div>
)}
</div>
<div className="group grid max-h-[420px] grid-cols-2 gap-2 overflow-y-auto pr-1 sm:grid-cols-3">
{pageRecords.map((r) => (
<LabelTile
key={r.filename}
dataset={dataset}
record={r}
onSaved={onSaved}
onDeleted={onDeleted}
/>
))}
</div>
</div>
);
}
// A tiny standalone control used by the panel to keep grid-refresh wiring in one place.
export function LabelingGridToggle({
count,
open,
onToggle,
}: {
count: number;
open: boolean;
onToggle: () => void;
}) {
return (
<Button
type="button"
variant="ghost"
size="sm"
className="h-7 w-fit px-1 text-xs text-muted-foreground hover:text-foreground"
onClick={onToggle}
aria-expanded={open}
>
{open ? "Hide captions" : `Review captions (${count} image${count === 1 ? "" : "s"})`}
</Button>
);
}

View file

@ -0,0 +1,127 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useEffect, useState } from "react";
import { Spinner } from "@/components/ui/spinner";
import {
diffusionDatasetImageUrl,
fetchGalleryObjectUrl,
listDiffusionDatasetImages,
} from "../api";
const MAX_TILES = 8;
// A single thumbnail tile: auth-fetches its object URL and revokes it on unmount.
function ShowcaseTile({ dataset, filename }: { dataset: string; filename: string }) {
const [url, setUrl] = useState<string | null>(null);
useEffect(() => {
let obj: string | null = null;
let cancelled = false;
fetchGalleryObjectUrl(diffusionDatasetImageUrl(dataset, filename, 256))
.then((u) => {
if (cancelled) {
URL.revokeObjectURL(u);
return;
}
obj = u;
setUrl(u);
})
.catch(() => {
/* a missing thumbnail just leaves the placeholder */
});
return () => {
cancelled = true;
if (obj) URL.revokeObjectURL(obj);
};
}, [dataset, filename]);
return (
<div className="size-14 shrink-0 overflow-hidden rounded-md bg-muted">
{url ? (
<img src={url} alt={filename} className="size-full object-cover" />
) : (
<div className="flex size-full items-center justify-center">
<Spinner className="size-3.5 text-muted-foreground" />
</div>
)}
</div>
);
}
// A compact preview strip of a dataset's images: up to 8 sampled thumbnails, plus a
// "+N more" tile that opens the full labeling grid. Refreshes when the dataset selection or
// `refreshKey` changes (e.g. after an upload/import). The whole strip is a button so a click
// anywhere jumps to Review captions.
export function DatasetShowcase({
dataset,
imageCount,
refreshKey = 0,
onBrowse,
}: {
dataset: string;
imageCount: number;
refreshKey?: number;
onBrowse: () => void;
}) {
const [names, setNames] = useState<string[] | null>(null);
useEffect(() => {
let cancelled = false;
setNames(null);
listDiffusionDatasetImages(dataset)
.then((r) => {
if (cancelled) return;
// Sample up to MAX_TILES evenly across the folder so the strip represents the whole
// set, not just the first few files.
const all = r.images.map((im) => im.filename);
if (all.length <= MAX_TILES) {
setNames(all);
return;
}
const stride = all.length / MAX_TILES;
const picked: string[] = [];
for (let i = 0; i < MAX_TILES; i++) picked.push(all[Math.floor(i * stride)]);
setNames(picked);
})
.catch(() => {
if (!cancelled) setNames([]);
});
return () => {
cancelled = true;
};
}, [dataset, refreshKey]);
if (names !== null && names.length === 0) return null;
const remaining = imageCount - (names?.length ?? 0);
return (
<button
type="button"
onClick={onBrowse}
title="Browse and caption these images"
className="flex w-full items-center gap-1.5 overflow-x-auto rounded-lg border border-border bg-muted/20 p-1.5 text-left transition-colors hover:border-foreground/20"
>
{names === null ? (
<div className="flex h-14 items-center gap-2 px-2 text-[11px] text-muted-foreground">
<Spinner className="size-3.5" /> Loading preview...
</div>
) : (
<>
{names.map((n) => (
<ShowcaseTile key={n} dataset={dataset} filename={n} />
))}
{remaining > 0 && (
<div className="flex size-14 shrink-0 flex-col items-center justify-center rounded-md bg-muted text-muted-foreground">
<span className="text-sm font-medium">+{remaining}</span>
<span className="text-[9px]">more</span>
</div>
)}
</>
)}
</button>
);
}

View file

@ -0,0 +1,143 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { type ReactElement, useMemo } from "react";
import type { TrainingSeriesPoint } from "@/features/training";
// The loss + LR cards are pure presentational (props only), so reuse them directly. We do
// NOT reuse ChartsSection/ChartsContent: those also render Grad Norm and an Eval Loss card,
// which are meaningless for diffusion LoRA training and showed as an empty card and an
// "Evaluation not configured" placeholder. This is a diffusion-only two-card layout.
// eslint-disable-next-line no-restricted-imports
import { LearningRateChartCard } from "@/features/studio/sections/charts/learning-rate-chart-card";
// eslint-disable-next-line no-restricted-imports
import { TrainingLossChartCard } from "@/features/studio/sections/charts/training-loss-chart-card";
// eslint-disable-next-line no-restricted-imports
import {
MAX_RENDER_POINTS,
buildStepTicks,
buildYDomain,
compressSeries,
ema,
} from "@/features/studio/sections/charts/utils";
// Fixed presentation defaults (the LLM tab exposes these via a settings sheet; here we pick
// sensible constants): EMA smoothing on, linear scale, raw + smoothed + average lines shown,
// no outlier trimming (diffusion loss is naturally noisy, not spiky-with-outliers).
const SMOOTHING = 0.8;
function toLossItems(series: TrainingSeriesPoint[]): { step: number; loss: number }[] {
return series
.filter((p) => Number.isFinite(p.value))
.map((p) => ({ step: p.step, loss: p.value }));
}
// The x-domain that spans all points (the LLM tab supports a scrollable window; a training
// run here is short enough to always show the whole thing).
function fullStepDomain(steps: number[]): [number, number] {
if (steps.length === 0) return [0, 1];
const min = steps[0];
const max = steps[steps.length - 1];
if (min === max) return [min, min + 4];
if (max - min < 6) return [Math.max(0, max - 6), max];
return [min, max];
}
// A diffusion-only metrics view: just Training Loss and Learning Rate, side by side, with a
// note under the loss card explaining why per-step loss looks noisy.
export function DiffusionCharts({
lossHistory,
lrHistory,
}: {
lossHistory: TrainingSeriesPoint[];
lrHistory: TrainingSeriesPoint[];
}): ReactElement | null {
const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]);
const smoothed = useMemo(
() => (lossItems.length > 0 ? ema(lossItems, SMOOTHING) : []),
[lossItems],
);
const reducedLoss = useMemo(
() => compressSeries(smoothed, MAX_RENDER_POINTS),
[smoothed],
);
const lossData = useMemo(
() =>
reducedLoss.map((p) => ({
...p,
displayLoss: p.loss,
displaySmoothed: p.smoothed,
})),
[reducedLoss],
);
const lrData = useMemo(
() =>
compressSeries(
lrHistory
.filter((p) => Number.isFinite(p.value))
.map((p) => ({ step: p.step, lr: p.value, displayLr: p.value })),
MAX_RENDER_POINTS,
),
[lrHistory],
);
const steps = useMemo(() => {
const set = new Set<number>();
for (const p of lossData) set.add(p.step);
for (const p of lrData) set.add(p.step);
return Array.from(set).sort((a, b) => a - b);
}, [lossData, lrData]);
const stepDomain = useMemo(() => fullStepDomain(steps), [steps]);
const xAxisTicks = useMemo(
() => buildStepTicks(stepDomain[0], stepDomain[1]),
[stepDomain],
);
const lossDomain = useMemo(
() => buildYDomain(lossData.flatMap((p) => [p.displayLoss, p.displaySmoothed])),
[lossData],
);
const lrDomain = useMemo(
() => buildYDomain(lrData.map((p) => p.displayLr)),
[lrData],
);
const avgRaw =
lossItems.length > 0
? +(lossItems.reduce((s, p) => s + p.loss, 0) / lossItems.length).toFixed(4)
: 0;
if (lossItems.length === 0 && lrData.length === 0) return null;
return (
<div className="grid grid-cols-1 gap-4 lg:grid-cols-2">
<div className="flex flex-col gap-1">
<TrainingLossChartCard
data={lossData}
domain={lossDomain}
visibleStepDomain={stepDomain}
xAxisTicks={xAxisTicks}
avgRaw={avgRaw}
avgDisplay={avgRaw}
showRaw={true}
showSmoothed={true}
showAvgLine={true}
scale="linear"
/>
<p className="px-1 text-[11px] leading-snug text-muted-foreground">
Per-step loss is noisy by design: every step samples a random noise level. Watch
the smoothed line for the trend, not the raw jitter.
</p>
</div>
<LearningRateChartCard
data={lrData}
domain={lrDomain}
visibleStepDomain={stepDomain}
xAxisTicks={xAxisTicks}
scale="linear"
/>
</div>
);
}

View file

@ -0,0 +1,875 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { ArrowDown01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import type { TrainingSeriesPoint } from "@/features/training";
// eslint-disable-next-line no-restricted-imports -- matches images-page.tsx's token access
import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store";
import { cn } from "@/lib/utils";
import { toast } from "@/lib/toast";
import {
type DiffusionDatasetExample,
type DiffusionTrainableFamily,
type DiffusionTrainingInfo,
type DiffusionTrainingStatus,
getDiffusionTrainingInfo,
getDiffusionTrainingStatus,
listDiffusionDatasetExamples,
startDiffusionTraining,
stopDiffusionTraining,
uploadDiffusionDataset,
} from "../api";
import { DatasetLabelingGrid, LabelingGridToggle } from "./dataset-labeling-grid";
import { DatasetShowcase } from "./dataset-showcase";
import { DiffusionCharts } from "./diffusion-charts";
import { ExampleDatasetCards, runExampleImport } from "./example-dataset-cards";
// The families the Train tab can train, in the popularity order the user asked for. This is
// the fallback used when the backend's /info does not yet report families (older backend);
// when it does, its list wins and these labels/notes fill any gaps.
type FamilyPreset = {
name: string;
label: string;
base_repos: string[];
defaults: { rank: number; lr: number; resolution: number };
vram_note: string;
gated?: boolean;
};
const FAMILY_PRESETS: FamilyPreset[] = [
{
name: "flux.1",
label: "FLUX.1-dev (12B)",
base_repos: ["black-forest-labs/FLUX.1-dev"],
defaults: { rank: 16, lr: 0.0001, resolution: 512 },
vram_note: "Gated repo - accept the license on Hugging Face and add your HF token. QLoRA (4-bit).",
gated: true,
},
{
name: "qwen-image",
label: "Qwen-Image (20B)",
base_repos: ["unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "Qwen/Qwen-Image"],
defaults: { rank: 16, lr: 0.00005, resolution: 512 },
vram_note: "Largest model - QLoRA (4-bit) on a big GPU. Start at 512px, batch 1.",
},
{
name: "z-image",
label: "Z-Image-Turbo (6B)",
base_repos: ["unsloth/Z-Image-Turbo-unsloth-bnb-4bit", "Tongyi-MAI/Z-Image-Turbo"],
defaults: { rank: 16, lr: 0.0001, resolution: 768 },
vram_note: "Lightest and fastest to train. bf16 only (fp16 is unstable for this family).",
},
{
name: "sdxl",
label: "SDXL (U-Net)",
base_repos: ["stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo"],
defaults: { rank: 16, lr: 0.0001, resolution: 1024 },
vram_note: "The classic text-to-image base. Trains comfortably at 1024px.",
},
];
const CUSTOM_BASE = "__custom__";
const UPLOAD_DATASET = "__upload__";
// Dataset-select option value prefix for a not-yet-imported example; picking it imports.
const EXAMPLE_PREFIX = "example:";
const DATASET_FILE_ACCEPT = ".png,.jpg,.jpeg,.webp,.bmp,.txt,.caption,.jsonl";
const selectClass = "h-8 w-full rounded-md border border-input bg-background px-2 text-xs";
// Merge the backend's reported families (if any) over the presets, keeping the preset
// ordering (popularity) and filling labels/notes/defaults the backend omits.
function mergeFamilies(reported?: DiffusionTrainableFamily[]): FamilyPreset[] {
if (!reported || reported.length === 0) return FAMILY_PRESETS;
const byName = new Map(reported.map((f) => [f.name, f]));
const merged: FamilyPreset[] = FAMILY_PRESETS.map((p) => {
const r = byName.get(p.name);
if (!r) return p;
byName.delete(p.name);
return {
name: p.name,
label: r.label || p.label,
base_repos: r.base_repos?.length ? r.base_repos : p.base_repos,
defaults: {
rank: r.defaults?.lora_rank ?? p.defaults.rank,
lr: r.defaults?.learning_rate ?? p.defaults.lr,
resolution: r.defaults?.resolution ?? p.defaults.resolution,
},
vram_note: r.vram_note || p.vram_note,
gated: r.gated ?? p.gated,
};
});
// Any backend family not in the presets goes last, so a newly added trainer still shows.
for (const r of byName.values()) {
merged.push({
name: r.name,
label: r.label || r.name,
base_repos: r.base_repos ?? [],
defaults: {
rank: r.defaults?.lora_rank ?? 16,
lr: r.defaults?.learning_rate ?? 0.0001,
resolution: r.defaults?.resolution ?? 768,
},
vram_note: r.vram_note ?? "",
gated: r.gated ?? false,
});
}
return merged;
}
// A full-page training workspace: left = configure (family, dataset, labeling, settings),
// right = live run (progress, loss/LR charts, completion + deploy). Kept mounted with the
// page so a long run survives Create/Train tab switches; polling is gated on `active`.
export function DiffusionTrainPanel({
active,
loadedFamily,
loadedBaseRepo,
onTrainingComplete,
onDeploy,
}: {
active: boolean;
// The currently loaded generation model's family / base repo, to preselect a matching
// training base when it is one we can train.
loadedFamily?: string | null;
loadedBaseRepo?: string | null;
// Bump the page's LoRA discovery so a freshly trained adapter appears in the picker.
onTrainingComplete?: () => void;
// Deploy a finished adapter into Create mode: load the base then preselect the adapter.
onDeploy?: (args: {
baseRepo: string;
family: string;
catalogPath: string;
trigger: string;
}) => void;
}) {
const [info, setInfo] = useState<DiffusionTrainingInfo | null>(null);
const families = useMemo(() => mergeFamilies(info?.families), [info?.families]);
const [familyName, setFamilyName] = useState(families[0]?.name ?? "flux.1");
const family = useMemo(
() => families.find((f) => f.name === familyName) ?? families[0],
[families, familyName],
);
const [baseChoice, setBaseChoice] = useState<string>(family?.base_repos[0] ?? "");
const [customBase, setCustomBase] = useState("");
const [dataset, setDataset] = useState<string>(UPLOAD_DATASET);
const [uploadName, setUploadName] = useState("my-images");
const [uploading, setUploading] = useState(false);
const fileInputRef = useRef<HTMLInputElement | null>(null);
const [gridOpen, setGridOpen] = useState(false);
const [gridRefresh, setGridRefresh] = useState(0);
const [examples, setExamples] = useState<DiffusionDatasetExample[]>([]);
const [importingId, setImportingId] = useState<string | null>(null);
const [outputDir, setOutputDir] = useState("");
const [instancePrompt, setInstancePrompt] = useState("");
const [showAdvanced, setShowAdvanced] = useState(false);
const [steps, setSteps] = useState(500);
const [learningRate, setLearningRate] = useState(family?.defaults.lr ?? 0.0001);
const [rank, setRank] = useState(family?.defaults.rank ?? 16);
const [resolution, setResolution] = useState(family?.defaults.resolution ?? 768);
const [batchSize, setBatchSize] = useState(1);
const [precision, setPrecision] = useState<"bf16" | "fp16" | "no">("bf16");
// Track whether the user hand-edited the numeric settings; if not, a family change
// re-seeds them from that family's defaults.
const settingsDirty = useRef(false);
const [starting, setStarting] = useState(false);
const [status, setStatus] = useState<DiffusionTrainingStatus | null>(null);
const refreshInfo = useCallback(async (): Promise<DiffusionTrainingInfo | null> => {
try {
const i = await getDiffusionTrainingInfo();
setInfo(i);
return i;
} catch {
return null;
}
}, []);
// On first activation, load the dataset list and preselect a base matching the loaded
// generation model when it is a trainable family.
useEffect(() => {
if (!active) return;
void refreshInfo().then((i) => {
setDataset((cur) => {
if (cur !== UPLOAD_DATASET && i?.datasets.some((d) => d.name === cur)) return cur;
return i && i.datasets.length > 0 ? i.datasets[0].name : UPLOAD_DATASET;
});
});
}, [active, refreshInfo]);
// Load the curated example list once (for the dropdown group + the cards). Best-effort:
// an older backend without the endpoint just yields no examples.
useEffect(() => {
if (!active) return;
let cancelled = false;
listDiffusionDatasetExamples()
.then((list) => {
if (!cancelled) setExamples(list);
})
.catch(() => {
if (!cancelled) setExamples([]);
});
return () => {
cancelled = true;
};
}, [active]);
// Examples whose folder is not on disk yet: shown in the dropdown's Examples group and as
// cards. An example imports into a folder named after its id, so a matching dataset name
// means it is already imported (and appears as a normal dataset instead).
const importedNames = useMemo(
() => new Set((info?.datasets ?? []).map((d) => d.name)),
[info?.datasets],
);
const pendingExamples = useMemo(
() => examples.filter((ex) => !importedNames.has(ex.id)),
[examples, importedNames],
);
// Import a curated example, then select the resulting folder. Seeds the trigger prompt from
// the example only when the field is meaningful (the import has no captions of its own).
const importExample = useCallback(
async (ex: DiffusionDatasetExample) => {
setImportingId(ex.id);
try {
const res = await runExampleImport(ex);
await refreshInfo();
setDataset(res.name);
setGridOpen(false);
setGridRefresh((k) => k + 1);
if (ex.suggested_trigger && res.caption_count === 0 && !instancePrompt.trim()) {
setInstancePrompt(ex.suggested_trigger);
}
} catch (e) {
toast.error(e instanceof Error ? e.message : "Import failed");
} finally {
setImportingId(null);
}
},
[refreshInfo, instancePrompt],
);
// If the loaded generation model is a trainable family, jump the family selector to it
// once (only when the panel first sees a loaded family).
const seededFromLoaded = useRef(false);
useEffect(() => {
if (seededFromLoaded.current) return;
if (!loadedFamily) return;
if (families.some((f) => f.name === loadedFamily)) {
setFamilyName(loadedFamily);
seededFromLoaded.current = true;
}
}, [loadedFamily, families]);
// Re-seed base + numeric settings from the family's defaults on family change (unless the
// user edited the numbers). Prefer the loaded base repo when it belongs to this family.
useEffect(() => {
if (!family) return;
const preferLoaded =
loadedBaseRepo && family.base_repos.includes(loadedBaseRepo)
? loadedBaseRepo
: family.base_repos[0] ?? CUSTOM_BASE;
setBaseChoice(preferLoaded);
if (!settingsDirty.current) {
setLearningRate(family.defaults.lr);
setRank(family.defaults.rank);
setResolution(family.defaults.resolution);
}
}, [family, loadedBaseRepo]);
// The base actually used everywhere (request, deploy, select value). baseChoice can
// briefly hold another family's repo between a family switch and the reseed effect
// (or if that effect is skipped); a raw <select value> would then DISPLAY the first
// option while the request still carried the stale repo -- the user saw FLUX's gated
// error while another family looked selected. Clamp to the current family's repos.
const effectiveBase =
baseChoice === CUSTOM_BASE || (family?.base_repos ?? []).includes(baseChoice)
? baseChoice
: family?.base_repos[0] ?? CUSTOM_BASE;
const poll = useCallback(async () => {
try {
setStatus(await getDiffusionTrainingStatus());
} catch {
/* best-effort; a failed poll should not surface an error while the tab is open */
}
}, []);
// Poll status while the panel is active.
useEffect(() => {
if (!active) return;
void poll();
const id = window.setInterval(() => void poll(), 1500);
return () => window.clearInterval(id);
}, [active, poll]);
// "Train another" dismisses the completed run's card locally (the backend keeps the
// terminal "completed" status until the next start, so we can't rely on it clearing).
const [dismissedJobId, setDismissedJobId] = useState<string | null>(null);
const running = Boolean(status?.active) || status?.status === "running";
const completed =
status?.status === "completed" && status.job_id !== dismissedJobId;
const pct =
status && status.total_steps > 0
? Math.min(100, Math.round((status.step / status.total_steps) * 100))
: 0;
// Notify the parent exactly once per completed run so it rescans the LoRA picker.
const notifiedComplete = useRef(false);
useEffect(() => {
if (status?.status === "completed" && !notifiedComplete.current) {
notifiedComplete.current = true;
onTrainingComplete?.();
} else if (status?.status === "running" && notifiedComplete.current) {
notifiedComplete.current = false;
}
}, [status?.status, onTrainingComplete]);
const selectedDataset =
dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined;
// A dataset where every image already ships a caption needs no trigger prompt; hide the
// field and explain why. Partial/no captions (or upload mode) still show it.
const fullyCaptioned = Boolean(
selectedDataset &&
selectedDataset.image_count > 0 &&
selectedDataset.caption_count >= selectedDataset.image_count,
);
// Map the backend's paired history arrays into the chart component's {step,value} series.
const lossHistory: TrainingSeriesPoint[] = useMemo(() => {
const h = status?.metric_history;
if (!h) return [];
return h.steps.map((step, i) => ({ step, value: h.loss[i] })).filter((p) => p.value != null);
}, [status?.metric_history]);
const lrHistory: TrainingSeriesPoint[] = useMemo(() => {
const h = status?.metric_history;
if (!h) return [];
return h.steps
.map((step, i) => ({ step, value: h.lr[i] }))
.filter((p): p is TrainingSeriesPoint => p.value != null);
}, [status?.metric_history]);
const onUpload = useCallback(async () => {
const files = Array.from(fileInputRef.current?.files ?? []);
if (files.length === 0) {
toast.error("Choose the images to upload first.");
return;
}
const name = uploadName.trim();
if (!name) {
toast.error("Give the dataset a folder name, e.g. my-style-photos.");
return;
}
setUploading(true);
try {
const res = await uploadDiffusionDataset(name, files);
toast.success(
`Uploaded ${res.uploaded} file${res.uploaded === 1 ? "" : "s"} - ` +
`"${res.name}" now has ${res.image_count} images`,
);
if (fileInputRef.current) fileInputRef.current.value = "";
await refreshInfo();
setDataset(res.name);
setGridRefresh((k) => k + 1);
} catch (e) {
toast.error(e instanceof Error ? e.message : "Upload failed");
} finally {
setUploading(false);
}
}, [uploadName, refreshInfo]);
const onStart = useCallback(async () => {
const baseModel = (effectiveBase === CUSTOM_BASE ? customBase : effectiveBase).trim();
if (!baseModel) {
toast.error("Pick a base model (or fill in the custom repo/path).");
return;
}
if (dataset === UPLOAD_DATASET) {
toast.error("Upload your training images first (or pick an existing dataset).");
return;
}
if (!outputDir.trim()) {
toast.error("Name the adapter (this becomes its folder under Studio outputs).");
return;
}
if (selectedDataset && selectedDataset.caption_count === 0 && !instancePrompt.trim()) {
toast.error(
"These images have no captions - add a trigger prompt so the trainer knows " +
"what to learn (it becomes the caption for every image).",
);
return;
}
if (steps < 1) return toast.error("Steps must be at least 1.");
if (rank < 1) return toast.error("LoRA rank must be at least 1.");
if (resolution < 64 || resolution % 8 !== 0) {
return toast.error("Resolution must be a multiple of 8 and at least 64.");
}
if (batchSize < 1) return toast.error("Batch size must be at least 1.");
if (learningRate <= 0) return toast.error("Learning rate must be greater than 0.");
setStarting(true);
try {
await startDiffusionTraining({
base_model: baseModel,
model_family: family?.name,
data_dir: dataset,
output_dir: outputDir.trim(),
instance_prompt: instancePrompt.trim() || undefined,
resolution,
train_steps: steps,
learning_rate: learningRate,
train_batch_size: batchSize,
lora_rank: rank,
mixed_precision: precision,
hf_token: hfApiToken(getHfToken()) || undefined,
});
toast.success("Training started");
void poll();
} catch (e) {
toast.error(e instanceof Error ? e.message : "Failed to start training");
} finally {
setStarting(false);
}
}, [
effectiveBase,
customBase,
family,
dataset,
selectedDataset,
outputDir,
instancePrompt,
resolution,
steps,
learningRate,
batchSize,
rank,
precision,
poll,
]);
const onStop = useCallback(async () => {
try {
await stopDiffusionTraining();
toast.success("Stop requested; finishing the current step.");
void poll();
} catch (e) {
toast.error(e instanceof Error ? e.message : "Failed to stop training");
}
}, [poll]);
const onDeployClick = useCallback(() => {
if (!status?.catalog_path) {
toast.error("The trained adapter is not available yet.");
return;
}
const baseRepo = status.base_model || (effectiveBase === CUSTOM_BASE ? customBase : effectiveBase);
if (!baseRepo) {
toast.error("Could not determine the base model to load for this adapter.");
return;
}
onDeploy?.({
baseRepo,
family: status.family || family?.name || "",
catalogPath: status.catalog_path,
trigger: instancePrompt.trim(),
});
}, [status, baseChoice, customBase, family, instancePrompt, onDeploy]);
const numberField = (
label: string,
value: number,
set: (n: number) => void,
fallback: number,
extra?: { min?: number; step?: number },
) => (
<div className="grid gap-1.5">
<Label className="text-xs">{label}</Label>
<Input
type="number"
min={extra?.min ?? 1}
step={extra?.step}
value={value}
onChange={(e) => {
settingsDirty.current = true;
set(Number(e.target.value) || fallback);
}}
className="h-8 text-xs"
/>
</div>
);
return (
<div className="flex min-h-0 min-w-0 flex-1 gap-4 overflow-hidden px-5 pb-8 sm:px-9">
{/* Left: configure */}
<div className="bg-card corner-squircle flex w-[380px] min-w-0 shrink-0 flex-col gap-4 overflow-y-auto overflow-x-hidden rounded-3xl p-5 ring-1 ring-foreground/10">
<div>
<h2 className="text-sm font-semibold">Train a LoRA</h2>
<p className="mt-1 text-[11px] leading-snug text-muted-foreground">
Teach an image model a style, character, or subject from your own images. The
finished adapter shows up in the Create tab&apos;s LoRA picker.
</p>
</div>
{/* Family + base */}
<div className="grid gap-1.5">
<Label className="text-xs">Model family</Label>
<select
value={familyName}
onChange={(e) => setFamilyName(e.target.value)}
className={selectClass}
aria-label="Model family"
>
{families.map((f) => (
<option key={f.name} value={f.name}>
{f.label}
</option>
))}
</select>
{family?.vram_note && (
<p className="text-[11px] leading-snug text-muted-foreground">{family.vram_note}</p>
)}
</div>
<div className="grid gap-1.5">
<Label className="text-xs">Base model</Label>
<select
value={effectiveBase}
onChange={(e) => setBaseChoice(e.target.value)}
className={selectClass}
aria-label="Base model"
>
{(family?.base_repos ?? []).map((repo) => (
<option key={repo} value={repo}>
{repo}
</option>
))}
<option value={CUSTOM_BASE}>Custom repo or local path...</option>
</select>
{effectiveBase === CUSTOM_BASE && (
<Input
value={customBase}
placeholder="my-org/my-base or /path/to/pipeline"
spellCheck={false}
onChange={(e) => setCustomBase(e.target.value)}
className="h-8 text-xs"
/>
)}
</div>
{/* Dataset */}
<div className="grid gap-1.5">
<Label className="text-xs">Training images</Label>
<select
value={dataset}
onChange={(e) => {
const v = e.target.value;
if (v.startsWith(EXAMPLE_PREFIX)) {
const ex = pendingExamples.find((x) => x.id === v.slice(EXAMPLE_PREFIX.length));
if (ex) void importExample(ex);
return; // controlled select snaps back to the current dataset while importing
}
setDataset(v);
setGridOpen(false);
}}
className={selectClass}
aria-label="Training images"
disabled={importingId !== null}
>
{(info?.datasets ?? []).map((d) => (
<option key={d.name} value={d.name}>
{d.name} ({d.image_count} image{d.image_count === 1 ? "" : "s"}
{d.caption_count > 0 ? `, ${d.caption_count} captions` : ""})
</option>
))}
{pendingExamples.length > 0 && (
<optgroup label="Examples (one-click import)">
{pendingExamples.map((ex) => (
<option key={ex.id} value={`${EXAMPLE_PREFIX}${ex.id}`}>
{ex.label} ({ex.image_cap} images, {ex.license})
</option>
))}
</optgroup>
)}
<option value={UPLOAD_DATASET}>Upload new images...</option>
</select>
{importingId && (
<p className="text-[11px] text-muted-foreground">
Importing {examples.find((e) => e.id === importingId)?.label ?? "example"}...
</p>
)}
{dataset === UPLOAD_DATASET ? (
<div className="grid gap-1.5 rounded-md border border-dashed border-border p-2">
<Input
value={uploadName}
placeholder="my-style-photos"
spellCheck={false}
onChange={(e) => setUploadName(e.target.value)}
className="h-8 text-xs"
aria-label="New dataset name"
/>
<div className="flex items-center gap-2">
<input
ref={fileInputRef}
type="file"
multiple
accept={DATASET_FILE_ACCEPT}
className="min-w-0 flex-1 text-xs file:mr-2 file:rounded-md file:border-0 file:bg-muted file:px-2 file:py-1 file:text-xs"
aria-label="Training image files"
/>
<Button
type="button"
size="sm"
variant="secondary"
className="h-8 shrink-0"
onClick={onUpload}
disabled={uploading}
>
{uploading ? "Uploading..." : "Upload"}
</Button>
</div>
<p className="text-[11px] text-muted-foreground">
10-50 images work well. Optional captions: a .txt per image or a
metadata.jsonl; without them the trigger prompt below captions every image.
</p>
</div>
) : (
selectedDataset && (
<>
{selectedDataset.image_count > 0 && !gridOpen && (
<DatasetShowcase
dataset={dataset}
imageCount={selectedDataset.image_count}
refreshKey={gridRefresh}
onBrowse={() => setGridOpen(true)}
/>
)}
<LabelingGridToggle
count={selectedDataset.image_count}
open={gridOpen}
onToggle={() => setGridOpen((o) => !o)}
/>
{gridOpen && (
<DatasetLabelingGrid
dataset={dataset}
refreshKey={gridRefresh}
onCountsChanged={() => void refreshInfo()}
/>
)}
{selectedDataset.caption_count === 0 && !gridOpen && (
<p className="text-[11px] text-muted-foreground">
No caption files - the trigger prompt below captions every image, or
open Review captions to label them.
</p>
)}
</>
)
)}
<ExampleDatasetCards
examples={pendingExamples}
busyId={importingId}
onImport={(ex) => void importExample(ex)}
/>
</div>
{/* Trigger + adapter name (trigger first: it describes the dataset, the name
just labels the output) */}
{fullyCaptioned ? (
<p className="text-[11px] leading-snug text-muted-foreground">
All {selectedDataset?.image_count} images have captions - no trigger prompt needed.
The style applies to any prompt after training.
</p>
) : (
<div className="grid gap-1.5">
<Label className="text-xs">
Trigger prompt (how you&apos;ll invoke the style later)
</Label>
<Input
value={instancePrompt}
placeholder="a photo in SKS style"
onChange={(e) => setInstancePrompt(e.target.value)}
className="h-8 text-xs"
/>
</div>
)}
<div className="grid gap-1.5">
<Label className="text-xs">Adapter name</Label>
<Input
value={outputDir}
placeholder="my-style-lora"
spellCheck={false}
onChange={(e) => setOutputDir(e.target.value)}
className="h-8 text-xs"
/>
</div>
{/* Collapsed training settings */}
<Button
type="button"
variant="ghost"
size="sm"
className="h-7 w-fit gap-1.5 px-1 text-xs text-muted-foreground hover:text-foreground"
onClick={() => setShowAdvanced((s) => !s)}
aria-expanded={showAdvanced}
>
<HugeiconsIcon
icon={ArrowDown01Icon}
className={cn("size-3.5 transition-transform", showAdvanced && "rotate-180")}
/>
{showAdvanced ? "Training settings" : "Training settings (defaults suit a first run)"}
</Button>
{showAdvanced && (
<>
<div className="grid grid-cols-2 gap-3 sm:grid-cols-4">
{numberField("Steps", steps, setSteps, 1)}
{numberField("LoRA rank", rank, setRank, 1)}
{numberField("Resolution", resolution, setResolution, 512, { min: 64, step: 64 })}
{numberField("Batch", batchSize, setBatchSize, 1)}
</div>
<div className="grid grid-cols-2 gap-3">
{numberField("Learning rate", learningRate, setLearningRate, 0.0001, {
min: 0,
step: 0.00001,
})}
<div className="grid gap-1.5">
<Label className="text-xs">Precision</Label>
<select
value={precision}
onChange={(e) => setPrecision(e.target.value as "bf16" | "fp16" | "no")}
className={selectClass}
>
<option value="bf16">bf16 (default)</option>
<option value="fp16">fp16 (older GPUs)</option>
<option value="no">fp32 (no mixed)</option>
</select>
</div>
</div>
</>
)}
<div className="mt-auto pt-2">
{running ? (
<Button type="button" variant="destructive" className="w-full" onClick={onStop}>
Stop training
</Button>
) : (
<Button
type="button"
className="w-full"
onClick={onStart}
disabled={starting || uploading}
>
{starting ? "Starting..." : "Start training"}
</Button>
)}
</div>
</div>
{/* Right: run view */}
<div className="flex min-w-0 flex-1 flex-col gap-4 overflow-y-auto">
{status &&
status.status !== "idle" &&
!(status.status === "completed" && status.job_id === dismissedJobId) ? (
<>
<div className="bg-card corner-squircle flex flex-col gap-3 rounded-3xl p-5 ring-1 ring-foreground/10">
<div className="flex items-center justify-between">
<span className="text-sm font-semibold capitalize">{status.status}</span>
<span className="text-xs text-muted-foreground">
{status.total_steps > 0 ? `${status.step}/${status.total_steps} steps` : ""}
</span>
</div>
<div className="h-2 w-full overflow-hidden rounded-full bg-border">
<div
className="h-full bg-primary transition-all"
style={{ width: `${pct}%` }}
/>
</div>
<div className="grid grid-cols-2 gap-3 sm:grid-cols-4">
<Stat label="Loss" value={status.loss != null ? status.loss.toFixed(4) : "-"} />
<Stat
label="Avg loss"
value={status.avg_loss != null ? status.avg_loss.toFixed(4) : "-"}
/>
<Stat
label="Speed"
value={
status.samples_per_second != null
? `${status.samples_per_second.toFixed(2)} img/s`
: "-"
}
/>
<Stat
label="Peak VRAM"
value={
status.peak_memory_gb != null ? `${status.peak_memory_gb.toFixed(1)} GB` : "-"
}
/>
</div>
{status.message && (
<p className="text-[11px] text-muted-foreground">{status.message}</p>
)}
</div>
<DiffusionCharts lossHistory={lossHistory} lrHistory={lrHistory} />
{completed && (
<div className="bg-card corner-squircle flex flex-col gap-2 rounded-3xl p-5 ring-1 ring-foreground/10">
<span className="text-sm font-semibold">Adapter ready</span>
<p className="text-[11px] text-muted-foreground">
Trained{status.family ? ` (${status.family})` : ""} and added to the LoRA
picker.
{status.lora_path && (
<span className="mt-1 block break-all">Saved: {status.lora_path}</span>
)}
</p>
<div className="mt-1 flex gap-2">
<Button type="button" size="sm" onClick={onDeployClick}>
Deploy to Create
</Button>
<Button
type="button"
size="sm"
variant="secondary"
onClick={() => setDismissedJobId(status.job_id)}
>
Train another
</Button>
</div>
</div>
)}
</>
) : (
<div className="flex flex-1 items-center justify-center">
<div className="max-w-sm text-center">
<p className="text-sm font-medium">No training run yet</p>
<p className="mt-1 text-xs text-muted-foreground">
Pick a family and dataset on the left, then Start training. The loss chart and
progress appear here live.
</p>
</div>
</div>
)}
</div>
</div>
);
}
function Stat({ label, value }: { label: string; value: string }) {
return (
<div className={cn("rounded-lg border border-border/60 bg-muted/20 px-2.5 py-1.5")}>
<div className="text-[10px] uppercase tracking-wide text-muted-foreground">{label}</div>
<div className="text-sm font-medium tabular-nums">{value}</div>
</div>
);
}

View file

@ -0,0 +1,151 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useEffect, useState } from "react";
import { Button } from "@/components/ui/button";
import { toast } from "@/lib/toast";
import {
type DiffusionDatasetExample,
type DiffusionDatasetImportResult,
importDiffusionDatasetExample,
} from "../api";
// Best-effort preview thumbnails pulled from the public HF datasets-server. Cached per repo
// (module-level) so re-renders and re-mounts do not refetch. A repo that the server cannot
// serve (e.g. diffusers/dog-example) resolves to an empty list and the card renders without
// previews - the import still works.
const _previewCache = new Map<string, Promise<string[]>>();
async function fetchPreviews(repo: string): Promise<string[]> {
const cached = _previewCache.get(repo);
if (cached) return cached;
const p = (async () => {
try {
const res = await fetch(
`https://datasets-server.huggingface.co/first-rows?dataset=${encodeURIComponent(
repo,
)}&config=default&split=train`,
);
if (!res.ok) return [];
const data = (await res.json()) as {
features?: { name: string; type?: { _type?: string } }[];
rows?: { row: Record<string, unknown> }[];
};
const imageCol = data.features?.find((f) => f.type?._type === "Image")?.name;
if (!imageCol || !data.rows) return [];
const urls: string[] = [];
for (const r of data.rows) {
const cell = r.row[imageCol] as { src?: string } | undefined;
if (cell?.src) urls.push(cell.src);
if (urls.length >= 3) break;
}
return urls;
} catch {
return [];
}
})();
_previewCache.set(repo, p);
return p;
}
function ExamplePreviews({ repo }: { repo: string }) {
const [urls, setUrls] = useState<string[] | null>(null);
useEffect(() => {
let cancelled = false;
void fetchPreviews(repo).then((u) => {
if (!cancelled) setUrls(u);
});
return () => {
cancelled = true;
};
}, [repo]);
if (!urls || urls.length === 0) return null;
return (
<div className="flex gap-1">
{urls.map((u) => (
<div key={u} className="size-10 shrink-0 overflow-hidden rounded-md bg-muted">
<img src={u} alt="" loading="lazy" className="size-full object-cover" />
</div>
))}
</div>
);
}
// One-click example-dataset importers. Each card shows the license so users see the terms
// before importing, plus a few preview thumbnails so the set is visible before download. On
// success the parent refreshes its dataset list and selects the imported folder (and can
// seed the trigger prompt from suggested_trigger).
//
// Layout: one card per row (the config column is narrow, so a two-column grid wrapped titles
// one word per line and let the long license text overrun into the next card).
export function ExampleDatasetCards({
examples,
busyId,
onImport,
}: {
examples: DiffusionDatasetExample[];
busyId: string | null;
onImport: (ex: DiffusionDatasetExample) => void;
}) {
if (examples.length === 0) return null;
return (
<div className="grid gap-2">
<span className="text-[11px] font-medium text-muted-foreground">
Or start from an example dataset
</span>
<div className="grid gap-2">
{examples.map((ex) => (
<div
key={ex.id}
className="flex min-w-0 flex-col gap-2 rounded-lg border border-border p-2.5"
>
<div className="flex min-w-0 items-center gap-3">
<div className="flex min-w-0 flex-1 flex-col gap-1">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 flex-1 truncate text-xs font-medium">{ex.label}</span>
<span
className="max-w-[110px] shrink truncate rounded-full bg-secondary px-2 py-0.5 text-[10px] font-normal text-secondary-foreground"
title={ex.license}
>
{ex.license}
</span>
</div>
<p className="line-clamp-2 text-[11px] leading-snug text-muted-foreground">
{ex.description}
</p>
</div>
<Button
type="button"
size="sm"
variant="secondary"
className="h-7 shrink-0 self-center px-3 text-xs"
onClick={() => onImport(ex)}
disabled={busyId !== null}
>
{busyId === ex.id ? "Importing..." : "Import"}
</Button>
</div>
<ExamplePreviews repo={ex.repo} />
</div>
))}
</div>
</div>
);
}
// Shared import helper so the panel's dropdown and the cards import identically.
export async function runExampleImport(
ex: DiffusionDatasetExample,
): Promise<DiffusionDatasetImportResult> {
const res = await importDiffusionDatasetExample(ex.id);
toast.success(
res.imported > 0
? `Imported ${res.image_count} images into "${res.name}"`
: `"${res.name}" already imported (${res.image_count} images)`,
);
return res;
}