unsloth/studio/backend
Daniel Han 07076430b3 Add FLUX.2 Klein and FLUX.2-dev DiT LoRA training
Register flux.2-klein and flux.2-dev in the DiT trainer following the
upstream DreamBooth references: latents train patchified and batch-norm
normalized from the VAE posterior mode, the packed forward reuses
step-invariant position ids, and the guidance vector (3.5) is gated on
the variant's guidance_embeds config. Conditioning stacks load per
variant (Mistral via Flux2Pipeline for dev, Qwen3 via Flux2KleinPipeline
for Klein) and are encoded and freed before the transformer lands on the
device. The fused single-stream to_qkv_mlp_proj joins the attention
projections in the LoRA targets; the single-stream out projection stays
dense because its to_out suffix would also match the double-stream
ModuleList container.

Wire both families through the training registry (family set, labels,
VRAM notes, rank 16 / lr 1e-4 defaults, bf16-only preflight), mark them
trainable with train base repos in the family registry, add FLUX.2-dev
to the gated-repo token check, and trust both official bases for
training downloads.

Verified on B200: 30-step klein int8 (19.6s) and nf4 (20.9s) and dev
int8 (52.0s) runs train with finite decreasing loss and the saved
adapters apply on the bf16 base pipeline (weight 0 reproduces the base
image exactly, weight 1 visibly restyles it).
2026-07-17 09:34:01 +00:00
..
assets Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning (#6908) 2026-07-07 06:13:43 -07:00
auth Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm (#6392) 2026-07-01 06:42:23 -07:00
core Add FLUX.2 Klein and FLUX.2-dev DiT LoRA training 2026-07-17 09:34:01 +00:00
hub Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename (#7031) 2026-07-09 06:46:00 -07:00
loggers Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
models Studio: preserve foreign gallery files, force safetensors on remote ControlNets, and close dataset/seed/GPU gaps 2026-07-13 10:02:42 +00:00
plugins Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
requirements Merge remote-tracking branch 'origin/main' into fold-integration 2026-07-07 05:52:16 +00:00
routes Studio: restore Reapply target on async image/video load errors; recheck training state before dataset commit; reclaim partial sd.cpp installs on retry 2026-07-13 15:27:46 +00:00
state Studio: shareable per-checkpoint preview links (#6486) 2026-06-24 06:31:53 -07:00
storage Sort chat recents by last activity (#6844) 2026-07-07 17:54:32 +01:00
tests Add FLUX.2 Klein and FLUX.2-dev DiT LoRA training 2026-07-17 09:34:01 +00:00
utils Studio: auto-detect completion masking markers, stop silent full-sequence training (#7054) 2026-07-11 05:13:45 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
_platform_compat.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
cloudflare_tunnel.py Reap Studio child processes when the parent dies abnormally (#6425) 2026-06-18 05:51:22 -07:00
colab.py Studio Colab: opt-in shareable Cloudflare tunnel link (#6684) 2026-06-26 00:56:23 -07:00
main.py Tighten comments in the diffusion training core and API models 2026-07-12 12:06:44 +00:00
run.py Studio: fix torch CUDA undefined-symbol errors from a conflicting LD_LIBRARY_PATH (#6905) 2026-07-06 23:48:23 -07:00
startup_banner.py Fix Windows Studio UTF-8 startup handling (#6614) 2026-07-01 13:47:33 +01:00