diff --git a/.gitignore b/.gitignore index 9f7d4b8c60..5017a114b8 100644 --- a/.gitignore +++ b/.gitignore @@ -237,3 +237,8 @@ package-lock.json llama.cpp/ # Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo. /~/ + +# Agent workspace runtime/scratch artifacts, never part of the repo. +async_task_outputs/ +temp/ +logs/ diff --git a/plans/controlnet-workflow.md b/plans/controlnet-workflow.md new file mode 100644 index 0000000000..9d02a1a0ab --- /dev/null +++ b/plans/controlnet-workflow.md @@ -0,0 +1,113 @@ +# Plan: ControlNet for the Studio Images workflow (stacked on #6769) + +## Context & research + +ControlNet is the **#2 most-used "beyond text-to-image" diffusion workflow** after LoRA (now shipped +in #6771). It conditions generation on a spatial control map so the output follows a structure. Web +research (this session) on ComfyUI / Forge / A1111 usage: + +- The dominant control types are **Depth, Canny (edges), and OpenPose (human pose)**. +- The biggest 2025 shift is toward **Union / all-in-one ControlNets** that bundle many control modes in + one model: **InstantX / Shakker-Labs `FLUX.1-dev-ControlNet-Union-Pro`** for FLUX, and **`xinsir/ + controlnet-union-sdxl-1.0`** for SDXL. SDXL has no official ControlNet (community: xinsir, TheMistoAI, + BRIA). SD1.5 has the original lllyasviel set. +- Sources: comfyui-wiki.com ControlNet collections (flux-1 / sdxl), stable-diffusion-art.com ControlNet + ComfyUI, education.civitai.com ControlNet guide, stablediffusiontutorials.com Qwen-Image ControlNets. + +**Both Studio backends can do ControlNet** (verified against the live tree): +- diffusers has `FluxControlNetPipeline` / `FluxControlNetModel`, `StableDiffusionXLControlNetPipeline` / + `ControlNetModel` / `ControlNetUnionModel`, `QwenImageControlNetPipeline`, `FluxControlNetInpaintPipeline` + (diffusers 0.38 in the studio venv). Also `FluxControlPipeline` (the Flux.1-Canny/Depth "Control" + in-model variants). +- native sd.cpp (`stable-diffusion.cpp`, the #6769 sd-cli base) has `--control-net `, + `--control-image `, `--control-strength `, `--control-net-cpu`, and a built-in + `preprocess_canny` (examples/cli/main.cpp:704, examples/common/common.cpp:422+). + +Studio has **no ControlNet wiring today**: `diffusion_families.py` has no controlnet fields and +`sd_cpp_args.py` has no `--control-*`. This adds it, mirroring the LoRA architecture (#6771) and the +existing img2img/inpaint/reference workflow patterns. + +## Scope (this PR = diffusers ControlNet, Union-first) + +Ship the highest-value slice first; keep it shippable and consistent with the shipped LoRA design. + +- **In scope:** diffusers ControlNet for the families with the strongest ecosystems and pipeline support: + **FLUX** (FluxControlNetPipeline + Union Pro), **SDXL** (StableDiffusionXLControlNetPipeline + xinsir + Union), **Qwen-Image** (QwenImageControlNetPipeline). Single ControlNet per generation. A `control_type` + hint (canny / depth / pose / tile / passthrough). **Canny preprocessing built in** (cheap, cv2/PIL) plus + **passthrough** for user-supplied control maps (depth/pose maps made elsewhere, matching ComfyUI where + preprocessing is separate). Strength + guidance start/end. Discovery endpoint + family-gated picker. +- **Out of scope (follow-ups):** native sd.cpp ControlNet (`--control-net`, needs GGUF ControlNet assets + + family support probe); heavy preprocessors (Depth-Anything, OpenPose detector) as optional server-side + auto-preprocess; multi-ControlNet stacking; ControlNet + inpaint combo. + +## Key facts (verified) + +- diffusers ControlNet pipelines are built with `Pipeline.from_pipe(base_pipe, controlnet=cn_model)` (or + `from_pretrained(base, controlnet=...)`), so the resident base modules are reused with **no reload** -- + same `from_pipe` machinery the img2img/inpaint/edit workflows already use (`diffusion.py` ~:1104-1130, + `_workflow_pipe`). The ControlNet model (`ControlNetModel` / `FluxControlNetModel` / `ControlNetUnionModel`) + is a small extra module loaded once and cached on `_LoadState`. +- ControlNet models are **family-specific** (SD1.5 CN != FLUX CN != SDXL CN). So discovery must be + **family-gated**, exactly like the LoRA picker's `supports_lora`/family filter. +- Generate-time contract mirrors reference/inpaint: a control image (b64) + params, threaded through + `routes/inference.py` into both backends (diffusers serves it; native rejects clearly until the + follow-up wires `--control-*`). +- Reuse: `diffusion_lora.py` discovery/resolve/family-gate patterns; the reference-image upload component + + the LoRA picker UI shape; `_workflow_pipe` from_pipe; `hf_hub_download_with_xet_fallback`. + +## Approach + +### Families (`core/inference/diffusion_families.py`) +- Add per-family ControlNet declaration: `controlnet_pipeline_class` (e.g. "FluxControlNetPipeline", + "StableDiffusionXLControlNetPipeline", "QwenImageControlNetPipeline"), `controlnet_model_class` + ("FluxControlNetModel" / "ControlNetModel" / a union class), and a small curated list of recommended + ControlNet repos tagged by control type. Expose a `controlnet: bool` capability (like `reference`). + +### Discovery (`core/inference/diffusion_controlnet.py`, new -- mirrors diffusion_lora.py) +- `list_controlnets(family)` = curated family-tagged repos + a local scan; `resolve_controlnet(id, family, + hf_token)` downloads via the xet-fallback helper; `preprocess_control(image, control_type)` (canny via + cv2/PIL; passthrough otherwise); `supports_controlnet(engine, family, model_kind, quant)` gate + (diffusers bf16 / bnb-4bit yes; GGUF-via-diffusers + torchao fp8/int8 dense = no, same rule as LoRA; + native = follow-up). + +### Backend -- diffusers (`core/inference/diffusion.py`) +- A ControlNet manager parallel to `_apply_loras`: load the requested `ControlNetModel` once (cache on + `_LoadState`, reset on unload/model change), build the CN pipeline via `from_pipe(base, controlnet=...)` + in `_workflow_pipe`, and pass `control_image` + `controlnet_conditioning_scale` + + `control_guidance_start/end` at generate time. Never fuse; CN model stays bf16. + +### Backend -- native (`core/inference/sd_cpp_backend.py`, sd_cpp_args.py) -- FOLLOW-UP +- Add `--control-net` / `--control-image` / `--control-strength` to the arg builder and a GGUF-ControlNet + resolve; gate to families sd.cpp supports. Deferred out of this PR. + +### Routes + request models (`models/inference.py`, `routes/inference.py`) +- Add optional `controlnet: ControlNetSpec` to `DiffusionGenerateRequest` (`{id, image, control_type, + strength (0..2, default 1), guidance_start (0..1), guidance_end (0..1)}`); thread into `backend.generate`; + surface `supports_controlnet` in status; persist the chosen CN + type in gallery recipe metadata. +- New `GET /api/models/diffusion-controlnets?family=` (mirror the LoRA discovery route). + +### Frontend (`features/images/images-page.tsx`, `api.ts`) +- A "ControlNet" control in the left rail (reuse the reference-image uploader + a control-type Select + + ControlNet-model Select gated by `supports_controlnet`/family + a strength SliderField). Show a small + preview of the preprocessed control map. Thread `controlnet` into `generateDiffusionImage`; omit when no + control image. + +## Verification +- **Unit:** request validation (optional/empty unchanged; bad strength rejected; unsupported family/quant + rejected); discovery (family filter, resolve, canny preprocess shape); diffusers manager (loads CN once, + from_pipe built, scale threaded, reset on model change) with a fake pipe; routes (no-CN path unchanged). +- **Live smoke (critical):** on GPU 4, drive the real diffusers backend with a real family + Union CN and a + canny control image; same prompt/seed at strength 0 vs 0.8; assert (a) output DIFFERS from no-control and + (b) the strong-control output structurally follows the control map (edge-overlap / SSIM vs the control). +- **Playwright (`unsloth_studio_workflow`):** upload a control image, pick type + model + strength, + generate; capture screenshots/GIF against the live secure studio. +- Full `pytest studio/backend/tests/` green; frontend `vite build` clean; ruff clean. + +## Delivery +- New branch `diffusion-controlnet` off `diffusion-image-workflows` (#6769 head) in an isolated worktree, + sibling to #6771 (LoRA) and #6772 (fp8 fix). PR on `unslothai/unsloth`, base = diffusion-image-workflows, + part of the single logical stack rooted on #6763 (continuation of #6658). Commit as Daniel Han; no AI/bot + mentions, no emojis, no em dashes. +- Follow-ups: native sd.cpp ControlNet; server-side Depth/OpenPose auto-preprocessors; multi-ControlNet; + ControlNet+inpaint. diff --git a/plans/diffusion-popularity-findings.md b/plans/diffusion-popularity-findings.md new file mode 100644 index 0000000000..630ef93291 --- /dev/null +++ b/plans/diffusion-popularity-findings.md @@ -0,0 +1,113 @@ +# Diffusion workflow popularity findings (HF download data) + +Read-only HF metadata pull (`scripts/investigate_popularity.py`), to ground the Studio +Images scope against what people actually download. Downloads are HF's 30-day count and +all-time count; pulled 2026-06-30. + +## Qwen-Image-Edit vs Qwen-Image-Layered (the explicit "determine popularity" question) + +| Model | dl / 30d | dl all-time | likes | pipeline | +|---|---:|---:|---:|---| +| Qwen/Qwen-Image-Edit-2509 | 511,996 | 2,942,966 | 1,185 | image-to-image | +| Qwen/Qwen-Image-Edit-2511 | 162,185 | 1,088,015 | 1,087 | image-to-image | +| Qwen/Qwen-Image-Edit (base) | 70,728 | 1,161,044 | 2,440 | image-to-image | +| **Qwen-Image-Edit (all variants)** | **~745,000** | **~5,192,000** | - | - | +| Qwen/Qwen-Image-Layered | 51,303 | 234,785 | 1,112 | image-text-to-image | +| unsloth/Qwen-Image-Edit-2511-GGUF | 218,313 | - | - | image-to-image | + +**Conclusion:** Qwen-Image-Edit is ~10-14x more downloaded than Layered (combined 745K/30d +vs 51K, 5.2M vs 235K all-time). Shipping Edit (2511 + the unsloth GGUF, which alone pulls +218K/30d) and rejecting/deferring Layered is the correct, data-backed call. Layered also +needs a dedicated pipeline (`additional_t_cond=True`) the standard QwenImagePipeline can't +drive, so it would be both niche AND extra engineering. Reject stands. + +## ControlNet is niche on the modern (diffusers/FLUX/Qwen) stack + +| Model | dl / 30d | dl all-time | likes | +|---|---:|---:|---:| +| InstantX/FLUX.1-dev-Controlnet-Canny | 2,891 | 136,727 | 194 | +| lllyasviel/ControlNet (SD1.5-era) | 0 | 14 | 3,820 | +| stabilityai/stable-diffusion-x4-upscaler | 10,040 | 2,976,405 | 725 | + +**Conclusion:** ControlNet's large user base lives in the older SD1.5 / A1111 ecosystem, not +the diffusers/FLUX/Qwen stack Studio targets (the modern FLUX ControlNet is ~3K/30d). It is +NOT part of the "most popular ~80%" for current-gen models, so deferring it is justified by +the data, not just by effort. The dedicated x4 upscaler is also low 30-day (10K) though high +all-time; our generic hires-fix upscale (img2img re-detail) covers the use case for any +loaded family without an extra model. + +## The shipped six cover the popular workflows + +Top text-to-image (HF list, 30-day): SD1.5 (1.78M), SDXL (1.32M), FLUX.1-dev (1.09M), +dreamshaper-7 (1.03M), **Tongyi-MAI/Z-Image-Turbo (886K)**, sd-turbo (684K), SD3.5-medium +(606K), sdxl-turbo (598K), Qwen-Image-Lightning (483K). All are plain txt2img -> our Create +tab; the GGUF/bnb families + Z-Image cover the modern ones. + +Top image-to-image (HF list, 30-day): Qwen-Image-Edit-2509 (512K) -> Edit tab; SDXL-refiner +(162K) -> Upscale/Transform; Kontext (150K) -> Edit tab. + +So Create / Transform / Inpaint / Extend / Upscale / Edit map onto the head of both +distributions. + +## SHIPPED: FLUX.2-klein image (reference) conditioning + +> Status: IMPLEMENTED + verified live (2026-06-30). `flux.2-klein` now has `reference=True`; +> the backend exposes a "reference" workflow that passes the image to the loaded +> Flux2KleinPipeline directly (no from_pipe, no strength, output at the requested size); the +> frontend has a "Reference" tab. Verified with `scripts/verify_reference_http.py` on +> `unsloth/FLUX.2-klein-4B-GGUF` (Q4_K_M): a reference-conditioned 1024x1024 result is +> non-blank, correctly sized, and DIFFERS from the identical-seed plain txt2img. +> +> FLUX.2-klein ALSO gained inpaint (`Flux2KleinInpaintPipeline` via from_pipe; verified with +> `scripts/verify_klein_inpaint.py`). It does NOT get outpaint/extend: FLUX.2 scales any >1MP +> input down to ~1MP, so a padded outpaint canvas shrinks back. "outpaint" is now a distinct +> capability advertised only for size-preserving inpaint families (`inpaint_preserves_size`). +> Multi-reference is shipped too (the pipeline accepts a list; the Reference tab has add/remove +> slots, backend caps at 3 extra; verified with `scripts/verify_multiref_http.py`: two +> references differ from one at the same seed). The analysis that motivated the work follows. + +The data surfaced this gap (now closed): + +| Model | dl / 30d | pipeline | +|---|---:|---| +| black-forest-labs/FLUX.2-klein-4B | 470,482 | image-to-image (#2 overall) | +| black-forest-labs/FLUX.2-dev | 271,037 | image-to-image | +| **unsloth/FLUX.2-klein-4B-GGUF** | **243,307** | image-to-image | +| black-forest-labs/FLUX.2-klein-9B | 178,964 | image-to-image | + +`flux.2-klein` is ALREADY a registered family in `diffusion_families.py` (txt2img only, +base `FLUX.2-klein-4B`, open repo). But `Flux2KleinPipeline.__call__` natively accepts an +`image` argument (verified in diffusers 0.38.0; params: image, prompt, height, width, +num_inference_steps, guidance_scale -- NOTE: no `strength`). FLUX.2 is a unified +text-to-image + reference/edit model: the SAME loaded pipe does both, depending on whether +`image` is passed. Today Studio exposes only txt2img for it, so the popular image-editing +mode of the #2 image-to-image model is unreachable. + +### Why it's a separate PR, not a tail-of-session add +FLUX.2 reference conditioning is a DIFFERENT semantic from the shipped workflows: +- No `strength` (it is reference-conditioning, not a denoise blend like img2img). +- Output size comes from width/height (txt2img-style), not from the input image size, so the + image-conditioned width/height rule we added for img2img/inpaint/upscale does NOT apply. +- FLUX.2 supports MULTIPLE reference images; single-image is the common case but the UX + should not preclude multi-ref. +This needs: read the Flux2KleinPipeline source for exact `image` semantics (list vs single, +how it is resized/tiled, recommended guidance), decide the UX (a "Reference" workflow that is +available alongside Create for `reference=True` families, distinct from the strength-based +Transform tab), then verify on the open FLUX.2-klein-4B base (and the unsloth GGUF) with a +reference image before/after. + +### Sketch (for the follow-up PR) +- `diffusion_families.py`: add `reference: bool = False`; set `reference=True` on flux.2-klein. +- `_family_workflows`: when `fam.reference`, expose `"reference"` (in addition to txt2img). +- `generate()`: a `reference` branch that passes `image` to `state.pipe` directly (no + from_pipe, no strength), with width/height = the requested size (NOT the input size). +- Frontend: a "Reference" tab (image dropzone + prompt), gated to `reference` families; + Create stays pure txt2img for the same model. +- Verify: load unsloth/FLUX.2-klein-4B-GGUF, pass a reference image, confirm the output is + conditioned on it and differs from a no-image run at the same seed. + +## Net +The seven shipped workflows (create, transform, inpaint, extend, upscale, reference, edit) +cover the popular ~80% across both the txt2img and image-to-image distributions, including the +#1 image-to-image model (Qwen-Image-Edit) and the #2 (FLUX.2-klein, now via the reference tab). +ControlNet / SD1.5-era ControlNet remain deferred with data backing (niche on the modern stack). diff --git a/plans/diffusion-workflows-pr-plan.md b/plans/diffusion-workflows-pr-plan.md new file mode 100644 index 0000000000..df07714cf4 --- /dev/null +++ b/plans/diffusion-workflows-pr-plan.md @@ -0,0 +1,156 @@ +# Stacked-PR plan: Studio diffusion workflows (Images redesign) + +Branch tip: `diffusion-eager-and-compile-cache` (latest commit "Phase 16 review fixes"). +Remote: `oobabooga/unsloth`. New PRs stack on top of the existing diffusion stack +(ultimately on top of unslothai/unsloth#6658), treated as one logical change. + +Nothing here is committed yet (commit/push only on explicit request). + +## CRITICAL: the working tree holds TWO uncommitted streams, and three core files INTERMINGLE them + +A full `git status` / marker audit (branch `diffusion-eager-and-compile-cache`, tip "Phase 16 +review fixes") shows the uncommitted tree is NOT a clean single feature. There are two streams: + +A) **The eager/compile-cache phase** (the branch's own in-progress work, NOT this session's — + zero of this feature's markers). Purely-its files, safe to NOT touch in the workflow PRs: + - new modules: `diffusion_arch_patches.py`, `diffusion_compile_cache.py`, + `diffusion_eager_patches.py`, `diffusion_gguf_compile.py`, `diffusion_patch_backend.py` + - new tests: `test_diffusion_arch_patches.py`, `test_diffusion_compile_cache.py`, + `test_diffusion_eager_patches.py`, `test_diffusion_gguf_compile.py` + - modified: `diffusion_speed.py`, `test_diffusion_speed.py`, `conftest.py`, + `scripts/diffusion_bench.py`, and ~25 untracked `scripts/*bench*/*probe*/*orchestrator*`. + +B) **The Images workflows feature** (this session): the workflow engine + frontend + installer. + +**The two streams INTERMINGLE inside three shared files and are NOT separable by file:** + - `studio/backend/core/inference/diffusion.py` — this feature's workflow hunks are interleaved + with the eager/compile wiring (imports at L67-75; `install_arch_patches`/`compile_cache.begin`/ + `.restore`/`.save` and the `eager_patched`/`compile_cache_ctx` state throughout + `load_pipeline`/`generate`/`unload`). A single `diffusion.py` cannot go into one PR without the + other stream's hunks. + - `studio/backend/models/inference.py` — this feature's `init_image`/`mask_image`/ + `reference_images`/`upscale`/`model_kind` fields sit next to the pre-existing `speed_mode`/ + `transformer_prequant_path` fields in the same request models. + - `studio/backend/tests/test_diffusion_backend.py` — this feature's workflow tests sit next to + the pre-existing `test_failed_load_rolls_back_eager_patches` (imports `diffusion_eager_patches`). + +**Implication / options (USER DECIDES — it is their branch + their eager/compile work):** + - CLEANLY separable now (purely this feature, can be committed/PR'd on their own anytime): + frontend `images-page.tsx` + `api.ts` + `pickers.tsx`, and the sd.cpp installer + `install_sd_cpp_prebuilt.py` + `test_sd_cpp_install.py`. (These are PR 2 and PR 3 below.) + - The backend engine (PR 1) CANNOT be cleanly split from the eager/compile phase via files. + Realistic paths: (a) finalize + commit the eager/compile phase first, then this feature's + backend lands as a clean diff on top; or (b) commit both streams together as the branch's + next chunk (consistent with treating the stack as one logical change); or (c) a manual + `git add -p` hunk split of the three shared files (tedious, risks a non-compiling + intermediate). NOT auto-doable safely without the owner's intent for the eager/compile work. + +## Proposed stack (3 PRs, bottom to top) + +### PR 1 - Backend: diffusion workflow engine (safetensors + image-conditioned + editing) +Files: +- `studio/backend/core/inference/diffusion.py` (the feature hunks: three load "kinds" + gguf/single_file/pipeline; `_workflow_pipe` via `from_pipe(torch_dtype=None)`; + `_align_vae_dtype`; `generate()` routing for reference/img2img/inpaint/upscale/edit; + image-conditioned width/height from the input image (but reference + txt2img use the slider + size); `upscale` (hires fix) branch on the img2img pipe; `reference` (FLUX.2) branch that + passes the image(s) to the loaded pipe directly (no from_pipe, no strength) incl. multi- + reference (`reference_images` combined into a list, capped at 3 extra); branch ORDER + inpaint/upscale before reference so a mask/upscale request on a reference family still routes + right; `_family_workflows` (adds "upscale" wherever img2img is supported, "reference" for + reference families, "outpaint" only for size-preserving inpaint families); `kind` on state + + `model_kind` in status; `load_progress` double-count fix). NOTE: this file ALSO carries + pre-existing speed hunks if any landed here - review per-hunk and exclude non-feature hunks. +- `studio/backend/core/inference/diffusion_families.py` (img2img/inpaint pipeline slots; + `edit` flag + `reference` flag + `inpaint_preserves_size` flag; `qwen-image-edit` + + `flux.1-kontext` families; flux.2-klein gains reference + inpaint (no outpaint: FLUX.2 + normalizes to ~1MP); `detect_family` longest-match + leftover-reject; `layered` reject). +- `studio/backend/core/inference/diffusion_engine_router.py` (model_kind -> diffusers for + non-gguf kinds). +- `studio/backend/core/inference/diffusion_memory.py` (`estimate_safetensors_dense_mib`). +- `studio/backend/core/inference/sd_cpp_backend.py` (model_kind passthrough; reject + img2img/inpaint on the native engine). +- `studio/backend/models/inference.py` (load request: optional gguf_filename, model_kind, + init/mask/strength, advanced knobs; status: workflows, model_kind). +- `studio/backend/routes/inference.py` (model_kind forwarding; ValueError -> 400; + exc_info logging). +- Tests: `test_diffusion_backend.py`, `test_diffusion_routes.py`. + +Title: `Studio diffusion: safetensors + image-conditioned + instruction-editing workflows` +Summary: Adds non-GGUF safetensors loading (full bnb-4bit pipelines + single-file fp8, +gated to unsloth/*), the image-conditioned workflows (img2img, inpaint, outpaint via the +inpaint path) built with `Pipeline.from_pipe` for zero-extra-VRAM component reuse, and +instruction editing as its own family kind (Qwen-Image-Edit-2511 + FLUX.1-Kontext-dev). +Fixes two bugs: `from_pipe` defaulting to a float32 recast that crashed torchao-quantized +transformers, and image-conditioned calls forcing the slider size onto the input image. + +### PR 2 - Frontend: redesigned Images page (workflow tabs + Advanced Options) +Files: +- `studio/frontend/src/features/images/images-page.tsx` (workflow tabs Create/Transform/ + Inpaint/Extend/Upscale/Reference/Edit; capability gating + auto-switch; `MaskCanvas`; + `buildOutpaint`; Upscale tab with Scale + Detail-strength sliders; Reference tab (FLUX.2, + reference image + add/remove extra references, no strength); Advanced Options accordion gated + to GGUF for transformer-quant; spinner-overlap fix). +- `studio/frontend/src/features/images/api.ts` (request/status types incl. model_kind, upscale, + reference_images). +- `studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx` (curated + safetensors + edit GGUF rows; `SUPPORTED_EDIT_KEYWORDS` un-hide; layered hide). + +Title: `Studio Images: workflow tabs (create/transform/inpaint/extend/upscale/reference/edit) + Advanced Options` +Summary: Redesigns the Images page around capability-gated workflow tabs with a brush mask +editor, client-side outpaint, a hires-fix upscale tab, a FLUX.2 reference tab, an instruction-edit +tab, and an Advanced Options panel (speed/quant/attention/memory/step-cache/offload), plus the +number-input spinner fix. + +### PR 3 - sd.cpp prebuilt installer hardening +Files: `studio/install_sd_cpp_prebuilt.py`, `studio/backend/tests/test_sd_cpp_install.py`. +Title: `Studio sd.cpp: pin release + verify sha256 + mirror-ready source` +Summary: Pins the stable-diffusion.cpp release (was tracking `latest`), verifies each +download's sha256 against GitHub's published asset digest before extract/execute, adds a +download timeout, and makes the source repo configurable (`UNSLOTH_SD_CPP_REPO`) so a +future unslothai mirror needs no code change. Cleanly separable from the rest. + +## Pre-PR review (done) +An independent 3-angle review (backend correctness, frontend/UX, security/robustness) ran over +the full session diff. No High findings; the load-gating to `unsloth/*` and the multi-reference +count caps were verified intact end-to-end. Fixes applied before the PRs: +- Frontend [Med]: multi-reference slots no longer renumber mid-edit (dropped the eager + `filter(Boolean)` in the per-slot onChange; empties dropped only at send). +- Backend [Med]: upscale now caps the ABSOLUTE output (longest side <= 2048), not just the + factor, so a large upload * 4x can't OOM. +- Backend/security [Med]: `_decode_b64_image` rejects images > 4096px/side (uniform guard for + init/mask/reference vs decompression-bomb / OOM inputs); base64 image fields capped at 32 MiB. +- Security [Low]: the native sd.cpp engine guard also rejects `reference_images`. +All covered by new tests (82 backend pass) and a post-fix e2e (all five workflows still pass). + +## Post-deploy user feedback fixes (done) +From live use of the deployed studio: +- Backend [Med]: image-conditioned workflows passed the raw upload size to pipelines that + require multiples of 16 (Z-Image/Qwen/FLUX), so an odd upload (e.g. 186px) failed with + "Height must be divisible by 16". Added `_snap_to_multiple` and auto-resize init (and the + matched mask) to the nearest /16 for img2img/inpaint/extend/edit. Verified live: a 186x250 + Transform and Inpaint now return 200 at 192x256. Tests added. +- Frontend [Med]: the Advanced options (FP8/INT8 quant, speed, attention, memory) were a + collapsed, muted accordion at the bottom of the left rail that users missed (HF screenshots + discussion #25). Moved them into a RIGHT-DOCKED panel mirroring Chat's settings panel. + Per follow-up (discussion #26): CLOSED by default, toggled by a SINGLE fixed top-bar button + using Chat's `LayoutAlignRightIcon` that stays in the exact same position in both states + (verified x/y identical open vs closed) and highlights when open. Controls extracted to a + render-local `advancedControls`; unused Accordion import removed. + +## Constraints for execution (when authorized) +- Write as the user; no AI/bot mentions, no emojis, no em dashes; well-formatted bodies. +- `gh auth status` first. Push to `oobabooga/unsloth`, stack on the current branch. +- Keep PR 3 independent; PR 2 depends on PR 1 (frontend needs the backend contract). +- Re-run `pytest studio/backend/tests/test_diffusion_*.py test_sd_cpp_install.py` + + frontend `tsc`/`build` before each PR. + +## Out of scope / follow-ups +Scope decisions are backed by HF download data in `plans/diffusion-popularity-findings.md`. +The seven shipped workflows are create, transform, inpaint, extend, upscale, reference, edit. +- Publish the unslothai/stable-diffusion.cpp mirror + macOS/Windows staging (#152). +- ControlNet / style-transfer: the goal's "most popular" set is covered by the seven shipped + workflows. ControlNet on the modern diffusers/FLUX/Qwen stack is niche by downloads + (~3K/30d), so deferring it is data-backed, not just an effort call. +- FLUX.2-klein inpaint and multi-reference are DONE (shipped). Outpaint is intentionally not + offered for FLUX.2 (it scales >1MP inputs to ~1MP). No further FLUX.2 follow-ups outstanding. diff --git a/plans/diffusion-workflows-studio.md b/plans/diffusion-workflows-studio.md new file mode 100644 index 0000000000..126bc79129 --- /dev/null +++ b/plans/diffusion-workflows-studio.md @@ -0,0 +1,125 @@ +# Plan: Unsloth Studio diffusion workflows + Images UI redesign + +Stacked as NEW PRs on top of the existing 14-PR diffusion stack above unslothai/unsloth#6658 +(treated as one logical base). Goal: cover ~80% of common real-world diffusion workflows, +across macOS/Linux/Windows/CPU, optimizing performance, accuracy, and memory. + +## Current state (verified by recon) + +- Backend is **text-to-image only** end-to-end. No image/mask/control plumbing in + `DiffusionGenerateRequest`, `/images/generate`, or `DiffusionBackend.generate()`. +- **diffusers 0.38.0 already imports every pipeline we need**: `*Img2ImgPipeline`, + `*InpaintPipeline`, `FluxFillPipeline`, `*ControlNetPipeline`/`ControlNetModel`, + `FluxKontextPipeline`, `QwenImageEditPipeline`/`QwenImageEditPlusPipeline`, + `StableDiffusion(Latent)UpscalePipeline`. No diffusers upgrade required. +- The native **sd.cpp engine already has dormant fields** (`init_img`, `strength`, `mask`, + `ref_images`) and a complete `upscale()` path — never wired to the request/route. +- **All advanced LOAD options are already wired** end-to-end (speed_mode, transformer_quant + fp8/int8/nvfp4/mxfp8, attention_backend, memory_mode, cpu_offload, transformer_cache, + vae_tiling status). The Advanced panel is mostly a FRONTEND surfacing job. +- Frontend has `Tabs` (`components/ui/tabs.tsx`) and `Accordion` ready. No dropzone and no + mask/brush canvas — both greenfield. `SliderField` is the customizer number input. +- `_EDIT_KEYWORDS = ("edit","kontext","inpaint","layered")` rejects edit/inpaint/Kontext/ + Layered repos at family detection. +- sd.cpp binary is downloaded prebuilt from **upstream leejet/stable-diffusion.cpp** (no + unsloth mirror, no checksum/manifest/version pin, not wired into setup.sh). llama.cpp uses + an `unslothai/llama.cpp` mirror with manifest+sha256+version pin+source fallback. +- Chat "Images" pill = provider-side (OpenAI/Gemini) hosted tool, separate from local diffusion. + +## Workflow popularity ranking (what to build for 80% coverage) + +1. txt2img (keep, polished) — done +2. img2img / variations — **P0** +3. inpainting (mask edit) — **P0** +4. upscaling / hires fix — **P0** +5. ControlNet (Canny/Depth/Pose/Lineart/Tile) — **P0/P1** +6. outpainting (canvas extend) — **P1** +7. instruction image editing (Qwen-Image-Edit, FLUX Kontext) — **P1** +8. style transfer / reference — **P1** (via img2img/edit/control) +9. batch generation/edit/upscale — **P1** +10. LoRA/style packs — **P2** + +## Editing-model decisions (researched) + +- **Qwen-Image-Edit / Edit-2511**: popular, best-in-class clean targeted edits + multilingual + text. **Support** (instruction edit, mask-optional). +- **FLUX.1 Kontext**: popular, character-consistent in-context editing. **Support** (note: + Kontext-dev is non-commercial/gated — surface license, don't block local custom models). +- **Qwen-Image-Layered**: newer, niche (Photoshop RGBA layer decomposition). Needs a dedicated + pipeline (`additional_t_cond`) — **defer** (keep rejected for now; optional later behind a + layered-specific view). This already crashed the standard path (the earlier bug). + +## UI design — workflow tabs (inside ImagesPage, route/nav unchanged) + +`Tabs` across the top of the controls area. Combine related workflows: +- **Create** — txt2img (current behavior preserved) +- **Transform** — img2img + style transfer (upload + strength/denoise + presets) +- **Edit** — inpaint (mask brush/upload/invert/feather, masked-vs-whole) + instruction edit + (Qwen-Image-Edit / FLUX Kontext, mask-optional) +- **Extend** — outpainting (directional handles, aspect presets, overlap/feather) +- **Control** — ControlNet (one control slot first: Canny/Depth/Pose/Lineart/Tile + preview) +- **Enhance** — upscaling (ESRGAN/RealESRGAN + latent/tiled) +- **Advanced Options** — Accordion surfacing existing load knobs (speed/compile/attention/ + quant fp8/int8/nvfp4/memory/offload/vae tiling/cache) with Auto defaults + resolved values. + +Capability gating: a workflow/control is shown enabled only when the selected engine+family+ +device+quant supports it; otherwise disabled with a plain-language "why". + +## Backend architecture + +- Extend `DiffusionGenerateRequest`: optional `workflow` (txt2img|img2img|inpaint|outpaint| + control|edit|upscale), `init_image` (b64), `mask_image` (b64), `control_image` (b64), + `strength`, `controlnet_conditioning_scale`, `control_start/end`, `upscale_factor`, + `ref_images`. Add an image-decode (b64→PIL) helper (none exists). +- `DiffusionFamily`: add optional pipeline-class slots (`img2img_pipeline_class`, + `inpaint_pipeline_class`, `edit_pipeline_class`, `controlnet_pipeline_class` + control repos). + Build the right pipeline around the already-loaded `transformer=` (reuse `_assemble_pipe` + shape); swap/cache pipeline class per workflow without reloading the transformer where + possible. +- `generate()` kwarg builder must branch: img2img/edit pipelines take `image=`/`strength=` and + reject `width/height`; inpaint adds `mask_image=`; control adds `control_image=`. Gate each + kwarg via `inspect.signature`. +- Capability resolver: maps engine+family+device+quant → supported workflows + reasons; echoed + in run metadata so the UI shows what actually ran. +- Memory planner must account for input/latent size, control models, VAE decode, upscale. + +## PR breakdown (stacked, small, capability-gated) + +- **PR-1 UI fixes + workflow shell**: fix number-input spinner overlap (DONE in tree), tab + scaffold (Create/Transform/Edit/Extend/Control/Enhance/Advanced), Advanced Options accordion + surfacing existing load knobs, capability banner, loading/empty/error states. +- **PR-2 Backend workflow contract + capability registry**: extend request/response, decode + helper, per-family pipeline slots, resolver. No new behavior yet beyond txt2img. +- **PR-3 img2img (Transform)**: backend + Transform tab + dropzone (adapt from + shared-composer `addFiles`/`PendingImageThumb`). Smoke test low vs high denoise. +- **PR-4 inpaint + instruction edit (Edit)**: mask canvas (greenfield), inpaint pipeline, + Qwen-Image-Edit/FLUX Kontext edit; relax `_EDIT_KEYWORDS` → route to edit family. +- **PR-5 outpaint (Extend)**: expanded-canvas inpaint, directional handles, feather/overlap. +- **PR-6 ControlNet (Control)**: one control slot + preprocessor preview + strength/start/end. +- **PR-7 upscaling (Enhance)**: wire dormant sd.cpp `upscale()` + diffusers upscale + `/images/upscale`. +- **PR-8 Advanced panel polish + FP8/INT8 verification matrix**. +- **PR-9 sd.cpp prebuilt packaging**: mirror to `unslothai/stable-diffusion.cpp`, manifest+ + sha256+version pin+`--published-repo`+source fallback, wire into setup.sh (ref + install_llama_prebuilt.py). +- **PR-10 cross-platform staging validation** (danielhanchen staging repos, small GGUFs). +- **PR-11 Playwright tests + screenshots/GIFs per tab** (studio_test_kit / unsloth_studio_workflow). +- **PR-12 batch + multi-control + reproducibility polish** (later). + +## Done so far + +- Fixed the customizer number-input spinner overlap (`SliderField` in images-page.tsx): native + spinners covered the value on the narrow field; now fully suppressed (webkit inner+outer + + Firefox `appearance:textfield`) and field widened to `w-14`. Frontend rebuilt clean. + +## Verification + +- Playwright (studio_test_kit) per tab: screenshots + GIFs, capability gating, upload/mask, + progress/cancel/error, gallery. +- B200 functional: load + generate one image per workflow per representative family. +- FP8 + INT8 verified (build matrix: SDXL/FLUX/Qwen-Image/Qwen-Image-Edit/GGUF; measure + black-image/NaN rate, peak VRAM, time-to-first-image, prompt adherence, source preservation). +- Cross-platform staging (Linux CUDA/CPU, Windows CUDA/CPU, macOS MPS) with small GGUFs. + +## Delivery + +New branch(es) off the current tip; new stacked PRs. Commit/push only when asked. diff --git a/plans/wobbly-jumping-narwhal.md b/plans/wobbly-jumping-narwhal.md new file mode 100644 index 0000000000..2261a2fe87 --- /dev/null +++ b/plans/wobbly-jumping-narwhal.md @@ -0,0 +1,141 @@ +# Plan: Publish unslothai/stable-diffusion.cpp mirror + our own CPU/Apple prebuilts + +## Context + +The Unsloth Studio native diffusion engine downloads a prebuilt `sd-cli` / `sd-server` +(stable-diffusion.cpp) via `studio/install_sd_cpp_prebuilt.py`. Today it pulls from +**leejet/stable-diffusion.cpp** upstream releases. We want to own this like we own +**unslothai/llama.cpp**: a fork that builds and publishes OUR OWN prebuilt binaries on a +schedule, so we control reproducibility, integrity, and the macOS load floor. + +**Why native is CPU/Apple-only.** On a GPU host, diffusers + our optimizations (regional +`torch.compile` ~2.2x, cuDNN/flash attention, FP8/INT8/NVFP4 quant, first-block-cache) is +faster than sd.cpp's CUDA path, which has none of those levers — so GPU hosts route to +diffusers. Native sd.cpp only wins where diffusers is weak: **CPU and Apple**. Therefore we +build native binaries ONLY for the platforms where native is actually the faster engine, and +skip CUDA/ROCm/Vulkan entirely (GPU = diffusers/torch). This also makes the CI far cheaper. + +The Studio side is already prepared: `install_sd_cpp_prebuilt.py` reads `UNSLOTH_SD_CPP_REPO` +(repo override) + `UNSLOTH_SD_CPP_TAG` (pinned tag) and verifies the GitHub asset `digest` +(`_verify_sha256`). So the bulk of the work is the mirror repo + release CI; the Studio change +is a small default flip + resolver tweak. + +## Coverage (user-confirmed): CPU / Apple ONLY + +| Platform | Arch | Build | Runner | Notes | +|---|---|---|---|---| +| macOS | arm64 | Metal (`-DSD_METAL=ON`) | macos-26, `OSX_DEPLOYMENT_TARGET=14.0` | Apple fast path (diffusers/MPS weak) | +| macOS | x86_64 | CPU | macos-15-intel, `OSX_DEPLOYMENT_TARGET=13.3` | Intel Macs | +| Linux | x86_64 | CPU | ubuntu-22.04 (glibc 2.35) | **also covers WSL** (WSL = Linux x64) | +| Linux | aarch64 | CPU | ubuntu-24.04-arm | ARM servers | +| Windows | x86_64 | CPU | windows-2022 (MSVC+Ninja) | | + +**Explicitly out of scope:** CUDA, ROCm, Vulkan native builds; GPU runners; cudart bundling; +per-gfx matrices. GPU stays on diffusers/torch. + +## Reference pattern (verified this session) + +`unslothai/llama.cpp` builds via `.github/workflows/unsloth-prebuilt.yml` (orchestrator) + six +per-accel children + `scripts/unsloth/` helpers (`assemble_metadata.py`, `package_bundle.py`, +`assert_macho_minos.sh`). Mechanisms to mirror: `resolve` (supply-chain aging — only build a +release public >=6h; stamp build-info + Unsloth fingerprint; upload ONE source artifact all +children extract) -> per-platform children (build from the source artifact, load-gate, package, +upload) -> `assemble` (fingerprint gate + manifest/sha256 index + coverage gate + **atomic +draft->publish**, no partial releases). Template files fetched to `workspace_81/temp/llamacpp_workflows/`. + +## Key facts (verified) + +- leejet builds both `sd-cli` and `sd-server` (`examples/cli`, `examples/server`) — the mirror + ships both (sd-server is used by PR #6768's persistent server). +- leejet naming: `sd--bin-.zip`. + leejet already ships macOS arm64, Linux x64 CPU, Windows CPU — we ADD macOS x86_64 and Linux + aarch64 (the gaps in our target set), and rebuild the rest under our own fingerprint/integrity. +- Studio resolver (`resolve_release_asset`, `install_sd_cpp_prebuilt.py:88`): filters to `.zip`; + macOS = darwin/macos + arch token; Linux = `linux` + arch + (no accel marker for auto/cpu); + Windows = `bin-win` + `avx2` else any. For a CPU-only mirror the resolver needs essentially NO + change — macOS x86_64 and Linux aarch64 already match by arch token; just confirm the Windows + CPU asset resolves (contains `bin-win`, falls back to the plain build). + +## Design + +### A. Mirror repo (fork of leejet/stable-diffusion.cpp) + +Fork so upstream C++ stays intact; add only `.github/workflows/` + `scripts/unsloth/`. Adapt the +llama.cpp orchestrator, heavily simplified (no CUDA/ROCm/Vulkan, no PR-mix): + +- **`resolve`**: pick the upstream leejet tag with the >=6h aging window; reuse leejet's + `master--` as the mirror tag (keeps `UNSLOTH_SD_CPP_TAG` comparable to upstream); + stamp a source tarball with build-info + the "Compiled by the Unsloth team" fingerprint; upload + the source artifact. Skip-if-already-published like llama.cpp. +- **Build children** (reusable `workflow_call`), each `cmake -DSD_BUILD_EXAMPLES=ON` (cli+server): + - `macos` (arm64 Metal + x64 CPU): pinned `CMAKE_OSX_DEPLOYMENT_TARGET`, `@loader_path` rpath, + load-gate via `assert_macho_minos.sh` (adapted for `sd-cli`/`sd-server`). + - `cpu-linux` (x64 + arm64) and `cpu-windows` (x64, MSVC+Ninja). +- **Asset naming = leejet-compatible**, all `.zip`: + `sd--bin-Darwin-macOS-arm64.zip`, `sd--bin-Darwin-macOS-x86_64.zip`, + `sd--bin-Linux-Ubuntu-24.04-x86_64.zip`, `sd--bin-Linux-Ubuntu-24.04-aarch64.zip`, + `sd--bin-win-cpu-x64.zip`. +- **`assemble`**: fingerprint gate (verify the mark in every archive), generate + `sd-prebuilt-manifest.json` + `sd-prebuilt-sha256.json`, coverage gate (all 5 assets present), + atomic draft->publish. GitHub sets each asset `digest`, which the Studio already verifies. +- **Signing/notarization:** none. The Studio downloads via `urllib` (not a browser), so no macOS + quarantine xattr is set and Gatekeeper does not block CLI-run binaries (matches llama.cpp). + +### B. Studio-side switch (PR on the diffusion stack, after the mirror's first green release) + +Small, in `studio/install_sd_cpp_prebuilt.py` + its test: +1. `DEFAULT_REPO = "unslothai/stable-diffusion.cpp"`; `DEFAULT_TAG` = the mirror's first tag. +2. Confirm `resolve_release_asset` picks correctly for all 5 CPU/Apple hosts (add a Windows CPU + token only if the plain-`bin-win` fallback proves insufficient; likely no change needed). + Keep the leejet fallback (env override still points back upstream). +3. Extend `test_sd_cpp_install.py` `_ASSETS` to the mirror's 5-asset set; assert host->pick for + macOS arm64/x64, Linux x64/arm64, Windows x64; assert GPU hosts are unaffected (still diffusers). + +## Critical files + +- New (mirror repo): `.github/workflows/unsloth-sd-prebuilt.yml` (+ `-macos.yml`, `-cpu-linux.yml`, + `-cpu-windows.yml`), `scripts/unsloth/{assemble_metadata.py,package_bundle.py,assert_macho_minos.sh}`. +- Studio: `studio/install_sd_cpp_prebuilt.py`, `studio/backend/tests/test_sd_cpp_install.py`. +- Local templates to adapt: `workspace_81/temp/llamacpp_workflows/{unsloth-prebuilt.yml,unsloth-prebuilt-macos.yml,unsloth-prebuilt-cpu.yml}`. + +## Sequencing (chicken-and-egg) + +1. Build the mirror repo + CI; `publish=false` dry run to validate the 5-way matrix (~10-20 min, + no GPU runners so cheap). +2. First green **published** release with all 5 assets + manifest/sha256. +3. THEN the Studio PR flips `DEFAULT_REPO`/`DEFAULT_TAG` + resolver test (on the diffusion stack). + +## Verification + +- **Resolver unit tests** (hermetic): feed the mirror's 5 asset names to `resolve_release_asset` + for macOS arm64/x64, Linux x64/arm64, Windows x64 -> correct pick; and CUDA/GPU host -> still + routes to diffusers (native not selected). +- **CI dry run**: `publish=false` artifact-only run; inspect the 5 archives each contain `sd-cli` + (+ `sd-server`) and carry the fingerprint. +- **Live install smoke** (this Linux box): `UNSLOTH_SD_CPP_REPO=unslothai/stable-diffusion.cpp + python studio/install_sd_cpp_prebuilt.py --print-asset` then real `install()`, confirm + `sd-cli --version` + `sd-server` launch, and drive one native CPU generation via the Studio. +- **Integrity**: each published archive matches its manifest sha256 and the GitHub asset digest. + +## Staging (user-confirmed): fork + push CI now + +Execution order: +1. **Preflight permissions**: `gh auth status`; confirm the token can create/fork under the + `unslothai` org and enable Actions. If it CANNOT, stop and report (fall back to scaffold-only, + or a private fork under danielhanchen), rather than pushing somewhere unintended. +2. **Fork** leejet/stable-diffusion.cpp -> `unslothai/stable-diffusion.cpp` (clone into the + workspace to add files). Keep upstream C++ intact. +3. **Add CI + scripts** on a branch: `.github/workflows/unsloth-sd-prebuilt.yml` + + `-macos.yml`/`-cpu-linux.yml`/`-cpu-windows.yml`, `scripts/unsloth/*`. Commit as Daniel Han + (no AI/bot mentions, no emojis, no em dashes). `unset GH_TOKEN`/use `gh` creds for pushes that + touch `.github/workflows/*` (needs `workflow` scope). +4. **Dry run**: trigger the orchestrator with `publish=false` (artifact-only), confirm all 5 + archives build + carry `sd-cli`/`sd-server` + the fingerprint. Iterate until green. +5. **First publish**: `publish=true` (or let the schedule run) -> a real release with the 5 + assets + manifest/sha256. +6. **Studio PR** (section B) on the diffusion stack once the release tag exists. + +## Follow-ups (not this task) + +- Nightly schedule + auto-bump of the Studio `DEFAULT_TAG` (PR bot), like llama.cpp. +- Add GPU native builds later ONLY if a real need appears (today: GPU = diffusers/torch). diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 54ff06d2ce..51f56af84e 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -10,19 +10,24 @@ the diffusers dreambooth scripts, form ``noisy = (1 - sigma) * latents + sigma * predict the velocity, and regress it onto ``target = noise - latents``. The per-family differences (latent normalisation + packing, the transformer forward -signature, and the LoRA save entrypoint) live in small ``_FamilySpec`` objects; the loop -itself is family-agnostic. Verified against diffusers 0.38.0. +signature, embedding collation, and the LoRA save entrypoint) live in small ``_FamilySpec`` +objects; the loop itself is family-agnostic. Verified against diffusers 0.38.0. Memory: the text encoder(s) are the largest module (T5-XXL ~9 GB for FLUX, Qwen2.5-VL ~7 GB for Qwen-Image, Qwen3 for Z-Image), so captions are encoded ONCE up front and the encoders -are freed before the loop. The transformer trains as a QLoRA (nf4) adapter by default with -gradient checkpointing and 8-bit AdamW, so only the (small) LoRA params + optimizer state -and the frozen 4-bit base sit in VRAM during the loop. +are freed before the loop. VAE latents are likewise precomputed into a small CPU cache +(``cache_latents``) and the VAE freed: the cache stores the posterior's affine parameters +(mean/std folded through the family's latent normalisation), so every step still draws a +fresh VAE sample -- distribution-identical to encoding in the loop, without keeping the VAE +resident or paying a per-step encode. The transformer trains as a QLoRA (nf4) adapter by +default with gradient checkpointing and 8-bit AdamW, so only the (small) LoRA params + +optimizer state and the frozen 4-bit base sit in VRAM during the loop. """ from __future__ import annotations import gc +import os import random import time from contextlib import nullcontext @@ -35,11 +40,19 @@ from core.training.diffusion_train_common import ( DEFAULT_LORA_TARGETS, DiffusionLoraConfig, EventCb, + LATENT_CACHE_OVER_BUDGET, StopCb, + _apply_perf_flags, _assert_trusted_base_model, _emit, + _latent_cache_forced, + _latent_cache_over_budget, + _plan_cache_variants, _publish_to_lora_catalog, + _restore_perf_flags, discover_image_caption_pairs, + has_functional_torchao, + repo_is_prequantized, ) # Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks @@ -81,14 +94,27 @@ class _FamilySpec: lora_targets: tuple[str, ...] # bf16 only (Z-Image overflows fp16 and its RoPE/embedder run in fp32). force_bf16: bool - # Builds (pipe, transformer, vae) with the transformer loaded as a trainable nf4 QLoRA - # when qlora=True. Returns the pipeline (for save_lora_weights + encode_prompt), the - # transformer to attach LoRA to, and the VAE (kept resident for latent encoding). - load: Callable[..., tuple[Any, Any, Any]] + # Approximate dense-bf16 transformer weight size, used by base_precision="auto" to + # decide which mode fits the free VRAM (with headroom for activations + optimizer). + dense_bf16_gb: float + # Phased load, so the (multi-GB) transformer never has to coexist with the text + # encoders + VAE: ``load_conditioners`` builds the pipeline WITHOUT its transformer + # (encode_prompt + VAE only) and returns (pipe, vae); ``load_transformer`` loads the + # transformer alone (as a trainable nf4 QLoRA when qlora=True) once the conditioning + # modules are freed. Roughly halves peak VRAM for the big DiTs. + load_conditioners: Callable[..., tuple[Any, Any]] + load_transformer: Callable[..., Any] # Encode a list of captions -> a per-caption tuple of CPU tensors (the family's embeds). encode_prompts: Callable[..., list[tuple]] # Encode a pixel tensor [B,3,H,W] in [-1,1] -> latents (family-normalised, on device). encode_latents: Callable[..., Any] + # Encode a pixel tensor -> (A, B) affine posterior parameters so a per-step sample is + # A + B * randn (family normalisation folded in). B is None for a deterministic + # (mode-based) family. Used by the latent cache. + encode_latent_stats: Callable[..., tuple] + # Collate a list of per-caption embed tuples -> one batched tuple on device. ``pad_to`` + # pins a fixed text length for families with variable-length embeds (compile). + collate: Callable[..., tuple] # One transformer forward: (transformer, noisy, timesteps, sigmas, embeds_batch, cfg, # device, weight_dtype) -> model_pred aligned with target = noise - latents. forward: Callable[..., Any] @@ -97,15 +123,12 @@ class _FamilySpec: # ── shared flow-matching helpers ────────────────────────────────────────────── -def _get_sigmas(scheduler, timesteps, device, dtype, n_dim): - """Gather per-sample sigmas for ``timesteps`` and broadcast to ``n_dim`` (matches the - diffusers dreambooth get_sigmas helper).""" - import torch - - sigmas = scheduler.sigmas.to(device = device, dtype = dtype) - schedule_timesteps = scheduler.timesteps.to(device) - step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] - sigma = sigmas[step_indices].flatten() +def _gather_sigmas(scheduler, indices, device, dtype, n_dim): + """Gather per-sample sigmas for schedule ``indices`` and broadcast to ``n_dim``. + Index-based (no per-item search): ``indices`` are the positions ``_sample_timesteps`` + drew from ``scheduler.timesteps``, and ``scheduler.sigmas`` is aligned with it, so this + returns exactly what the diffusers ``get_sigmas`` timestep-matching helper would.""" + sigma = scheduler.sigmas[indices].to(device = device, dtype = dtype).flatten() while sigma.ndim < n_dim: sigma = sigma.unsqueeze(-1) return sigma @@ -114,7 +137,6 @@ def _get_sigmas(scheduler, timesteps, device, dtype, n_dim): def _sample_timesteps(scheduler, batch_size, device): """Logit-normal density timestep sampling (weighting_scheme='logit_normal'), returning (timesteps, indices) into the scheduler's schedule.""" - import torch from diffusers.training_utils import compute_density_for_timestep_sampling u = compute_density_for_timestep_sampling( @@ -127,7 +149,7 @@ def _sample_timesteps(scheduler, batch_size, device): num_train = scheduler.config.num_train_timesteps indices = (u * num_train).long().clamp(0, num_train - 1) timesteps = scheduler.timesteps.to(device)[indices].to(device) - return timesteps + return timesteps, indices def _encoders_to_device(pipe, device) -> None: @@ -156,12 +178,9 @@ def _bnb_4bit_config(): ) -def _repo_is_prequantized(base_model: str) -> bool: - """Heuristic: a repo whose name marks a bitsandbytes 4-bit build already ships a - quantized transformer, so we load it as-is rather than re-quantizing on the fly. A - dense (bf16) base instead gets on-the-fly nf4 quantization for QLoRA.""" - name = str(base_model or "").lower() - return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name +# Kept as a module name for existing callers/tests; the heuristic itself moved to +# diffusion_train_common so config validation can use it without importing this module. +_repo_is_prequantized = repo_is_prequantized def _load_quantized_transformer(transformer_cls, cfg): @@ -176,32 +195,204 @@ def _load_quantized_transformer(transformer_cls, cfg): ) -# ── FLUX.1-dev ──────────────────────────────────────────────────────────────── -def _flux_load(cfg, device, weight_dtype, qlora): +def _load_pipe_without_transformer(pipe_cls, cfg, device): + """Load a pipeline for conditioning only: ``transformer = None`` skips the multi-GB + denoiser entirely (the documented diffusers pattern), leaving just the text encoders + + tokenizers + VAE + scheduler. The transformer loads later, after these are freed.""" import torch - from diffusers import FluxPipeline, FluxTransformer2DModel - if qlora: - transformer = FluxTransformer2DModel.from_pretrained( + pipe = pipe_cls.from_pretrained( + cfg.base_model, + transformer = None, + torch_dtype = torch.bfloat16, + token = cfg.hf_token, + ) + pipe.vae.to(device, dtype = torch.float32) + return pipe, pipe.vae + + +def _load_dit_transformer(transformer_cls, cfg, device, base_precision): + """Load the transformer alone in the resolved ``base_precision``: + + - nf4: a prequant (bnb-4bit) repo carries its quantization config and loads 4-bit + as-is; a dense base is quantized to nf4 on the fly. The memory floor. + - bf16 / fp8: the dense transformer (fp8 converts its frozen linears to float8 + training compute AFTER the LoRA attaches; storage stays bf16). + - int8: the dense transformer quantized in place to torchao weight-only int8 (the + PEFT-attachable scheme), roughly halving the bf16 weight footprint.""" + import torch + + if base_precision == "nf4": + if not repo_is_prequantized(cfg.base_model): + return _load_quantized_transformer(transformer_cls, cfg) + transformer = transformer_cls.from_pretrained( cfg.base_model, subfolder = "transformer", - quantization_config = _bnb_4bit_config(), torch_dtype = torch.bfloat16, token = cfg.hf_token, ) - pipe = FluxPipeline.from_pretrained( - cfg.base_model, - transformer = transformer, - torch_dtype = torch.bfloat16, - token = cfg.hf_token, + # A prequant load is already device-placed by bitsandbytes. + if not getattr(transformer, "is_loaded_in_4bit", False): + transformer = transformer.to(device) + return transformer + + # Dense load for bf16 / fp8 / int8. int8 quantizes AFTER the LoRA attaches (see + # _int8_quantize_base): quantizing first makes peft dispatch its TorchaoLoraLinear + # wrapper, whose peft-0.18 constructor is incompatible with the torchao-0.16 config API + # (missing get_apply_tensor_subclass). + return transformer_cls.from_pretrained( + cfg.base_model, + subfolder = "transformer", + torch_dtype = torch.bfloat16, + token = cfg.hf_token, + ).to(device) + + +def _int8_quantize_base(transformer) -> None: + """torchao weight-only int8 on the big frozen linears, applied after add_adapter so + the base_layer inside each LoRA wrapper quantizes while the adapters stay high + precision. ``make_filter_fn`` (shared with the inference quant layer) keeps only + Linears with >= 512 features -- which also naturally skips the rank-sized LoRA + matrices -- and drops the M=1 modulation projections int8 kernels reject.""" + from core.inference.diffusion_transformer_quant import exclude_tokens_for_scheme, make_filter_fn + from torchao.quantization import Int8WeightOnlyConfig, quantize_ + + quantize_( + transformer, + Int8WeightOnlyConfig(), + filter_fn = make_filter_fn(512, exclude_name_tokens = exclude_tokens_for_scheme("int8")), + ) + + +def _fp8_module_filter(mod, fqn: str) -> bool: + """Which frozen linears get float8 training compute: skip anything LoRA-owned (the + adapters must stay high precision -- PEFT has no float8 base support), the output + projection, and shapes float8 kernels reject (dims not divisible by 16), matching the + diffusers FLUX2 reference filter.""" + import torch.nn as nn + + if not isinstance(mod, nn.Linear): + return False + if "lora_" in fqn: + return False + if fqn.endswith("proj_out") or ".proj_out." in fqn: + return False + return mod.in_features % 16 == 0 and mod.out_features % 16 == 0 + + +def _apply_fp8_training(transformer, on_event) -> bool: + """Convert the frozen base linears to torchao float8 training compute (dynamic scaling; + weights stay bf16 in memory). Applied AFTER add_adapter so the filter can exclude the + LoRA modules. Never fatal: on any failure the run continues in bf16 with a warning.""" + try: + from torchao.float8 import Float8LinearConfig, convert_to_float8_training + convert_to_float8_training( + transformer, + module_filter_fn = _fp8_module_filter, + config = Float8LinearConfig(pad_inner_dim = True), ) - else: - pipe = FluxPipeline.from_pretrained( - cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token - ) - transformer = pipe.transformer - pipe.vae.to(device, dtype = torch.float32) - return pipe, transformer, pipe.vae + return True + except Exception as exc: # noqa: BLE001 -- fp8 is an optimisation, never fatal + _emit(on_event, "warning", message = f"fp8 training unavailable, using bf16 compute: {exc}") + return False + + +def _pick_auto_precision( + prequant, + device, + free_gb, + dense_gb, + capability, + has_fp8, + has_torchao = True, +) -> str: + """Pure policy for base_precision="auto": nf4 for a prequant base or no CUDA; else the + fastest dense mode whose weights + headroom (activations, optimizer, cache) fit the + free VRAM at decision time. bf16 + regional compile is the measured speed winner + (2.3-2.6x over nf4 on B200); fp8 stays an explicit opt-in because torchao float8's + dynamic-scaling overhead made it SLOWER than compiled bf16 at LoRA-training shapes on + the same hardware. int8 must still materialise the full bf16 transformer before + ``quantize_`` shrinks it module-by-module, so its band requires the dense-load + transient (1.15x dense) to fit -- what int8 buys in that band is steady-state + headroom for activations and the latent cache, not load-time memory. int8 also needs + torchao at runtime (``_int8_quantize_base`` has no fallback, unlike fp8), so auto only + picks it when torchao is importable and drops to nf4 otherwise. ``capability``/``has_fp8`` + remain parameters so the policy can be revisited per GPU generation without changing + callers.""" + _ = capability, has_fp8 + if prequant or device != "cuda" or not free_gb or not dense_gb: + return "nf4" + if free_gb > dense_gb * 1.5: + return "bf16" + if free_gb > dense_gb * 1.15: + return "int8" if has_torchao else "nf4" + return "nf4" + + +def _resolve_base_precision(cfg, spec, device) -> str: + """Resolve "auto" against the live GPU (free VRAM measured BEFORE anything loads); + explicit modes pass through (normalized() already validated them against the repo and + compute dtype) but are re-checked against the live device here: the dense modes are + CUDA-only, and /info never advertises them on a host without a GPU, so an explicit + request from a stale or direct client fails fast instead of loading a full dense + transformer onto the CPU.""" + mode = (cfg.base_precision or "nf4").strip().lower() + if mode != "auto": + if mode in ("bf16", "int8", "fp8") and device != "cuda": + raise ValueError( + f"base_precision={mode!r} needs a CUDA GPU; this host has none. " + f"Use base_precision='nf4' or 'auto'." + ) + # int8 has no runtime fallback (_int8_quantize_base imports torchao unconditionally), + # so an explicit int8 against a missing torchao or the Windows-ROCm stub would leave + # the transformer dense with compile disabled as if it were int8 -- the memory saving + # silently gone and a likely OOM. The auto pick and /info already gate on a FUNCTIONAL + # torchao; apply the same gate to the explicit request so it fails fast with a clear + # message. fp8 keeps its own graceful fallback (_apply_fp8_training), so this is int8-only. + if mode == "int8" and not has_functional_torchao(): + raise ValueError( + "base_precision='int8' needs a functional torchao install; this host's " + "torchao is missing or the non-functional Windows-ROCm stub. Use " + "base_precision='nf4', 'bf16', or 'auto'." + ) + return mode + # auto may only resolve to the dense modes when the run uses bf16 compute, mirroring + # the normalized() rule for explicit dense modes; otherwise stay on the nf4 floor. + if getattr(cfg, "mixed_precision", "bf16") != "bf16": + return "nf4" + prequant = repo_is_prequantized(cfg.base_model) + free_gb = None + capability = None + has_fp8 = False + # int8 quantization has no runtime fallback, so gate the auto pick on a FUNCTIONAL + # torchao: a plain find_spec("torchao") is satisfied by the Windows-ROCm import stub, + # whose quantize_ is a no-op that would leave the transformer dense while compile is + # disabled as if it were int8. has_functional_torchao imports the exact symbols + # _int8_quantize_base uses and rejects the stub. + has_torchao = has_functional_torchao() + if device == "cuda": + try: + import torch + + free_gb = torch.cuda.mem_get_info()[0] / 1e9 + capability = torch.cuda.get_device_capability() + has_fp8 = hasattr(torch, "float8_e4m3fn") + except Exception: # noqa: BLE001 -- probe failure -> the safe mode + pass + return _pick_auto_precision( + prequant, device, free_gb, spec.dense_bf16_gb, capability, has_fp8, has_torchao + ) + + +# ── FLUX.1-dev ──────────────────────────────────────────────────────────────── +def _flux_load_conditioners(cfg, device, weight_dtype): + from diffusers import FluxPipeline + return _load_pipe_without_transformer(FluxPipeline, cfg, device) + + +def _flux_load_transformer(cfg, device, weight_dtype, base_precision): + from diffusers import FluxTransformer2DModel + return _load_dit_transformer(FluxTransformer2DModel, cfg, device, base_precision) def _flux_encode_prompts(pipe, captions, device): @@ -231,24 +422,66 @@ def _flux_encode_latents(vae, pixel_values): return lat -def _flux_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): +def _flux_encode_latent_stats(vae, pixel_values): import torch + + with torch.no_grad(): + dist = vae.encode(pixel_values.to(torch.float32)).latent_dist + scale = vae.config.scaling_factor + return (dist.mean - vae.config.shift_factor) * scale, dist.std * scale + + +def _flux_collate( + entries, + device, + weight_dtype, + pad_to = None, +): + import torch + + # FLUX embeds are fixed-length (encode_prompt pads to max_sequence_length), so a plain + # cat batches them; text_ids are shared position ids, identical across prompts. + pe = torch.cat([e[0] for e in entries]).to(device = device, dtype = weight_dtype) + pooled = torch.cat([e[1] for e in entries]).to(device = device, dtype = weight_dtype) + text_ids = entries[0][2].to(device = device, dtype = torch.float32) + return (pe, pooled, text_ids) + + +# Per-run cache of the step-invariant FLUX conditioning tensors (RoPE image ids + the +# guidance vector): their shapes are fixed once resolution/batch are, so rebuilding them +# every step is pure allocator churn. Cleared at run start (subprocess-local anyway). +_FLUX_STATIC: dict[tuple, tuple] = {} + + +def _flux_static_inputs(bsz, h, w, device): + import torch + from diffusers import FluxPipeline + + key = (bsz, h, w, str(device)) + hit = _FLUX_STATIC.get(key) + if hit is None: + # Position ids drive RoPE and are indices, not activations -- keep them float32 (the + # dtype diffusers' own pipeline builds) regardless of the bf16 training dtype. + img_ids = FluxPipeline._prepare_latent_image_ids(bsz, h // 2, w // 2, device, torch.float32) + guidance = torch.full((bsz,), 1.0, device = device, dtype = torch.float32) + hit = _FLUX_STATIC[key] = (img_ids, guidance) + return hit + + +def _flux_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): from diffusers import FluxPipeline pe, pooled, text_ids = embeds_batch bsz, c, h, w = noisy.shape packed = FluxPipeline._pack_latents(noisy, bsz, c, h, w) - # Position ids drive RoPE and are indices, not activations -- keep them float32 (the - # dtype diffusers' own pipeline builds) regardless of the bf16 training dtype. - img_ids = FluxPipeline._prepare_latent_image_ids(bsz, h // 2, w // 2, device, torch.float32) - guidance = torch.full((bsz,), 1.0, device = device, dtype = torch.float32) + img_ids, guidance = _flux_static_inputs(bsz, h, w, device) model_pred = transformer( hidden_states = packed, timestep = timesteps / 1000, guidance = guidance, - pooled_projections = pooled.to(weight_dtype), - encoder_hidden_states = pe.to(weight_dtype), - txt_ids = text_ids.to(torch.float32), + pooled_projections = pooled, + encoder_hidden_states = pe, + txt_ids = text_ids, img_ids = img_ids, return_dict = False, )[0] @@ -265,19 +498,17 @@ def _flux_save(pipe_cls, out_dir, transformer_lora_layers): # ── Qwen-Image ──────────────────────────────────────────────────────────────── -def _qwen_load(cfg, device, weight_dtype, qlora): - import torch - from diffusers import QwenImagePipeline, QwenImageTransformer2DModel +def _qwen_load_conditioners(cfg, device, weight_dtype): + from diffusers import QwenImagePipeline + return _load_pipe_without_transformer(QwenImagePipeline, cfg, device) + +def _qwen_load_transformer(cfg, device, weight_dtype, base_precision): # The prequant default (unsloth/Qwen-Image-2512-unsloth-bnb-4bit) ships the transformer - # 4-bit, so from_pretrained loads it trainable as-is. A dense (bf16) base -- the 20B - # Qwen/Qwen-Image -- is quantized to nf4 on the fly so QLoRA still fits. - kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token} - if qlora and not _repo_is_prequantized(cfg.base_model): - kwargs["transformer"] = _load_quantized_transformer(QwenImageTransformer2DModel, cfg) - pipe = QwenImagePipeline.from_pretrained(cfg.base_model, **kwargs) - pipe.vae.to(device, dtype = torch.float32) - return pipe, pipe.transformer, pipe.vae + # 4-bit and loads trainable as-is under nf4; the dense modes need the 20B + # Qwen/Qwen-Image base. + from diffusers import QwenImageTransformer2DModel + return _load_dit_transformer(QwenImageTransformer2DModel, cfg, device, base_precision) def _qwen_encode_prompts(pipe, captions, device): @@ -297,6 +528,15 @@ def _qwen_encode_prompts(pipe, captions, device): return out +def _qwen_latent_affine(vae, ref): + import torch + + z = vae.config.z_dim + mean = torch.tensor(vae.config.latents_mean, device = ref.device, dtype = ref.dtype) + std = torch.tensor(vae.config.latents_std, device = ref.device, dtype = ref.dtype) + return mean.view(1, z, 1, 1, 1), std.view(1, z, 1, 1, 1) + + def _qwen_encode_latents(vae, pixel_values): import torch @@ -305,16 +545,53 @@ def _qwen_encode_latents(vae, pixel_values): px = pixel_values.to(torch.float32).unsqueeze(2) # [B,3,1,H,W] with torch.no_grad(): lat = vae.encode(px).latent_dist.sample() # [B,16,1,h,w] - z = vae.config.z_dim - mean = torch.tensor(vae.config.latents_mean, device = lat.device, dtype = lat.dtype) - std = torch.tensor(vae.config.latents_std, device = lat.device, dtype = lat.dtype) - mean = mean.view(1, z, 1, 1, 1) - std = std.view(1, z, 1, 1, 1) + mean, std = _qwen_latent_affine(vae, lat) return (lat - mean) / std -def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): +def _qwen_encode_latent_stats(vae, pixel_values): import torch + + px = pixel_values.to(torch.float32).unsqueeze(2) + with torch.no_grad(): + dist = vae.encode(px).latent_dist + mean, std = _qwen_latent_affine(vae, dist.mean) + return (dist.mean - mean) / std, dist.std / std + + +def _qwen_collate( + entries, + device, + weight_dtype, + pad_to = None, +): + import torch + import torch.nn.functional as F + + # Qwen embeds are variable-length: pad to the batch max (or a pinned ``pad_to`` bucket + # under compile so the graph shape stays fixed) and batch the validity mask with them. + seqs = [e[0].shape[1] for e in entries] + target = max(pad_to or 0, max(seqs)) + pes, masks = [], [] + for pe, mask in entries: + s = pe.shape[1] + if mask is None: + mask = torch.ones((1, s), dtype = torch.int64) + if s < target: + pe = F.pad(pe, (0, 0, 0, target - s)) + mask = F.pad(mask, (0, target - s)) + pes.append(pe) + masks.append(mask) + pe_b = torch.cat(pes).to(device = device, dtype = weight_dtype) + mask_b = torch.cat(masks).to(device) + # A single unpadded sample keeps the legacy None mask (identical math; avoids any + # behaviour delta for existing single-image runs whose pipeline returned None). + if len(entries) == 1 and entries[0][1] is None and target == seqs[0]: + mask_b = None + return (pe_b, mask_b) + + +def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): from diffusers import QwenImagePipeline pe, mask = embeds_batch @@ -325,8 +602,8 @@ def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, devi img_shapes = [[(1, h // 2, w // 2)]] * bsz pred = transformer( hidden_states = packed, - encoder_hidden_states = pe.to(weight_dtype), - encoder_hidden_states_mask = mask.to(device) if mask is not None else None, + encoder_hidden_states = pe, + encoder_hidden_states_mask = mask, timestep = timesteps / 1000, img_shapes = img_shapes, return_dict = False, @@ -344,18 +621,16 @@ def _qwen_save(pipe_cls, out_dir, transformer_lora_layers): # ── Z-Image ─────────────────────────────────────────────────────────────────── -def _zimage_load(cfg, device, weight_dtype, qlora): - import torch - from diffusers import ZImagePipeline, ZImageTransformer2DModel +def _zimage_load_conditioners(cfg, device, weight_dtype): + from diffusers import ZImagePipeline + return _load_pipe_without_transformer(ZImagePipeline, cfg, device) - # Prequant default loads 4-bit as-is; the dense bf16 Tongyi-MAI base is quantized to nf4 - # on the fly. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows). - kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token} - if qlora and not _repo_is_prequantized(cfg.base_model): - kwargs["transformer"] = _load_quantized_transformer(ZImageTransformer2DModel, cfg) - pipe = ZImagePipeline.from_pretrained(cfg.base_model, **kwargs) - pipe.vae.to(device, dtype = torch.float32) - return pipe, pipe.transformer, pipe.vae + +def _zimage_load_transformer(cfg, device, weight_dtype, base_precision): + # Prequant default loads 4-bit as-is under nf4; the dense modes use the bf16 Tongyi-MAI + # base. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows). + from diffusers import ZImageTransformer2DModel + return _load_dit_transformer(ZImageTransformer2DModel, cfg, device, base_precision) def _zimage_encode_prompts(pipe, captions, device): @@ -384,16 +659,31 @@ def _zimage_encode_latents(vae, pixel_values): return (lat - vae.config.shift_factor) * vae.config.scaling_factor +def _zimage_encode_latent_stats(vae, pixel_values): + # Z-Image trains from the posterior mode (deterministic), so the cached entry is the + # final latent itself: B is None and the loop skips the per-step sampling draw. + return _zimage_encode_latents(vae, pixel_values), None + + +def _zimage_collate( + entries, + device, + weight_dtype, + pad_to = None, +): + caps = [e[0].to(device = device, dtype = weight_dtype) for e in entries] + return (caps,) + + def _zimage_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): import torch - (emb,) = embeds_batch + (caps,) = embeds_batch # List I/O: one [C,1,H,W] latent + one [seq,2560] caption per sample. The timestep # convention is REVERSED ((1000 - t) / 1000) and the prediction is NEGATED. x_list = list(noisy.unsqueeze(2).unbind(dim = 0)) - cap_list = [emb.to(device = device, dtype = weight_dtype)] t_norm = (1000 - timesteps) / 1000 - out = transformer(x_list, t_norm, cap_list, return_dict = False)[0] + out = transformer(x_list, t_norm, list(caps), return_dict = False)[0] return -torch.stack(out, dim = 0).squeeze(2) @@ -411,9 +701,13 @@ _SPECS: dict[str, _FamilySpec] = { family = "flux.1", lora_targets = _FLUX_TARGETS, force_bf16 = False, - load = _flux_load, + dense_bf16_gb = 23.8, + load_conditioners = _flux_load_conditioners, + load_transformer = _flux_load_transformer, encode_prompts = _flux_encode_prompts, encode_latents = _flux_encode_latents, + encode_latent_stats = _flux_encode_latent_stats, + collate = _flux_collate, forward = _flux_forward, save = _flux_save, ), @@ -421,9 +715,13 @@ _SPECS: dict[str, _FamilySpec] = { family = "qwen-image", lora_targets = _QWEN_TARGETS, force_bf16 = True, - load = _qwen_load, + dense_bf16_gb = 41.0, + load_conditioners = _qwen_load_conditioners, + load_transformer = _qwen_load_transformer, encode_prompts = _qwen_encode_prompts, encode_latents = _qwen_encode_latents, + encode_latent_stats = _qwen_encode_latent_stats, + collate = _qwen_collate, forward = _qwen_forward, save = _qwen_save, ), @@ -431,9 +729,13 @@ _SPECS: dict[str, _FamilySpec] = { family = "z-image", lora_targets = _ZIMAGE_TARGETS, force_bf16 = True, - load = _zimage_load, + dense_bf16_gb = 12.3, + load_conditioners = _zimage_load_conditioners, + load_transformer = _zimage_load_transformer, encode_prompts = _zimage_encode_prompts, encode_latents = _zimage_encode_latents, + encode_latent_stats = _zimage_encode_latent_stats, + collate = _zimage_collate, forward = _zimage_forward, save = _zimage_save, ), @@ -456,18 +758,32 @@ def _assert_gated_access(base_model: str, hf_token: Optional[str]) -> None: ) -def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng): - """Load an image -> a normalised [3,H,W] tensor in [-1,1]. Same geometry as the SDXL - loader but without the SDXL time-ids (DiT families don't use them).""" - import numpy as np - import torch +def _open_resized(path, resolution): + """Open + EXIF-orient + short-side resize to ``resolution`` (same geometry as the SDXL + loader). Returns the resized PIL image and its (rw, rh).""" from PIL import Image, ImageOps img = ImageOps.exif_transpose(Image.open(path)).convert("RGB") w0, h0 = img.size scale = resolution / min(w0, h0) rw, rh = max(resolution, round(w0 * scale)), max(resolution, round(h0 * scale)) - img = img.resize((rw, rh), Image.LANCZOS) + return img.resize((rw, rh), Image.LANCZOS), rw, rh + + +def _to_unit_tensor(img): + import numpy as np + import torch + + arr = np.asarray(img, dtype = np.float32) / 255.0 + return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0 + + +def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng): + """Load an image -> a normalised [3,H,W] tensor in [-1,1]. Same geometry as the SDXL + loader but without the SDXL time-ids (DiT families don't use them).""" + from PIL import Image + + img, rw, rh = _open_resized(path, resolution) if center_crop: left, top = (rw - resolution) // 2, (rh - resolution) // 2 else: @@ -476,8 +792,192 @@ def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng): img = img.crop((left, top, left + resolution, top + resolution)) if random_flip and rng.random() < 0.5: img = img.transpose(Image.FLIP_LEFT_RIGHT) - arr = np.asarray(img, dtype = np.float32) / 255.0 - return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0 + return _to_unit_tensor(img) + + +def _load_pixel_tensor_planned(path, resolution, center_crop, u_left, u_top, flip): + """Deterministic variant of ``_load_pixel_tensor`` for the latent cache: the crop comes + as unit fractions (mapped uniformly over the same inclusive integer range ``randint`` + draws from) and the flip as a bool. ``center_crop`` reproduces the exact legacy + floor-div center so a cached center-crop run matches the uncached one bit-for-bit.""" + from PIL import Image + + img, rw, rh = _open_resized(path, resolution) + if center_crop: + left, top = (rw - resolution) // 2, (rh - resolution) // 2 + else: + left = min(int(u_left * (rw - resolution + 1)), max(0, rw - resolution)) + top = min(int(u_top * (rh - resolution + 1)), max(0, rh - resolution)) + img = img.crop((left, top, left + resolution, top + resolution)) + if flip: + img = img.transpose(Image.FLIP_LEFT_RIGHT) + return _to_unit_tensor(img) + + +def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_event, check_stop): + """Precompute the per-image latent posterior cache: for each planned crop/flip variant, + encode once and store the affine (A, B) pair on CPU (pinned when possible) in fp32. The + stats stay fp32 so the per-step sample happens in fp32 and only the RESULT is cast to + weight_dtype, matching the in-loop path (encode fp32 -> sample/normalise fp32 -> + .to(weight_dtype)); fp32 doubles the cache RAM over bf16 but the cache is tiny (a handful + of latents per image). Returns None if the build was interrupted by a stop request.""" + + plan = _plan_cache_variants( + len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed + ) + + def _hold(t): + if t is None: + return None + import torch + + t = t.to(torch.float32).cpu() + if device == "cuda": + try: + t = t.pin_memory() + except RuntimeError: + pass + return t + + cache: list[list[tuple]] = [] + total = len(image_paths) + total_variants = sum(len(v) for v in plan) + forced = _latent_cache_forced() + gated = False + for i, path in enumerate(image_paths): + variants = [] + for u_left, u_top, flip in plan[i]: + px = ( + _load_pixel_tensor_planned( + path, cfg.resolution, cfg.center_crop, u_left, u_top, flip + ) + .unsqueeze(0) + .to(device) + ) + a, b = spec.encode_latent_stats(vae, px) + a, b = _hold(a), _hold(b) + if not forced and not gated: + # Size-gate the automatic cache off the first REAL encoded variant, before + # building the rest: packed 16-channel DiT latents x variants x images of two + # fp32 tensors can exhaust host/pinned RAM. Over budget we bail with the VAE + # still resident so the loop encodes latents per step instead. ``b`` is None + # for a deterministic-latent family, so only ``a`` contributes bytes there. + per_variant = a.numel() * a.element_size() + if b is not None: + per_variant += b.numel() * b.element_size() + if _latent_cache_over_budget(per_variant, total_variants): + _emit( + on_event, + "warning", + message = ( + "Latent cache disabled: estimated " + f"{per_variant * total_variants / 1024 ** 3:.1f} GiB over the " + "budget; encoding latents per step instead. Set " + "UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE=1 to keep it." + ), + ) + return LATENT_CACHE_OVER_BUDGET + gated = True + variants.append((a, b)) + cache.append(variants) + if (i + 1) % 4 == 0 or i + 1 == total: + _emit(on_event, "preparing", stage = "cache_latents", done = i + 1, total = total) + if check_stop(): + return None + return cache + + +def _sample_cached_latents(cache, idxs, variant_rng, device, weight_dtype): + """Draw one latent per index from the cache: pick a variant, then sample the posterior + (A + B * randn) when the family is stochastic. Fresh noise per step, exactly like an + in-loop ``latent_dist.sample()``. The cached stats are fp32, so the sample is drawn in + fp32 and only the RESULT is cast to weight_dtype (matching the in-loop path's + ``encode_latents(...).to(weight_dtype)``).""" + import torch + + parts_a, parts_b = [], [] + for i in idxs: + variants = cache[i] + a, b = variants[variant_rng.randrange(len(variants))] if len(variants) > 1 else variants[0] + parts_a.append(a) + parts_b.append(b) + lat_a = torch.cat(parts_a).to(device, non_blocking = True) + if parts_b[0] is None: + return lat_a.to(dtype = weight_dtype) + lat_b = torch.cat(parts_b).to(device, non_blocking = True) + return (lat_a + lat_b * torch.randn_like(lat_a)).to(dtype = weight_dtype) + + +def _should_compile( + cfg, + base_is_bnb, + device, + base_precision = "nf4", +) -> bool: + mode = (cfg.compile_transformer or "auto").strip().lower() + if device != "cuda" or mode == "off": + return False + # torch.compile cannot trace the torchao int8 subclass in training (inductor rejects + # the aliased subclass graph outputs), so int8 always runs eager. + if base_precision == "int8": + return False + if mode == "on": + return True + # auto: regional compile is the whole point of the dense modes (measured 2.6x on + # Z-Image bf16) but fragile over bitsandbytes 4-bit modules (graph breaks in the + # dequant path), so it stays off for QLoRA. fp8 is only competitive compiled. + return base_precision in ("bf16", "fp8") + + +def _maybe_compile_transformer( + transformer, + cfg, + base_is_bnb, + device, + on_event, + base_precision = "nf4", +) -> bool: + """Regionally compile the transformer blocks (diffusers compile_repeated_blocks) after + the LoRA is attached. Never fatal: a wrap failure falls back to eager with a warning + event, and dynamo's suppress_errors keeps a frame that fails to COMPILE at the first + step running eager instead of raising mid-run.""" + if not _should_compile(cfg, base_is_bnb, device, base_precision): + if base_precision == "fp8": + _emit( + on_event, + "warning", + message = "fp8 training without torch.compile is slow; enable compile for the speedup.", + ) + return False + import torch + + fn = getattr(transformer, "compile_repeated_blocks", None) + if not callable(fn): + _emit( + on_event, "warning", message = "torch.compile unavailable for this model; running eager." + ) + return False + try: + dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None) + if dynamo_cfg is not None: + # Heterogeneous-block DiTs (Z-Image: ~11 distinct block shapes) exceed dynamo's + # default recompile limit of 8; bump it the same way the inference speed layer + # does (diffusers' documented regional-compile fix). + for attr in ("recompile_limit", "cache_size_limit"): + if hasattr(dynamo_cfg, attr): + setattr(dynamo_cfg, attr, max(getattr(dynamo_cfg, attr) or 0, 64)) + if hasattr(dynamo_cfg, "suppress_errors"): + dynamo_cfg.suppress_errors = True + # dynamic=True matches the inference speed layer's proven default: on torch 2.10 / + # B200 the dynamic=False specialisation fused a gemm_and_bias epilogue that failed + # with CUBLAS_STATUS_EXECUTION_FAILED then an illegal memory access on the FLUX + # training graph. fullgraph only on a dense base: bnb 4-bit layers graph-break by + # design. + fn(fullgraph = not base_is_bnb, dynamic = True) + return True + except Exception as exc: # noqa: BLE001 -- optimisation only, never fatal + _emit(on_event, "warning", message = f"torch.compile disabled (eager fallback): {exc}") + return False def run_dit_lora_training( @@ -503,14 +1003,10 @@ def run_dit_lora_training( ) import torch - import torch.nn.functional as F - from diffusers import FlowMatchEulerDiscreteScheduler - from diffusers.training_utils import cast_training_params - from peft import LoraConfig - from peft.utils import get_peft_model_state_dict rng = random.Random(cfg.seed) torch.manual_seed(cfg.seed) + _FLUX_STATIC.clear() save_on_stop = True @@ -528,15 +1024,14 @@ def run_dit_lora_training( device = "cuda" if torch.cuda.is_available() else "cpu" # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is # unsupported for real runs but keeps import/unit tests architecture-agnostic). - # Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the - # run would otherwise die deep in model load with an opaque dtype error. + # Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the run + # would otherwise die deep in model load with an opaque dtype error. if device == "cuda" and not torch.cuda.is_bf16_supported(): raise ValueError( "This trainer requires a bfloat16-capable GPU (Ampere or newer); " "this CUDA device does not support bf16." ) weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32 - use_lora_targets = _select_lora_targets(cfg.lora_target_modules, spec.lora_targets) _assert_trusted_base_model(cfg.base_model) _assert_gated_access(cfg.base_model, cfg.hf_token) @@ -551,9 +1046,42 @@ def run_dit_lora_training( ) return str(out_dir) - # QLoRA by default for the big DiTs (nf4 transformer). The prequant Qwen/Z-Image repos - # are already 4-bit; FLUX quantizes its transformer on the fly. - pipe, transformer, vae = spec.load(cfg, device, weight_dtype, qlora = True) + # TF32 / cudnn.benchmark for the run, restored on the way out (the trainer subprocess is + # disposable, but restoring keeps in-process callers -- tests, notebooks -- clean). + perf_snap = _apply_perf_flags(cfg, device) + try: + return _train_dit( + cfg, + spec, + pairs, + rng, + device, + weight_dtype, + on_event, + _check_stop, + lambda: save_on_stop, + ) + finally: + _restore_perf_flags(perf_snap) + + +def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_stop, _save_on_stop): + """The body of ``run_dit_lora_training``, split out so the backend perf flags are + snapshot/restored around it in exactly one place.""" + import torch + import torch.nn.functional as F + from diffusers import FlowMatchEulerDiscreteScheduler + from diffusers.optimization import get_scheduler + from diffusers.training_utils import cast_training_params + from peft import LoraConfig + from peft.utils import get_peft_model_state_dict + + use_lora_targets = _select_lora_targets(cfg.lora_target_modules, spec.lora_targets) + out_dir = Path(cfg.output_dir).expanduser() + + # Phase 1: conditioning only. The pipeline loads WITHOUT its transformer, so the text + # encoders + VAE never share VRAM with the multi-GB denoiser. + pipe, vae = spec.load_conditioners(cfg, device, weight_dtype) # Precompute all caption embeddings, then free the (large) text encoder(s): captions are # constant and the encoders are frozen, so this is exact and the biggest memory win. @@ -567,6 +1095,51 @@ def run_dit_lora_training( if device == "cuda": torch.cuda.empty_cache() + # Phase 2: the VAE latent cache, then free the VAE too (see module docstring: the cache + # keeps the posterior affine parameters, so per-step sampling noise is preserved). + use_cache = cfg.cache_latents and os.environ.get( + "UNSLOTH_DIFFUSION_NO_LATENT_CACHE", "" + ) not in ("1", "true") + latent_cache = None + if use_cache: + latent_cache = _build_latent_cache( + spec, vae, image_paths, cfg, device, weight_dtype, on_event, _check_stop + ) + if latent_cache is LATENT_CACHE_OVER_BUDGET: + # The estimated cache exceeded the host-memory budget; keep the VAE resident and + # fall through to the in-loop encode path (latent_cache stays None). + latent_cache = None + elif latent_cache is None: # stopped during the cache build; nothing trained yet + _emit( + on_event, + "complete", + output_dir = str(out_dir), + lora_path = None, + stopped = True, + steps_run = 0, + ) + return str(out_dir) + else: + try: + pipe.vae = None + except Exception: # noqa: BLE001 -- a pipeline without a settable vae keeps it + pass + del vae + vae = None + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + # Variant picks use their own stream so the training loop's index/noise draws stay on + # the same seed-deterministic sequence whether or not the cache is enabled. + variant_rng = random.Random(cfg.seed + 1) + + # Phase 3: only now load the transformer, in the resolved base precision (nf4 QLoRA by + # default; bf16 / int8 / fp8 are the dense speed modes; "auto" picks from free VRAM + # measured before the load). + base_precision = _resolve_base_precision(cfg, spec, device) + transformer = spec.load_transformer(cfg, device, weight_dtype, base_precision) + base_is_bnb = base_precision == "nf4" + # Freeze the base; attach the trainable LoRA to the transformer. transformer.requires_grad_(False) transformer.add_adapter( @@ -591,77 +1164,95 @@ def run_dit_lora_training( cast_training_params(transformer, dtype = torch.float32) lora_params = [p for p in transformer.parameters() if p.requires_grad] - optimizer = _make_optimizer(lora_params, cfg.learning_rate) - # One lr_sched.step() per optimizer update (cfg.train_steps total), matching the - # SDXL trainer: counting micro-steps instead would stretch warmup past the run. - from diffusers.optimization import get_scheduler + # int8 / fp8 convert the frozen base linears AFTER the LoRA attaches, so the adapter + # modules are excluded and stay high precision. + if base_precision == "int8": + _int8_quantize_base(transformer) + if base_precision == "fp8" and not _apply_fp8_training(transformer, on_event): + base_precision = "bf16" + compiled = _maybe_compile_transformer( + transformer, cfg, base_is_bnb, device, on_event, base_precision + ) + # Compiled Qwen graphs need one fixed text length across steps: pin the pad bucket to + # the dataset's longest caption (encode_prompt already caps it at 1024 tokens). + qwen_pad_to = None + if compiled and spec.family == "qwen-image": + qwen_pad_to = max(e[0].shape[1] for e in caption_embeds.values()) + + optimizer = _make_optimizer(lora_params, cfg.learning_rate) + scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + cfg.base_model, subfolder = "scheduler", token = cfg.hf_token + ) + # The LR schedule advances once per optimizer update, so warmup/decay are counted in + # optimizer steps (matching the SDXL trainer; multiplying by the accumulation factor + # would stretch warmup past the run and never reach the decay). lr_sched = get_scheduler( cfg.lr_scheduler, optimizer = optimizer, num_warmup_steps = cfg.lr_warmup_steps, num_training_steps = cfg.train_steps, ) - scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - cfg.base_model, subfolder = "scheduler", token = cfg.hf_token - ) - _emit(on_event, "model_load_completed") + _emit(on_event, "model_load_completed", compiled = compiled, base_precision = base_precision) transformer.train() + n_images = len(image_paths) + batch_size = cfg.train_batch_size stopped = False running_loss = 0.0 peak_gb = 0.0 t_start = time.time() + t_steady = None done = 0 - # Honor train_batch_size by folding it into the micro-step count: averaging the - # gradient over batch * accum single-image passes is mathematically identical to - # true batching with a mean loss, and keeps the QLoRA memory profile flat (one - # image's activations at a time). Previously batch_size > 1 silently trained at 1. - micro_steps = cfg.gradient_accumulation_steps * cfg.train_batch_size + # bf16 autocast around the forward + loss, matching the diffusers dreambooth scripts' + # accelerator.autocast: it reconciles the fp32 LoRA params with the bnb 4-bit base + # matmuls in one compute dtype. Without it the 4-bit backward on FLUX dies with an + # illegal-address / CUBLAS failure. + autocast = ( + torch.autocast(device_type = "cuda", dtype = torch.bfloat16) + if device == "cuda" + else nullcontext() + ) for opt_step in range(cfg.train_steps): optimizer.zero_grad(set_to_none = True) step_loss = 0.0 - for _ in range(micro_steps): - i = rng.randrange(len(image_paths)) - px = ( - _load_pixel_tensor( - image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng + for _ in range(cfg.gradient_accumulation_steps): + idxs = [rng.randrange(n_images) for _ in range(batch_size)] + if latent_cache is not None: + latents = _sample_cached_latents( + latent_cache, idxs, variant_rng, device, weight_dtype ) - .unsqueeze(0) - .to(device) - ) - latents = spec.encode_latents(vae, px).to(weight_dtype) + else: + px = torch.stack( + [ + _load_pixel_tensor( + image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng + ) + for i in idxs + ] + ).to(device) + latents = spec.encode_latents(vae, px).to(weight_dtype) noise = torch.randn_like(latents) - timesteps = _sample_timesteps(scheduler, latents.shape[0], device) - sigmas = _get_sigmas(scheduler, timesteps, device, weight_dtype, latents.ndim) + timesteps, t_indices = _sample_timesteps(scheduler, latents.shape[0], device) + sigmas = _gather_sigmas(scheduler, t_indices, device, weight_dtype, latents.ndim) noisy = (1.0 - sigmas) * latents + sigmas * noise - emb = caption_embeds[captions[i]] - emb_dev = tuple( - t.to(device = device, dtype = weight_dtype) - if (t is not None and t.is_floating_point()) - else (t.to(device) if t is not None else None) - for t in emb - ) - # bf16 autocast around the forward + loss, matching the diffusers dreambooth - # scripts' accelerator.autocast: it reconciles the fp32 LoRA params with the - # bnb 4-bit base matmuls in one compute dtype. Without it the 4-bit backward - # on FLUX dies with an illegal-address / CUBLAS failure. - autocast = ( - torch.autocast(device_type = "cuda", dtype = torch.bfloat16) - if device == "cuda" - else nullcontext() + embeds = spec.collate( + [caption_embeds[captions[i]] for i in idxs], + device, + weight_dtype, + pad_to = qwen_pad_to, ) with autocast: model_pred = spec.forward( - transformer, noisy, timesteps, sigmas, emb_dev, cfg, device, weight_dtype + transformer, noisy, timesteps, sigmas, embeds, cfg, device, weight_dtype ) target = noise - latents loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean") - (loss / micro_steps).backward() - step_loss += float(loss.detach()) / micro_steps + (loss / cfg.gradient_accumulation_steps).backward() + step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps grad_norm: Optional[float] = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: @@ -673,14 +1264,19 @@ def run_dit_lora_training( running_loss += step_loss done = opt_step + 1 + now = time.time() + if done == 1: + # Step 1 pays the one-time costs (cudnn autotune, torch.compile warmup), so the + # reported rate starts after it and reflects the steady state. + t_steady = now if done % cfg.log_every == 0 or done == cfg.train_steps: if device == "cuda": peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) - sps = round( - (done * cfg.train_batch_size * cfg.gradient_accumulation_steps) - / max(time.time() - t_start, 1e-6), - 3, - ) + per_step = batch_size * cfg.gradient_accumulation_steps + if t_steady is not None and done > 1: + sps = round((done - 1) * per_step / max(now - t_steady, 1e-6), 3) + else: + sps = round(done * per_step / max(now - t_start, 1e-6), 3) _emit( on_event, "progress", @@ -697,10 +1293,9 @@ def run_dit_lora_training( stopped = True break - out_dir = Path(cfg.output_dir).expanduser() lora_path: Optional[str] = None catalog_path: Optional[str] = None - if not (stopped and not save_on_stop): + if not (stopped and not _save_on_stop()): out_dir.mkdir(parents = True, exist_ok = True) layers = get_peft_model_state_dict(transformer) spec.save(pipe, str(out_dir), layers) @@ -716,19 +1311,28 @@ def run_dit_lora_training( base_model = cfg.base_model, stopped = stopped, steps_run = done if cfg.train_steps else 0, + wall_seconds = round(time.time() - t_start, 1), ) return str(out_dir) def _make_optimizer(params, lr): """8-bit AdamW (bitsandbytes) when available -- half the optimizer state, no accuracy - regression for LoRA -- else the torch AdamW fallback.""" + regression for LoRA -- else torch AdamW, fused on CUDA (with a fallback when this + build/device lacks the fused kernel).""" import torch + try: import bitsandbytes as bnb return bnb.optim.AdamW8bit(params, lr = lr) except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW - return torch.optim.AdamW(params, lr = lr) + pass + if torch.cuda.is_available(): + try: + return torch.optim.AdamW(params, lr = lr, fused = True) + except Exception: # noqa: BLE001 -- fused unsupported on this build/device + pass + return torch.optim.AdamW(params, lr = lr) def _free_text_encoders(pipe) -> None: diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index db2eca1011..8c3a79afaf 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -20,6 +20,13 @@ Design: - ``run_diffusion_training_process`` is the thin mp.Queue adapter; it dispatches to the trainer registered for the resolved family (SDXL here, DiT families in a follow-up). ``main`` is a CLI. + +Memory/perf: captions are encoded once up front and the CLIP text encoders freed; VAE +latents are likewise precomputed into a small CPU cache (``cache_latents``) and the VAE +freed. The cache stores the posterior's affine pair (mean/std, scale folded in), so every +step still draws a fresh VAE sample -- distribution-identical to encoding in the loop, +without keeping the VAE resident or paying a per-step encode. TF32 matmuls + cudnn +autotuning are enabled for the run under ``cfg.enable_tf32``. """ from __future__ import annotations @@ -40,12 +47,18 @@ from core.training.diffusion_train_common import ( # noqa: F401 EventCb, StopCb, DiffusionLoraConfig, + LATENT_CACHE_OVER_BUDGET, + _apply_perf_flags, _assert_trusted_base_model, _coerce_gradient_checkpointing, _config_from_dict, _CONFIG_ALIASES, _emit, + _latent_cache_forced, + _latent_cache_over_budget, + _plan_cache_variants, _publish_to_lora_catalog, + _restore_perf_flags, discover_image_caption_pairs, get_trainer, ) @@ -99,6 +112,43 @@ def _load_image_tensor( return tensor, time_ids +def _load_image_tensor_planned( + path: str, resolution: int, center_crop: bool, u_left: float, u_top: float, flip: bool +) -> tuple[Any, tuple[int, int, int, int, int, int]]: + """Deterministic variant of ``_load_image_tensor`` for the latent cache: the crop comes + as unit fractions (mapped uniformly over the same inclusive integer range ``randint`` + draws from) and the flip as a bool. Geometry (EXIF transpose, LANCZOS short-side resize, + the SDXL ``add_time_ids`` from the original size + actual crop offset) matches + ``_load_image_tensor`` exactly; ``center_crop`` reproduces the legacy floor-div center + bit-for-bit. The flip does not change time_ids (only the mirrored crop_left does).""" + import numpy as np + import torch + from PIL import Image, ImageOps + + img = ImageOps.exif_transpose(Image.open(path)).convert("RGB") + original_w, original_h = img.size + scale = resolution / min(original_w, original_h) + resized_w = max(resolution, round(original_w * scale)) + resized_h = max(resolution, round(original_h * scale)) + img = img.resize((resized_w, resized_h), Image.LANCZOS) + if center_crop: + left, top = (resized_w - resolution) // 2, (resized_h - resolution) // 2 + else: + left = min(int(u_left * (resized_w - resolution + 1)), max(0, resized_w - resolution)) + top = min(int(u_top * (resized_h - resolution + 1)), max(0, resized_h - resolution)) + img = img.crop((left, top, left + resolution, top + resolution)) + crop_left = left + if flip: + img = img.transpose(Image.FLIP_LEFT_RIGHT) + # Mirror the crop's left origin so the conditioning matches the flipped pixels, the + # same mirroring ``_load_image_tensor`` applies on a random flip. + crop_left = max(0, resized_w - resolution - left) + arr = np.asarray(img, dtype = np.float32) / 255.0 + tensor = torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0 + time_ids = (original_h, original_w, top, crop_left, resolution, resolution) + return tensor, time_ids + + def _encode_sdxl_prompts( prompts: list[str], tokenizers: list, text_encoders: list, device: Any ) -> tuple: @@ -126,6 +176,102 @@ def _encode_sdxl_prompts( return prompt_embeds, pooled +def _build_sdxl_latent_cache( + vae, vae_scale, image_paths, cfg, device, weight_dtype, on_event, check_stop +): + """Precompute the per-image latent posterior cache: for each planned crop/flip variant, + encode once and store ``(A, B, time_ids)`` on CPU in fp32. ``A`` and ``B`` are the affine + posterior parameters (mean/std with the VAE scale folded in) so a per-step sample is + ``A + B * randn`` -- distribution-identical to an in-loop ``latent_dist.sample()`` -- and + ``time_ids`` is the SDXL micro-conditioning for the crop. The stats stay fp32 so the + per-step sample happens in fp32 and only the RESULT is cast to weight_dtype, matching the + in-loop path (encode fp32 -> sample fp32 -> scale -> .to(weight_dtype)); fp32 doubles the + cache RAM over bf16 but the cache is tiny (a handful of latents per image). Returns None if + the build was interrupted by a stop request. ``vae_scale`` is read before the VAE is freed.""" + import torch + + plan = _plan_cache_variants( + len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed + ) + + def _hold(t): + t = t.to(torch.float32).cpu() + if device == "cuda": + try: + t = t.pin_memory() + except RuntimeError: + pass + return t + + cache: list[list[tuple]] = [] + total = len(image_paths) + total_variants = sum(len(v) for v in plan) + forced = _latent_cache_forced() + gated = False + for i, path in enumerate(image_paths): + variants = [] + for u_left, u_top, flip in plan[i]: + tensor, time_ids = _load_image_tensor_planned( + path, cfg.resolution, cfg.center_crop, u_left, u_top, flip + ) + pixel_values = tensor.unsqueeze(0).to(device, dtype = torch.float32) + with torch.no_grad(): + dist = vae.encode(pixel_values).latent_dist + a = _hold(dist.mean * vae_scale) + b = _hold(dist.std * vae_scale) + if not forced and not gated: + # Size-gate the automatic cache off the first REAL encoded variant, before + # building the rest: thousands of images x variants of two fp32 tensors can + # exhaust host/pinned RAM with no fallback. Over budget we bail with the VAE + # still resident so the loop encodes latents per step instead. + per_variant = a.numel() * a.element_size() + b.numel() * b.element_size() + if _latent_cache_over_budget(per_variant, total_variants): + _emit( + on_event, + "warning", + message = ( + "Latent cache disabled: estimated " + f"{per_variant * total_variants / 1024 ** 3:.1f} GiB over the " + "budget; encoding latents per step instead. Set " + "UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE=1 to keep it." + ), + ) + return LATENT_CACHE_OVER_BUDGET + gated = True + variants.append((a, b, tuple(time_ids))) + cache.append(variants) + if (i + 1) % 4 == 0 or i + 1 == total: + _emit(on_event, "preparing", stage = "cache_latents", done = i + 1, total = total) + if check_stop(): + return None + return cache + + +def _sample_sdxl_cached_latents(cache, idxs, variant_rng, device, weight_dtype): + """Draw one latent + its time_ids per index from the cache: pick a variant, then sample + the posterior (A + B * randn) with fresh noise per step, exactly like an in-loop + ``latent_dist.sample() * vae_scale``. The cached stats are fp32, so the sample is drawn in + fp32 and only the RESULT is cast to weight_dtype (matching the in-loop path). Returns + ``(latents, batch_time_ids)`` already on ``device`` in the training dtype (scale is folded + into the cache).""" + import torch + + parts_a, parts_b, tid_rows = [], [], [] + for i in idxs: + variants = cache[i] + a, b, time_ids = ( + variants[variant_rng.randrange(len(variants))] if len(variants) > 1 else variants[0] + ) + parts_a.append(a) + parts_b.append(b) + tid_rows.append(time_ids) + lat_a = torch.cat(parts_a).to(device, non_blocking = True) + lat_b = torch.cat(parts_b).to(device, non_blocking = True) + latents = (lat_a + lat_b * torch.randn_like(lat_a)).to(dtype = weight_dtype) + batch_time_ids = torch.tensor(tid_rows, device = device, dtype = weight_dtype) + return latents, batch_time_ids + + def run_diffusion_lora_training( config: DiffusionLoraConfig, *, @@ -174,250 +320,320 @@ def run_diffusion_lora_training( precision = "fp16" weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision] - # Preflight the base model against the same trust gate as inference, before any fetch. - _assert_trusted_base_model(cfg.base_model) + # TF32 / cudnn.benchmark for the run, restored on the way out (the trainer subprocess is + # disposable, but restoring keeps in-process callers -- tests, notebooks -- clean). Wraps + # the whole body so every return (early stop and normal) restores the backend flags. + snap = _apply_perf_flags(cfg, device) + try: + # Preflight the base model against the same trust gate as inference, before any fetch. + _assert_trusted_base_model(cfg.base_model) - pairs = discover_image_caption_pairs( - cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column - ) - _emit(on_event, "model_load_started", num_images = len(pairs)) + pairs = discover_image_caption_pairs( + cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column + ) + _emit(on_event, "model_load_started", num_images = len(pairs)) - # Honour a stop requested before the (potentially large / slow) base model loads, the - # same way the LLM training worker checks its stop thread around model load. - if _check_stop(): + # Honour a stop requested before the (potentially large / slow) base model loads, the + # same way the LLM training worker checks its stop thread around model load. + if _check_stop(): + out_dir = Path(cfg.output_dir).expanduser() + _emit( + on_event, + "complete", + output_dir = str(out_dir), + lora_path = None, + stopped = True, + steps_run = 0, + ) + return str(out_dir) + + pipe = StableDiffusionXLPipeline.from_pretrained( + cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token, add_watermarker = False + ) + unet, vae = pipe.unet, pipe.vae + tokenizers = [pipe.tokenizer, pipe.tokenizer_2] + text_encoders = [pipe.text_encoder, pipe.text_encoder_2] + noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config) + + # Freeze the base; only the LoRA trains. The SDXL VAE overflows fp16, so keep it fp32. + for m in (unet, vae, *text_encoders): + m.requires_grad_(False) + vae.to(device, dtype = torch.float32) + for m in (unet, *text_encoders): + m.to(device, dtype = weight_dtype) + + unet.add_adapter( + LoraConfig( + r = cfg.lora_rank, + lora_alpha = cfg.lora_alpha, + lora_dropout = cfg.lora_dropout, + init_lora_weights = "gaussian", + target_modules = list(cfg.lora_target_modules), + ) + ) + if cfg.gradient_checkpointing: + unet.enable_gradient_checkpointing() + # LoRA params must be fp32 for a stable optimizer under mixed precision. + if weight_dtype != torch.float32: + cast_training_params(unet, dtype = torch.float32) + + lora_params = [p for p in unet.parameters() if p.requires_grad] + optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate) + # The scheduler advances once per optimizer update: lr_sched.step() runs a single + # time per outer opt_step (after the accumulation inner loop), for cfg.train_steps + # total. Count warmup/decay in those optimizer steps -- multiplying by the + # accumulation factor would stretch warmup past the run and never reach the decay. + lr_sched = get_scheduler( + cfg.lr_scheduler, + optimizer = optimizer, + num_warmup_steps = cfg.lr_warmup_steps, + num_training_steps = cfg.train_steps, + ) + + vae_scale = vae.config.scaling_factor + prediction_type = noise_scheduler.config.prediction_type + + # Precompute text embeddings once per unique caption, then free the CLIP text encoders. + # SDXL re-encoded captions every step (pure waste: captions are constant) and kept both + # text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no + # torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and + # lighter. The env toggle exists purely so the accuracy guard can A/B the two paths. + precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true") + caption_embeds: dict[str, tuple] = {} + if precompute: + for cap in sorted({c for _, c in pairs}): + pe, pooled_c = _encode_sdxl_prompts([cap], tokenizers, text_encoders, device) + caption_embeds[cap] = (pe.cpu(), pooled_c.cpu()) + for te in text_encoders: + te.to("cpu") + text_encoders = [] + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + + # Precompute the VAE latent cache, then free the VAE: the cache holds the posterior + # affine pair (mean/std, scale folded in) so per-step sampling noise is preserved. The + # env toggle lets the accuracy guard A/B the cached vs in-loop encode paths. + use_cache = cfg.cache_latents and os.environ.get( + "UNSLOTH_DIFFUSION_NO_LATENT_CACHE", "" + ) not in ("1", "true") + latent_cache = None + if use_cache: + latent_cache = _build_sdxl_latent_cache( + vae, + vae_scale, + [p for p, _ in pairs], + cfg, + device, + weight_dtype, + on_event, + _check_stop, + ) + if latent_cache is LATENT_CACHE_OVER_BUDGET: + # The estimated cache exceeded the host-memory budget; keep the VAE resident + # and fall through to the in-loop encode path (latent_cache stays None). + latent_cache = None + elif latent_cache is None: # stopped during the cache build; nothing trained yet + out_dir = Path(cfg.output_dir).expanduser() + _emit( + on_event, + "complete", + output_dir = str(out_dir), + lora_path = None, + stopped = True, + steps_run = 0, + ) + return str(out_dir) + else: + try: + pipe.vae = None + except Exception: # noqa: BLE001 -- a pipeline without a settable vae keeps it + pass + del vae + vae = None + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + # Variant picks use their own stream so the loop's index/noise draws stay on the same + # seed-deterministic sequence whether or not the cache is enabled. + variant_rng = random.Random(cfg.seed + 1) + + _emit(on_event, "model_load_completed") + + def _next_batch() -> tuple[list[int], list[str], list[str]]: + idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs))) + chosen = [pairs[i] for i in idx] + return idx, [c[0] for c in chosen], [c[1] for c in chosen] + + unet.train() + stopped = False + micro = 0 + running_loss = 0.0 + peak_gb = 0.0 + t_start = time.time() + done = 0 + for opt_step in range(cfg.train_steps): + optimizer.zero_grad(set_to_none = True) + step_loss = 0.0 + for _ in range(cfg.gradient_accumulation_steps): + idx, img_paths, captions = _next_batch() + if latent_cache is not None: + # Scale is folded into the cache; the sampler draws in fp32 and casts the + # result to weight_dtype (matching the in-loop path below). + latents, batch_time_ids = _sample_sdxl_cached_latents( + latent_cache, idx, variant_rng, device, weight_dtype + ) + else: + loaded = [ + _load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng) + for p in img_paths + ] + pixel_values = torch.stack([t for t, _ in loaded]).to( + device, dtype = torch.float32 + ) + # Per-sample SDXL micro-conditioning from the actual crop (original size + offset). + batch_time_ids = torch.tensor( + [tid for _, tid in loaded], device = device, dtype = weight_dtype + ) + + with torch.no_grad(): + latents = vae.encode(pixel_values).latent_dist.sample() * vae_scale + latents = latents.to(dtype = weight_dtype) + + noise = torch.randn_like(latents) + bsz = latents.shape[0] + timesteps = torch.randint( + 0, noise_scheduler.config.num_train_timesteps, (bsz,), device = device + ).long() + noisy = noise_scheduler.add_noise(latents, noise, timesteps) + + if precompute: + prompt_embeds = torch.cat([caption_embeds[c][0] for c in captions]).to(device) + pooled = torch.cat([caption_embeds[c][1] for c in captions]).to(device) + else: + prompt_embeds, pooled = _encode_sdxl_prompts( + captions, tokenizers, text_encoders, device + ) + prompt_embeds = prompt_embeds.to(dtype = weight_dtype) + pooled = pooled.to(dtype = weight_dtype) + added = {"text_embeds": pooled, "time_ids": batch_time_ids} + + model_pred = unet( + noisy, timesteps, prompt_embeds, added_cond_kwargs = added, return_dict = False + )[0] + + if prediction_type == "v_prediction": + target = noise_scheduler.get_velocity(latents, noise, timesteps) + else: + target = noise + + if cfg.snr_gamma is not None: + snr = compute_snr(noise_scheduler, timesteps) + w = torch.stack([snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim = 1).min( + dim = 1 + )[0] + w = w / snr if prediction_type != "v_prediction" else w / (snr + 1) + loss = F.mse_loss(model_pred.float(), target.float(), reduction = "none") + loss = loss.mean(dim = list(range(1, loss.ndim))) * w + loss = loss.mean() + else: + loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean") + + (loss / cfg.gradient_accumulation_steps).backward() + step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps + micro += 1 + + # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that); + # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning). + grad_norm: Optional[float] = None + if cfg.max_grad_norm and cfg.max_grad_norm > 0: + # clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal). + grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) + optimizer.step() + lr_sched.step() + + running_loss += step_loss + done = opt_step + 1 + if done % cfg.log_every == 0 or done == cfg.train_steps: + # ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so + # these progress events are directly consumable by the existing training + # status/SSE machinery when the diffusion trainer is wired into the worker. + if device == "cuda": + peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) + samples_per_second = round( + (done * cfg.train_batch_size * cfg.gradient_accumulation_steps) + / max(time.time() - t_start, 1e-6), + 3, + ) + _emit( + on_event, + "progress", + step = done, + total_steps = cfg.train_steps, + loss = round(step_loss, 5), + avg_loss = round(running_loss / done, 5), + learning_rate = lr_sched.get_last_lr()[0], + grad_norm = round(grad_norm, 5) if grad_norm is not None else None, + samples_per_second = samples_per_second, + peak_memory_gb = peak_gb or None, + ) + + if _check_stop(): + stopped = True + break + + # Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless + # the run was cancelled with save disabled -- then leave no partial adapter behind. out_dir = Path(cfg.output_dir).expanduser() + lora_path: Optional[str] = None + catalog_path: Optional[str] = None + if not (stopped and not save_on_stop): + out_dir.mkdir(parents = True, exist_ok = True) + unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) + StableDiffusionXLPipeline.save_lora_weights( + save_directory = str(out_dir), + unet_lora_layers = unet_lora, + safe_serialization = True, + weight_name = DEFAULT_LORA_FILENAME, + ) + lora_path = str(out_dir / DEFAULT_LORA_FILENAME) + # Mirror into the Studio diffusion LoRA directory so the Images picker discovers it + # (its scan lists only files directly under loras/diffusion, not subdirectories). + catalog_path = _publish_to_lora_catalog(lora_path, cfg) _emit( on_event, "complete", output_dir = str(out_dir), - lora_path = None, - stopped = True, - steps_run = 0, + lora_path = lora_path, + catalog_path = catalog_path, + family = cfg.resolved_family, + base_model = cfg.base_model, + stopped = stopped, + steps_run = done if cfg.train_steps else 0, ) return str(out_dir) - - pipe = StableDiffusionXLPipeline.from_pretrained( - cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token, add_watermarker = False - ) - unet, vae = pipe.unet, pipe.vae - tokenizers = [pipe.tokenizer, pipe.tokenizer_2] - text_encoders = [pipe.text_encoder, pipe.text_encoder_2] - noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config) - - # Freeze the base; only the LoRA trains. The SDXL VAE overflows fp16, so keep it fp32. - for m in (unet, vae, *text_encoders): - m.requires_grad_(False) - vae.to(device, dtype = torch.float32) - for m in (unet, *text_encoders): - m.to(device, dtype = weight_dtype) - - unet.add_adapter( - LoraConfig( - r = cfg.lora_rank, - lora_alpha = cfg.lora_alpha, - lora_dropout = cfg.lora_dropout, - init_lora_weights = "gaussian", - target_modules = list(cfg.lora_target_modules), - ) - ) - if cfg.gradient_checkpointing: - unet.enable_gradient_checkpointing() - # LoRA params must be fp32 for a stable optimizer under mixed precision. - if weight_dtype != torch.float32: - cast_training_params(unet, dtype = torch.float32) - - lora_params = [p for p in unet.parameters() if p.requires_grad] - optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate) - # The scheduler advances once per optimizer update: lr_sched.step() runs a single - # time per outer opt_step (after the accumulation inner loop), for cfg.train_steps - # total. Count warmup/decay in those optimizer steps -- multiplying by the - # accumulation factor would stretch warmup past the run and never reach the decay. - lr_sched = get_scheduler( - cfg.lr_scheduler, - optimizer = optimizer, - num_warmup_steps = cfg.lr_warmup_steps, - num_training_steps = cfg.train_steps, - ) - - vae_scale = vae.config.scaling_factor - prediction_type = noise_scheduler.config.prediction_type - - # Precompute text embeddings once per unique caption, then free the CLIP text encoders. - # SDXL re-encoded captions every step (pure waste: captions are constant) and kept both - # text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no - # torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and - # lighter. The env toggle exists purely so the accuracy guard can A/B the two paths. - precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true") - caption_embeds: dict[str, tuple] = {} - if precompute: - for cap in sorted({c for _, c in pairs}): - pe, pooled_c = _encode_sdxl_prompts([cap], tokenizers, text_encoders, device) - caption_embeds[cap] = (pe.cpu(), pooled_c.cpu()) - for te in text_encoders: - te.to("cpu") - text_encoders = [] - gc.collect() - if device == "cuda": - torch.cuda.empty_cache() - - _emit(on_event, "model_load_completed") - - def _next_batch() -> tuple[list[str], list[str]]: - idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs))) - chosen = [pairs[i] for i in idx] - return [c[0] for c in chosen], [c[1] for c in chosen] - - unet.train() - stopped = False - micro = 0 - running_loss = 0.0 - peak_gb = 0.0 - t_start = time.time() - done = 0 - for opt_step in range(cfg.train_steps): - optimizer.zero_grad(set_to_none = True) - step_loss = 0.0 - for _ in range(cfg.gradient_accumulation_steps): - img_paths, captions = _next_batch() - loaded = [ - _load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng) - for p in img_paths - ] - pixel_values = torch.stack([t for t, _ in loaded]).to(device, dtype = torch.float32) - # Per-sample SDXL micro-conditioning from the actual crop (original size + offset). - batch_time_ids = torch.tensor( - [tid for _, tid in loaded], device = device, dtype = weight_dtype - ) - - with torch.no_grad(): - latents = vae.encode(pixel_values).latent_dist.sample() * vae_scale - latents = latents.to(dtype = weight_dtype) - - noise = torch.randn_like(latents) - bsz = latents.shape[0] - timesteps = torch.randint( - 0, noise_scheduler.config.num_train_timesteps, (bsz,), device = device - ).long() - noisy = noise_scheduler.add_noise(latents, noise, timesteps) - - if precompute: - prompt_embeds = torch.cat([caption_embeds[c][0] for c in captions]).to(device) - pooled = torch.cat([caption_embeds[c][1] for c in captions]).to(device) - else: - prompt_embeds, pooled = _encode_sdxl_prompts( - captions, tokenizers, text_encoders, device - ) - prompt_embeds = prompt_embeds.to(dtype = weight_dtype) - pooled = pooled.to(dtype = weight_dtype) - added = {"text_embeds": pooled, "time_ids": batch_time_ids} - - model_pred = unet( - noisy, timesteps, prompt_embeds, added_cond_kwargs = added, return_dict = False - )[0] - - if prediction_type == "v_prediction": - target = noise_scheduler.get_velocity(latents, noise, timesteps) - else: - target = noise - - if cfg.snr_gamma is not None: - snr = compute_snr(noise_scheduler, timesteps) - w = torch.stack([snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim = 1).min( - dim = 1 - )[0] - w = w / snr if prediction_type != "v_prediction" else w / (snr + 1) - loss = F.mse_loss(model_pred.float(), target.float(), reduction = "none") - loss = loss.mean(dim = list(range(1, loss.ndim))) * w - loss = loss.mean() - else: - loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean") - - (loss / cfg.gradient_accumulation_steps).backward() - step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps - micro += 1 - - # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that); - # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning). - grad_norm: Optional[float] = None - if cfg.max_grad_norm and cfg.max_grad_norm > 0: - # clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal). - grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) - optimizer.step() - lr_sched.step() - - running_loss += step_loss - done = opt_step + 1 - if done % cfg.log_every == 0 or done == cfg.train_steps: - # ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so - # these progress events are directly consumable by the existing training - # status/SSE machinery when the diffusion trainer is wired into the worker. - if device == "cuda": - peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) - samples_per_second = round( - (done * cfg.train_batch_size * cfg.gradient_accumulation_steps) - / max(time.time() - t_start, 1e-6), - 3, - ) - _emit( - on_event, - "progress", - step = done, - total_steps = cfg.train_steps, - loss = round(step_loss, 5), - avg_loss = round(running_loss / done, 5), - learning_rate = lr_sched.get_last_lr()[0], - grad_norm = round(grad_norm, 5) if grad_norm is not None else None, - samples_per_second = samples_per_second, - peak_memory_gb = peak_gb or None, - ) - - if _check_stop(): - stopped = True - break - - # Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless - # the run was cancelled with save disabled -- then leave no partial adapter behind. - out_dir = Path(cfg.output_dir).expanduser() - lora_path: Optional[str] = None - catalog_path: Optional[str] = None - if not (stopped and not save_on_stop): - out_dir.mkdir(parents = True, exist_ok = True) - unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) - StableDiffusionXLPipeline.save_lora_weights( - save_directory = str(out_dir), - unet_lora_layers = unet_lora, - safe_serialization = True, - weight_name = DEFAULT_LORA_FILENAME, - ) - lora_path = str(out_dir / DEFAULT_LORA_FILENAME) - # Mirror into the Studio diffusion LoRA directory so the Images picker discovers it - # (its scan lists only files directly under loras/diffusion, not subdirectories). - catalog_path = _publish_to_lora_catalog(lora_path, cfg) - _emit( - on_event, - "complete", - output_dir = str(out_dir), - lora_path = lora_path, - catalog_path = catalog_path, - family = cfg.resolved_family, - base_model = cfg.base_model, - stopped = stopped, - steps_run = done if cfg.train_steps else 0, - ) - return str(out_dir) + finally: + _restore_perf_flags(snap) def _make_lora_optimizer(params: list, lr: float) -> Any: """8-bit AdamW (bitsandbytes) by default -- half the optimizer state, no meaningful - quality cost for LoRA -- falling back to fp32 AdamW when unavailable or when - UNSLOTH_DIFFUSION_FP32_OPTIM is set (used by the accuracy guard).""" + quality cost for LoRA -- falling back to torch AdamW (fused on CUDA) when unavailable. + UNSLOTH_DIFFUSION_FP32_OPTIM forces plain (non-fused) AdamW: the accuracy guard wants the + reference optimizer, so it must not take the fused path.""" import torch - if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") not in ("1", "true"): + if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") in ("1", "true"): + return torch.optim.AdamW(params, lr = lr) + try: + import bitsandbytes as bnb + return bnb.optim.AdamW8bit(params, lr = lr) + except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW + pass + if torch.cuda.is_available(): try: - import bitsandbytes as bnb - return bnb.optim.AdamW8bit(params, lr = lr) - except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW + return torch.optim.AdamW(params, lr = lr, fused = True) + except Exception: # noqa: BLE001 -- fused unsupported on this build/device pass return torch.optim.AdamW(params, lr = lr) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index d70f9bce00..5d9140317f 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -17,6 +17,7 @@ from __future__ import annotations import json import os +import random import re import time from dataclasses import dataclass, field, replace @@ -123,6 +124,82 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None return "sdxl" +def repo_is_prequantized(base_model: str) -> bool: + """Heuristic: a repo whose name marks a bitsandbytes 4-bit build already ships a + quantized transformer, so it loads as-is for nf4 and cannot serve the dense + (bf16/int8/fp8) base precisions.""" + name = str(base_model or "").lower() + return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name + + +def _module_is_torchao_stub(module: Any) -> bool: + """True iff ``module`` is the Unsloth Windows-ROCm torchao import stub rather than the + real package. The stub (core/_torchao_stub.py) satisfies find_spec and even lets + ``from torchao.quantization import quantize_`` succeed -- but the imported symbols are + no-op stub types, so the quantization never happens. Every stub module carries the + ``_unsloth_stub`` sentinel, so match on it (comparing against the stub module's own + sentinel object, not identity of a re-created one).""" + if module is None: + return False + sentinel = getattr(module, "_unsloth_stub", None) + if sentinel is None: + return False + try: + from core._torchao_stub import _STUB_SENTINEL + except Exception: # noqa: BLE001 -- stub module absent -> nothing to compare against + return False + return sentinel is _STUB_SENTINEL + + +def has_functional_torchao() -> bool: + """True iff the real torchao quantization API is importable (not the Windows-ROCm stub). + + ``_int8_quantize_base`` needs ``Int8WeightOnlyConfig`` + ``quantize_`` from + ``torchao.quantization`` and has no runtime fallback, so gate both the auto int8 pick + and the advertised int8 mode on a FUNCTIONAL import: a plain ``find_spec("torchao")`` + is satisfied by the stub, whose quantize_ is a no-op that leaves the transformer dense + while compile is disabled as if it were int8. Import the exact symbols the int8 path + uses and reject the stub module. Never raises.""" + try: + import importlib + + quant = importlib.import_module("torchao.quantization") + if _module_is_torchao_stub(quant): + return False + # The symbols the int8 path actually imports must exist on the real module. + return hasattr(quant, "Int8WeightOnlyConfig") and hasattr(quant, "quantize_") + except Exception: # noqa: BLE001 -- torchao absent / broken build -> treat as unavailable + return False + + +def train_precision_modes() -> tuple[list[str], str]: + """(supported base_precision modes, recommended pick) for the current machine: nf4 + always works; bf16/auto need a bf16-capable CUDA GPU (Ampere+); int8/fp8 additionally + need a FUNCTIONAL torchao (their explicit paths import torchao with no fallback, and the + Windows-ROCm stub only looks installed). fp8 also needs an fp8-capable GPU (sm89+). The + dense modes all train in bf16 compute, which the DiT trainer requires, so a non-bf16 CUDA + GPU (T4/V100/RTX 20xx) is offered only nf4 -- otherwise /info would advertise a start that + evicts resident models and then fails the trainer's bf16 guard. Used by the /info endpoint + so the UI can gate the precision selector. Never raises.""" + modes = ["nf4"] + recommended = "nf4" + try: + import torch + if torch.cuda.is_available() and torch.cuda.is_bf16_supported(): + modes.append("bf16") + torchao_ok = has_functional_torchao() + if torchao_ok: + modes.append("int8") + major, minor = torch.cuda.get_device_capability() + if torchao_ok and (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"): + modes.append("fp8") + modes.append("auto") + recommended = "auto" + except Exception: # noqa: BLE001 -- no torch / probe failure -> nf4 only + pass + return modes, recommended + + def get_trainer(family: str) -> Callable[..., str]: """Return the training entrypoint for ``family``. Imports the trainer module lazily so this shared module stays free of the heavy trainer imports (and any import cycle).""" @@ -170,19 +247,96 @@ _FAMILY_VRAM_NOTES = { "z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.", } +# The flow-matching DiT families (run by diffusion_dit_trainer). They expose the +# base_precision / compile levers and require bf16 compute on CUDA; SDXL is absent because +# it uses its own mixed_precision path. Kept as a set so the UI gate, the bf16 preflight, +# and any future dispatch stay in sync. +_DIT_TRAIN_FAMILIES = frozenset({"flux.1", "qwen-image", "z-image"}) + + +def bf16_unsupported_reason(resolved_family: str) -> Optional[str]: + """Return a user-facing error string if ``resolved_family`` needs bf16 compute that the + live GPU cannot provide, else None. The DiT trainer requires a bf16-capable GPU (Ampere + or newer) and otherwise raises deep in model load; the start route uses this to fail fast + BEFORE evicting resident GPU workloads. CPU-only hosts (which fall back to fp32 for + import/unit tests) and SDXL (its own mixed_precision path) are exempt. Never raises.""" + if (resolved_family or "").strip().lower() not in _DIT_TRAIN_FAMILIES: + return None + try: + import torch + if torch.cuda.is_available() and not torch.cuda.is_bf16_supported(): + return ( + "This trainer requires a bfloat16-capable GPU (Ampere or newer); this CUDA " + "device does not support bf16. Train the DiT families on a newer GPU." + ) + except Exception: # noqa: BLE001 -- torch probe failure must not block a start + return None + return None + + +def training_precision_preflight_error(resolved_family: str, base_precision: str) -> Optional[str]: + """Reason the requested DiT precision cannot run on this host, else None -- checked by the + start route BEFORE evicting resident GPU workloads (the trainer's own checks fire only in the + child, after eviction). Three gates, all mirroring _resolve_base_precision so a doomed run is + rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); the dense + precisions (bf16/int8/fp8) requiring a CUDA GPU; and an explicit int8 needing a FUNCTIONAL + torchao (its _int8_quantize_base has no fallback). Never raises.""" + reason = bf16_unsupported_reason(resolved_family) + if reason: + return reason + fam = (resolved_family or "").strip().lower() + mode = (base_precision or "").strip().lower() + if fam in _DIT_TRAIN_FAMILIES and mode in ("bf16", "int8", "fp8"): + # The DiT trainer's dense precisions all require CUDA (_resolve_base_precision rejects + # bf16/int8/fp8 on device != "cuda"). bf16_unsupported_reason exempts a CPU-only host (the + # fp32 fallback for import/unit tests), so without this a dense request on a GPU-less host + # would pass the preflight, evict resident workloads, then raise only in the child. + try: + import torch + has_cuda = torch.cuda.is_available() + except Exception: # noqa: BLE001 -- no torch / probe failure -> treat as no CUDA + has_cuda = False + if not has_cuda: + return ( + f"base_precision={mode!r} needs a CUDA GPU; this host has none. " + "Use base_precision='nf4' or 'auto'." + ) + if mode == "int8" and not has_functional_torchao(): + return ( + "base_precision='int8' needs a functional torchao install; this host's torchao is " + "missing or the non-functional Windows-ROCm stub. Use 'nf4', 'bf16', or 'auto'." + ) + return None + def family_train_infos() -> list[dict[str, Any]]: """Describe every trainable family for the Train UI: name, label, the default + allowed base repos, the recommended starting hyperparameters, and a VRAM/access note. Built from the family registry so it stays in sync with what the trainers actually support.""" from core.inference.diffusion_families import detect_family + from core.inference.diffusion_transformer_quant import _family_denied + dit_modes, dit_recommended = train_precision_modes() infos: list[dict[str, Any]] = [] for name in trainable_family_names(): fam = detect_family("", override = name) if fam is None: continue repos = list(fam.train_base_repos) or [fam.base_repo] + # base_precision / compile apply to the DiT trainer only; SDXL keeps its + # mixed_precision lever, so the UI hides the selector for it. + is_dit = name in _DIT_TRAIN_FAMILIES + # On a non-bf16 CUDA GPU the start route's preflight rejects EVERY DiT family (even nf4, + # since the DiT trainer requires bf16 unconditionally on CUDA), so advertise no precision + # for it -- otherwise /info offers an nf4 DiT option that always 400s. Otherwise drop any + # scheme this family's DiT corrupts (fp8 on Qwen-Image: activation outliers exceed fp8's + # range; the inference path denies the same set), so the UI never offers a mode + # normalized() would then reject. + dit_block = bf16_unsupported_reason(name) if is_dit else None + if not is_dit or dit_block: + fam_modes: list[str] = [] + else: + fam_modes = [m for m in dit_modes if not _family_denied(name, m)] infos.append( { "name": name, @@ -190,7 +344,10 @@ def family_train_infos() -> list[dict[str, Any]]: "default_base": repos[0], "base_repos": repos, "defaults": train_defaults(name), - "vram_note": _FAMILY_VRAM_NOTES.get(name, ""), + "vram_note": dit_block or _FAMILY_VRAM_NOTES.get(name, ""), + "precision_modes": fam_modes, + "recommended_precision": "nf4" if (not is_dit or dit_block) else dit_recommended, + "supports_compile": bool(is_dit and not dit_block), } ) return infos @@ -228,6 +385,23 @@ class DiffusionLoraConfig: caption_column: str = "text" # column in metadata.jsonl adapter_name: str = "default" hf_token: Optional[str] = None + # Precompute the VAE latents once (freeing the VAE for the whole run) instead of + # re-encoding every step. ``cache_variants`` crop/flip draws are frozen per image; + # the per-step VAE sampling noise itself is preserved (see the DiT trainer docstring). + cache_latents: bool = True + cache_variants: int = 4 + # Regional torch.compile of the transformer blocks: "off" | "on" | "auto" (auto turns + # it on only for a dense, non-bitsandbytes base where it is a clean win). + compile_transformer: str = "auto" + # TF32 matmuls + high fp32 matmul precision + cudnn autotuning for the run. Near-lossless; + # disable for strict bit-reproducibility A/Bs. + enable_tf32: bool = True + # DiT base transformer precision: "nf4" (bitsandbytes QLoRA, the memory floor and the + # default), "bf16" (dense, fastest eager, compile-friendly), "int8" (torchao + # weight-only, half of bf16), "fp8" (torchao float8 training compute on the frozen + # linears, Ada/Hopper/Blackwell + compile), or "auto" (pick by free VRAM + GPU class). + # Non-nf4 modes need a dense base repo (not a prequant bnb-4bit one). SDXL ignores it. + base_precision: str = "nf4" # How often to emit a progress event (in optimizer steps). log_every: int = 1 # Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from @@ -259,6 +433,43 @@ class DiffusionLoraConfig: raise ValueError("resolution must be a multiple of 8 and >= 64") if self.mixed_precision not in ("bf16", "fp16", "no"): raise ValueError("mixed_precision must be one of bf16 / fp16 / no") + if not 1 <= int(self.cache_variants) <= 16: + raise ValueError("cache_variants must be between 1 and 16") + compile_transformer = str(self.compile_transformer or "auto").strip().lower() + if compile_transformer not in ("off", "on", "auto"): + raise ValueError("compile_transformer must be one of off / on / auto") + base_precision = str(self.base_precision or "nf4").strip().lower() + if base_precision not in ("nf4", "bf16", "int8", "fp8", "auto"): + raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / auto") + # base_precision is a DiT-only lever (nf4/bf16/int8/fp8/auto for the transformer + # load); SDXL uses its own mixed_precision path and ignores base_precision entirely, + # so the dense-mode gates (prequant base / non-bf16 compute) apply only to the DiT + # families. The mode-name validity check above still runs for every family. + if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8"): + if repo_is_prequantized(self.base_model): + raise ValueError( + f"base_precision={base_precision!r} needs a dense base repo, but " + f"'{self.base_model}' is already bitsandbytes-quantized. Pick the " + f"family's dense (bf16) base repo for this mode, or use nf4/auto." + ) + if self.mixed_precision != "bf16": + raise ValueError( + f"base_precision={base_precision!r} trains in bf16 compute; set " + f"mixed_precision to bf16." + ) + # Some DiT families are corrupted by fp8's activation range: outliers exceed even + # per-row fp8's dynamic range, so the frozen linears' float8 training compute + # learns against a garbage forward pass. The inference path already denies these + # schemes; mirror that deny here so the run fails fast instead of silently + # producing a broken adapter. int8 (per-token) is unaffected and stays allowed. + from core.inference.diffusion_transformer_quant import _family_denied + + if _family_denied(resolved_family, base_precision): + raise ValueError( + f"base_precision={base_precision!r} is not supported for " + f"{resolved_family}: its activations exceed fp8's range and corrupt the " + f"trained result. Use 'nf4', 'int8', 'bf16', or 'auto'." + ) # A zero/negative gamma would zero out (or invert) the min-SNR weight and # silently train on a degenerate loss; None is the documented disable. if self.snr_gamma is not None and float(self.snr_gamma) <= 0: @@ -283,6 +494,9 @@ class DiffusionLoraConfig: lora_target_modules = targets, max_grad_norm = float(self.max_grad_norm), hf_token = token or None, + cache_variants = int(self.cache_variants), + compile_transformer = compile_transformer, + base_precision = base_precision, resolved_family = resolved_family, ) @@ -367,6 +581,145 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None: on_event({"type": type_, "ts": time.time(), **kw}) +def _plan_cache_variants( + num_images: int, cache_variants: int, center_crop: bool, random_flip: bool, seed: int +) -> list[list[tuple[float, float, bool]]]: + """Seed-deterministic crop/flip plan for the latent cache: per image, up to + ``cache_variants`` draws of (u_left, u_top, flip) with the crop as unit fractions the + loader maps onto its integer crop range. Uses its own rng stream so the training + loop's draws are untouched. Center-crop / no-flip collapse duplicate variants (a + center crop without flip is one variant no matter how many draws), so callers encode + each distinct variant exactly once. Pure (no torch) for CPU unit tests.""" + crop_rng = random.Random(seed) + plan: list[list[tuple[float, float, bool]]] = [] + for _ in range(max(0, num_images)): + variants: list[tuple[float, float, bool]] = [] + for _ in range(max(1, cache_variants)): + u_left, u_top = crop_rng.random(), crop_rng.random() + flip = bool(random_flip and crop_rng.random() < 0.5) + if center_crop: + u_left = u_top = 0.5 # loader ignores the fractions for a center crop + key = (u_left, u_top, flip) + if key not in variants: + variants.append(key) + plan.append(variants) + return plan + + +# Host-memory budget for the AUTOMATIC latent cache. The cache holds two fp32 posterior +# tensors (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA +# host. At 1024px an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times +# that, so a few thousand images x cache_variants can exhaust host or pinned RAM with no +# fallback. Over this budget the default falls back to per-step VAE encoding. A fixed +# constant (rather than a psutil RAM fraction) keeps the gate dependency-free and identical +# across hosts; it is deliberately conservative, well under a typical training host's RAM. +_LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3 # 4 GiB + +# Returned by the cache builders when the estimated cache exceeds the budget: the caller +# keeps the VAE resident and encodes each step's latents in-loop. A distinct sentinel from +# ``None`` (which means a stop was requested mid-build) so the two are not conflated. +LATENT_CACHE_OVER_BUDGET: Any = object() + + +def _latent_cache_forced() -> bool: + """The user explicitly forced the latent cache on, bypassing the size gate. This is the + explicit opt-in counterpart to ``UNSLOTH_DIFFUSION_NO_LATENT_CACHE`` (the explicit + opt-out); only the automatic default is size-gated, so an explicit choice is honoured + verbatim in either direction.""" + return os.environ.get("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "") in ("1", "true") + + +def _latent_cache_over_budget( + per_variant_bytes: int, + total_variants: int, + budget_bytes: Optional[int] = None, +) -> bool: + """True when a cache of ``total_variants`` entries, each two fp32 tensors totalling + ``per_variant_bytes``, is estimated to exceed ``budget_bytes``. ``per_variant_bytes`` is + measured from a real encoded latent, so the estimate tracks the actual per-family tensor + shape (SDXL 4-channel vs. a packed 16-channel DiT latent) rather than a guess. The budget + is read from the module constant at call time when not given, so tests can override it.""" + if budget_bytes is None: + budget_bytes = _LATENT_CACHE_BUDGET_BYTES + return per_variant_bytes * max(0, total_variants) > budget_bytes + + +def _apply_perf_flags( + cfg: "DiffusionLoraConfig", + device: str, + cudnn_benchmark: bool = False, +) -> dict: + """Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision + when ``cfg.enable_tf32`` is on, strict fp32 (all TF32 flags cleared) when it is off, + plus cudnn autotuning when the caller opts in. Autotune is + for the conv-heavy SDXL U-Net only: measured on B200, it DOUBLES peak VRAM (fp32 VAE + conv workspaces) while the DiT loop -- pure matmuls once the latent cache is built -- + gains nothing from it. Returns a snapshot for ``_restore_perf_flags``. Best-effort: + missing attributes on a CPU/other-vendor build are skipped.""" + from core.inference.diffusion_speed import snapshot_backend_flags + + snap: dict[str, Any] = {"flags": snapshot_backend_flags(), "matmul_precision": None} + if device != "cuda": + return snap + try: + import torch + + snap["matmul_precision"] = torch.get_float32_matmul_precision() + if cfg.enable_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision("high") + else: + # The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather + # than inherit ambient state (cudnn TF32 defaults to ON in torch). + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + if cudnn_benchmark: + torch.backends.cudnn.benchmark = True + # The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes + # on torch 2.10 + cu130 (B200): mha_graph.execute fails, then poisons the context + # into illegal memory accesses. Flash / mem-efficient SDPA are mathematically + # equivalent, so pin those for the run (restored on exit). + cuda_backends = getattr(torch.backends, "cuda", None) + if cuda_backends is not None and hasattr(cuda_backends, "enable_cudnn_sdp"): + try: + snap["cudnn_sdp"] = bool(cuda_backends.cudnn_sdp_enabled()) + except Exception: # noqa: BLE001 -- flag unreadable: skip the tweak entirely + snap["cudnn_sdp"] = None + if snap["cudnn_sdp"]: + cuda_backends.enable_cudnn_sdp(False) + except Exception: # noqa: BLE001 -- perf flags are never fatal + pass + return snap + + +def _restore_perf_flags(snap: Optional[dict]) -> None: + """Undo ``_apply_perf_flags`` (the trainer subprocess is disposable, but in-process + callers -- tests, notebooks -- must not inherit mutated globals).""" + if not snap: + return + from core.inference.diffusion_speed import restore_backend_flags + + restore_backend_flags(snap.get("flags")) + try: + import torch + + if snap.get("matmul_precision"): + torch.set_float32_matmul_precision(snap["matmul_precision"]) + # Restore the exact pre-run cudnn SDPA state; None means the flag was unreadable + # (or absent) at apply time and was never touched. + cuda_backends = getattr(torch.backends, "cuda", None) + if ( + snap.get("cudnn_sdp") is not None + and cuda_backends is not None + and hasattr(cuda_backends, "enable_cudnn_sdp") + ): + cuda_backends.enable_cudnn_sdp(bool(snap["cudnn_sdp"])) + except Exception: # noqa: BLE001 -- best-effort restore + pass + + def _assert_trusted_base_model(base_model: str) -> None: """Gate the training base model the same way the inference backend gates non-GGUF loads: a local path or a trusted repo (``unsloth/*`` or an allowlisted official base). This runs @@ -446,6 +799,15 @@ def _coerce_gradient_checkpointing(value: Any) -> bool: return bool(value) +def _coerce_bool(value: Any) -> bool: + """Coerce a flag that may arrive as a string through the generic Studio config path + (e.g. "false" / "0" / "off"). A non-empty string like "false" is otherwise truthy, so + an opt-out would silently no-op. A real bool passes through.""" + if isinstance(value, str): + return value.strip().lower() not in ("", "none", "false", "0", "no", "off") + return bool(value) + + def _config_from_dict(config: dict) -> DiffusionLoraConfig: """Build a DiffusionLoraConfig from a plain dict. Unknown keys are ignored so a richer request payload (UI form) does not break construction; a small set of generic Studio @@ -465,4 +827,7 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig: kwargs["gradient_checkpointing"] = _coerce_gradient_checkpointing( kwargs["gradient_checkpointing"] ) + for flag in ("cache_latents", "enable_tf32"): + if flag in kwargs: + kwargs[flag] = _coerce_bool(kwargs[flag]) return DiffusionLoraConfig(**kwargs) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 5c6a184a86..da28f2d588 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -17,6 +17,7 @@ runs a scripted target on a thread. from __future__ import annotations +import math import multiprocessing as mp import threading import time @@ -31,6 +32,21 @@ _CTX = mp.get_context("spawn") _TERMINAL = ("complete", "error") +def _finite_or_none(value: Any) -> Optional[float]: + """Coerce a numeric progress field to a finite float, or None. A divergent run (or a + grad clip that returns inf) can push loss / grad_norm to NaN or +/-Infinity, and those + are invalid in strict JSON -- FastAPI's encoder would emit the JS-only NaN/Infinity + tokens that break a strict client parse. Nulling them here (the single service ingestion + point both trainers feed) keeps every status snapshot and persisted record JSON-safe.""" + if value is None: + return None + try: + f = float(value) + except (TypeError, ValueError): + return None + return f if math.isfinite(f) else None + + def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) -> None: # Imported lazily so this module (and the route layer) stays torch-free at import. from .diffusion_lora_trainer import run_diffusion_training_process @@ -108,21 +124,13 @@ def _append_metric( return if istep <= 0 or loss is None: return - try: - floss = float(loss) - except (TypeError, ValueError): + floss = _finite_or_none(loss) + if floss is None: # non-numeric or non-finite (NaN/Inf): skip, keep the curve JSON-safe return - if floss != floss: # NaN guard - return - - def _opt_float(v: Any) -> Optional[float]: - try: - return float(v) if v is not None else None - except (TypeError, ValueError): - return None - - flr = _opt_float(lr) - fgn = _opt_float(grad_norm) + # lr / grad_norm may be None (sparse series) or non-finite; a non-finite value is + # nulled, not dropped, so a bad point never taints the (loss-driven) history. + flr = _finite_or_none(lr) + fgn = _finite_or_none(grad_norm) steps = state["metric_steps"] losses = state["metric_loss"] lrs = state["metric_lr"] @@ -228,17 +236,25 @@ class DiffusionTrainingService: self._pump.start() return job_id - def stop(self) -> bool: - """Request a clean stop (the trainer finishes the current step and saves a partial - adapter). Returns True if a stop was signalled, False if nothing was running.""" + def stop(self, save: bool = True) -> bool: + """Request a clean stop: the trainer finishes the current step, then either saves + a partial adapter (``save=True``, the default) or discards the run (``save=False``, + matching the LLM trainer's cancel). Returns True if a stop was signalled, False if + nothing was running.""" with self._lock: if self._proc is None or not self._proc.is_alive() or self._stop_queue is None: return False try: - self._stop_queue.put(True) + # Bare True keeps the wire format older trainers expect; the dict form + # carries the no-save cancel flag the trainer's _check_stop understands. + self._stop_queue.put(True if save else {"save": False}) except Exception: # noqa: BLE001 return False - self._state["message"] = "Stop requested; finishing the current step..." + self._state["message"] = ( + "Stop requested; finishing the current step and saving a partial adapter..." + if save + else "Cancel requested; finishing the current step (no adapter will be saved)..." + ) self._state["updated_at"] = time.time() return True @@ -302,15 +318,47 @@ class DiffusionTrainingService: s["num_images"] = ev.get("num_images") elif etype == "model_load_completed": s.update(in_model_load = False, message = "Training...") + elif etype == "preparing": + # A long precompute phase (e.g. the VAE latent cache) between model load and + # the first step; surfaced so the UI shows visible progress instead of a + # silent "Loading base model..." stall. + done, total = ev.get("done"), ev.get("total") + stage = str(ev.get("stage", "prepare")).replace("_", " ") + s.update( + status = "running", + in_model_load = True, + message = ( + f"Preparing ({stage} {done}/{total})..." + if done is not None and total is not None + else f"Preparing ({stage})..." + ), + ) + elif etype == "warning": + # Non-fatal trainer notes (e.g. torch.compile falling back to eager); keep + # training state, surface the text. + s["message"] = str(ev.get("message", "warning")) elif etype == "progress": + # Null any non-finite float (NaN/Inf from a divergent step or an inf grad + # norm) so the JSON status stays strict-parseable; a missing key keeps the + # last value, a present-but-non-finite one becomes None. + loss = _finite_or_none(ev["loss"]) if "loss" in ev else s["loss"] + avg_loss = _finite_or_none(ev["avg_loss"]) if "avg_loss" in ev else s["avg_loss"] + learning_rate = ( + _finite_or_none(ev["learning_rate"]) + if "learning_rate" in ev + else s["learning_rate"] + ) + grad_norm = ( + _finite_or_none(ev["grad_norm"]) if "grad_norm" in ev else s["grad_norm"] + ) s.update( status = "running", step = ev.get("step", s["step"]), total_steps = ev.get("total_steps", s["total_steps"]), - loss = ev.get("loss", s["loss"]), - avg_loss = ev.get("avg_loss", s["avg_loss"]), - learning_rate = ev.get("learning_rate", s["learning_rate"]), - grad_norm = ev.get("grad_norm", s["grad_norm"]), + loss = loss, + avg_loss = avg_loss, + learning_rate = learning_rate, + grad_norm = grad_norm, message = "Training...", ) # Fold optional perf fields (emitted by the trainers) so the UI can show @@ -337,7 +385,11 @@ class DiffusionTrainingService: status = "stopped" if ev.get("stopped") else "completed", output_dir = ev.get("output_dir"), lora_path = ev.get("lora_path"), - message = "Stopped (partial adapter saved)." + message = ( + "Stopped (partial adapter saved)." + if ev.get("lora_path") + else "Stopped (no adapter saved)." + ) if ev.get("stopped") else "Training complete.", ) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 3256dfa8e7..61d5b4ad3d 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -719,6 +719,34 @@ class DiffusionTrainingStartRequest(BaseModel): random_flip: bool = Field(True) caption_column: str = Field("text") hf_token: Optional[str] = Field(None) + cache_latents: bool = Field( + True, description = "Precompute VAE latents once and free the VAE for the run" + ) + cache_variants: int = Field( + 4, ge = 1, le = 16, description = "Frozen crop/flip variants per image in the latent cache" + ) + compile_transformer: Literal["off", "on", "auto"] = Field( + "auto", description = "Regional torch.compile of the transformer blocks" + ) + enable_tf32: bool = Field( + True, description = "TF32 matmuls + cudnn autotuning (near-lossless speedup)" + ) + base_precision: Literal["nf4", "bf16", "int8", "fp8", "auto"] = Field( + "nf4", + description = ( + "DiT base transformer precision: nf4 QLoRA (memory floor, default), bf16 dense, " + "int8 torchao weight-only, fp8 float8 training compute (Ada/Hopper/Blackwell), " + "or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base." + ), + ) + + +class DiffusionTrainingStopRequest(BaseModel): + """Optional body for stopping a diffusion training job. ``save`` mirrors the LLM + trainer's stop: True (default) exports the partial adapter, False cancels without + leaving one behind.""" + + save: bool = Field(True) class DiffusionTrainingStartResponse(BaseModel): @@ -790,6 +818,12 @@ class DiffusionTrainableFamily(BaseModel): base_repos: List[str] = Field(default_factory = list) defaults: dict = Field(default_factory = dict) vram_note: str = "" + # base_precision modes this machine supports for the family (empty = the family has no + # precision selector, e.g. SDXL), plus the recommended pick and whether regional + # torch.compile applies. Defaults keep older backends' payloads valid. + precision_modes: List[str] = Field(default_factory = list) + recommended_precision: str = "nf4" + supports_compile: bool = False class DiffusionTrainingInfoResponse(BaseModel): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 07dbd5c812..b08efa5753 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -73,6 +73,7 @@ from models.training import ( DiffusionTrainingStartRequest, DiffusionTrainingStartResponse, DiffusionTrainingStatusResponse, + DiffusionTrainingStopRequest, ) from models.responses import TrainingStopResponse, TrainingMetricsResponse from pydantic import BaseModel as PydanticBaseModel @@ -1233,10 +1234,23 @@ async def start_diffusion_training( from core.training.diffusion_lora_trainer import _config_from_dict try: - _config_from_dict(config).normalized() + normalized_cfg = _config_from_dict(config).normalized() except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) + # Preflight the requested DiT precision BEFORE freeing GPU residents: the DiT trainer's own + # checks (a bf16-capable GPU is required; an explicit int8 needs a functional torchao) fire + # only in the child, AFTER _free_gpu_for_diffusion_training() already evicted the user's + # chat/Images model. Fail fast (400) so a pre-Ampere GPU (T4 / V100 / RTX 20xx) or a + # stub-torchao host never tears down resident models for a run that cannot start. + from core.training.diffusion_train_common import training_precision_preflight_error + + _precision_reason = training_precision_preflight_error( + normalized_cfg.resolved_family, normalized_cfg.base_precision + ) + if _precision_reason: + raise HTTPException(status_code = 400, detail = _precision_reason) + # Run the trainers' trust gate here too (both assert the same predicate before # from_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents # instead of tearing down the user's chat/Images model and failing in the child. @@ -1291,11 +1305,17 @@ async def start_diffusion_training( @router.post("/diffusion/stop") -async def stop_diffusion_training(current_subject: str = Depends(get_current_subject)): - """Request a clean stop of the running diffusion training job (partial adapter saved).""" +async def stop_diffusion_training( + body: Optional[DiffusionTrainingStopRequest] = None, + current_subject: str = Depends(get_current_subject), +): + """Request a clean stop of the running diffusion training job. The optional body's + ``save`` mirrors the LLM /stop: true (default, also for an empty POST) exports the + partial adapter, false cancels without saving one.""" from core.training.diffusion_training_service import get_diffusion_training_service - stopped = get_diffusion_training_service().stop() + save = body.save if body is not None else True + stopped = get_diffusion_training_service().stop(save = save) return {"status": "stopping" if stopped else "idle"} diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py new file mode 100644 index 0000000000..66bebb28c6 --- /dev/null +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -0,0 +1,547 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""CPU-only unit tests for the DiT base_precision work. + +Covers the new precision plumbing the precision PR adds: the ``base_precision`` +config validation (dense-vs-prequant + mixed-precision gating), the prequant-repo +heuristic and its trainer alias, the pure ``auto`` precision policy table, the +explicit-mode passthrough of ``_resolve_base_precision``, the fp8 module filter, the +fp8 branch of the compile policy, the ``train_precision_modes`` machine probe, the +family-info precision fields, and the request-model ``base_precision`` field. No GPU / +model load: every helper here is pure or name-based, so the config validation runs on +name matching (``resolve_trainable_family`` is offline) and the torch probe is monkeypatched. +""" + +from __future__ import annotations + +import pytest +import torch.nn as nn + +import core.training.diffusion_train_common as common +from core.training import diffusion_dit_trainer as dit +from core.training.diffusion_train_common import ( + DiffusionLoraConfig, + _config_from_dict, + repo_is_prequantized, + train_precision_modes, +) +from models.training import DiffusionTrainingStartRequest + +# A dense (non-prequant) DiT base and a prequant bnb-4bit base. Both resolve a trainer +# family from their names alone, so normalized() runs without a network call. +_FLUX_DENSE = "black-forest-labs/FLUX.1-dev" +_Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit" +# An SDXL base whose name LOOKS prequant (bnb-4bit): SDXL ignores base_precision, so the +# dense-mode gates must not fire for it even with a dense mode + fp16 compute. +_SDXL_PREQUANT_NAME = "some/sdxl-model-bnb-4bit" +# A dense Qwen-Image base: its DiT is corrupted by fp8 (activation outliers), so fp8 is +# denied for training the same way the inference path denies it. +_QWEN_DENSE = "Qwen/Qwen-Image" + + +def _cfg(base_model = _FLUX_DENSE, **kw) -> DiffusionLoraConfig: + return DiffusionLoraConfig(base_model = base_model, data_dir = "d", output_dir = "o", **kw) + + +# ── base_precision validation ───────────────────────────────────────────────── +def test_base_precision_validation(): + # Default normalizes to the nf4 memory floor. + assert _cfg().normalized().base_precision == "nf4" + + # An unknown mode is rejected by name. + with pytest.raises(ValueError, match = "base_precision"): + _cfg(base_precision = "banana").normalized() + + # A dense mode is case/space-insensitive and stored lowered: " FP8 " on a dense base + # with bf16 compute normalizes cleanly to "fp8". + norm = _cfg(base_precision = " FP8 ", mixed_precision = "bf16").normalized() + assert norm.base_precision == "fp8" + + # A dense mode against a prequant (bnb-4bit) base is refused: the repo already ships a + # 4-bit transformer and cannot serve the dense precisions. + with pytest.raises(ValueError, match = "dense base repo"): + _cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized() + + # A dense mode with non-bf16 compute is refused: these modes train in bf16 compute. + with pytest.raises(ValueError, match = "bf16 compute"): + _cfg(base_precision = "int8", mixed_precision = "fp16").normalized() + + # "auto" is ACCEPTED by normalized() even on a prequant base: the concrete mode is + # resolved at runtime against the live GPU, not at config validation. + assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto" + + +def test_base_precision_denies_fp8_for_corrupted_family(): + # fp8 corrupts the Qwen-Image DiT (activation outliers exceed fp8's range), so a dense + # Qwen base with base_precision="fp8" is refused up front -- mirroring the inference deny. + with pytest.raises(ValueError, match = "fp8"): + _cfg(base_model = _QWEN_DENSE, base_precision = "fp8", mixed_precision = "bf16").normalized() + + # The deny is fp8-specific: int8 (per-token, unaffected) and the other dense modes stay + # allowed for the same Qwen base. + for mode in ("nf4", "bf16", "int8", "auto"): + norm = _cfg( + base_model = _QWEN_DENSE, base_precision = mode, mixed_precision = "bf16" + ).normalized() + assert norm.resolved_family == "qwen-image" + assert norm.base_precision == mode + + # A family the deny does not cover (FLUX) still accepts fp8. + flux = _cfg(base_model = _FLUX_DENSE, base_precision = "fp8", mixed_precision = "bf16").normalized() + assert flux.resolved_family == "flux.1" + assert flux.base_precision == "fp8" + + +def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch): + # /info advertises the machine's DiT modes per family, but a family whose DiT the mode + # corrupts must not offer it: with fp8 in the machine list, Qwen-Image drops fp8 while + # FLUX keeps it, so the UI never surfaces a mode normalized() would reject. + monkeypatch.setattr( + common, "train_precision_modes", lambda: (["nf4", "bf16", "int8", "fp8", "auto"], "auto") + ) + # family_train_infos reads the live GPU via bf16_unsupported_reason; pin it to "bf16 OK" so + # this positive-path assertion is deterministic across GPU types (a non-bf16 CUDA box would + # otherwise empty every DiT family's modes). The empty-on-non-bf16 path is covered separately. + monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None) + infos = {i["name"]: i for i in common.family_train_infos()} + assert "fp8" not in infos["qwen-image"]["precision_modes"] + assert "int8" in infos["qwen-image"]["precision_modes"] # int8 is fine on Qwen + assert "fp8" in infos["flux.1"]["precision_modes"] + + +def test_resolve_base_precision_explicit_int8_gates_on_torchao(monkeypatch): + # Explicit int8 has no runtime fallback, so a missing/stub torchao must fail fast here + # rather than load dense with compile disabled. Gate the explicit request the same way + # auto + /info already gate it. + spec = dit._SPECS["flux.1"] + cfg = _cfg(base_precision = "int8") + + monkeypatch.setattr(dit, "has_functional_torchao", lambda: False) # torchao absent / stub + with pytest.raises(ValueError, match = "torchao"): + dit._resolve_base_precision(cfg, spec, "cuda") + + # With a functional torchao the explicit int8 passes straight through. + monkeypatch.setattr(dit, "has_functional_torchao", lambda: True) + assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8" + + # The gate is int8-specific: explicit bf16/fp8 pass through regardless of torchao (fp8 has + # its own graceful fallback; bf16 needs no torchao). + monkeypatch.setattr(dit, "has_functional_torchao", lambda: False) + assert dit._resolve_base_precision(_cfg(base_precision = "bf16"), spec, "cuda") == "bf16" + assert dit._resolve_base_precision(_cfg(base_precision = "fp8"), spec, "cuda") == "fp8" + + +def test_bf16_unsupported_reason(monkeypatch): + # The route uses this to fail fast on a non-bf16 GPU BEFORE evicting resident workloads. + import torch + + from core.training.diffusion_train_common import bf16_unsupported_reason + + # SDXL (own mixed_precision path) and unknown families are always exempt. + assert bf16_unsupported_reason("sdxl") is None + assert bf16_unsupported_reason("") is None + + # A DiT family on a CUDA GPU without bf16 -> a clear reason. + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False) + assert "bfloat16" in (bf16_unsupported_reason("flux.1") or "") + + # A bf16-capable GPU -> no reason. + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) + assert bf16_unsupported_reason("qwen-image") is None + + # A CPU-only host (fp32 fallback for import/unit tests) -> no reason even for a DiT family. + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert bf16_unsupported_reason("z-image") is None + + +def test_training_precision_preflight_error(monkeypatch): + # The start route calls this BEFORE evicting resident GPU workloads: it folds the bf16-GPU + # requirement together with the explicit-int8 torchao requirement, so both fail fast instead + # of only surfacing in the trainer child after the GPU has already been freed. + import torch + + from core.training.diffusion_train_common import training_precision_preflight_error + + # Present a bf16-capable CUDA GPU so the int8 gate (not the bf16 gate) is what we exercise. + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) + + # The bf16 gate takes precedence: a non-bf16 GPU rejects any DiT precision first. + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False) + assert "bfloat16" in (training_precision_preflight_error("flux.1", "int8") or "") + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) + + # Explicit int8 on a DiT family with a NON-functional torchao -> a clear int8 reason + # (its _int8_quantize_base has no fallback, so the child would otherwise raise post-eviction). + monkeypatch.setattr(common, "has_functional_torchao", lambda: False) + reason = training_precision_preflight_error("qwen-image", "int8") + assert reason is not None and "int8" in reason and "torchao" in reason + + # The same int8 request is fine once torchao is functional. + monkeypatch.setattr(common, "has_functional_torchao", lambda: True) + assert training_precision_preflight_error("qwen-image", "int8") is None + + # With a broken torchao, only EXPLICIT int8 is gated -- nf4/bf16/auto pass, and the int8 + # gate never applies to a non-DiT (SDXL) or unknown family. + monkeypatch.setattr(common, "has_functional_torchao", lambda: False) + assert training_precision_preflight_error("flux.1", "nf4") is None + assert training_precision_preflight_error("flux.1", "auto") is None + assert training_precision_preflight_error("sdxl", "int8") is None + assert training_precision_preflight_error("", "int8") is None + + # On a CUDA-ABSENT host, bf16_unsupported_reason exempts CPU-only, but the DiT trainer's dense + # precisions still require CUDA (mirroring _resolve_base_precision), so bf16/int8/fp8 for a DiT + # family are rejected UP FRONT rather than after eviction. nf4/auto (and SDXL) still pass. + monkeypatch.setattr(common, "has_functional_torchao", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + for dense in ("bf16", "int8", "fp8"): + reason = training_precision_preflight_error("flux.1", dense) + assert reason is not None and "CUDA" in reason + assert training_precision_preflight_error("flux.1", "nf4") is None + assert training_precision_preflight_error("flux.1", "auto") is None + assert training_precision_preflight_error("sdxl", "bf16") is None + + +def test_family_train_infos_empties_dit_modes_on_non_bf16(monkeypatch): + # On a non-bf16 GPU the start route rejects EVERY DiT family (even nf4), so /info must not + # advertise a DiT precision option that always 400s: the modes empty, the reason surfaces in + # vram_note, compile is off, and the recommendation degrades to nf4. SDXL (non-DiT) is exempt. + from core.training.diffusion_train_common import _DIT_TRAIN_FAMILIES, family_train_infos + + monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: "no bfloat16 on this GPU") + + infos = {info["name"]: info for info in family_train_infos()} + dit_seen = False + for name, info in infos.items(): + if name in _DIT_TRAIN_FAMILIES: + dit_seen = True + assert info["precision_modes"] == [] + assert info["vram_note"] == "no bfloat16 on this GPU" + assert info["recommended_precision"] == "nf4" + assert info["supports_compile"] is False + assert dit_seen # the registry must still expose at least one DiT family to have covered it + + +def test_base_precision_gates_skip_sdxl(): + # SDXL ignores base_precision, so the dense-mode gates (prequant base / non-bf16 compute) + # must not fire for it: a prequant-looking SDXL name with base_precision="bf16" does not + # raise, and the mode is still stored lowered. + norm = _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "bf16").normalized() + assert norm.resolved_family == "sdxl" + assert norm.base_precision == "bf16" + + # The non-bf16-compute gate is also skipped for SDXL (fp16 is a valid SDXL mixed + # precision), even with a dense base_precision requested. + norm2 = _cfg( + base_model = "stabilityai/stable-diffusion-xl-base-1.0", + base_precision = "int8", + mixed_precision = "fp16", + ).normalized() + assert norm2.resolved_family == "sdxl" + + # The mode-name validity check still runs for SDXL: an unknown mode is rejected. + with pytest.raises(ValueError, match = "base_precision"): + _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "banana").normalized() + + # The gates STILL fire for a DiT family: a prequant DiT base with a dense mode raises. + with pytest.raises(ValueError, match = "dense base repo"): + _cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized() + + +# ── repo_is_prequantized heuristic + trainer alias ──────────────────────────── +@pytest.mark.parametrize( + "repo, expected", + [ + ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", True), + ("some/model-4bit", True), + ("some/model-int4", True), + ("some/model-nf4", True), + ("black-forest-labs/FLUX.1-dev", False), + ("Tongyi-MAI/Z-Image-Turbo", False), + ], +) +def test_repo_is_prequantized_cases(repo, expected): + assert repo_is_prequantized(repo) is expected + + +def test_repo_is_prequantized_alias_is_same_object(): + # The trainer keeps a module-level alias for callers/tests; it must be the exact same + # function object as the common heuristic (moved there for config validation). + assert dit._repo_is_prequantized is repo_is_prequantized + + +# ── _pick_auto_precision policy table (pure) ────────────────────────────────── +def test_pick_auto_precision_policy_table(): + p = dit._pick_auto_precision + + # A prequant base always resolves to nf4 (it can only serve 4-bit). + assert p(True, "cuda", 140, 23.8, (10, 0), True) == "nf4" + # No CUDA -> nf4 (the dense modes need a GPU). + assert p(False, "cpu", 140, 23.8, (10, 0), True) == "nf4" + # Missing free-VRAM number -> the safe nf4 mode. + assert p(False, "cuda", None, 23.8, (10, 0), True) == "nf4" + + # Plenty of free VRAM -> bf16 regardless of fp8 capability: compiled bf16 measured + # FASTER than torchao float8 at LoRA-training shapes, so fp8 is opt-in only. + assert p(False, "cuda", 140, 23.8, (10, 0), True) == "bf16" + assert p(False, "cuda", 140, 23.8, (8, 0), True) == "bf16" + assert p(False, "cuda", 140, 23.8, (10, 0), False) == "bf16" + + # Middle band (30 > 23.8 * 1.15 = 27.4, but not > 23.8 * 1.5 = 35.7) -> int8. + assert p(False, "cuda", 30, 23.8, (10, 0), True) == "int8" + # int8 needs torchao at runtime (no fallback), so the int8 band drops to nf4 when + # torchao is not importable while the bf16 band is unaffected. + assert p(False, "cuda", 30, 23.8, (10, 0), True, False) == "nf4" + assert p(False, "cuda", 140, 23.8, (10, 0), True, False) == "bf16" + # int8 still materialises the full bf16 transformer before quantize_ shrinks it, so + # free VRAM below the dense-load transient (25 < 27.4) must fall back to nf4 even + # though the QUANTIZED weights would have fit. + assert p(False, "cuda", 25, 23.8, (10, 0), True) == "nf4" + # Too little free VRAM for any dense load -> nf4. + assert p(False, "cuda", 10, 23.8, (10, 0), True) == "nf4" + + +# ── _resolve_base_precision passthrough ─────────────────────────────────────── +def test_resolve_base_precision_passes_explicit_through(): + # An explicit mode passes straight through without probing the GPU (normalized() already + # validated it); the spec is only consulted for "auto". + spec = dit._SPECS["flux.1"] + cfg = _cfg(base_precision = "bf16") + assert dit._resolve_base_precision(cfg, spec, "cuda") == "bf16" + + # The dense modes are CUDA-only: an explicit request on a GPU-less host fails fast + # (before any model load) instead of silently proceeding; /info never advertised it. + with pytest.raises(ValueError, match = "CUDA"): + dit._resolve_base_precision(cfg, spec, "cpu") + # nf4 stays a passthrough on any device (the bnb load path owns its own errors). + assert dit._resolve_base_precision(_cfg(base_precision = "nf4"), spec, "cpu") == "nf4" + + +def test_resolve_auto_requires_bf16_compute(): + # auto may resolve to bf16/int8 which train in bf16 compute, so a non-bf16 + # mixed_precision pins auto to the nf4 floor BEFORE any GPU probe (pure, no CUDA + # needed here) -- mirroring the normalized() rule for explicit dense modes. + spec = dit._SPECS["flux.1"] + cfg = _cfg(base_precision = "auto", mixed_precision = "fp16") + assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" + + +def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch): + # The int8 auto band needs a FUNCTIONAL torchao at runtime; when torchao is not + # importable _resolve_base_precision must fall to nf4 instead of picking an int8 that + # would crash in _int8_quantize_base. Drive the probe into the int8 band and toggle the + # functional-torchao probe (shared with train_precision_modes, imported into the trainer). + import torch + + spec = dit._SPECS["flux.1"] # dense_bf16_gb = 23.8 + cfg = _cfg(base_precision = "auto", mixed_precision = "bf16") + + class _FakeCuda: + # Free VRAM in the int8 band (30 > 23.8 * 1.15) but below the bf16 band. + @staticmethod + def mem_get_info(): + return (int(30 * 1e9), int(80 * 1e9)) + + @staticmethod + def get_device_capability(): + return (10, 0) + + monkeypatch.setattr(torch, "cuda", _FakeCuda) + + monkeypatch.setattr(dit, "has_functional_torchao", lambda: False) # torchao absent / stub + assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" + + # With a functional torchao the same band picks int8. + monkeypatch.setattr(dit, "has_functional_torchao", lambda: True) + assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8" + + +def test_resolve_auto_int8_band_treats_stub_as_absent(monkeypatch): + # Simulate the Windows-ROCm torchao STUB: has_functional_torchao returns False (the + # stub satisfies find_spec but its quantize_ is a no-op), so the int8 band must fall to + # nf4 rather than pick an int8 whose quantization silently does nothing. + import torch + + spec = dit._SPECS["flux.1"] + cfg = _cfg(base_precision = "auto", mixed_precision = "bf16") + + class _FakeCuda: + @staticmethod + def mem_get_info(): + return (int(30 * 1e9), int(80 * 1e9)) + + @staticmethod + def get_device_capability(): + return (10, 0) + + monkeypatch.setattr(torch, "cuda", _FakeCuda) + # The stub scenario: the probe reports no functional torchao. + monkeypatch.setattr(dit, "has_functional_torchao", lambda: False) + assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" + + +def test_has_functional_torchao_rejects_stub(monkeypatch): + # has_functional_torchao must reject the Unsloth import stub: even though + # `from torchao.quantization import quantize_` would succeed against the stub, the + # symbols are no-op stub types. Simulate a stub torchao.quantization module carrying the + # stub sentinel and assert the probe returns False. + import importlib + import types + + from core._torchao_stub import _STUB_SENTINEL + + real_import_module = importlib.import_module + + stub_quant = types.ModuleType("torchao.quantization") + stub_quant._unsloth_stub = _STUB_SENTINEL + + def _fake_import(name, *args, **kwargs): + if name == "torchao.quantization": + return stub_quant + return real_import_module(name, *args, **kwargs) + + monkeypatch.setattr(importlib, "import_module", _fake_import) + assert common.has_functional_torchao() is False + + # A real module exposing the int8 symbols (no stub sentinel) probes True. + real_like = types.ModuleType("torchao.quantization") + real_like.Int8WeightOnlyConfig = object + real_like.quantize_ = lambda *a, **k: None + + def _fake_import_real(name, *args, **kwargs): + if name == "torchao.quantization": + return real_like + return real_import_module(name, *args, **kwargs) + + monkeypatch.setattr(importlib, "import_module", _fake_import_real) + assert common.has_functional_torchao() is True + + +# ── _fp8_module_filter ──────────────────────────────────────────────────────── +def test_fp8_module_filter(): + lin = nn.Linear(64, 64) + # A plain feed-forward Linear with divisible dims gets float8 training compute. + assert dit._fp8_module_filter(lin, "transformer_blocks.0.ff.net.0") is True + # A LoRA-owned module is skipped (adapters stay high precision). + assert dit._fp8_module_filter(lin, "transformer_blocks.0.attn.to_q.lora_A.default") is False + # The output projection is skipped. + assert dit._fp8_module_filter(lin, "proj_out") is False + # An in_features not divisible by 16 is rejected (float8 kernels reject the shape). + assert dit._fp8_module_filter(nn.Linear(30, 64), "transformer_blocks.0.ff.net.0") is False + # A non-Linear module is never float8. + assert dit._fp8_module_filter(nn.LayerNorm(64), "transformer_blocks.0.norm") is False + + +# ── _should_compile fp8 branch ──────────────────────────────────────────────── +def test_should_compile_fp8_branch(): + # fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb) + # cuda base. + cfg = _cfg(compile_transformer = "auto") + assert dit._should_compile(cfg, False, "cuda", "fp8") is True + # fp8 forces compile under auto even when the base is (hypothetically) reported as bnb. + assert dit._should_compile(cfg, True, "cuda", "fp8") is True + # An explicit "off" still wins over fp8 -- compile stays off. + assert dit._should_compile(_cfg(compile_transformer = "off"), False, "cuda", "fp8") is False + + +# ── train_precision_modes machine probe ─────────────────────────────────────── +def test_train_precision_modes_no_cuda(monkeypatch): + # Patch the torch module attribute the function imports so it observes a CPU-only box: + # no CUDA -> the nf4-only floor with nf4 recommended, and it never raises. + import torch + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert train_precision_modes() == (["nf4"], "nf4") + + +def test_train_precision_modes_gates_int8_fp8_on_torchao(monkeypatch): + # int8/fp8 are only advertised when torchao is FUNCTIONAL: on a CUDA host WITHOUT a real + # torchao (or with only the Windows-ROCm stub) /info must not offer int8/fp8, since their + # explicit paths import torchao with no fallback. bf16 + auto stay advertised. + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 0)) + + # No functional torchao (absent or stub): bf16 + auto only, int8/fp8 dropped. + monkeypatch.setattr(common, "has_functional_torchao", lambda: False) + modes, recommended = train_precision_modes() + assert modes == ["nf4", "bf16", "auto"] + assert "int8" not in modes and "fp8" not in modes + assert recommended == "auto" + + # With a functional torchao on an fp8-capable GPU, int8 + fp8 are advertised again. + monkeypatch.setattr(common, "has_functional_torchao", lambda: True) + modes2, _ = train_precision_modes() + assert "int8" in modes2 and "fp8" in modes2 + + +def test_train_precision_modes_gates_dense_on_bf16_support(monkeypatch): + # The dense modes (bf16/int8/fp8/auto) all train in bf16 compute, which the DiT trainer + # requires. On a CUDA GPU that cannot do bf16 (T4/V100/RTX 20xx), /info must offer ONLY + # nf4 -- otherwise the UI advertises a start that evicts resident models and then fails the + # trainer's bf16 guard. + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5)) # Turing, no bf16 + monkeypatch.setattr(common, "has_functional_torchao", lambda: True) + modes, recommended = train_precision_modes() + assert modes == ["nf4"] + assert recommended == "nf4" + + +# ── family_train_infos precision fields ─────────────────────────────────────── +def test_family_train_infos_carries_precision_fields(monkeypatch): + # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL + # (no precision selector) stays empty regardless of the probe. + monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto")) + # Also pin bf16_unsupported_reason (family_train_infos reads the live GPU through it): "bf16 OK" + # so this positive-path assertion is deterministic across GPU types, not just on CPU-only CI. + monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None) + infos = {i["name"]: i for i in common.family_train_infos()} + + flux = infos["flux.1"] + assert flux["precision_modes"] == ["nf4", "bf16"] + assert flux["recommended_precision"] == "auto" + assert flux["supports_compile"] is True + + sdxl = infos["sdxl"] + assert sdxl["precision_modes"] == [] + assert sdxl["recommended_precision"] == "nf4" + assert sdxl["supports_compile"] is False + + +# ── request model base_precision field ──────────────────────────────────────── +def test_request_model_base_precision(): + # The request defaults to the nf4 memory floor. + req = DiffusionTrainingStartRequest(base_model = "x", data_dir = "d", output_dir = "o") + assert req.base_precision == "nf4" + + # An allowed dense mode is accepted. + assert ( + DiffusionTrainingStartRequest( + base_model = "x", data_dir = "d", output_dir = "o", base_precision = "fp8" + ).base_precision + == "fp8" + ) + + # An out-of-Literal value is rejected by pydantic. + with pytest.raises(Exception): + DiffusionTrainingStartRequest( + base_model = "x", data_dir = "d", output_dir = "o", base_precision = "int4" + ) + + # The generic Studio dict path carries base_precision through onto DiffusionLoraConfig. + cfg = _config_from_dict( + { + "base_model": _FLUX_DENSE, + "data_dir": "d", + "output_dir": "o", + "base_precision": "bf16", + } + ) + assert cfg.base_precision == "bf16" diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py new file mode 100644 index 0000000000..162822ec57 --- /dev/null +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""CPU-only unit tests for the diffusion training performance work. + +Covers the new pure helpers and small policy functions that the perf PR adds: +the seed-deterministic latent-cache crop/flip plan, the per-family collate fns, the +index-based sigma gather, the new config validation + request-model fields, the +torch.compile policy, the stop save/cancel flag, and the ``preparing`` / ``warning`` +service events. No GPU / model load: the collates and gathers run on CPU tensors, the +scheduler is default-initialised (no ``from_pretrained``), and the route/service tests +inject in-thread fakes exactly like ``test_diffusion_training.py``. +""" + +from __future__ import annotations + +import itertools + +import pytest +import torch +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from auth.authentication import get_current_subject +from core.training.diffusion_dit_trainer import ( + _flux_collate, + _gather_sigmas, + _qwen_collate, + _sample_timesteps, + _should_compile, + _zimage_collate, +) +from core.training.diffusion_train_common import ( + DiffusionLoraConfig, + LATENT_CACHE_OVER_BUDGET, + _apply_perf_flags, + _config_from_dict, + _latent_cache_forced, + _latent_cache_over_budget, + _plan_cache_variants, + _restore_perf_flags, +) +import core.training.diffusion_lora_trainer as sdxl_trainer +import core.training.diffusion_train_common as train_common +from core.training.diffusion_training_service import DiffusionTrainingService +from models.training import DiffusionTrainingStartRequest, DiffusionTrainingStopRequest +from routes.training import router as training_router + +# A trainable SDXL base so DiffusionLoraConfig.normalized() resolves a family without a +# network call (resolve_trainable_family is pure name matching for this repo). +_SDXL = "stabilityai/stable-diffusion-xl-base-1.0" + + +def _cfg(**kw) -> DiffusionLoraConfig: + return DiffusionLoraConfig(base_model = _SDXL, data_dir = "d", output_dir = "o", **kw) + + +# ── _plan_cache_variants (pure, seed-deterministic) ─────────────────────────── +def test_plan_cache_variants_deterministic_and_deduped(): + # Same seed -> byte-identical plan (its own rng stream, so it is fully reproducible). + p1 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123) + p2 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123) + assert p1 == p2 + assert len(p1) == 3 + + # cache_variants=1 -> exactly one variant per image. + p_one = _plan_cache_variants(3, 1, center_crop = False, random_flip = True, seed = 7) + assert [len(v) for v in p_one] == [1, 1, 1] + + # A center crop with no flip collapses to a single distinct variant no matter how many + # draws are requested, and that variant is the fixed (0.5, 0.5, False) center. + p_cc = _plan_cache_variants(2, 8, center_crop = True, random_flip = False, seed = 7) + assert [len(v) for v in p_cc] == [1, 1] + assert p_cc[0][0] == (0.5, 0.5, False) + + # A center crop WITH flip has at most two distinct variants (flip on/off; crop is fixed). + p_cf = _plan_cache_variants(2, 8, center_crop = True, random_flip = True, seed = 7) + assert all(len(v) <= 2 for v in p_cf) + + # Every crop fraction is a valid unit fraction the loader can map onto its crop range. + for u_left, u_top, flip in itertools.chain.from_iterable(p1): + assert 0.0 <= u_left < 1.0 + assert 0.0 <= u_top < 1.0 + assert isinstance(flip, bool) + + +# ── per-family collate fns ──────────────────────────────────────────────────── +def test_flux_collate_shapes(): + # FLUX embeds are fixed length: 3 entries batch by a plain cat; text_ids are shared. + entries = [(torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)] + pe, pooled, text_ids = _flux_collate(entries, "cpu", torch.float32) + assert pe.shape == (3, 512, 32) + assert pooled.shape == (3, 16) + assert text_ids.shape == (512, 3) + # Position ids stay float32 regardless of the requested weight dtype. + assert pe.dtype == torch.float32 + assert pooled.dtype == torch.float32 + assert text_ids.dtype == torch.float32 + + +def test_qwen_collate_pads_and_masks(): + dim = 8 + # A short (mask=None) and a long (mask=ones) entry -> pad to the batch max and build the + # validity mask, with the padded tail of the short sample masked out. + short = (torch.randn(1, 5, dim), None) + long = (torch.randn(1, 9, dim), torch.ones(1, 9, dtype = torch.int64)) + pe, mask = _qwen_collate([short, long], "cpu", torch.float32) + assert pe.shape == (2, 9, dim) + assert mask.shape == (2, 9) + assert torch.equal(mask[0, 5:], torch.zeros(4, dtype = mask.dtype)) + + # A single unpadded sample with a None mask keeps the legacy None mask (no behaviour delta). + pe1, mask1 = _qwen_collate([(torch.randn(1, 5, dim), None)], "cpu", torch.float32) + assert pe1.shape == (1, 5, dim) + assert mask1 is None + + # A single sample pinned to a compile pad bucket must pad AND expose a mask so the padded + # positions are attended to as invalid. + pe2, mask2 = _qwen_collate([(torch.randn(1, 5, dim), None)], "cpu", torch.float32, pad_to = 16) + assert pe2.shape == (1, 16, dim) + assert mask2 is not None + assert torch.equal(mask2[0, 5:], torch.zeros(11, dtype = mask2.dtype)) + + +def test_zimage_collate_list(): + # Z-Image uses list I/O: the batch is one tuple carrying a list of per-sample tensors, each + # cast to the requested dtype. + entries = [(torch.randn(7, 2560),), (torch.randn(9, 2560),)] + out = _zimage_collate(entries, "cpu", torch.float32) + assert isinstance(out, tuple) and len(out) == 1 + (caps,) = out + assert isinstance(caps, list) and len(caps) == 2 + assert all(t.dtype == torch.float32 for t in caps) + + +# ── index-based sigma gather ────────────────────────────────────────────────── +def test_gather_sigmas_matches_search_based_gather(): + # CI installs the backend test deps without diffusers; the scheduler math is what we + # are checking, so skip rather than fail there. + pytest.importorskip("diffusers") + from diffusers import FlowMatchEulerDiscreteScheduler + + torch.manual_seed(0) + sched = FlowMatchEulerDiscreteScheduler() # default init, no from_pretrained / no network + timesteps, indices = _sample_timesteps(sched, 16, "cpu") + + # The index path must return exactly what the old per-item timestep-matching search did. + schedule_timesteps = sched.timesteps.to("cpu") + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + assert step_indices == indices.tolist() + + sigma = _gather_sigmas(sched, indices, "cpu", torch.float32, 4) + assert sigma.ndim == 4 + expected = sched.sigmas[step_indices].flatten() + while expected.ndim < 4: + expected = expected.unsqueeze(-1) + assert torch.equal(sigma, expected) + + +# ── config validation of the new perf fields ────────────────────────────────── +def test_config_validates_new_fields(): + # Defaults normalize cleanly and carry the new perf fields through. + norm = _cfg().normalized() + assert norm.cache_variants == 4 + assert norm.compile_transformer == "auto" + assert norm.enable_tf32 is True + assert norm.cache_latents is True + + # cache_variants is bounded to 1..16 inclusive. + for bad in (0, 17): + with pytest.raises(ValueError): + _cfg(cache_variants = bad).normalized() + + # An unknown compile mode is rejected. + with pytest.raises(ValueError): + _cfg(compile_transformer = "banana").normalized() + + # compile_transformer is case/space-insensitive and stored lowered. + assert _cfg(compile_transformer = " ON ").normalized().compile_transformer == "on" + + # The generic Studio dict path preserves the flags without inventing defaults. + cfg = _config_from_dict( + { + "base_model": _SDXL, + "data_dir": "d", + "output_dir": "o", + "enable_tf32": False, + "cache_latents": False, + } + ) + assert cfg.enable_tf32 is False + assert cfg.cache_latents is False + + # String flags from the generic Studio dict path are coerced: "false" is otherwise a + # non-empty (truthy) string, so an opt-out would silently no-op. + cfg = _config_from_dict( + { + "base_model": _SDXL, + "data_dir": "d", + "output_dir": "o", + "enable_tf32": "false", + "cache_latents": "0", + } + ) + assert cfg.enable_tf32 is False + assert cfg.cache_latents is False + + +# ── torch.compile policy ────────────────────────────────────────────────────── +def test_should_compile_policy(): + # off never compiles, even on cuda. + assert _should_compile(_cfg(compile_transformer = "off"), False, "cuda") is False + # on always compiles on cuda. + assert _should_compile(_cfg(compile_transformer = "on"), False, "cuda") is True + # auto stays off over a bitsandbytes base (graph breaks in the dequant path). + assert _should_compile(_cfg(compile_transformer = "auto"), True, "cuda") is False + # auto turns on for the dense bf16 base precision on cuda. + assert ( + _should_compile(_cfg(compile_transformer = "auto"), False, "cuda", base_precision = "bf16") + is True + ) + # Any mode is a no-op on cpu. + for mode in ("off", "on", "auto"): + assert _should_compile(_cfg(compile_transformer = mode), False, "cpu") is False + + +# ── service stop save/cancel flag ───────────────────────────────────────────── +class _StopQueue: + """Records what stop() puts on the wire (put-only for these tests).""" + + def __init__(self) -> None: + self.items: list = [] + + def put(self, x) -> None: + self.items.append(x) + + +class _AliveProc: + def is_alive(self) -> bool: + return True + + +def test_service_stop_save_flag(): + svc = DiffusionTrainingService() + # Nothing running -> stop is a no-op and returns False. + assert svc.stop() is False + + # Attach a fake live proc + stop queue so stop() has a target. + svc._proc = _AliveProc() + q = _StopQueue() + svc._stop_queue = q + + # save=False is the cancel path: the dict form {"save": False} goes on the queue. + assert svc.stop(save = False) is True + assert q.items[-1] == {"save": False} + + # The default (save) path keeps the bare-True wire format. + assert svc.stop() is True + assert q.items[-1] is True + + +# ── preparing / warning events + stopped completion messages ────────────────── +def test_apply_event_preparing_and_warning(): + svc = DiffusionTrainingService() + svc._apply_event({"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8}) + st = svc.status() + assert st["status"] == "running" + assert st["in_model_load"] is True + assert "4/8" in st["message"] + + svc._apply_event({"type": "warning", "message": "compile disabled"}) + assert svc.status()["message"] == "compile disabled" + + # A stop with no saved adapter reports the no-adapter message and the stopped status. + svc_no = DiffusionTrainingService() + svc_no._apply_event({"type": "complete", "stopped": True, "lora_path": None}) + st_no = svc_no.status() + assert st_no["status"] == "stopped" + assert st_no["message"] == "Stopped (no adapter saved)." + + # A stop that DID save a partial adapter reports the partial-adapter message. + svc_partial = DiffusionTrainingService() + svc_partial._apply_event( + {"type": "complete", "stopped": True, "lora_path": "/o/pytorch_lora_weights.safetensors"} + ) + assert svc_partial.status()["message"] == "Stopped (partial adapter saved)." + + +# ── route: stop body forwards the save flag ─────────────────────────────────── +class _FakeService: + """Records the save flag the /diffusion/stop route forwards. A local copy of the + test_diffusion_training.py pattern so the two suites stay decoupled.""" + + def __init__(self) -> None: + self._running = True + self.stopped_with_save = None + + def stop(self, save = True): + self.stopped_with_save = save + was = self._running + self._running = False + return was + + +@pytest.fixture +def client(monkeypatch): + fake = _FakeService() + monkeypatch.setattr( + "core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake + ) + app = FastAPI() + app.include_router(training_router, prefix = "/api/train") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + c = TestClient(app) + c._fake = fake # type: ignore[attr-defined] + return c + + +def test_route_stop_save_body(client): + # An explicit {"save": false} body forwards save=False to the service. + r = client.post("/api/train/diffusion/stop", json = {"save": False}) + assert r.status_code == 200, r.text + assert client._fake.stopped_with_save is False + + # A body-less POST defaults to save=True. + r2 = client.post("/api/train/diffusion/stop") + assert r2.status_code == 200, r2.text + assert client._fake.stopped_with_save is True + + +# ── request models: new perf fields + stop schema ───────────────────────────── +def test_request_models_new_fields(): + req = DiffusionTrainingStartRequest(base_model = "b", data_dir = "d", output_dir = "o") + assert req.cache_latents is True + assert req.cache_variants == 4 + assert req.compile_transformer == "auto" + assert req.enable_tf32 is True + + # cache_variants is validated against its 1..16 bound by pydantic. + with pytest.raises(Exception): + DiffusionTrainingStartRequest( + base_model = "b", data_dir = "d", output_dir = "o", cache_variants = 32 + ) + + # The stop request defaults to saving a partial adapter. + assert DiffusionTrainingStopRequest().save is True + + +# ── perf flags round-trip on cpu ────────────────────────────────────────────── +def test_perf_flags_cpu_roundtrip(): + # On a cpu device (or a torch build without cuda), applying the perf flags is a no-op + # snapshot path and restoring it must not raise. + snap = _apply_perf_flags(_cfg(), "cpu") + assert isinstance(snap, dict) + _restore_perf_flags(snap) # no exception + + +def test_perf_flags_tf32_off_clears_flags(): + # enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags + # (cudnn TF32 defaults ON in torch) rather than inherit ambient state, and restore must + # put the ambient values back. The flag attributes are plain Python state, present and + # settable on CPU-only torch builds, so this runs without a GPU. + import torch + + before = ( + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.get_float32_matmul_precision(), + ) + snap = _apply_perf_flags(_cfg(enable_tf32 = False), "cuda") + try: + assert torch.backends.cuda.matmul.allow_tf32 is False + assert torch.backends.cudnn.allow_tf32 is False + assert torch.get_float32_matmul_precision() == "highest" + finally: + _restore_perf_flags(snap) + after = ( + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.get_float32_matmul_precision(), + ) + assert after == before + + +# ── latent cache size gate ──────────────────────────────────────────────────── +class _FakeLatentDist: + def __init__(self, shape): + self.mean = torch.zeros(shape, dtype = torch.float32) + self.std = torch.ones(shape, dtype = torch.float32) + + +class _FakeEncoded: + def __init__(self, shape): + self.latent_dist = _FakeLatentDist(shape) + + +class _FakeVae: + # Minimal VAE stand-in: encode() returns a posterior of the requested latent shape so the + # builder measures a real per-variant byte size without a model load or image files. + def __init__(self, shape): + self._shape = shape + + def encode(self, pixel_values): + return _FakeEncoded(self._shape) + + +def _fake_planned_loader(path, resolution, center_crop, u_left, u_top, flip): + # The fake VAE ignores pixels; return a valid tensor + square SDXL time_ids. + tensor = torch.zeros(3, resolution, resolution, dtype = torch.float32) + return tensor, (resolution, resolution, 0, 0, resolution, resolution) + + +def _build_fake_sdxl_cache(monkeypatch, num_images, latent_shape): + # center_crop + no flip collapses to one variant per image, so total_variants == num_images. + monkeypatch.setattr(sdxl_trainer, "_load_image_tensor_planned", _fake_planned_loader) + cfg = _cfg(cache_variants = 1, center_crop = True, random_flip = False).normalized() + return sdxl_trainer._build_sdxl_latent_cache( + _FakeVae(latent_shape), + 1.0, + [f"img{i}.png" for i in range(num_images)], + cfg, + "cpu", + torch.float32, + None, + lambda: False, + ) + + +def test_latent_cache_over_budget_boundary(): + # 32 bytes per variant x 4 variants = 128 bytes; exactly at budget is not "over". + assert _latent_cache_over_budget(32, 4, budget_bytes = 200) is False + assert _latent_cache_over_budget(32, 4, budget_bytes = 128) is False + assert _latent_cache_over_budget(32, 4, budget_bytes = 127) is True + # An empty plan can never overflow. + assert _latent_cache_over_budget(1_000_000, 0, budget_bytes = 1) is False + + +def test_latent_cache_forced_env(monkeypatch): + monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False) + assert _latent_cache_forced() is False + monkeypatch.setenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "1") + assert _latent_cache_forced() is True + + +def test_sdxl_cache_built_under_budget(monkeypatch): + # Default (4 GiB) budget: a handful of tiny latents fits, so the full cache is returned. + monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False) + cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8)) + assert cache is not LATENT_CACHE_OVER_BUDGET and cache is not None + assert len(cache) == 3 + assert all(len(variants) == 1 for variants in cache) + + +def test_sdxl_cache_gated_over_budget(monkeypatch): + # A budget below one variant forces the gate on the first encode: the sentinel is returned + # so the caller keeps the VAE resident and encodes per step. + monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False) + monkeypatch.setattr(train_common, "_LATENT_CACHE_BUDGET_BYTES", 8) + cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8)) + assert cache is LATENT_CACHE_OVER_BUDGET + + +def test_sdxl_cache_force_bypasses_gate(monkeypatch): + # An explicit force-on must be honoured verbatim even when the estimate is over budget. + monkeypatch.setenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "1") + monkeypatch.setattr(train_common, "_LATENT_CACHE_BUDGET_BYTES", 8) + cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8)) + assert cache is not LATENT_CACHE_OVER_BUDGET and cache is not None + assert len(cache) == 3 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index b138de3980..156803778a 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -190,6 +190,45 @@ def test_apply_event_transitions(): assert svc.status()["status"] == "error" and svc.status()["message"] == "boom" +def test_progress_nulls_non_finite_floats_for_strict_json(): + # A divergent step (or an inf grad norm) can push loss / avg_loss / learning_rate to + # NaN or Infinity, which strict JSON forbids. The service must null those so the status + # snapshot and the metric history stay strict-JSON serializable. + import json + import math + + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + svc._apply_event( + { + "type": "progress", + "step": 3, + "total_steps": 10, + "loss": float("nan"), + "avg_loss": float("inf"), + "learning_rate": float("-inf"), + } + ) + snap = svc.status() + assert snap["loss"] is None + assert snap["avg_loss"] is None + assert snap["learning_rate"] is None + # The non-finite point is skipped in the history, so the loss series stays clean. + assert snap["metric_loss"] == [] + assert snap["metric_steps"] == [] + # strict JSON (allow_nan=False) round-trips without a ValueError from NaN/Infinity. + json.dumps(snap, allow_nan = False) + + # A finite point after the bad one is recorded and preserved verbatim. + svc._apply_event( + {"type": "progress", "step": 4, "total_steps": 10, "loss": 0.5, "learning_rate": 1e-4} + ) + snap2 = svc.status() + assert snap2["loss"] == 0.5 + assert snap2["metric_loss"] == [0.5] and snap2["metric_steps"] == [4] + assert math.isfinite(snap2["learning_rate"]) + json.dumps(snap2, allow_nan = False) + + def test_terminal_events_clear_model_load_flag(): # A stop or error during model load emits complete/error WITHOUT a preceding # model_load_completed, so the terminal update must reset in_model_load or the @@ -211,6 +250,7 @@ class _FakeService: def __init__(self): self._running = False self.started_with = None + self.stopped_with_save = None # Extra keys merged into status() so a test can inject metric history / perf fields. self.status_extra: dict = {} @@ -219,7 +259,8 @@ class _FakeService: self._running = True return "job-123" - def stop(self): + def stop(self, save = True): + self.stopped_with_save = save was = self._running self._running = False return was @@ -564,6 +605,38 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat assert client._fake.started_with is None +def test_route_start_refuses_non_bf16_gpu_without_freeing_gpu(client, monkeypatch): + # A DiT precision the host cannot run (no bf16 GPU, or explicit int8 without a functional + # torchao) must 400 BEFORE resident GPU workloads are freed: otherwise the host tears down the + # user's chat/Images model and the run then dies in the trainer child. The route imports + # training_precision_preflight_error locally, so patch it on its home module. + import routes.training as tr + + freed = [] + monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1)) + monkeypatch.setattr( + "core.training.diffusion_train_common.training_precision_preflight_error", + lambda fam, prec: ( + "This trainer requires a bfloat16-capable GPU (Ampere or newer)." + if fam != "sdxl" + else None + ), + ) + r = client.post( + "/api/train/diffusion/start", + json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"}, + ) + assert r.status_code == 400 + assert "bfloat16" in r.json()["detail"] + assert freed == [] + assert client._fake.started_with is None + + # SDXL (its own mixed_precision path) is exempt: the same probe returns None, so an SDXL + # start proceeds normally past the preflight. + r2 = client.post("/api/train/diffusion/start", json = _BODY) + assert r2.status_code == 200, r2.text + + # ── metric history + perf/family fields (PR A platform) ────────────────────── def test_apply_event_records_metric_history_and_perf(): svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)