From 55043acf501a63fc6a9f6645eff5f4a030c660a2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 18:35:28 -0700 Subject: [PATCH] Diffusion training base precision modes: bf16 speed mode (2.3-2.6x), int8, fp8, auto (#6839) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add SDXL diffusion family (U-Net pipeline support) SDXL is the first U-Net family in the diffusion backend: its denoiser is pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file .safetensors is the whole pipeline rather than a transformer-only file. The backend previously assumed a DiT transformer everywhere, so add the two hooks a U-Net family needs and register SDXL. DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file with the base repo as config, instead of transformer_class.from_single_file plus a companion assembly). _align_vae_dtype now reads the denoiser generically so img2img and inpaint keep the VAE and U-Net dtypes aligned. The non-GGUF trust gate is extended with a short, exact-match, safetensors-only allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay trusted as before; a random repo, even one that detects as SDXL, is still rejected. The image-conditioned and ControlNet workflows are the standard SDXL pipelines, built around the resident modules via from_pipe like every other family, so SDXL gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no native sd.cpp mapping yet, so the no-GPU route falls back to diffusers. Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG). Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads both as a pipeline and as a single file and generates coherent txt2img + img2img. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Images: LoRA free-text Hugging Face entry + recipe round-trip The backend has always accepted a bare Hugging Face repo id (owner/name, or owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the picker only rendered when the curated catalog had entries, and the catalog is empty, so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded model supports LoRA, and replace the curated-only dropdown with a text input: type a Hub repo id, or pick a discovered adapter from a datalist of suggestions when the catalog is populated. Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's "id:weight" strings (splitting on the last colon, since the id itself may contain one for a specific weight file) back into the selection, so replaying a saved image reproduces its adapters. The generate payload trims hand-typed ids and drops empty / zero-weight rows, and a model swap clears the selection (a LoRA is family-specific) without discarding a free-text pick that is not in the curated list. * Add diffusion LoRA training (SDXL text-to-image) First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an image + caption dataset and export it as a diffusers .safetensors that the existing diffusion LoRA loader (and any diffusers pipeline) can load. core/training/diffusion_lora_trainer.py: - DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed, mixed precision). - discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image .txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested). - run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise + timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction), AdamW + get_scheduler + grad accumulation + grad clipping, then export via save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and polls should_stop for a clean stop with a partial save. - run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue), so the training worker can spawn it; plus a CLI entry point. Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are follow-ups. The trainer is decoupled and worker-ready. Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar / instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from baseline (mean abs pixel diff ~55/255). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * diffusion trainer: emit learning_rate in progress events (Studio pump compatibility) The Studio training pump reads 'learning_rate' from progress events; the diffusion trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are directly consumable by the existing training status/SSE machinery when it is wired into the worker, without a translation shim. * Wire diffusion LoRA training into the Studio API Make the SDXL LoRA trainer reachable from the app with a small, self-contained job service and JSON routes, deliberately separate from the LLM TrainingBackend (whose lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat- inference -- is text-training specific and would mis-handle a diffusion run). core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a time -- validate the config cheaply (before any spawn), spawn the trainer subprocess (spawn context, parent-lifetime bound), pump its events (model_load_* / progress / complete / error) into an in-memory status snapshot, and support a clean stop. The subprocess context and target are injectable so the full start -> pump -> status -> complete path is unit-tested without real multiprocessing or torch. routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status (JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas mirroring DiffusionLoraConfig, so model_dump() passes straight through. Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn, concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions; plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop) with a mocked service. The diffusion trainer's progress events already use the field names this path expects. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Images: add a Train LoRA (SDXL) dialog Surface the diffusion training API in the Images page. A "Train LoRA" button in the top bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick the base model, dataset folder, output folder, an optional instance prompt, and the core hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog polls the training status while open and shows a progress bar, step count, live loss, and the saved adapter path, with a Stop button for a clean stop. The dialog is independent of the loaded generation model (training runs in its own subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining / getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes. * Import diffusion training schemas from models.training directly The import-hoist lint flags newly re-exported names in the models/__init__.py hub as unused (it does not treat __all__ membership as a use). Import the three diffusion training schemas straight from models.training in routes/training.py, where they are used in the route annotations and calls, and drop the __init__ re-export. * Remove stray async task scratch outputs committed by mistake * ControlNet: reject filesystem-like ids and do not cache a model past an unload race Two review findings on the ControlNet path: - resolve_controlnet's bare-repo fallback accepted any id with a slash, so a path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory. Restrict the fallback to a strict owner/name HF repo id shape. - _controlnet_pipe now re-checks the cancel event after the blocking from_pretrained: an unload that raced the download had already cleared the caches, so caching the late module would pin it past the unload. * Pipeline prefetch: fetch only the default torch weights A full-pipeline prefetch kept every repo file outside assets/, so an official repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO, Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never loads. Skip non-torch exports and dtype-variant twins in _pipeline_file_downloaded, and drop a component .bin when the same directory carries a picked safetensors weight (diffusers' own preference). * Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16 The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 / RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported() and drop to fp16 there. * Diffusion training service: join the old pump outside the lock start() joined a finished job's pump thread while holding the service lock, but the pump's final state writes need that same lock, so the join always burned its full timeout and a stale pump could then overwrite the new job's state. Join outside the lock (with a re-check after), and fence _apply_event and the exit handler by process identity so a superseded pump can never touch the current job's state. Adds regression tests for both. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety Addresses review findings on the SDXL LoRA trainer: - Gate the base model with the same trust check as inference (unsloth/*, allowlisted official bases, or a local path) before from_pretrained, so an untrusted remote repo is never fetched or deserialised. - Check the stop signal before the (slow) model load, not only between steps, so a cancel during download is honoured; a stop may carry save=False to cancel without leaving a partial adapter. - Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor. - Apply EXIF orientation before resize/crop so rotated photos train upright. - Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead of scaling every gradient to zero. - Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative lora_alpha or learning_rate. - Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/ lr_scheduler_type/random_seed) onto the diffusion field names. - Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it. - Report worker exceptions in both message and error keys so the failure is not lost. Adds regression tests for the config coercion/validation and aliasing. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ControlNet: address review findings on the diffusers path - resolve_controlnet enforces catalog family compatibility so a direct API call cannot load a ControlNet built for another family through the wrong pipeline. - Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError to ValueError) instead of a generic 500. - strength 0 disables ControlNet entirely, so a no-op selection never pays the download / VRAM cost; the control image is decoded and validated BEFORE the ControlNet is resolved or built, so a malformed image fails fast for the same reason. - ControlNet loads use the base compute dtype (state.dtype is a display string, not a torch.dtype, so it silently fell back to float32) and honor the base offload policy via group offloading instead of forcing the module resident. - Empty/malformed HF token coerced to anonymous access. - Flux Union ControlNet control_mode mapped from the selected control type. - resolve_controlnet drops the unused hf_token/cancel_event params. - ControlNetSpec validates guidance_start <= guidance_end (clean 422). - Images UI ControlNet Select shows its placeholder when nothing is selected. Adds regression tests for family enforcement and the union control-mode map. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs Four review findings on the diffusion training start path: - It spawned the SDXL trainer without checking the LLM TrainingBackend, so a start while an LLM run was active put two trainers on the same GPU. Add a symmetric interlock: diffusion start returns 409 when LLM training is active, and LLM start refuses while a diffusion job is active. - It went straight to service.start() without freeing GPU residents. Add a pre-spawn free of the export subprocess, the resident Images pipeline (with an arbiter release), and chat models, mirroring the LLM start path. - data_dir / output_dir were passed through unresolved, so Studio-relative names failed and absolute paths bypassed containment. Resolve them with resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained path). - The request model dropped max_grad_norm and lora_target_modules, so runs that set them trained with defaults. Add both fields. The gemini pump-join deadlock was already fixed earlier (join outside the lock + proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop change owned by the diffusion training engine PR (should_stop polled before the first optimizer step). Adds route + model regression tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown Address review findings on the LoRA path: - resolve_one: normalise a blank/whitespace hf_token to None (anonymous access) and reject a client-supplied weight file with traversal / absolute path. - resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError so the route returns 400 instead of a generic 500. - _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs foo.gguf) so each is uniquely addressable. - inject_prompt_tags: the backend-validated weight now wins over a user-typed for a selected adapter; unselected user tags are left alone. - diffusers _apply_loras: reject a .gguf adapter with a clear error before touching the pipe (diffusers loads safetensors only). - _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is dropped wholesale (freeing adapters), so the previous call could race an in-flight denoise on the same pipe. - Images page: use a stable LoRA key and clear the selection (not just the options) when the catalog refresh fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Diffusion: guard trust check against OSError and validate conditioning inputs - _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid characters (or a bare owner/name id) can't raise OSError; treat any failure as not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request still raises the clear FileNotFoundError for a genuinely missing local pick. - generate(): reject mask_image / upscale / reference_images supplied without an input image, and reject reference_images on a family that does not support reference conditioning, instead of silently degrading to txt2img / img2img. * SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers Addresses review findings on the SDXL family: - Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request, before the route evicts the current model; SDXL has no transformer-only GGUF variant. - Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file (config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no longer triggers a multi-GB base download. - Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline but this backend loads every sdxl repo as the base txt2img pipeline. - Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch degrades to anonymous instead of erroring on a malformed token. - Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and access state.family.denoiser_attr directly. Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter. * Images: preserve restored LoRAs through model load and never send hidden LoRAs - The LoRA effect cleared the selection on every load->capable transition, which wiped adapters restored from a gallery recipe before the model finished loading. Track the previously-loaded family in a ref and clear only on a real family swap; keep the selection on the initial load and on unload. - Gate the generate payload's loras on loraCapable so a restored selection that is hidden (loaded model does not support LoRA) is never sent to the backend. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh Nine review findings on the SDXL training dialog: - Forward the saved Hub token so a gated/private SDXL base can be trained (the image load flow already sends it). - Re-seed the base-model field from the current default each time the dialog opens; the keep-alive dialog otherwise kept its mount-time default after a model loaded. - Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open. - Add client-side validation of steps/rank/resolution/batch/learning-rate before the request. - Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from the UI, not only the API. - Gate the dialog on the active Images route (active && trainOpen) so switching tabs closes it and stops its polling. - Rescan the LoRA picker when a run completes, so a freshly-trained adapter appears without a model reload. - Cap the dialog height and scroll the body so the Start/Stop footer stays reachable on short viewports. - Correct the copy to not over-promise picker auto-discovery. Freeing the resident Images pipeline before training is handled backend-side in the diffusion training start route. * Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch The catalog-refresh .catch from the lower branch clears the selected adapters too, which is right for its catalog-only picker but wrong here: this picker holds free-text HF repo ids that are valid without being in the catalog, so a transient refresh failure must not wipe them. Family swaps still clear the selection and hidden LoRAs are never sent. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Train LoRA dialog: stop suggesting absolute paths the backend rejects The dataset and output placeholders showed /path/to/... examples, but the training routes resolve those fields inside the Studio home and reject absolute paths outside the approved roots, so following the placeholder produced a 400. Use folder-name placeholders and say in the labels and the dialog description where each folder resolves. * Align the VAE to the denoiser's first FLOATING dtype, not its first parameter A GGUF-quantized transformer's leading parameters are packed uint8 storage, so reading next(parameters()).dtype handed nn.Module.to() an integer dtype and every image-conditioned generation on a GGUF model (Qwen-Image-Edit) failed with a 500. Probe the parameters for the first floating dtype, treat an all-integer module as a no-op, and also catch TypeError so an unexpected dtype can never break generation. Regression test included. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Count LR scheduler warmup/decay in optimizer steps, not micro-steps lr_sched.step() runs once per outer optimizer step (after the gradient accumulation inner loop), for train_steps total. The scheduler was configured with num_warmup_steps and num_training_steps multiplied by gradient_accumulation_steps, so with accumulation > 1 a warmup or non-constant schedule stretched past the run and never reached the intended decay. Count both in optimizer steps. * Address Codex review findings on the image-workflows PR Keep diffusion.py importable without torch: the compile/arch patch modules import torch at module level, so import them lazily at their load/unload call sites instead of at module load. This restores the torchless contract so get_diffusion_backend() works on a CPU/native sd.cpp install. Match family reject keywords and aliases as whole path/name segments, not raw substrings, so an unrelated word like edited, edition, or kontextual no longer misroutes or hides a valid base image model, while supported edit families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same segment matching in the picker task filter. Route FLUX.2-dev native guidance through --guidance like the other FLUX families rather than --cfg-scale. Reject native upscale requests that have no input image. Read image header dimensions and reject over-limit inputs before decoding pixels, so a crafted small-payload image cannot spike memory. Reject an upscale that would shrink the source below its input size. Validate the model_kind against the filename extension before the GPU handoff. Estimate a local diffusers pipeline's size from its on-disk weights so auto memory planning does not skip offload and OOM. Report workflows: [txt2img] from the native backend status so the Create tab stays enabled for a loaded native model. Clamp the outpaint canvas to the backend's 4096px decode limit. Adds regression tests for segment matching and kind/extension validation. * Guard inference loads and worker lifetime against diffusion training Teach the chat and image load guards about an active diffusion (SDXL) LoRA job: a chat load is refused (its footprint cannot be fit-checked against the trainer) and an image load is refused outright, mirroring the existing LLM training guards, so a load can no longer allocate GPU memory alongside the trainer and undo the pre-start cleanup. Bind the diffusion trainer subprocess to the parent's lifetime and scrub the native path lease secret from it by running the child through run_without_native_path_secret, matching the inference/export/LLM workers, so a Studio crash or kill no longer leaves the trainer holding the GPU. Reset in_model_load on the complete and error terminal events: a stop or failure during model loading otherwise leaves the status reporting a stale loading indicator after the job has ended. * Harden diffusion LoRA handling on the diffusers and native paths Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max): diffusers requires the adapter loaded before compilation, so applying one to the already-compiled module fails with adapter-key mismatches. The status gate now hides the picker and generate raises a clear message instead. Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the diffusion cancellation sentinel in resolve_specs, so an unload/superseding load during resolution maps to a 409 instead of a generic server error. Drop weight-0 LoRA rows before the native support gate so a request carrying only disabled adapters stays a no-op on families where native LoRA is unsupported, matching the diffusers path. Reject duplicate LoRA ids in the request model: both apply paths suffix colliding names, so a repeated id would stack the same adapter past its per-adapter weight bound. Strip all user-typed prompt tags on the native path (only the selected adapters are materialized in the managed lora-model-dir, so an unselected tag can never resolve), and restore saved LoRA selections from a gallery recipe so restore reproduces a LoRA image. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden ControlNet resolve, gallery metadata, and the control-type picker Check cancellation immediately after a ControlNet from_pretrained and before any device placement, so an unload/eviction that raced the download does not allocate several GB onto the GPU after the load was already cleared. Require a loadable weight or shard index (not just config.json) before a local ControlNet folder is advertised, so an interrupted copy is hidden instead of failing deep in from_pretrained as a generic 500. Do not record a strength-0 ControlNet in the gallery recipe: it is treated as disabled and skipped, so the image is unconditioned and the metadata must not claim a ControlNet was applied. Build the control-type picker from the selected ControlNet's advertised control_types instead of a hardcoded passthrough/canny pair, so a union model with a precomputed depth or pose map sends the correct control_mode. * Address further Codex findings on the image-workflows PR - Persist the actual output image size in the gallery recipe instead of the request sliders: Transform/Inpaint/Edit derive the size from the uploaded image, Extend grows the canvas, and Upscale resizes it, so the sliders recorded (and later restored) the wrong dimensions for those workflows. - Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file name) in validate_load_request, so the unloadable pick fails before chat is evicted rather than deep in from_pretrained. - Only publish an image-conditioned from_pipe wrapper to the shared aux cache when the load is still current: from_pipe runs under the generate lock but not the state lock, so an unload racing its construction could otherwise cache a wrapper over torn-down modules that a later load would reuse. - Verify the Windows CUDA runtime archive checksum before extracting it, like the main sd-cli archive, so a corrupt or tampered runtime is rejected rather than extracted next to the binary. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refuse non-SDXL base models at diffusion training start The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image repo or a GGUF filename passed as base_model was accepted and then failed minutes later inside StableDiffusionXLPipeline.from_pretrained with an unrelated-looking error. Add a name-based guard in normalized() so known DiT-family names and .gguf checkpoints are rejected up front, which the API start route surfaces as an immediate 400 with a message that says exactly which bases are trainable. Unrecognisable names still pass through so custom local SDXL checkpoints keep working. * Add diffusion dataset upload and training info endpoints Training an image LoRA required knowing the Studio home layout and copying files onto the server by hand, which is the most confusing step of the whole flow. Two small endpoints fix that: - GET /api/train/diffusion/info reports the datasets and outputs roots plus every dataset folder that contains images (with image/caption counts), so the UI can offer a picker instead of a blind free-text path. - POST /api/train/diffusion/dataset uploads images and optional caption .txt / metadata.jsonl files into a named folder under the datasets root, creating it on first use and accumulating on repeat uploads so large sets can arrive in batches. Names are validated to a single path component and files stream to disk under the same per-upload size cap as LLM dataset uploads. The returned name is a valid data_dir for /diffusion/start. * Rework the Train LoRA dialog into a guided SDXL flow The dialog assumed users knew the Studio home layout and that only SDXL is trainable, and hid both facts behind free-text fields. Restructure it around the three real decisions: - Base model is a dropdown of the trainable SDXL picks (Base 1.0, Turbo, the loaded SDXL pipeline when there is one) with a custom repo/path escape hatch, instead of a bare text field defaulting to a repo id. - Training images come from an in-browser upload (new dataset endpoints) or a picker over existing dataset folders with image/caption counts. No shell access or knowledge of the datasets root is needed any more, and the captioning rules are explained inline. - The output field is now Adapter name and the instance prompt is labelled as the trigger prompt, with a no-captions warning wired to the selected dataset's actual caption count. Hyperparameters collapse behind a training settings toggle since the defaults suit a first run. A completed run says where the adapter went and offers Done / Train another, and the top-bar button gets an icon and a plainer description. The dialog title states the SDXL-only scope and that other families load LoRAs but cannot train them yet. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Validate diffusion training config before freeing the GPU The start route freed resident GPU workloads (export, Images pipeline, chat) before the service validated the config, so a start that was then refused, now including a non-SDXL base model, tore down the user's loaded model for nothing. Run the same cheap normalise pass first; the LLM path already follows this rule via its before_spawn hook. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactor diffusion LoRA training into a family-aware platform Split the SDXL trainer into a shared, architecture-agnostic layer so more model families can be trained without duplicating the plumbing: - New core/training/diffusion_train_common.py holds the config + validation, dataset discovery, event emission, stop protocol, adapter publishing, and a lazy trainer registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and re-exports the moved names so existing imports are unchanged. - The SDXL-only base-model blocklist becomes a positive check: the family is resolved from the base model (or an explicit model_family) via the diffusion family registry, and a known-but-not-yet-trainable family is refused with a clear message. Unknown custom names still default to the SDXL trainer. - DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable. DiT families flip on when their trainers land. - Trained adapters now write a .json metadata sidecar (family, base model, rank, trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the picker instead of showing it as unknown for every model. - The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo, and Qwen-Image repos (safetensors-only, no remote code). * Retain diffusion training loss history and expose it in status The training service kept only the latest loss, so a live loss chart could show a single point. Fold each progress event into bounded (step, loss, lr) history arrays (capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM, and record the family / base model / catalog path on completion. The status endpoint returns these as a nested metric_history object the UI can chart directly, and the start request accepts an optional model_family override. * Tests for the diffusion training platform Cover the trainer registry (get_trainer resolves SDXL, unknown family raises), family resolution (explicit model_family validation, resolved_family on the config), the metadata sidecar write + scan read with family gating, and the service loss-history folding (append, bad-point skipping, decimation at cap, family/perf fields) plus the status route nesting metric_history. * Add diffusion dataset labeling and example-import endpoints The Train tab needs to let users caption small datasets in the browser and pull in a ready-made set to see training work end to end, neither of which the upload-only endpoint supported. Add, under /api/train/diffusion/dataset: - GET {name}/images lists every image with its resolved caption (metadata beats a per-image sidecar, matching the trainer's discovery order) so uncaptioned images are visible and flaggable. - GET {name}/image/{filename} serves an image, with ?thumb= returning a cached downscaled JPEG kept in a hidden .thumbs subdir (regenerated when the source is newer) so the labeling grid stays light. - PUT {name}/caption/{filename} writes, or when blank clears, the .txt sidecar; DELETE {name}/image/{filename} removes the image plus its sidecars and thumbnails. - GET dataset-examples lists a curated, license-labelled registry, and POST dataset/import-example materializes one into a dataset folder as numbered images + .txt captions. Two loaders cover the shapes seen in the wild: streaming rows from datasets.load_dataset (dog-example, Tuxemon) and a snapshot + jsonl walk for imagefolder repos whose captions live in a non-standard *.jsonl (the public-domain tarot set). Imports are idempotent and cap the image count. Filenames and dataset names are validated against path traversal and pinned inside the datasets root. * Test diffusion dataset labeling and example-import endpoints Cover caption precedence, thumbnail generation and .thumbs exclusion, caption write/clear, image delete cleanup, path-traversal rejection on names and filenames, and example import with a mocked datasets.load_dataset (files plus sidecars written, idempotent second call, cap respected, load failure mapped to 502). * Add flow-matching DiT LoRA trainers (FLUX.1-dev, Qwen-Image, Z-Image) Extends diffusion LoRA training beyond SDXL to the three popular DiT families via a single shared flow-matching loop parameterised by small per-family specs (loading, prompt/latent encoding, transformer forward, save). Verified against diffusers 0.38.0: - FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the user's HF token). - Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std, img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16 base). - Z-Image: list I/O with the reversed timestep convention and a negated prediction, bf16 only. The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now route these families to the DiT trainer; the SDXL blocklist guard is replaced by a positive family resolution that also rejects GGUF repos (inference-only) and still-unsupported families. Per-family defaults + labels + VRAM notes are exposed via family_train_infos for the Train UI. Memory: caption embeddings are precomputed once and the text encoders freed before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit) and 8-bit AdamW are on by default. * Speed up + shrink SDXL LoRA training (precompute text embeds, 8-bit AdamW) SDXL re-encoded every caption with both CLIP text encoders on every step (pure waste, since captions are constant) and kept the encoders resident. Precompute each unique caption's embeddings once, then free the text encoders before the loop: numerically identical (embeddings are deterministic and this consumes no torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32 fallback, halving optimizer state with no meaningful LoRA quality cost. Env toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard A/B the paths. * Expose trainable families in /diffusion/info and preflight gated bases The training info endpoint now returns the trainable model families (name, label, default + allowed base repos, recommended defaults, and a VRAM/access note) so the Train UI can offer a base picker with realistic guidance. The start route preflights a gated base repo (HEAD model_index.json with the user's token) BEFORE freeing resident GPU workloads, so a missing FLUX.1-dev license/token fails fast with an actionable 400 instead of evicting the loaded model and then hitting a confusing mid-load 401. * Tests for DiT trainers, family resolution, info families, gated preflight Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image are trainable (and GGUF repos are rejected as inference-only), the families list in /diffusion/info, and the gated-base 400 preflight that leaves the GPU untouched. * Add diffusion training API client: metrics, families, dataset labeling, examples Extends the Images training client for the Train tab: the status type now carries metric_history (step/loss/lr) plus catalog_path/family/base_model/samples_per_second/ peak_memory_gb; the start request gains model_family; and info gains an optional families list (per-family bases + defaults). Adds typed calls for the dataset labeling and one-click example endpoints: list images with captions, thumbnail URL, write/clear a caption, delete an image, list example datasets, and import an example. * Add diffusion Train panel: config, dataset labeling, live charts, deploy New full-page training workspace for the Images tab. Left column configures the run: model family (FLUX.1-dev, Qwen-Image, Z-Image, SDXL in popularity order, with per-family VRAM/license notes and defaults, backfilled from the backend families list when present), base repo, dataset (existing folder, browser upload, or one-click example import), an in-browser caption labeling grid (per-image thumbnail + caption saved on blur, delete, uncaptioned highlight), adapter name, trigger prompt, and collapsed training settings. Right column shows the live run: progress + loss/avg/speed/peak-VRAM readouts, the reused training loss/LR charts fed from metric_history, and a completion card that deploys the adapter into Create or starts another run. * Wire Create/Train tab switch into the Images page and deploy flow Replaces the Train LoRA dialog with a top-bar Create | Train segmented control next to the model selector. Create renders the existing generation workspace unchanged; Train renders the full-page training panel (unmounted in Create so its polling stops while the backend run and its retained metric history survive a tab switch). Adds a deploy handler: loading the trained adapter's base as a pipeline, queueing the adapter so the LoRA discovery effect applies it once the base is loaded and LoRA-capable for the matching family (with a mismatch warning), seeding the prompt with the trigger, and switching back to Create. Removes the now-unused dialog. * Wrap the DiT training forward in bf16 autocast The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single compute dtype during the forward, exactly like the diffusers dreambooth scripts run under accelerator.autocast. Without it the 4-bit backward on FLUX.1-dev fails with an illegal-address CUBLAS error partway into the first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL path (its own trainer) is untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Train tab example cards and Create/Train tab layout The example-dataset cards used a two-column grid in the ~340px config column, which wrapped titles one word per line and let the long license text overrun into the neighbouring card. Switch to one card per row with a horizontal layout: title with a compact truncated license badge (full text in the tooltip), a two-line clamped description, and the Import button on the right. The Create/Train switch had an icon inside the Train trigger that overhung the pill corner. Drop the icon, make both triggers a fixed equal width so the active pill sits flush in the top bar. * Show only loss and learning-rate charts for diffusion training The Train tab reused the LLM charts section, which also rendered an empty Grad Norm card and an Eval Loss card showing an Evaluation not configured placeholder with a red smear. Neither applies to diffusion LoRA training. Add a diffusion-only two-card view that reuses the loss and learning-rate cards directly with fixed presentation defaults, and note under the loss chart that per-step loss is noisy by design so users read the smoothed line for the trend rather than the raw jitter. * Add a dataset preview strip to the Train tab When a dataset with images is selected, show a strip of up to 8 sampled thumbnails with a +N more tile, so users can see what is in the folder before training. Clicking the strip opens the existing caption review grid. Samples are drawn evenly across the folder and refresh on dataset change or after an upload/import. * Stop example cards from overflowing the Train config column The example-dataset cards still overran the ~340px config column: the license used the Badge component whose baked-in w-fit and whitespace-nowrap ignored the max-width and truncate, and the grid children had the default min-width auto so wide content pushed past the column edge and clipped the Import buttons. Replace the badge with a plain truncating pill span, and give the config column min-w-0 with overflow-x-hidden so nothing escapes its width. * Add Smithsonian Butterflies and Nouns example datasets Two permissive ~100-image sets for the Train tab: huggan/smithsonian_butterflies_subset (CC0, the classic diffusers-docs training set, imported as a subject set with a trigger prompt since its metadata columns are species names not captions) and m1guelpf/nouns (CC0, captioned pixel-art avatars via the text column). Both cap at 100 images. * Paginate the Train tab caption grid with prev/next controls Large example datasets (100+ images) rendered every tile at once, so the caption review grid grew unbounded. Show 24 images per page with < > chevrons and an x-y of N indicator; a new dataset or refresh resets to the first page. * Offer example datasets in the Train dropdown with previews Add an Examples group to the training-images dropdown that imports a curated dataset in one pick, alongside the existing cards. Cards now show up to three preview thumbnails pulled from the public HF datasets-server so the set is visible before download. Hide the trigger prompt when every image already has a caption (a captioned style set needs no trigger), and turn the training-settings toggle into a ghost button with a rotating chevron. * Clamp the training base repo to the selected family The base-model select's state could briefly hold the previous family's repo after a family switch (the reseed effect runs a beat later, and a value with no matching option makes the browser display the first option anyway). The request then carried the stale repo: picking Qwen or Z-Image still sent black-forest-labs/FLUX.1-dev and surfaced FLUX's gated-repo error under the wrong family. Derive an effectiveBase clamped to the current family's repos and use it for the select value, the start request, and the deploy fallback. Also move the Trigger prompt above Adapter name: the trigger describes the dataset, the name only labels the output. * Speed up diffusion LoRA training and cut DiT peak VRAM by a third Perf core for the diffusion trainers, defaults preserving the training math: - Phased model loading: the pipeline now loads without its transformer (conditioning only), captions are encoded and the text encoders freed, the VAE latent cache is built and the VAE freed, and only then does the transformer load. The multi-GB denoiser never shares VRAM with the encoders, cutting measured peak VRAM on B200: FLUX 17.1 -> 10.4 GB, Qwen-Image 19.1 -> 12.8 GB, Z-Image 7.3 -> 4.7 GB. - Latent cache (cache_latents, default on): per-image crop/flip variants (cache_variants, default 4 vs the single frozen variant of the diffusers --cache_latents) store the VAE posterior's affine parameters, so every step still draws a fresh VAE sample; a cached center-crop Z-Image run matches the uncached one at the bf16 nondeterminism floor. - True batching: train_batch_size now actually batches the transformer forward (it was silently 1). nf4 dequant dominates the step cost, so batch 4 lands near batch-1 step time: 4.0x samples/s on Qwen-Image, 3.1x on FLUX, 2.1x on Z-Image, with multi-seed loss envelopes overlapping batch-1. - LR scheduler support in the DiT loop (lr_scheduler / lr_warmup_steps were accepted but ignored); progress events now report the real per-step LR. - TF32 + high fp32 matmul precision under enable_tf32 (default on), snapshot/restored around the run. cudnn.benchmark is scoped to a caller opt-in only: autotuning the fp32 VAE convs doubled peak VRAM on the DiT families for zero steady-state gain. - Vectorized sigma gathering (drops a per-step Python search loop), cached FLUX img_ids/guidance, fused torch AdamW fallback, steady-state samples_per_second (excludes the first-step warmup). - Regional torch.compile plumbing (compile_transformer off/on/auto with eager fallback): auto stays off over a bitsandbytes base where compile is a net loss (27 s warmup, slightly slower steady on Z-Image); it arms automatically for the dense/quantized speed modes that follow. - Stop parity with the LLM trainer: /api/train/diffusion/stop accepts an optional {save} body and the service forwards save=False as a no-save cancel; a new preparing event surfaces cache-build progress. - SDXL trainer gets the same latent cache, perf flags, and fused fallback; its batching, LR schedule, and min-SNR stay as they were. Verified: 83 backend tests green; per-family 30-40 step runs with adapter round-trip generation through the normal LoRA path (FLUX, Qwen-Image, Z-Image all pass). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add base_precision speed modes to DiT training: bf16 2.3-2.6x, int8, fp8 New base_precision config for the DiT trainers: nf4 (unchanged default) | bf16 | int8 | fp8 | auto, advertised per family + per machine through /api/train/diffusion/info (precision_modes, recommended_precision, supports_compile) so the UI can gate the selector. - bf16: dense transformer + regional torch.compile (auto-armed). The measured speed mode: 2.3x nf4 on FLUX (1.81 -> 4.12 steps/s), 2.6x on Z-Image (2.5 -> 6.38 steps/s) on B200, at dense-weight VRAM (FLUX 24.7 GB / Z-Image 13.6 GB peak vs 10.4 / 4.7 for nf4). - int8: torchao weight-only int8 on the frozen base, quantized AFTER add_adapter (quantizing first trips peft 0.18's TorchaoLoraLinear, which is incompatible with the torchao 0.16 config API). Runs eager: inductor rejects the int8 subclass training graph (aliased subclass outputs), so compile is force-disabled for it. - fp8: torchao convert_to_float8_training on the frozen linears (filter skips lora_ modules, proj_out, non-divisible-by-16 dims, pad_inner_dim), applied after add_adapter, compile auto-armed. Works and round-trips, but measured SLOWER than compiled bf16 at LoRA-training shapes (FLUX 3.15 vs 4.12 steps/s; Z-Image similar), so it is an explicit opt-in and auto never picks it. - auto: free VRAM (measured before load) + dense-size table -> bf16 when it fits with headroom, int8 in the middle band, else nf4. Prequant bnb repos always resolve to nf4; dense modes on them are rejected at validation with a pointer to the family's dense base. Two crashes found and fixed along the way: - The cuDNN SDPA backend's training graph fails on the FLUX attention shapes (torch 2.10 + cu130, B200): mha_graph.execute errors, then the context degrades into illegal memory accesses. The perf-flag guard now pins flash/mem-efficient SDPA for the run (mathematically equivalent, snapshot/restored). nf4 escaped it by routing attention differently. - Regional compile now uses dynamic=True (the inference layer's proven default): dynamic=False specialisation fused a gemm_and_bias epilogue that failed with CUBLAS_STATUS_EXECUTION_FAILED on the FLUX training graph; dynamic=True is also faster (Z-Image 3.84 -> 6.38 steps/s). Verified: 98 backend tests green (new test_diffusion_base_precision.py: validation, auto policy table, fp8 filter, compile gating, /info fields); per-mode 40-step runs on FLUX + Z-Image with loss means inside the nf4 envelope and adapter round-trip generation through the normal LoRA path for bf16-, fp8-, and int8-trained adapters. * Clear TF32 flags when enable_tf32 is off so the opt-out is strict fp32 * Address review: auto int8 requires the dense-load transient to fit, dense modes are CUDA-only, auto respects bf16 compute, exact cudnn SDPA restore * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: fp32 latent cache stats + strict-JSON-safe progress floats - Latent caches (DiT + SDXL) now hold the posterior mean/std in fp32 and draw the per-step sample in fp32, casting only the result to the training dtype. This matches the in-loop path (encode fp32 -> sample fp32 -> cast) exactly instead of sampling in bf16; the cache is tiny so the doubled RAM is negligible. - The training service nulls non-finite floats (NaN/Inf loss, avg_loss, learning_rate) at its single ingestion point so status snapshots and persisted run records stay strict-JSON serializable; the metric history skips non-finite loss points. Test covers NaN/Inf progress followed by a finite point. * Address review: gate auto int8 on torchao, scope dense validation to DiT - base_precision="auto" only picks int8 when torchao is importable (the int8 quantize has no runtime fallback, unlike fp8); otherwise the middle band falls back to nf4. Threaded as a parameter so the policy stays pure. - The dense-mode validation (prequant base / bf16 compute) now applies only to DiT families: sdxl ignores base_precision entirely, so a leftover value can no longer fail an SDXL run. The mode-name validity check still runs everywhere. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate int8 and fp8 on a functional torchao import, not find_spec The Windows ROCm torchao import stub satisfies find_spec and even lets from torchao.quantization import quantize_ succeed, but its quantize_ is a no-op: auto would pick int8, leave the transformer dense, and disable compile as if it were quantized. has_functional_torchao imports the exact symbols the int8 path uses and rejects the stub via its sentinel; both the auto picker and the /info advertised modes now use it * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip the sigma-gather test when diffusers is not installed CI runs the backend suite without diffusers; the test checks our index math against the scheduler's own gather, so it skips rather than fails there. * Coerce cache_latents and enable_tf32 string flags in the config dict path The generic Studio config dict path can deliver these flags as strings, and a non-empty string like "false" is truthy, so an opt-out silently no-ops (the latent cache still builds, TF32 stays on). Coerce them the same way gradient_checkpointing already is. * Remove committed runtime scratch artifacts and ignore their dirs logs/ (a 1.3 MB ComfyUI object_info dump plus stale PID files), temp/ (PR body and commit message scratch), and async_task_outputs/ (agent task transcripts) are environment specific runtime artifacts that were committed by accident and carry stale local state into every checkout. Remove them and gitignore the directories so they cannot be re-added. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore pre-Ampere bf16 fail-fast in the DiT trainer The perf rewrite dropped the bf16 capability guard, so a pre-Ampere CUDA device (T4/V100/RTX 20xx) would die deep in model load with an opaque dtype error instead of a clear message. Restores parity with the SDXL trainer. * Size-gate the automatic diffusion latent cache The latent cache holds two fp32 posterior tensors per crop/flip variant per image, pinned on CUDA hosts, so datasets with thousands of images can exhaust host or pinned memory with no fallback. Estimate the cache size from the first real encoded latent and fall back to per-step VAE encoding when it exceeds a 4 GiB budget. UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE bypasses the gate; the existing UNSLOTH_DIFFUSION_NO_LATENT_CACHE opt-out is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate DiT training precision: deny fp8 for Qwen, gate explicit int8 on torchao, gate advertised dense modes + route on bf16 - normalized() + family_train_infos() mirror the inference fp8 deny for Qwen-Image (activation outliers exceed fp8's range and corrupt the trained result); int8 stays allowed and the UI no longer advertises fp8 for it. - _resolve_base_precision() gates an explicit int8 on a FUNCTIONAL torchao, the same gate auto and /info already apply, so a missing/stub torchao fails fast instead of silently loading dense with compile disabled. - train_precision_modes() gates the dense modes (bf16/int8/fp8/auto) on torch.cuda.is_bf16_supported(), so a non-bf16 CUDA GPU (T4/V100/RTX 20xx) is offered only nf4 instead of a start that evicts resident models and then fails. - start_diffusion_training preflights bf16 support for the DiT families BEFORE _free_gpu_for_diffusion_training(), so any DiT start (nf4 included, since the trainer requires bf16 unconditionally on CUDA) fails fast without eviction. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate DiT training on functional torchao for explicit int8; hide always-400 DiT modes on non-bf16 GPUs The start route preflight only rejected non-bf16 GPUs; an explicit int8 request on a host with a missing or stub torchao passed the preflight, evicted resident GPU workloads, then died in the trainer child (its int8 base quantizer has no fallback). Fold both gates into training_precision_preflight_error so int8-without-torchao fails fast before eviction. Also empty the advertised DiT precision_modes (and surface the reason in vram_note, drop compile) whenever the bf16 preflight would reject the family, so /info never offers an nf4 DiT option the route always 400s. * Reject dense DiT precisions on a CUDA-absent host before eviction; stabilize family-info tests The start-route preflight caught the bf16-GPU and int8-torchao requirements but not the dense precisions' CUDA requirement: on a GPU-less host bf16_unsupported_reason exempts CPU-only, so a bf16/fp8 (or int8-with-torchao) DiT request passed the preflight, evicted resident workloads, then raised only in the trainer child. Add the dense-mode CUDA gate mirroring _resolve_base_precision so the doomed run is rejected up front. Also pin bf16_unsupported_reason in the two positive-path family-info tests so they are deterministic across GPU types (a non-bf16 CUDA box would otherwise empty every DiT family's advertised modes). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .gitignore | 5 + plans/controlnet-workflow.md | 113 +++ plans/diffusion-popularity-findings.md | 113 +++ plans/diffusion-workflows-pr-plan.md | 156 +++ plans/diffusion-workflows-studio.md | 125 +++ plans/wobbly-jumping-narwhal.md | 141 +++ .../core/training/diffusion_dit_trainer.py | 920 +++++++++++++++--- .../core/training/diffusion_lora_trainer.py | 672 ++++++++----- .../core/training/diffusion_train_common.py | 367 ++++++- .../training/diffusion_training_service.py | 100 +- studio/backend/models/training.py | 34 + studio/backend/routes/training.py | 28 +- .../tests/test_diffusion_base_precision.py | 547 +++++++++++ .../tests/test_diffusion_train_perf.py | 469 +++++++++ .../backend/tests/test_diffusion_training.py | 75 +- 15 files changed, 3449 insertions(+), 416 deletions(-) create mode 100644 plans/controlnet-workflow.md create mode 100644 plans/diffusion-popularity-findings.md create mode 100644 plans/diffusion-workflows-pr-plan.md create mode 100644 plans/diffusion-workflows-studio.md create mode 100644 plans/wobbly-jumping-narwhal.md create mode 100644 studio/backend/tests/test_diffusion_base_precision.py create mode 100644 studio/backend/tests/test_diffusion_train_perf.py 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)