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:
parent
78e3cb237f
commit
17e55db92b
7 changed files with 1886 additions and 72 deletions
|
|
@ -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 }),
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
127
studio/frontend/src/features/images/train/dataset-showcase.tsx
Normal file
127
studio/frontend/src/features/images/train/dataset-showcase.tsx
Normal 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>
|
||||
);
|
||||
}
|
||||
143
studio/frontend/src/features/images/train/diffusion-charts.tsx
Normal file
143
studio/frontend/src/features/images/train/diffusion-charts.tsx
Normal 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>
|
||||
);
|
||||
}
|
||||
|
|
@ -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'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'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>
|
||||
);
|
||||
}
|
||||
|
|
@ -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;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue