Commit graph

327 commits

Author SHA1 Message Date
Daniel Han
37185e651d Reject doomed local picks before the GPU handoff; load bare local safetensors
Several image/video/training preflights ran before the route acquires the GPU or
frees resident models, but let a doomed local pick through and only failed deep in
the background load, after the user's chat/Images/Video model was already evicted.

- Local base_repo / base_model: _is_trusted_diffusion_repo accepts any existing
  local path, but the base loads via from_pretrained (needs model_index.json). A
  local dir that is not a diffusers pipeline passed the trust gate, evicted the
  resident model, then failed. Add a shared _assert_local_base_is_pipeline check
  and call it in the image, video, and training preflights.
- Dataset images: discover_image_caption_pairs only checked filenames, so a
  corrupt or zero-byte upload passed the start-route preflight, freed the GPU, then
  crashed the spawned trainer in PIL. Add an opt-in verify_images decode probe
  (cheap PIL header check) that the start route enables; the trainers leave it off
  since they decode every image anyway.
- Local single-file safetensors: the On-Device scanner advertises a bare
  .safetensors directory (no model_index.json) as a text-to-image model, but the
  picker starts it as a pipeline with no filename, so every click 400s. Reinterpret
  such a pick as a single_file load of the sole checkpoint (resolve_local_single_file)
  so the advertised model is actually loadable.

Regression tests for each: local non-pipeline base (image/video/training), the
verify_images decode gate, and resolve_local_single_file.
2026-07-07 08:06:54 +00:00
Daniel Han
22dad1df73 Gate mxfp8 DiT training precision before evicting resident GPU models
The start route's precision preflight folded bf16/int8/fp8 into the CUDA
requirement but omitted mxfp8, so an mxfp8 request on a GPU-less host (or an
older CUDA GPU without Blackwell) passed the preflight, evicted resident image
and chat models, then raised only in the spawned trainer child. Mirror
_resolve_base_precision: require CUDA for mxfp8 and re-check the Blackwell
(sm100+) capability up front, so a doomed run is rejected before teardown.
2026-07-07 06:10:28 +00:00
Daniel Han
e1dd2dda6b Merge remote-tracking branch 'origin/diffusion-train-perf2' into fold-integration
# Conflicts:
#	studio/backend/core/training/diffusion_dit_trainer.py
#	studio/backend/core/training/diffusion_train_common.py
2026-07-07 01:06:42 +00:00
Daniel Han
2dbfd3c4a9 Merge remote-tracking branch 'origin/diffusion-krea2' into fold-integration
# Conflicts:
#	studio/backend/core/training/diffusion_train_common.py
2026-07-07 01:03:21 +00:00
Daniel Han
b52e7a5cc2 Merge remote-tracking branch 'origin/diffusion-train-tab-2' into fold-integration 2026-07-07 01:01:00 +00:00
pre-commit-ci[bot]
a8494e25f3 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-06 15:24:16 +00:00
Daniel Han
0f2f4334e2 Reject dense DiT precisions on a CUDA-absent host before eviction; stabilize family-info tests
The start-route preflight caught the bf16-GPU and int8-torchao requirements but not the dense
precisions' CUDA requirement: on a GPU-less host bf16_unsupported_reason exempts CPU-only, so a
bf16/fp8 (or int8-with-torchao) DiT request passed the preflight, evicted resident workloads, then
raised only in the trainer child. Add the dense-mode CUDA gate mirroring _resolve_base_precision so
the doomed run is rejected up front. Also pin bf16_unsupported_reason in the two positive-path
family-info tests so they are deterministic across GPU types (a non-bf16 CUDA box would otherwise
empty every DiT family's advertised modes).
2026-07-06 15:22:58 +00:00
Daniel Han
aa54a062ec Gate DiT training on functional torchao for explicit int8; hide always-400 DiT modes on non-bf16 GPUs
The start route preflight only rejected non-bf16 GPUs; an explicit int8 request on
a host with a missing or stub torchao passed the preflight, evicted resident GPU
workloads, then died in the trainer child (its int8 base quantizer has no fallback).
Fold both gates into training_precision_preflight_error so int8-without-torchao fails
fast before eviction. Also empty the advertised DiT precision_modes (and surface the
reason in vram_note, drop compile) whenever the bf16 preflight would reject the family,
so /info never offers an nf4 DiT option the route always 400s.
2026-07-06 13:35:43 +00:00
Daniel Han
2d974219bf Deploy Krea adapters on Turbo and use its distilled recipe over the API
- DiffusionFamily gains deploy_base_repo (krea/Krea-2-Turbo): deploying a LoRA
  trained on Raw now previews it on Turbo, not the non-distilled Raw checkpoint.
  Scoped to a same-precision override so it never turns an nf4 train base into a
  larger bf16 deploy load; exposed through family_train_infos -> the Train UI's
  onDeployClick / historical-run deploy resolve the deploy base.
- _GENERATION_DEFAULTS gains a Krea entry (8 steps, 0 CFG) so the OpenAI
  /v1/images/generations route matches the Create UI's documented distilled recipe
  instead of falling through to the generic (9, 0.0).
2026-07-06 11:16:25 +00:00
pre-commit-ci[bot]
cc6d7c96a6 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-06 11:04:54 +00:00
Daniel Han
6bd3e87c6f Gate DiT training precision: deny fp8 for Qwen, gate explicit int8 on torchao, gate advertised dense modes + route on bf16
- normalized() + family_train_infos() mirror the inference fp8 deny for
  Qwen-Image (activation outliers exceed fp8's range and corrupt the trained
  result); int8 stays allowed and the UI no longer advertises fp8 for it.
- _resolve_base_precision() gates an explicit int8 on a FUNCTIONAL torchao, the
  same gate auto and /info already apply, so a missing/stub torchao fails fast
  instead of silently loading dense with compile disabled.
- train_precision_modes() gates the dense modes (bf16/int8/fp8/auto) on
  torch.cuda.is_bf16_supported(), so a non-bf16 CUDA GPU (T4/V100/RTX 20xx) is
  offered only nf4 instead of a start that evicts resident models and then fails.
- start_diffusion_training preflights bf16 support for the DiT families BEFORE
  _free_gpu_for_diffusion_training(), so any DiT start (nf4 included, since the
  trainer requires bf16 unconditionally on CUDA) fails fast without eviction.
2026-07-06 11:04:07 +00:00
pre-commit-ci[bot]
477757eac7 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-06 10:42:55 +00:00
Daniel Han
6384fea272 dit trainer: preserve biases under mxfp8, gate explicit mxfp8 to Blackwell
- The torchao 0.17 MX training path swaps a matched frozen Linear's weight for a wrapper tensor
  whose linear override computes input @ weight_t and drops the bias, so mxfp8'ing a biased frozen
  linear silently loses its bias and corrupts the base output the LoRA regresses against (verified
  on Blackwell: the bias term is fully dropped). Skip biased linears in _mx_module_filter.
- _resolve_base_precision re-checked explicit dense modes against the live device but only rejected
  CPU, so an explicit mxfp8 request on a non-Blackwell CUDA GPU passed and then crashed at the first
  MX GEMM after a full dense-transformer load. /info only advertises mxfp8 on sm100+; mirror that
  gate here and fail fast for a stale or direct client below Blackwell.
2026-07-06 10:41:02 +00:00
pre-commit-ci[bot]
1c40fa855f [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-05 07:59:24 +00:00
pre-commit-ci[bot]
d0f7dad7ec [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-05 07:58:51 +00:00
pre-commit-ci[bot]
b1fdefb43d [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-05 07:57:51 +00:00
pre-commit-ci[bot]
f343eadcd7 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-05 07:57:16 +00:00
Daniel Han
466f853f14 Merge branch 'diffusion-krea2' into diffusion-train-perf2 2026-07-05 07:56:51 +00:00
Daniel Han
4b93558ad4 Merge branch 'diffusion-train-tab-2' into diffusion-krea2 2026-07-05 07:56:47 +00:00
Daniel Han
69f32b9586 Merge branch 'diffusion-train-precision' into diffusion-train-tab-2 2026-07-05 07:56:20 +00:00
Daniel Han
3925aea07f Merge branch 'diffusion-train-perf' into diffusion-train-precision 2026-07-05 07:56:14 +00:00
Daniel Han
8c00f81a5d Size-gate the automatic diffusion latent cache
The latent cache holds two fp32 posterior tensors per crop/flip variant per
image, pinned on CUDA hosts, so datasets with thousands of images can exhaust
host or pinned memory with no fallback. Estimate the cache size from the first
real encoded latent and fall back to per-step VAE encoding when it exceeds a
4 GiB budget. UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE bypasses the gate; the
existing UNSLOTH_DIFFUSION_NO_LATENT_CACHE opt-out is unchanged.
2026-07-05 07:53:12 +00:00
Daniel Han
da3a79468e Use permutation-cycle index sampling in diffusion trainers and guard non-object run records
Replace the with-replacement per-batch index draw in the SDXL and DiT LoRA
trainers with a shared PermutationBatchSampler that visits every image once per
cycle before repeating, so short runs cover the whole dataset. The sampler
reshuffles from the run's rng so the index stream stays seed-deterministic.

Guard the diffusion run detail route against a valid-JSON non-object record,
which previously raised TypeError and returned a 500; it now 404s like the list
path's shape check.

Add regression tests for both.
2026-07-05 07:49:30 +00:00
Daniel Han
6a9db8017b Merge branch 'diffusion-krea2' into diffusion-train-perf2 2026-07-05 05:35:55 +00:00
Daniel Han
e0786a98be Merge branch 'diffusion-train-tab-2' into diffusion-krea2 2026-07-05 05:35:54 +00:00
Daniel Han
1978b2a844 Merge branch 'diffusion-train-precision' into diffusion-train-tab-2 2026-07-05 05:35:53 +00:00
Daniel Han
1919491665 Merge branch 'diffusion-train-perf' into diffusion-train-precision 2026-07-05 05:35:52 +00:00
Daniel Han
d6795ed077 Restore pre-Ampere bf16 fail-fast in the DiT trainer
The perf rewrite dropped the bf16 capability guard, so a pre-Ampere CUDA
device (T4/V100/RTX 20xx) would die deep in model load with an opaque dtype
error instead of a clear message. Restores parity with the SDXL trainer.
2026-07-05 05:27:58 +00:00
Daniel Han
551c38bd4a Merge branch 'diffusion-krea2' into diffusion-train-perf2 2026-07-05 02:11:26 +00:00
Daniel Han
dc36983562 Merge branch 'diffusion-train-tab-2' into diffusion-krea2 2026-07-05 02:11:25 +00:00
Daniel Han
a99b951c33 Merge branch 'diffusion-train-precision' into diffusion-train-tab-2
# Conflicts:
#	studio/backend/tests/test_diffusion_training.py
2026-07-05 02:11:15 +00:00
Daniel Han
32a77623ba Merge branch 'diffusion-train-perf' into diffusion-train-precision
# Conflicts:
#	studio/backend/core/training/diffusion_train_common.py
2026-07-05 02:10:23 +00:00
Daniel Han
c2ab1a0e61 Merge branch 'image-generation' into diffusion-train-perf
# Conflicts:
#	studio/backend/core/training/diffusion_dit_trainer.py
#	studio/backend/core/training/diffusion_train_common.py
2026-07-05 02:09:39 +00:00
Daniel Han
e605075508 Fix diffusion training validation and honor lr_scheduler and batch size in the DiT trainer 2026-07-05 01:51:34 +00:00
Daniel Han
6a8b0b47e7 Fix review findings on image generation: failed-load VRAM, API defaults, preflights
- Free reserved VRAM in the diffusion load worker's failure path: a load-time OOM
  never commits _state and the next load's _unload_locked early-returns, so nothing
  else reclaimed the half-built pipeline's memory
- Use a monotonic clock for the denoise ETA rate
- Sync _GENERATION_DEFAULTS with the UI table: kontext, flux.2-dev, sdxl-turbo and
  SDXL base rows so /v1/images/generations stops falling back to 9 steps / CFG 0
- 400 (not sanitized 500) when /v1/images/generations hits an edit-only model
- Fail fast on pre-Ampere CUDA in the DiT trainer instead of dying in model load
- Run the trainer trust gate in the diffusion training route before freeing GPU
  residents so an untrusted base cannot tear down loaded chat/Images models
- Protect native sd.cpp companion VAE/text-encoder repos from cache deletion while
  a load is downloading them
- Exempt the task-scoped Images picker from the chat-only GGUF/MLX format gate so
  local diffusers pipelines stay selectable on no-GPU hosts
2026-07-05 01:00:47 +00:00
Daniel Han
08674a297e Merge branch 'diffusion-krea2' into diffusion-train-perf2 2026-07-05 00:31:34 +00:00
Daniel Han
99eb248607 Merge branch 'diffusion-train-tab-2' into diffusion-krea2 2026-07-05 00:31:33 +00:00
Daniel Han
71c20ded19 Merge branch 'diffusion-train-precision' into diffusion-train-tab-2 2026-07-05 00:31:31 +00:00
Daniel Han
79b97e9ad0 Merge branch 'diffusion-train-perf' into diffusion-train-precision 2026-07-05 00:31:30 +00:00
Daniel Han
ad7d5c6827 Coerce cache_latents and enable_tf32 string flags in the config dict path
The generic Studio config dict path can deliver these flags as strings, and a
non-empty string like "false" is truthy, so an opt-out silently no-ops (the
latent cache still builds, TF32 stays on). Coerce them the same way
gradient_checkpointing already is.
2026-07-05 00:12:47 +00:00
Daniel Han
8e589f8b27 Merge diffusion-krea2: CI test fixes (diffusers import order, arbiter device pin, sigma-gather skip) 2026-07-04 05:08:51 +00:00
pre-commit-ci[bot]
e6d775d0fc [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-04 05:08:43 +00:00
Daniel Han
8324cdc407 Merge diffusion-train-tab-2: CI test fixes (diffusers import order, arbiter device pin, sigma-gather skip) 2026-07-04 05:08:18 +00:00
Daniel Han
3fcf218614 Merge diffusion-train-precision: CI test fixes (diffusers import order, arbiter device pin, sigma-gather skip) 2026-07-04 05:07:46 +00:00
Daniel Han
f5d5b09ae5 Merge diffusion-train-perf: CI test fixes (diffusers import order, arbiter device pin, sigma-gather skip) 2026-07-04 05:06:59 +00:00
Daniel Han
c63df7d918 Merge image-generation: CI test fixes (diffusers import order, arbiter device pin) 2026-07-04 05:05:54 +00:00
Daniel Han
1d3aa53d1f Validate the training config before importing diffusers and pin the arbiter test's device
The fp16-on-bf16-family refusal in run_dit_lora_training now fires before the heavy
imports, so a host without diffusers gets the real validation error instead of
ModuleNotFoundError. test_in_progress_returns_409_after_validation_passes pins the
resolved device to cuda because the load route only takes the GPU arbiter for non-CPU
loads, which made the ownership assert host-dependent.
2026-07-04 05:01:58 +00:00
pre-commit-ci[bot]
3204874f7e [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-04 04:42:28 +00:00
pre-commit-ci[bot]
6ecbb2b8a3 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-04 04:41:56 +00:00
pre-commit-ci[bot]
96cecf9c47 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-04 04:41:23 +00:00