ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Invalidate latents on a VAE swap, keep a cut-off generation, surface the EMA adapter

Browse source 
- source_revision() scanned the checkpoint root plus text_encoder/tokenizer but
  not vae, so swapping or fine-tuning the VAE in place left the conditioning
  cache namespace unchanged and a warm run trained against latents from the old
  checkpoint. Include the vae directory, like any other component the cached
  tensors come from.
- /images/generate answers only when the images are saved, and secure mode's
  tunnel caps an origin response near 100 seconds, which a native CPU or a
  high-step run passes routinely. The page reported failure while the work kept
  running, and a retry would duplicate it. A lost response (fetch rejection or
  a gateway status the origin never answered) is now told apart from a refusal:
  the page waits out generate-progress and reloads the gallery, so the run it
  started still lands.
- The trainer emits the EMA adapter's path with the terminal event, but the
  state update dropped it, so neither the run history nor either response
  schema carried it and an enabled EMA left nothing discoverable. Keep it, and
  show it next to the primary adapter.
- weighting_scheme advertised a choice of timestep sampling; sampling is always
  logit-normal and the flag only selects the bell loss weights. Describe what
  it does.
This commit is contained in:
Daniel Han 2026-07-26 20:47:46 +00:00
commit 68912bfaf2
9 changed files with 517 additions and 458 deletions
Download patch file Download diff file Expand all files Collapse all files

374
studio/backend/core/training/diffusion_dit_trainer.py
View file

@ -64,9 +64,9 @@ from core.training.diffusion_train_extras import (
save_ema_adapter,
)
# Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks
# also carry added-kv projections; Z-Image is single-stream. Architecture-specific, so kept
# here rather than in the generic DEFAULT_LORA_TARGETS.
# Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks also
# carry added-kv projections; Z-Image is single-stream. Architecture-specific, so kept here
# rather than in the generic DEFAULT_LORA_TARGETS.
_FLUX_TARGETS = (
"to_q",
"to_k",
@ -79,9 +79,9 @@ _FLUX_TARGETS = (
)
_QWEN_TARGETS = _FLUX_TARGETS
_ZIMAGE_TARGETS = ("to_q", "to_k", "to_v", "to_out.0")
# The Krea 2 authors' recommended default targets (their DreamBooth reference script):
# attention + SwiGLU + text-fusion projector + conditioning embedders. For long runs they
# suggest narrowing to the attention layers so prompt adherence doesn't drop.
# The Krea 2 authors' recommended default targets (their DreamBooth reference script): attention
# + SwiGLU + text-fusion projector + conditioning embedders. For long runs they suggest narrowing
# to the attention layers so prompt adherence doesn't drop.
_KREA2_TARGETS = (
"img_in",
"final_layer.linear",
@ -123,29 +123,27 @@ class _FamilySpec:
lora_targets: tuple[str, ...]
# bf16 only (Z-Image overflows fp16 and its RoPE/embedder run in fp32).
force_bf16: bool
# Approximate dense-bf16 transformer weight size, used by base_precision="auto" to
# decide which mode fits the free VRAM (with headroom for activations + optimizer).
# 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 coexists with the text encoders + VAE:
# ``load_conditioners`` builds the pipeline WITHOUT its transformer (encode_prompt + VAE
# only) and returns (pipe, vae); ``load_transformer`` then loads the transformer alone (as
# trainable nf4 QLoRA when qlora=True) once the conditioners are freed. Roughly halves peak
# VRAM for the big DiTs.
# ``load_conditioners`` builds the pipeline WITHOUT its transformer and returns (pipe, vae);
# ``load_transformer`` then loads the transformer alone (trainable nf4 QLoRA when qlora=True)
# once the conditioners 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 params 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 a pixel tensor to (A, B) affine posterior params so a per-step sample is A + B * randn
# (family normalisation folded in). B is None for a deterministic 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 a list of per-caption embed tuples into 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.
# One transformer forward: (transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device,
# weight_dtype) to model_pred aligned with target = noise - latents.
forward: Callable[..., Any]
# Save the LoRA in diffusers format via the family pipeline's save_lora_weights.
save: Callable[..., None]
@ -186,8 +184,8 @@ def _training_sigma_table(scheduler, flow_shift):
mu = float(getattr(sc, "max_shift", None) or math.log(3.0))
shifted = scheduler.time_shift(mu, 1.0, sigmas)
if getattr(sc, "shift_terminal", None):
# The stretch scales off the table's own final sigma, so it must see the full
# descending schedule (as it does at inference set_timesteps), never a batch.
# The stretch scales off the table's own final sigma, so it must see the full descending schedule
# (as it does at inference set_timesteps), never a batch.
shifted = scheduler.stretch_shift_to_terminal(shifted)
return shifted
s = float(flow_shift)
@ -250,10 +248,9 @@ def _bnb_4bit_config():
load_in_4bit = True,
bnb_4bit_quant_type = "nf4",
bnb_4bit_compute_dtype = torch.bfloat16,
# Double quantization compresses the per-block absmax scales with a second 8-bit
# pass: ~0.4 bits/param off the frozen base at no fidelity cost (the scales are
# metadata, not weights), which matters most on the 20B+ DiTs where the nf4 base
# dominates VRAM.
# Double quantization compresses the per-block absmax scales with a second 8-bit pass: ~0.4
# bits/param off the frozen base at no fidelity cost (the scales are metadata, not weights),
# which matters most on the 20B+ DiTs where the nf4 base dominates VRAM.
bnb_4bit_use_double_quant = True,
)
@ -317,9 +314,8 @@ def _load_dit_transformer(transformer_cls, cfg, device, base_precision):
return transformer
# Dense load for bf16 / fp8 / mxfp8 / 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).
# _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.
return transformer_cls.from_pretrained(
cfg.base_model,
subfolder = "transformer",
@ -407,9 +403,8 @@ def _mx_module_filter(mod, fqn: str) -> bool:
if fqn.endswith("proj_out") or ".proj_out." in fqn:
return False
# Skip biased linears: the torchao 0.17 MX training path swaps the weight for a wrapper whose
# linear override computes input @ weight_t and drops the bias entirely, so an mxfp8'd FROZEN
# base linear loses its bias and changes the output the LoRA regresses against (verified on
# Blackwell: bias fully dropped). Keep biased linears in bf16.
# linear override computes input @ weight_t and drops the bias entirely, so an mxfp8'd frozen
# base linear changes the output the LoRA regresses against. Keep biased linears in bf16.
if getattr(mod, "bias", None) is not None:
return False
return mod.in_features % 32 == 0 and mod.out_features % 32 == 0
@ -498,21 +493,19 @@ def _resolve_base_precision(cfg, spec, device) -> str:
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 -- memory saving gone, 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. fp8 keeps its own fallback (_apply_fp8_training), so int8-only.
# 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 -- memory saving gone, likely OOM. The auto pick and /info already
# gate on a FUNCTIONAL torchao; apply the same gate here. fp8 keeps its own fallback.
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'."
)
# mxfp8 needs Blackwell (sm100+): its MX GEMM has no kernel below sm100 and raises at the
# first training step, AFTER a full dense load. /info advertises mxfp8 only on sm100+
# (train_precision_modes), so re-check here to fail fast for a stale/direct client on an
# older GPU instead of crashing mid-run.
# mxfp8 needs Blackwell (sm100+): its MX GEMM has no kernel below sm100 and raises at the first
# training step, AFTER a full dense load. /info advertises mxfp8 only on sm100+, so re-check here
# to fail fast for a stale/direct client instead of crashing mid-run.
if mode == "mxfp8" and device == "cuda":
try:
import torch
@ -525,8 +518,8 @@ def _resolve_base_precision(cfg, spec, device) -> str:
"Use base_precision='bf16', 'int8', 'nf4', 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.
# 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)
@ -534,9 +527,8 @@ def _resolve_base_precision(cfg, spec, device) -> str:
capability = None
has_fp8 = False
# int8 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 leaves the transformer dense with compile disabled. has_functional_torchao
# imports the exact symbols _int8_quantize_base uses and rejects the stub.
# find_spec("torchao") is satisfied by the Windows-ROCm import stub whose quantize_ is a no-op.
# has_functional_torchao imports the exact symbols _int8_quantize_base uses and rejects the stub.
has_torchao = has_functional_torchao()
if device == "cuda":
try:
@ -607,8 +599,8 @@ def _flux_collate(
):
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.
# 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)
@ -617,7 +609,7 @@ def _flux_collate(
# Per-run cache of the step-invariant FLUX conditioning tensors (RoPE image ids + guidance
# vector): shapes are fixed once resolution/batch are, so rebuilding them every step is pure
# allocator churn. Cleared at run start (subprocess-local anyway).
# allocator churn. Cleared at run start.
_FLUX_STATIC: dict[tuple, tuple] = {}
@ -628,8 +620,8 @@ def _flux_static_inputs(bsz, h, w, device):
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.
# Position ids drive RoPE and are indices, not activations, so 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)
@ -672,8 +664,8 @@ def _qwen_load_conditioners(cfg, device, weight_dtype):
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 and loads trainable as-is under nf4; the dense modes need the 20B Qwen/Qwen-Image base.
# The prequant default ships the transformer 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)
@ -707,8 +699,8 @@ def _qwen_latent_affine(vae, ref):
def _qwen_encode_latents(vae, pixel_values):
import torch
# AutoencoderKLQwenImage is a 3D (video) VAE: add a temporal dim, encode, drop it back
# into a [B,16,1,H,W] latent normalised by the per-channel latents_mean / latents_std.
# AutoencoderKLQwenImage is a 3D (video) VAE: add a temporal dim, encode, drop it back into a
# [B,16,1,H,W] latent normalised by the per-channel latents_mean / latents_std.
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]
@ -735,8 +727,8 @@ def _qwen_collate(
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.
# 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 = [], []
@ -751,8 +743,8 @@ def _qwen_collate(
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).
# A single unpadded sample keeps the legacy None mask (identical math; avoids a behaviour delta
# for existing single-image runs).
if len(entries) == 1 and entries[0][1] is None and target == seqs[0]:
mask_b = None
return (pe_b, mask_b)
@ -764,8 +756,8 @@ def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, devi
pe, mask = embeds_batch
bsz, c, f, h, w = noisy.shape
packed = QwenImagePipeline._pack_latents(noisy, bsz, c, h, w)
# Each batch entry is a LIST of one (frame, h/2, w/2) tuple: the transformer indexes
# sample[0] / sample[1:] per entry (transformer_qwenimage.py), so a flat list breaks it.
# Each batch entry is a LIST of one (frame, h/2, w/2) tuple: the transformer indexes sample[0] /
# sample[1:] per entry (transformer_qwenimage.py), so a flat list breaks it.
img_shapes = [[(1, h // 2, w // 2)]] * bsz
pred = transformer(
hidden_states = packed,
@ -794,8 +786,8 @@ def _zimage_load_conditioners(cfg, device, weight_dtype):
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).
# 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)
@ -827,8 +819,8 @@ def _zimage_encode_latents(vae, pixel_values):
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.
# 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
@ -846,8 +838,8 @@ def _zimage_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, de
import torch
(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.
# 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))
t_norm = (1000 - timesteps) / 1000
out = transformer(x_list, t_norm, list(caps), return_dict = False)[0]
@ -866,8 +858,8 @@ def _zimage_save(pipe_cls, out_dir, transformer_lora_layers):
# ── Krea 2 ────────────────────────────────────────────────────────────────────
def _krea2_load_conditioners(cfg, device, weight_dtype):
# The krea repo ships transformers-5.x style configs the pinned 4.x line can't parse, so the
# conditioning pipeline is assembled per-component (with tokenizer and rope compat) rather
# than from_pretrained(transformer = None); see diffusion_krea2.py for the compat story.
# conditioning pipeline is assembled per-component rather than from_pretrained(transformer=None);
# see diffusion_krea2.py.
import torch
from core.inference.diffusion_krea2 import load_krea2_pipeline
@ -879,8 +871,8 @@ def _krea2_load_conditioners(cfg, device, weight_dtype):
def _krea2_load_transformer(cfg, device, weight_dtype, base_precision):
# The transformer subfolder is diffusers-format (no transformers compat needed).
# There is no prequant repo yet, so nf4 quantizes the 12B transformer on the fly.
# The transformer subfolder is diffusers-format. No prequant repo yet, so nf4 quantizes the 12B
# transformer on the fly.
from diffusers import Krea2Transformer2DModel
return _load_dit_transformer(Krea2Transformer2DModel, cfg, device, base_precision)
@ -893,9 +885,8 @@ def _krea2_encode_prompts(pipe, captions, device):
with torch.no_grad():
for cap in captions:
# encode_prompt pads/truncates to the fixed max_sequence_length, so every embed is
# [1, 512, num_text_layers, 2560] with a [1, 512] validity mask -- static shapes
# (padding sits mid-template, BEFORE the assistant suffix, matching how the model was
# sampled at training time).
# [1, 512, num_text_layers, 2560] with a [1, 512] validity mask (static shapes; padding sits
# mid-template, BEFORE the assistant suffix, matching how the model was sampled at training).
pe, mask = pipe.encode_prompt(
prompt = cap,
device = device,
@ -906,8 +897,8 @@ def _krea2_encode_prompts(pipe, captions, device):
return out
# Krea 2 conditions on the Qwen-Image VAE (AutoencoderKLQwenImage) with the same
# per-channel latents_mean / latents_std normalisation, so latent encoding is shared.
# Krea 2 conditions on the Qwen-Image VAE with the same per-channel latents_mean / latents_std
# normalisation, so latent encoding is shared.
_krea2_encode_latents = _qwen_encode_latents
_krea2_encode_latent_stats = _qwen_encode_latent_stats
@ -920,8 +911,8 @@ def _krea2_collate(
):
import torch
# Fixed-length embeds (see _krea2_encode_prompts), so collation is a plain concat
# with the mask riding along; ``pad_to`` is moot because the shapes are static.
# Fixed-length embeds (see _krea2_encode_prompts), so collation is a plain concat with the mask
# riding along; ``pad_to`` is moot because the shapes are static.
pe_b = torch.cat([e[0] for e in entries]).to(device = device, dtype = weight_dtype)
mask_b = torch.cat([e[1] for e in entries]).to(device)
return (pe_b, mask_b)
@ -931,9 +922,9 @@ def _krea2_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, dev
from diffusers import Krea2Pipeline
pe, mask = embeds_batch
# [B,16,1,H,W] -> [B, (H/2)*(W/2), 64] 2x2 patches. Krea2Pipeline._pack_latents /
# _unpack_latents are instance methods (they read self.patch_size), so the packing is inlined
# here like the reference DreamBooth script (patch_size = 2).
# [B,16,1,H,W] to [B, (H/2)*(W/2), 64] 2x2 patches. Krea2Pipeline._pack_latents / _unpack_latents
# are instance methods (they read self.patch_size), so the packing is inlined here like the
# reference DreamBooth script (patch_size = 2).
bsz, c, _f, h, w = noisy.shape
packed = noisy.reshape(bsz, c, h // 2, 2, w // 2, 2)
packed = packed.permute(0, 2, 4, 1, 3, 5).reshape(bsz, (h // 2) * (w // 2), c * 4)
@ -962,27 +953,25 @@ def _krea2_save(pipe_cls, out_dir, transformer_lora_layers):
# ── FLUX.2 (dev + Klein) ──────────────────────────────────────────────────────
# Both variants share Flux2Transformer2DModel and the packing/forward conventions of the
# upstream DreamBooth references (train_dreambooth_lora_flux2[_klein].py); they differ only
# in the conditioning stack (dev: Mistral-3-Small via Flux2Pipeline; Klein: Qwen3 via
# Flux2KleinPipeline) and size. Latents train in the PATCHIFIED, batch-norm-normalised space
# the references use (space-to-depth then (x - bn_mean) / bn_std from the VAE's running
# stats), from the posterior MODE (deterministic, so the latent cache stores the final
# latent and skips the per-step draw).
# Both variants share Flux2Transformer2DModel and the packing/forward conventions of the upstream
# DreamBooth references; they differ only in the conditioning stack (dev: Mistral-3-Small via
# Flux2Pipeline; Klein: Qwen3 via Flux2KleinPipeline) and size. Latents train in the PATCHIFIED,
# batch-norm-normalised space the references use, from the posterior MODE (deterministic, so the
# latent cache stores the final latent and skips the per-step draw).
_FLUX2_TARGETS = (
# Double-stream blocks: separate q/k/v plus the ModuleList out proj.
"to_k",
"to_q",
"to_v",
"to_out.0",
# Single-stream blocks: the fused qkv+mlp input projection carries the bulk of the
# capacity. Their out proj is a PLAIN Linear named to_out whose suffix would also match
# the double-stream ModuleList container (which peft cannot wrap), and the upstream
# scripts enumerate it per block index (a variant-dependent count), so it stays dense.
# Single-stream blocks: the fused qkv+mlp input projection carries the bulk of the capacity.
# Their out proj is a PLAIN Linear named to_out whose suffix would also match the double-stream
# ModuleList container (which peft cannot wrap), and the upstream scripts enumerate it per block
# index, so it stays dense.
"to_qkv_mlp_proj",
)
# The references train dev with its guidance-distillation vector at 3.5 (Klein applies it
# only when the variant's config carries guidance_embeds).
# The references train dev with its guidance-distillation vector at 3.5 (Klein applies it only
# when the variant's config carries guidance_embeds).
_FLUX2_TRAIN_GUIDANCE = 3.5
@ -1008,10 +997,9 @@ def _flux2_encode_prompts(pipe, captions, device):
out = []
with torch.no_grad():
for cap in captions:
# encode_prompt pads to max_sequence_length (padding="max_length"), so the
# embeds are fixed-length and text_ids are per-caption [1, txt_len, 4]. The
# per-class text_encoder_out_layers defaults differ (dev (10,20,30) vs Klein
# (9,18,27)) and are left to the pipeline.
# encode_prompt pads to max_sequence_length, so the embeds are fixed-length and text_ids are
# per-caption [1, txt_len, 4]. The per-class text_encoder_out_layers defaults differ (dev
# (10,20,30) vs Klein (9,18,27)) and are left to the pipeline.
pe, text_ids = pipe.encode_prompt(
prompt = cap,
device = device,
@ -1029,16 +1017,15 @@ def _flux2_encode_latents(vae, pixel_values):
with torch.no_grad():
lat = vae.encode(pixel_values.to(torch.float32)).latent_dist.mode()
lat = Flux2Pipeline._patchify_latents(lat)
# The FLUX.2 VAE normalises latents with its BatchNorm running stats, not
# shift/scaling_factor (mirrors the upstream reference).
# The FLUX.2 VAE normalises latents with its BatchNorm running stats, not shift/scaling_factor.
mean = vae.bn.running_mean.view(1, -1, 1, 1).to(lat)
std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(lat)
return (lat - mean) / std
def _flux2_encode_latent_stats(vae, pixel_values):
# FLUX.2 trains from the posterior mode (deterministic): the cached entry is the final
# latent itself; B is None and the loop skips the per-step sampling draw.
# FLUX.2 trains from the posterior mode (deterministic): the cached entry is the final latent
# itself; B is None and the loop skips the per-step sampling draw.
return _flux2_encode_latents(vae, pixel_values), None
@ -1050,9 +1037,8 @@ def _flux2_collate(
):
import torch
# Fixed-length embeds (padding="max_length" in encode_prompt) -> plain concat;
# text_ids are PER-SAMPLE [1, txt_len, 4] (unlike FLUX.1's shared grid), so they batch
# by concat too. ``pad_to`` is moot because the shapes are static.
# Fixed-length embeds give a plain concat; text_ids are PER-SAMPLE [1, txt_len, 4] (unlike
# FLUX.1's shared grid), so they batch by concat too. ``pad_to`` is moot with static shapes.
pe = torch.cat([e[0] for e in entries]).to(device = device, dtype = weight_dtype)
text_ids = torch.cat([e[1] for e in entries]).to(device = device, dtype = torch.float32)
return (pe, text_ids)
@ -1078,8 +1064,8 @@ def _flux2_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, dev
from diffusers import Flux2Pipeline
pe, text_ids = embeds_batch
# ``noisy`` is already in the patchified normalised space (_flux2_encode_latents), so
# packing is the [B,C,H,W] -> [B, H*W, C] flatten the pipeline uses.
# ``noisy`` is already in the patchified normalised space, so packing is the [B,C,H,W] to
# [B, H*W, C] flatten the pipeline uses.
packed = Flux2Pipeline._pack_latents(noisy)
img_ids = _flux2_static_img_ids(noisy, device)
guidance = None
@ -1097,9 +1083,9 @@ def _flux2_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, dev
return_dict = False,
)[0]
pred = pred[:, : packed.size(1)]
# _unpack_latents_with_ids scatters per sample; diffusers 0.39 stacks the per-sample
# tensors itself (its list annotation is stale), older/newer builds may hand back the
# list. Same-resolution training batches align either way with target = noise - latents.
# _unpack_latents_with_ids scatters per sample; diffusers 0.39 stacks the per-sample tensors
# itself (its list annotation is stale), other builds may hand back the list. Same-resolution
# batches align either way with target = noise - latents.
unpacked = Flux2Pipeline._unpack_latents_with_ids(pred, img_ids)
if isinstance(unpacked, (list, tuple)):
unpacked = torch.stack(unpacked)
@ -1200,8 +1186,8 @@ _SPECS: dict[str, _FamilySpec] = {
family = "flux.2-dev",
lora_targets = _FLUX2_TARGETS,
force_bf16 = True,
# 32B DiT; the Mistral conditioning stack (~46 GB bf16) is loaded, encoded, and
# freed BEFORE this lands on the device (the shared phased load).
# 32B DiT; the Mistral conditioning stack (~46 GB bf16) is loaded, encoded and freed BEFORE this
# lands on the device (the shared phased load).
dense_bf16_gb = 64.5,
load_conditioners = _flux2_load_conditioners,
load_transformer = _flux2_load_transformer,
@ -1215,9 +1201,9 @@ _SPECS: dict[str, _FamilySpec] = {
}
# HF repos that gate access behind a license acceptance: training needs a token whose account
# accepted the license. Checked by name (no network) so a missing token fails fast with an
# actionable message instead of a confusing 401 mid-load.
# HF repos gating access behind a license acceptance: training needs a token whose account
# accepted it. Checked by name (no network) so a missing token fails fast with an actionable
# message instead of a confusing 401 mid-load.
_GATED_TRAIN_REPOS = frozenset({"black-forest-labs/flux.1-dev", "black-forest-labs/flux.2-dev"})
@ -1354,11 +1340,10 @@ def _build_latent_cache(
pass
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 resident so
# the loop encodes per step. ``b`` is None for a deterministic-latent family, so
# only ``a`` contributes bytes there.
# 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 resident so the loop encodes per step. ``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()
@ -1489,16 +1474,16 @@ def _should_compile(
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.
# 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/mxfp8 are only competitive compiled (eager, their per-matmul
# dynamic casts run 4-5x slower than bf16).
# 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/mxfp8 are only competitive compiled (eager, their per-matmul dynamic casts run
# 4-5x slower than bf16).
return base_precision in ("bf16", "fp8", "mxfp8")
@ -1536,16 +1521,15 @@ def _maybe_compile_transformer(
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 like the inference speed layer does
# (diffusers' documented regional-compile fix).
# Heterogeneous-block DiTs (Z-Image: ~11 distinct block shapes) exceed dynamo's default recompile
# limit of 8; bump it like the inference speed layer does.
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
# 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)
@ -1567,10 +1551,9 @@ def run_dit_lora_training(
if spec is None:
raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
# DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). An explicit fp16
# request on a bf16-only family is refused, not silently upgraded, so the choice is never
# misrepresented. Validation runs before the heavy imports so a host without diffusers still
# sees the real error.
# DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). An explicit fp16 request
# on a bf16-only family is refused, not silently upgraded, so the choice is never misrepresented.
# Validation runs before the heavy imports so a host without diffusers still sees the real error.
if cfg.mixed_precision == "fp16" and spec.force_bf16:
raise ValueError(
f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
@ -1597,12 +1580,10 @@ def run_dit_lora_training(
return True
device = "cuda" if torch.cuda.is_available() else "cpu"
# The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box: unsupported for
# real runs but keeps import/unit tests architecture-agnostic). Fail fast on pre-Ampere CUDA
# (T4/V100/RTX 20xx): bf16 is required and the run would otherwise die deep in model load with
# an opaque dtype error. Gate on NATIVE bf16 (capability major >= 8) -- is_bf16_supported()
# counts pre-Ampere emulation as supported, which this guard rejects; shared with /info modes
# + start preflight.
# The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box: unsupported for real
# runs but keeps import/unit tests architecture-agnostic). Fail fast on pre-Ampere CUDA: bf16 is
# required and the run would otherwise die deep in model load with an opaque dtype error. Gate on
# NATIVE bf16 (capability major >= 8), since is_bf16_supported() counts pre-Ampere emulation.
if device == "cuda" and not native_bf16_supported():
raise ValueError(
"This trainer requires a bfloat16-capable GPU (Ampere or newer); "
@ -1615,9 +1596,8 @@ def run_dit_lora_training(
pairs = discover_image_caption_pairs(
cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
)
# Resolve num_epochs -> a concrete train_steps now the dataset size is known, and rebind cfg
# so every downstream read (scheduler length, loop range, progress total_steps, steps_run)
# sees the same value.
# Resolve num_epochs into a concrete train_steps now the dataset size is known, and rebind cfg so
# every downstream read (scheduler length, loop range, progress totals) sees the same value.
cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0)
_emit(on_event, "model_load_started", num_images = len(pairs))
if _check_stop():
@ -1628,7 +1608,7 @@ def run_dit_lora_training(
return str(out_dir)
# 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).
# disposable, but restoring keeps in-process callers clean).
perf_snap = _apply_perf_flags(cfg, device)
try:
return _train_dit(
@ -1660,14 +1640,14 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
use_lora_targets = _select_lora_targets(cfg.lora_target_modules, spec.lora_targets)
out_dir = Path(cfg.output_dir).expanduser()
# Phase 0: the persistent conditioning cache (opt-in via cond_cache_dir). When every
# planned latent variant AND every caption embedding is already on disk, the run is
# "warm": the VAE and the multi-GB text encoders are never loaded at all.
# Phase 0: the persistent conditioning cache (opt-in via cond_cache_dir). When every planned
# latent variant AND every caption embedding is already on disk, the run is "warm": the VAE and
# the multi-GB text encoders are never loaded at all.
image_paths = [p for p, _ in pairs]
captions = [c for _, c in pairs]
uniq = sorted(set(captions))
# CFG dropout swaps a sample's conditioning for the empty prompt, so its embedding must
# be precomputed alongside the captions (the encoders are freed right after encoding).
# CFG dropout swaps a sample's conditioning for the empty prompt, so its embedding must be
# precomputed alongside the captions (the encoders are freed right after encoding).
cfg_dropout = float(getattr(cfg, "cfg_dropout", 0.0) or 0.0)
to_encode = uniq + ([""] if cfg_dropout > 0 and "" not in uniq else [])
use_cache = cfg.cache_latents and os.environ.get(
@ -1676,18 +1656,16 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
pcache = None
if getattr(cfg, "cond_cache_dir", None):
try:
# Namespace on the CHECKPOINT, not just the family: keys below carry only the
# caption/image content and crop variant, so one cache dir reused for two
# checkpoints (or the same repo at a new revision) would let a warm run train
# on embeddings and latent stats produced by the other model -- incompatible
# shapes at best, silently wrong conditioning at worst.
# Namespace on the CHECKPOINT, not just the family: the keys below carry only caption/image
# content and crop variant, so one cache dir reused for two checkpoints (or a new revision) would
# let a warm run train on the other model's embeddings and latent stats.
from .diffusion_train_extras import source_revision # noqa: PLC0415
namespace = f"{spec.family}_{cfg.base_model}_{source_revision(cfg.base_model)}"
pcache = PersistentConditioningCache(cfg.cond_cache_dir, namespace, cfg.resolution)
except Exception as exc: # noqa: BLE001 -- the cache is an optimisation, never fatal
_emit(on_event, "warning", message = f"conditioning cache disabled: {exc}")
# The crop/flip variant plan is seed-deterministic, so the persistent keys are stable
# across runs of the same config and the warm check can run before anything loads.
# The crop/flip variant plan is seed-deterministic, so the persistent keys are stable across runs
# of the same config and the warm check can run before anything loads.
plan = _plan_cache_variants(
len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed
)
@ -1709,10 +1687,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
total = len(image_paths),
)
if caption_embeds is None:
# Phase 1 (cold): conditioning only. The pipeline loads WITHOUT its transformer, so
# the text encoders + VAE never share VRAM with the multi-GB denoiser. Caption
# embeddings are precomputed once and the (large) encoders freed: captions are
# constant and the encoders frozen, so this is exact and the biggest memory win.
# Phase 1 (cold): conditioning only. The pipeline loads WITHOUT its transformer, so the text
# encoders + VAE never share VRAM with the multi-GB denoiser. Caption embeddings are precomputed
# once and the encoders freed: captions are constant and the encoders frozen, so this is exact.
latent_cache = None
pipe, vae = spec.load_conditioners(cfg, device, weight_dtype)
encoded = _encode_prompts_cached(spec, pipe, to_encode, device, pcache)
@ -1722,9 +1699,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
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).
# Phase 2: the VAE latent cache, then free the VAE too (the cache keeps the posterior affine
# parameters, so per-step sampling noise is preserved).
if use_cache:
latent_cache = _build_latent_cache(
spec,
@ -1739,8 +1715,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
plan = plan,
)
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).
# 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(
@ -1762,13 +1738,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
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 picks use their own stream so the training loop's index/noise draws stay
# seed-deterministic 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 / mxfp8 are the dense speed modes; "auto" picks from free VRAM
# measured before the load).
# Phase 3: only now load the transformer, in the resolved base precision (nf4 QLoRA by default;
# bf16 / int8 / fp8 / mxfp8 are the dense speed modes; "auto" picks from free VRAM).
base_precision = _resolve_base_precision(cfg, spec, device)
transformer = spec.load_transformer(cfg, device, weight_dtype, base_precision)
base_is_bnb = base_precision == "nf4"
@ -1785,10 +1760,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
)
)
if cfg.gradient_checkpointing:
# Non-reentrant checkpointing: reentrant recompute of a bnb 4-bit LoRA linear can trip
# an illegal memory access on the larger FLUX transformer, and non-reentrant is the
# recommended mode anyway (it also handles a checkpointed segment whose inputs don't
# require grad, which happens with a frozen 4-bit base).
# Non-reentrant checkpointing: reentrant recompute of a bnb 4-bit LoRA linear can trip an illegal
# memory access on the larger FLUX transformer, and non-reentrant is the recommended mode anyway
# (it also handles a checkpointed segment whose inputs don't require grad).
import functools
import torch.utils.checkpoint as _ckpt
transformer.enable_gradient_checkpointing(
@ -1797,8 +1771,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
cast_training_params(transformer, dtype = torch.float32)
lora_params = [p for p in transformer.parameters() if p.requires_grad]
# int8 / fp8 / mxfp8 convert the frozen base linears AFTER the LoRA attaches, so the
# adapter modules are excluded and stay high precision.
# int8 / fp8 / mxfp8 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):
@ -1806,17 +1780,16 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
if base_precision == "mxfp8" and not _apply_mxfp8_training(transformer, on_event):
base_precision = "bf16"
# LoRA EMA (opt-in via cfg.ema_decay > 0): shadows ONLY the trainable adapter params
# (megabytes, not the base's gigabytes), initialised after the precision conversions so
# the shadows track the final fp32 param objects. Exported as a second adapter under
# <output_dir>/ema at save time.
# LoRA EMA (opt-in via cfg.ema_decay above 0): shadows ONLY the trainable adapter params
# (megabytes, not the base's gigabytes), initialised after the precision conversions so the
# shadows track the final fp32 param objects. Exported as a second adapter under the ema subdir.
ema = LoRAEMA(transformer, decay = cfg.ema_decay) if getattr(cfg, "ema_decay", 0.0) else None
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).
# 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())
@ -1825,8 +1798,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
cfg.base_model, subfolder = "scheduler", token = cfg.hf_token
)
# Timestep-shift + loss-weighting setup (see _training_sigma_table). getattr defaults
# keep an un-normalized config (tests, direct callers) on the historical behavior.
# Timestep-shift + loss-weighting setup (see _training_sigma_table). getattr defaults keep an
# un-normalized config (tests, direct callers) on the historical behavior.
flow_shift = getattr(cfg, "flow_shift", None)
if flow_shift is None:
flow_shift = "auto" if spec.family == "qwen-image" else 1.0
@ -1838,9 +1811,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
if str(getattr(cfg, "weighting_scheme", "none") or "none") == "bell"
else None
)
# 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).
# The LR schedule advances once per optimizer update, so warmup/decay are counted in optimizer
# steps; multiplying by the accumulation factor would stretch warmup past the run.
lr_sched = get_scheduler(
cfg.lr_scheduler,
optimizer = optimizer,
@ -1854,8 +1826,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
n_images = len(image_paths)
batch_size = cfg.train_batch_size
# Permutation-cycle index sampler (shared with the SDXL trainer): visits every image once per
# cycle before repeating, so a short run covers the whole dataset instead of the old
# with-replacement draw. Uses the loop's own rng to stay seed-deterministic.
# cycle before repeating, so a short run covers the whole dataset. Uses the loop's own rng to
# stay seed-deterministic.
index_sampler = PermutationBatchSampler(n_images, rng)
stopped = False
running_loss = 0.0
@ -1863,10 +1835,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
t_start = time.time()
t_steady = None
done = 0
# bf16 autocast around the forward + loss, matching the diffusers dreambooth scripts'
# accelerator.autocast: 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.
# bf16 autocast around the forward + loss, matching the diffusers dreambooth scripts: 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"
@ -1896,9 +1867,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
timesteps, t_indices = _sample_timesteps(scheduler, latents.shape[0], device)
sigmas = _gather_sigmas(sigma_table, t_indices, device, weight_dtype, latents.ndim)
if shift_active:
# The model's timestep conditioning must follow the shifted sigma (the
# scheduler convention is timestep = sigma * num_train_timesteps). Gather
# in fp32 from the table so bf16 rounding never skews the conditioning.
# The model's timestep conditioning must follow the shifted sigma (timestep = sigma *
# num_train_timesteps). Gather in fp32 from the table so bf16 rounding never skews it.
timesteps = (
sigma_table[t_indices].to(device = device, dtype = torch.float32).flatten()
* num_train_ts
@ -1935,8 +1905,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
grad_norm = None
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
# clip_grad_norm_ returns the total PRE-clip norm: the health signal the UI
# charts (an exploding norm shows up here even while the clip caps the update).
# clip_grad_norm_ returns the total PRE-clip norm: the health signal the UI charts (an exploding
# norm shows up even while the clip caps the update).
grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
optimizer.step()
lr_sched.step()
@ -1947,8 +1917,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
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.
# 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":
@ -1985,7 +1955,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
catalog_path = _publish_to_lora_catalog(lora_path, cfg)
if ema is not None and ema.updates > 0:
try:
ema_path = save_ema_adapter(ema, transformer, spec.save, str(out_dir))
# Report the adapter FILE, like lora_path, so a caller can load it directly.
ema_dir = save_ema_adapter(ema, transformer, spec.save, str(out_dir))
ema_path = str(Path(ema_dir) / DEFAULT_LORA_FILENAME)
except Exception as exc: # noqa: BLE001 -- the primary adapter is already saved
_emit(on_event, "warning", message = f"EMA adapter save failed: {exc}")
_emit(

23
studio/backend/core/training/diffusion_train_extras.py
View file

@ -43,10 +43,9 @@ from typing import Any, Iterable, Optional
# ── LoRA EMA ──────────────────────────────────────────────────────────────────
# Warmup horizon for the EMA decay ramp: effective decay is
# min(decay, (1 + updates) / (WARMUP_OFFSET + updates)), the standard inverse
# ramp. With the offset at 10, step 1 averages aggressively (~0.18) and the
# ramp reaches 0.99 after ~1000 updates, so a 300-step run still ends with a
# shadow that has absorbed most of the trajectory instead of the init.
# min(decay, (1 + updates) / (WARMUP_OFFSET + updates)), the standard inverse ramp. With the
# offset at 10, step 1 averages aggressively (~0.18) and the ramp reaches 0.99 after ~1000
# updates, so a 300-step run still ends with a shadow that absorbed most of the trajectory.
_EMA_WARMUP_OFFSET = 10.0
@ -177,7 +176,7 @@ def source_revision(ref: Any) -> str:
A repo id or directory path is not a version: a Hub repo that advances to a new
revision, or a directory edited in place, keeps the same string while its encoders
change, so cached conditioning from the old encoder would be reused. Shared by the
or VAE change, so cached conditioning from the old ones would be reused. Shared by the
trainer's cache namespace and the inference-side wrapper; deliberately cheap and
never touches the encoders, since the point of the cache is that a warm run does
not load them.
@ -194,7 +193,9 @@ def source_revision(ref: Any) -> str:
roots += [
e.path
for e in it
if e.is_dir() and e.name.startswith(("text_encoder", "tokenizer"))
# vae too: cached latents come from it, so an in-place VAE swap must
# invalidate them just like an encoder change.
if e.is_dir() and e.name.startswith(("text_encoder", "tokenizer", "vae"))
]
for root in roots:
with os.scandir(root) as it:
@ -309,13 +310,13 @@ class PersistentConditioningCache:
# ── aspect-ratio bucketing ────────────────────────────────────────────────────
# Pixel-dimension divisor for bucket shapes. The DiT families divide by 8 in
# the VAE and 2 again in latent patching, and regional torch.compile prefers a
# small set of distinct shapes, so buckets snap to multiples of 64 pixels.
# Pixel-dimension divisor for bucket shapes. The DiT families divide by 8 in the VAE and 2 again
# in latent patching, and regional torch.compile prefers few distinct shapes, so buckets snap to
# multiples of 64 pixels.
BUCKET_DIVISOR = 64
# Widest aspect ratio a bucket may take; anything more extreme clamps to it
# (matching the common multi-tier bucketing practice of capping panoramas).
# Widest aspect ratio a bucket may take; anything more extreme clamps to it (matching the common
# practice of capping panoramas).
MAX_BUCKET_RATIO = 2.0

40
studio/backend/core/training/diffusion_training_service.py
View file

@ -56,22 +56,23 @@ def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) ->
def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
# First thing in the child (before torch): self-bind to parent death and scrub the native
# path secret, like the other workers (multiprocessing children can't be given a preexec_fn).
# First thing in the child (before torch): self-bind to parent death and scrub the native path
# secret, like the other workers (multiprocessing children can't be given a preexec_fn).
from utils.native_path_leases import run_without_native_path_secret
run_without_native_path_secret(
_run_diffusion_child, event_queue = event_queue, stop_queue = stop_queue, config = config
)
# Cap on retained metric points; over it, arrays are decimated (every other point) so a
# long run stays bounded while the loss chart keeps its shape.
# Cap on retained metric points; over it, arrays are decimated (every other point) so a long run
# stays bounded while the loss chart keeps its shape.
_METRIC_CAP = 4000
# ── persisted run history ──────────────────────────────────────────────────────
# Every terminal run is recorded as one JSON file (summary + scrubbed config + metric logs)
# so the Train tab can show history. JSON (not LLM sqlite) keeps diffusion runs off the LLM Runs page.
# Every terminal run is recorded as one JSON file (summary + scrubbed config + metric logs) so
# the Train tab can show history. JSON, not the LLM sqlite, so diffusion runs stay off the LLM
# Runs page.
def _runs_dir() -> Path:
from utils.paths.storage_roots import studio_root
@ -93,8 +94,8 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]:
rec = json.loads(p.read_text(encoding = "utf-8"))
except Exception: # noqa: BLE001 -- a corrupt record never breaks the listing
continue
# Skip a wrong-shape record (not a dict, or missing string job_id/status) so one bad
# file can't blow up the route's DiffusionTrainingRunSummary(**r) or the whole panel.
# Skip a wrong-shape record (not a dict, or missing string job_id/status) so one bad file can't
# blow up the route's DiffusionTrainingRunSummary(**r) or the whole panel.
if not isinstance(rec, dict):
continue
if not (isinstance(rec.get("job_id"), str) and isinstance(rec.get("status"), str)):
@ -133,6 +134,8 @@ def _idle_state() -> dict[str, Any]:
"in_model_load": False,
"output_dir": None,
"lora_path": None,
# The optional second (EMA-averaged) adapter, when ema_decay was enabled.
"ema_path": None,
"catalog_path": None,
"family": None,
"base_model": None,
@ -172,7 +175,8 @@ def _append_metric(
floss = _finite_or_none(loss)
if floss is None: # non-numeric or non-finite: skip, keep the curve JSON-safe
return
# lr / grad_norm may be None or non-finite; non-finite is nulled (not dropped) to stay index-aligned.
# lr / grad_norm may be None or non-finite; non-finite is nulled (not dropped) to stay
# index-aligned.
flr = _finite_or_none(lr)
fgn = _finite_or_none(grad_norm)
steps = state["metric_steps"]
@ -208,8 +212,8 @@ class DiffusionTrainingService:
self._ctx = ctx if ctx is not None else _CTX
self._target = target if target is not None else _default_target
self._lock = threading.Lock()
# Set by reserve() while a start is in flight (before the route frees GPU models) so the
# load guards refuse a concurrent load during the free-then-spawn window. Cleared by unreserve().
# Set by reserve() while a start is in flight (before the route frees GPU models) so the load
# guards refuse a concurrent load during the free-then-spawn window. Cleared by unreserve().
self._reserved = False
self._proc: Any = None
self._stop_queue: Any = None
@ -258,8 +262,8 @@ class DiffusionTrainingService:
_config_from_dict(config).normalized()
# Join a finished job's pump OUTSIDE the lock: its final state writes take this lock, so
# joining under it would stall the start and let the stale pump overwrite the new state.
# Join a finished job's pump OUTSIDE the lock: its final state writes take this lock, so joining
# under it would stall the start and let the stale pump overwrite the new state.
with self._lock:
if self._proc is not None and self._proc.is_alive():
raise RuntimeError("A diffusion training job is already running.")
@ -403,6 +407,7 @@ class DiffusionTrainingService:
"started_at": s.get("started_at"),
"ended_at": s.get("updated_at"),
"lora_path": s.get("lora_path"),
"ema_path": s.get("ema_path"),
"catalog_path": s.get("catalog_path"),
"saved": bool(s.get("lora_path")),
"config": cfg,
@ -440,8 +445,8 @@ class DiffusionTrainingService:
elif etype == "model_load_completed":
s.update(in_model_load = False, message = "Training...")
elif etype == "preparing":
# A long precompute phase (e.g. VAE latent cache) before the first step; surfaced
# so the UI shows progress instead of a silent "Loading base model..." stall.
# A long precompute phase (e.g. VAE latent cache) before the first step; surfaced so the UI shows
# 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(
@ -457,8 +462,8 @@ class DiffusionTrainingService:
# Non-fatal trainer notes; keep training state, surface the text.
s["message"] = str(ev.get("message", "warning"))
elif etype == "progress":
# Null any non-finite float so the JSON stays strict-parseable; a missing key
# keeps the last value, a present-but-non-finite one becomes None.
# Null any non-finite float so the JSON 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 = (
@ -501,6 +506,7 @@ class DiffusionTrainingService:
status = "stopped" if ev.get("stopped") else "completed",
output_dir = ev.get("output_dir"),
lora_path = ev.get("lora_path"),
ema_path = ev.get("ema_path"),
message = (
"Stopped (partial adapter saved)."
if ev.get("lora_path")

45
studio/backend/models/training.py
View file

@ -747,23 +747,22 @@ class DiffusionTrainingStartRequest(BaseModel):
"ceil(N / (batch x grad_accum)) optimizer steps over the N-image dataset"
),
)
# Upper bound as well as positive, matching the LLM schema's _parse_lr: JSON accepts
# 1e309, which floats to inf and satisfies a gt-only constraint, and the route then
# evicts the resident models and starts AdamW with an infinite rate -- the first step
# destroys the adapter while the run reports normal progress and saves the result.
# (NaN already fails gt.) Values at or above 1.0 always diverge anyway.
# Upper bound as well as positive, matching the LLM schema's _parse_lr: JSON accepts 1e309, which
# floats to inf and satisfies a gt-only constraint, and the route would then evict the resident
# models and start AdamW with an infinite rate, destroying the adapter on the first step while
# reporting normal progress. (NaN already fails gt.) Values at or above 1.0 always diverge.
learning_rate: float = Field(1e-4, gt = 0, lt = 1.0)
train_batch_size: int = Field(1, ge = 1, le = 64)
gradient_accumulation_steps: int = Field(1, ge = 1, le = 256)
lora_rank: int = Field(16, ge = 1, le = 320)
lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank")
# Strictly below 1: PEFT turns lora_dropout into nn.Dropout(p=...), so 1.0 zeroes every input
# to the LoRA branch -- lora_A/lora_B receive no gradient and the run saves an untrained
# adapter while reporting normal progress. Matches TrainingStartRequest's [0, 1) validator.
# Strictly below 1: PEFT turns lora_dropout into nn.Dropout(p=...), so 1.0 zeroes every input to
# the LoRA branch and the run saves an untrained adapter while reporting normal progress. Matches
# TrainingStartRequest's [0, 1) validator.
lora_dropout: float = Field(0.0, ge = 0.0, lt = 1.0)
# Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that sets
# them isn't silently trained with defaults. Default targets to the SDXL attention
# projections (the trainer's DEFAULT_LORA_TARGETS) so it is never None.
# Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that sets them
# isn't silently trained with defaults. Default targets to the SDXL attention projections (the
# trainer's DEFAULT_LORA_TARGETS) so it is never None.
lora_target_modules: List[str] = Field(
default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"],
description = "U-Net modules to attach LoRA to",
@ -811,8 +810,8 @@ class DiffusionTrainingStartRequest(BaseModel):
"or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base."
),
)
# DiT-only levers the trainer implements. Undeclared, they were silently dropped by
# model_dump(); the defaults match DiffusionLoraConfig so an existing caller is unaffected.
# DiT-only levers the trainer implements. Undeclared, they were silently dropped by model_dump();
# the defaults match DiffusionLoraConfig so an existing caller is unaffected.
ema_decay: float = Field(
0.0, ge = 0.0, lt = 1.0, description = "EMA of the LoRA weights; 0 disables it"
)
@ -820,7 +819,11 @@ class DiffusionTrainingStartRequest(BaseModel):
0.0, ge = 0.0, le = 1.0, description = "Chance of dropping the caption to an empty prompt"
)
weighting_scheme: Literal["none", "bell"] = Field(
"none", description = "Flow-matching timestep sampling: uniform, or logit-normal (bell)"
"none",
description = (
"Per-sample loss weighting over the drawn timestep: none (unweighted MSE) or bell "
"(Gaussian bell centered mid-schedule). Timesteps are always logit-normal sampled."
),
)
flow_shift: Optional[float | Literal["auto"]] = Field(
None,
@ -869,15 +872,16 @@ class DiffusionTrainingStatusResponse(BaseModel):
loss: Optional[float] = None
avg_loss: Optional[float] = None
learning_rate: Optional[float] = None
# Total pre-clip gradient norm from the last optimizer step (health signal the UI charts
# alongside the loss).
# Total pre-clip gradient norm from the last optimizer step (health signal the UI charts).
grad_norm: Optional[float] = None
num_images: Optional[int] = None
in_model_load: bool = False
output_dir: Optional[str] = None
lora_path: Optional[str] = None
# The second, EMA-averaged adapter written in the run's ema subdir when ema_decay was enabled.
ema_path: Optional[str] = None
# Where the adapter was mirrored into the Studio LoRA catalog, and what family / base it was
# trained from -- lets the UI deploy it onto the right base.
# trained from, so the UI can deploy it onto the right base.
catalog_path: Optional[str] = None
family: Optional[str] = None
base_model: Optional[str] = None
@ -919,6 +923,7 @@ class DiffusionTrainingRunDetail(DiffusionTrainingRunSummary):
peak_memory_gb: Optional[float] = None
num_images: Optional[int] = None
lora_path: Optional[str] = None
ema_path: Optional[str] = None
config: Optional[dict] = None
metric_history: Optional[DiffusionMetricHistory] = None
@ -945,9 +950,9 @@ 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 = no precision selector,
# e.g. SDXL), plus the recommended pick and whether regional torch.compile applies. Defaults
# keep older backends' payloads valid.
# base_precision modes this machine supports for the family (empty = 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

17
studio/backend/tests/test_diffusion_train_extras.py
View file

@ -61,7 +61,7 @@ def test_ema_fixed_decay_math():
def test_ema_warmup_ramp_is_responsive_early_and_capped_late():
m = _TinyLoRAish()
ema = LoRAEMA(m, decay = 0.99, warmup = True)
# First update: decay = min(0.99, 1/11) -- the shadow mostly adopts the new value.
# First update: decay = min(0.99, 1/11), so the shadow mostly adopts the new value.
assert ema.effective_decay() == pytest.approx(1 / 11)
with torch.no_grad():
m.lora_A.fill_(2.0)
@ -269,8 +269,8 @@ def test_config_blank_cond_cache_dir_means_off():
def test_source_revision_marks_a_dir_update_and_never_raises(tmp_path):
# The trainer namespaces its conditioning cache on this, so an in-place checkpoint
# update must change the marker or a warm run trains on the old encoder's embeddings.
# The trainer namespaces its conditioning cache on this, so an in-place checkpoint update must
# change the marker or a warm run trains on the old encoder's embeddings.
from core.training.diffusion_train_extras import source_revision
d = tmp_path / "ckpt"
@ -280,6 +280,15 @@ def test_source_revision_marks_a_dir_update_and_never_raises(tmp_path):
first = source_revision(str(d))
assert first == source_revision(str(d)) # stable while untouched
w.write_bytes(b"v2-longer")
assert source_revision(str(d)) != first
second = source_revision(str(d))
assert second != first
# The VAE produces the cached latents, so an in-place VAE swap must invalidate them too.
(d / "vae").mkdir()
v = d / "vae" / "diffusion_pytorch_model.safetensors"
v.write_bytes(b"vae-v1")
third = source_revision(str(d))
assert third != second
v.write_bytes(b"vae-v2-longer")
assert source_revision(str(d)) != third
for ref in (None, "", "no/such/repo-xyz", "/does/not/exist", 7):
assert isinstance(source_revision(ref), str)

301
studio/backend/tests/test_diffusion_training.py
View file

@ -121,7 +121,7 @@ def _stoppable_target(*, event_queue, stop_queue, config):
def _crashing_target(*, event_queue, stop_queue, config):
event_queue.put({"type": "model_load_started"})
# Exits without a terminal event -> the pump must mark it as an error.
# Exits without a terminal event, so the pump must mark it as an error.
_CFG = {"base_model": "b", "data_dir": "d", "output_dir": "/tmp/out", "train_steps": 2}
@ -203,9 +203,9 @@ def test_apply_event_transitions():
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.
# 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
@ -246,8 +246,8 @@ def test_progress_nulls_non_finite_floats_for_strict_json():
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
# client shows a stale loading indicator after the job ended.
# model_load_completed, so the terminal update must reset in_model_load or the client shows a
# stale loading indicator after the job ended.
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
svc._apply_event({"type": "model_load_started"})
assert svc.status()["in_model_load"] is True
@ -260,6 +260,22 @@ def test_terminal_events_clear_model_load_flag():
assert svc2.status()["in_model_load"] is False and svc2.status()["status"] == "error"
def test_complete_event_keeps_the_ema_adapter_path():
# A DiT run with ema_decay writes a SECOND adapter and reports it as ema_path. Dropping the
# field on the way into the snapshot leaves that adapter undiscoverable by any client.
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
svc._apply_event(
{"type": "complete", "output_dir": "/o", "lora_path": "/o/a.safetensors",
"ema_path": "/o/ema/a.safetensors"}
)
snap = svc.status()
assert snap["lora_path"] == "/o/a.safetensors"
assert snap["ema_path"] == "/o/ema/a.safetensors"
# A run without EMA leaves it null rather than carrying the previous run's path.
svc._apply_event({"type": "complete", "output_dir": "/o2", "lora_path": "/o2/a.safetensors"})
assert svc.status()["ema_path"] is None
# ── route wiring (mocked service) ─────────────────────────────────────────────
class _FakeService:
def __init__(self):
@ -330,15 +346,14 @@ def client(monkeypatch):
monkeypatch.setattr(
"core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake
)
# Neutralize the LLM interlock + GPU-free for the wiring tests (their own tests below
# exercise those behaviors). The route imports get_training_backend at module scope.
# Neutralize the LLM interlock + GPU-free for the wiring tests (their own tests exercise those).
# The route imports get_training_backend at module scope.
import routes.training as tr
monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False))
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
# The dataset preflight runs the trainer's discovery against _BODY's fake
# data_dir; stub it here so wiring tests pass, and let the dedicated preflight
# tests below re-point it at a real tmp dataset.
# The dataset preflight runs the trainer's discovery against _BODY's fake data_dir; stub it here
# so wiring tests pass, and let the dedicated preflight tests re-point it at a real tmp dataset.
monkeypatch.setattr(
"core.training.diffusion_train_common.discover_image_caption_pairs",
lambda data_dir, **kw: [("img.png", "caption")],
@ -346,8 +361,8 @@ def client(monkeypatch):
app = FastAPI()
app.include_router(training_router, prefix = "/api/train")
app.dependency_overrides[get_current_subject] = lambda: "test-user"
# Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there,
# so the wiring tests below behave as before. The guard test flips this override.
# Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there, so the
# wiring tests behave as before. The guard test flips this override.
app.dependency_overrides[authenticated_via_api_key] = lambda: False
c = TestClient(app)
c._fake = fake # type: ignore[attr-defined]
@ -377,11 +392,10 @@ def test_route_start_ok(client):
def test_route_start_frees_gpu_off_the_coroutine_thread(client, monkeypatch):
# The GPU cleanup can block for seconds (engine unload waits on generation locks; the
# export subprocess join can take seconds), so the async start route must offload it via
# asyncio.to_thread rather than run it inline and freeze the event loop for concurrent
# status/progress/cancel requests. Assert the cleanup runs on a DIFFERENT thread than the
# inline coroutine body (service.start), which an inline (un-offloaded) call could not.
# The GPU cleanup can block for seconds (engine unload waits on generation locks; the export
# subprocess join can take seconds), so the async start route must offload it via asyncio.to_thread
# rather than freeze the event loop. Assert the cleanup runs on a DIFFERENT thread than the inline
# coroutine body, which an inline call could not.
import threading
import routes.training as tr
@ -409,8 +423,7 @@ def test_route_start_frees_gpu_off_the_coroutine_thread(client, monkeypatch):
def test_route_start_reserves_before_freeing_gpu(client, monkeypatch):
# The training slot must be reserved (is_active -> true) BEFORE the route frees resident GPU
# models, so a concurrent /images/load or /video/load guard refuses during the free-then-spawn
# window instead of double-allocating the GPU. Assert the ordering: reserve is logged before
# the GPU free runs, and the service reports active while the free is in flight.
# window. Assert reserve is logged before the GPU free, and the service reports active meanwhile.
import routes.training as tr
order: list = []
@ -432,11 +445,9 @@ def test_route_start_reserves_before_freeing_gpu(client, monkeypatch):
def test_route_start_reserves_before_scanning_dataset(client, monkeypatch):
# The dataset preflight scan decode-probes every image and can take noticeable time on a large
# folder. It must run AFTER the slot is reserved (is_active -> true), so a concurrent
# upload/caption/delete guarded by _require_diffusion_dataset_mutable() can't mutate the dataset
# the trainer is about to read during that window. Assert reserve precedes the scan and the
# service is active while the scan is in flight.
# The dataset preflight scan decode-probes every image and can take a while, so it must run AFTER
# the slot is reserved, or a concurrent upload/caption/delete could mutate the dataset the trainer
# is about to read. Assert reserve precedes the scan and the service is active during it.
order: list = []
def _record_scan(data_dir, **kw):
@ -456,8 +467,8 @@ def test_route_start_reserves_before_scanning_dataset(client, monkeypatch):
def test_route_start_unreserves_when_dataset_preflight_fails(client, monkeypatch):
# A dataset preflight failure AFTER the reservation must roll it back (unreserve), or a rejected
# start would leave training permanently "active" and keep blocking loads and dataset edits.
# A dataset preflight failure AFTER the reservation must roll it back, or a rejected start would
# leave training permanently "active" and keep blocking loads and dataset edits.
def _bad_scan(data_dir, **kw):
raise ValueError("no captioned images found")
@ -475,9 +486,8 @@ def test_route_start_unreserves_when_dataset_preflight_fails(client, monkeypatch
def test_service_reserve_marks_active_and_rolls_back():
# The real service: reserve() flips is_active true before any proc exists (so a load guard
# refuses during the free window), and unreserve() clears it without a live proc, so a failed
# start is not left permanently "active".
# The real service: reserve() flips is_active true before any proc exists (so a load guard refuses
# during the free window), and unreserve() clears it without a live proc.
from core.training.diffusion_training_service import DiffusionTrainingService
svc = DiffusionTrainingService()
@ -490,11 +500,9 @@ def test_service_reserve_marks_active_and_rolls_back():
def test_service_reserve_is_compare_and_set():
# reserve() is the concurrency gate: two /diffusion/start requests can interleave between the
# is_active() check and the reservation, so reserve() itself must reject a second reservation
# atomically. Without the compare-and-set, both callers would reserve, both would free the
# GPU's resident chat/image model, and the loser would only 409 AFTER the eviction -- exactly
# the evict-then-fail the reservation exists to prevent. A second reserve must raise; the first
# stays reserved; after unreserve the slot is claimable again.
# is_active() check and the reservation, so reserve() itself must reject a second one atomically.
# Without the compare-and-set both would free the GPU's resident model and the loser would 409
# only AFTER the eviction. A second reserve must raise; after unreserve the slot is claimable.
from core.training.diffusion_training_service import DiffusionTrainingService
svc = DiffusionTrainingService()
@ -508,11 +516,9 @@ def test_service_reserve_is_compare_and_set():
def test_route_start_preflights_gated_base_off_the_coroutine_thread(client, monkeypatch):
# _preflight_gated_base does a blocking urlopen HEAD (up to a 5s timeout) to Hugging Face, so
# the async start route must offload it via asyncio.to_thread rather than run it inline and
# freeze the event loop for concurrent status/progress/cancel requests. Assert it runs on a
# DIFFERENT thread than the inline coroutine body (service.start), which an inline call could
# not.
# _preflight_gated_base does a blocking urlopen HEAD (up to 5s) to Hugging Face, so the async
# start route must offload it via asyncio.to_thread rather than freeze the event loop. Assert it
# runs on a DIFFERENT thread than the inline coroutine body.
import threading
import routes.training as tr
@ -547,8 +553,8 @@ def test_route_start_forwards_extra_training_knobs(client):
def test_route_start_forwards_num_epochs(client):
# Epochs mode: the frontend omits train_steps and sends num_epochs; it must reach the
# service so the trainer can resolve it against the dataset size.
# Epochs mode: the frontend omits train_steps and sends num_epochs; it must reach the service so
# the trainer can resolve it against the dataset size.
body = {k: v for k, v in _BODY.items() if k != "train_steps"}
r = client.post("/api/train/diffusion/start", json = {**body, "num_epochs": 8})
assert r.status_code == 200, r.text
@ -556,8 +562,8 @@ def test_route_start_forwards_num_epochs(client):
def test_route_start_forwards_dit_loss_knobs(client):
# The trainer implements these, but the request schema did not declare them, so
# model_dump() dropped them and the run silently used the defaults.
# The trainer implements these, but the request schema did not declare them, so model_dump()
# dropped them and the run silently used the defaults.
body = {
**_BODY,
"ema_decay": 0.99,
@ -590,10 +596,9 @@ def test_request_model_dit_loss_knob_bounds():
def test_request_model_rejects_lora_dropout_of_one():
# lora_dropout = 1.0 makes PEFT build nn.Dropout(p=1.0), which zeroes every input to the LoRA
# branch: the adapter output is identically the frozen base and both lora_A/lora_B get exactly
# zero gradient, so the run reports normal progress and saves an untrained adapter. The generic
# training schema already requires < 1 (TrainingStartRequest._check_lora_dropout); the diffusion
# schema must match instead of accepting the degenerate boundary.
# branch: the adapter output is identically the frozen base and both lora_A/lora_B get zero
# gradient, so the run saves an untrained adapter while reporting normal progress. The generic
# training schema already requires below 1; the diffusion schema must match.
from pydantic import ValidationError
from models.training import DiffusionTrainingStartRequest
@ -621,8 +626,8 @@ def test_request_model_num_epochs_bounds():
def test_request_model_base_precision_accepts_mxfp8():
# The base_precision Literal now includes mxfp8 (the DiT dense speed mode); a bogus
# mode is still rejected.
# The base_precision Literal now includes mxfp8 (the DiT dense speed mode); a bogus mode is still
# rejected.
from pydantic import ValidationError
from models.training import DiffusionTrainingStartRequest
@ -635,10 +640,9 @@ def test_request_model_base_precision_accepts_mxfp8():
def test_config_from_dict_epoch_mode_drops_max_steps_sentinel():
# The generic Studio epoch-mode payload sends max_steps: 0 as the "use epochs" sentinel.
# The max_steps -> train_steps alias would copy that 0 and normalized() would reject
# train_steps < 1 before epochs are resolved; _config_from_dict must drop the falsy
# value so the default train_steps stands in until resolve_train_steps applies num_epochs.
# The generic Studio epoch-mode payload sends max_steps: 0 as the "use epochs" sentinel. The
# max_steps alias would copy that 0 and normalized() would reject train_steps below 1 before
# epochs are resolved, so _config_from_dict must drop the falsy value.
from core.training.diffusion_train_common import DiffusionLoraConfig, _config_from_dict
cfg = _config_from_dict(
@ -657,9 +661,8 @@ def test_config_from_dict_epoch_mode_drops_max_steps_sentinel():
norm = cfg.normalized()
assert norm.num_epochs == 2
# An explicit non-zero max_steps in epochs mode is still honored (only the 0 sentinel is
# dropped), and a plain steps payload (no num_epochs) keeps max_steps: 0 -> train_steps 0
# so normalized() surfaces the invalid value as before.
# An explicit non-zero max_steps in epochs mode is still honored (only the 0 sentinel is dropped),
# and a plain steps payload keeps max_steps: 0 -> train_steps 0 so normalized() still surfaces it.
cfg_explicit = _config_from_dict(
{
"base_model": "stabilityai/stable-diffusion-xl-base-1.0",
@ -673,9 +676,9 @@ def test_config_from_dict_epoch_mode_drops_max_steps_sentinel():
def test_permutation_sampler_covers_dataset_once_per_cycle():
# Every index must appear exactly once per cycle before any repeat (epoch-style pass),
# so a short run over a small dataset never leaves images unseen the way the old
# with-replacement draw did. Consecutive cycles must be reshuffled (differ).
# Every index must appear exactly once per cycle before any repeat (an epoch-style pass), so a
# short run over a small dataset never leaves images unseen the way the old with-replacement draw
# did. Consecutive cycles must be reshuffled.
import random
from core.training.diffusion_train_common import PermutationBatchSampler
@ -683,8 +686,8 @@ def test_permutation_sampler_covers_dataset_once_per_cycle():
n = 100
sampler = PermutationBatchSampler(n, random.Random(0))
# Draw exactly one cycle in batches of 3 (n not divisible by the batch, so a batch spans
# the cycle boundary); the first n indices must be a permutation of range(n).
# Draw exactly one cycle in batches of 3 (n not divisible by the batch, so a batch spans the cycle
# boundary); the first n indices must be a permutation of range(n).
drawn: list[int] = []
while len(drawn) < n:
drawn.extend(sampler.next_batch(3))
@ -703,8 +706,8 @@ def test_permutation_sampler_covers_dataset_once_per_cycle():
replay = PermutationBatchSampler(n, random.Random(0))
assert replay.next_batch(n) == cycle_a
# A batch larger than the dataset refills across cycles so it never shrinks (batch shape
# preserved), even though it must then repeat indices within the batch.
# A batch larger than the dataset refills across cycles so it never shrinks, even though it must
# then repeat indices within the batch.
big = PermutationBatchSampler(4, random.Random(1))
batch = big.next_batch(10)
assert len(batch) == 10
@ -712,10 +715,9 @@ def test_permutation_sampler_covers_dataset_once_per_cycle():
def test_permutation_sampler_honors_batch_on_tiny_dataset():
# Regression for the SDXL LoRA trainer clamp: a dataset smaller than train_batch_size must
# still yield exactly train_batch_size indices (refilled across cycles), so a tiny dataset
# trains at the configured/reported effective batch instead of silently shrinking it -- the
# contract the SDXL _next_batch now relies on (matching the DiT trainer, which never clamps).
# Regression for the SDXL LoRA trainer clamp: a dataset smaller than train_batch_size must still
# yield exactly train_batch_size indices, so a tiny dataset trains at the configured batch instead
# of silently shrinking it -- the contract SDXL's _next_batch now relies on.
import random
from core.training.diffusion_train_common import PermutationBatchSampler
@ -730,15 +732,15 @@ def test_permutation_sampler_honors_batch_on_tiny_dataset():
def test_route_start_accepts_zero_max_grad_norm(client):
# 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_);
# the request model must not reject it.
# 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_), so the request
# model must not reject it.
r = client.post("/api/train/diffusion/start", json = {**_BODY, "max_grad_norm": 0.0})
assert r.status_code == 200, r.text
assert client._fake.started_with["max_grad_norm"] == 0.0
def test_route_start_rejects_nonpositive_snr_gamma(client):
# gamma <= 0 zeroes/inverts the min-SNR loss weight; null is the disable value.
# A gamma at or below 0 zeroes/inverts the min-SNR loss weight; null is the disable value.
r = client.post("/api/train/diffusion/start", json = {**_BODY, "snr_gamma": 0})
assert r.status_code == 422
@ -750,12 +752,10 @@ def test_route_start_rejects_uncontained_paths(client):
def test_route_start_resolves_bare_name_under_image_dataset_root(client, monkeypatch, tmp_path):
# The upload/labeling routes manage image datasets directly under datasets_root() and
# the UI passes the bare folder name back as data_dir. The generic resolve_dataset_path
# searches the LLM uploads and recipe roots FIRST, so an unrelated upload file or recipe
# folder sharing the name would shadow the just-uploaded image dataset (preflight 400
# "not a directory", or training the wrong data). The route must prefer the image
# dataset root for a bare name that exists there.
# The upload/labeling routes manage image datasets directly under datasets_root() and the UI
# passes the bare folder name back as data_dir. The generic resolve_dataset_path searches the LLM
# uploads and recipe roots FIRST, so an unrelated file sharing the name would shadow the
# just-uploaded image dataset. The route must prefer the image dataset root for a bare name.
import utils.paths as up
ds_root = tmp_path / "assets" / "datasets"
@ -809,8 +809,8 @@ def test_route_start_conflict_maps_to_409(client):
def test_route_start_over_api_with_inference_in_flight_is_409(client, monkeypatch):
# An API-key client must not start diffusion training (which frees VRAM by unloading
# chat) while an inference request is streaming; it should 409 instead of killing it.
# An API-key client must not start diffusion training (which frees VRAM by unloading chat) while
# an inference request is streaming; it should 409 instead of killing it.
client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
monkeypatch.setattr(
"core.inference.llama_keepwarm.other_inference_request_count",
@ -851,8 +851,8 @@ def test_route_status_and_stop(client):
def test_service_restart_after_completion():
# A finished job's pump is joined OUTSIDE the lock (it needs the lock for its
# final state writes), so a second start neither stalls nor deadlocks.
# A finished job's pump is joined OUTSIDE the lock (it needs the lock for its final state writes),
# so a second start neither stalls nor deadlocks.
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
svc.start(dict(_CFG))
_wait_status(svc, "completed")
@ -865,8 +865,8 @@ def test_service_restart_after_completion():
def test_stale_pump_events_cannot_corrupt_new_job():
# An event carrying a superseded job's proc identity must be dropped, so a
# straggler pump can never overwrite the state of a newly started job.
# An event carrying a superseded job's proc identity must be dropped, so a straggler pump can
# never overwrite the state of a newly started job.
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
svc.start(dict(_CFG))
_wait_status(svc, "completed")
@ -914,9 +914,9 @@ def test_diffusion_info_lists_image_dataset_folders(client, dataset_roots):
def test_diffusion_info_counts_metadata_captions(client, dataset_roots):
# A dataset captioned via metadata.jsonl must not report caption_count=0 (which the
# Train UI treats as uncaptioned); metadata rows count like sidecars, without
# double-counting an image that has both.
# A dataset captioned via metadata.jsonl must not report caption_count=0 (which the Train UI
# treats as uncaptioned); metadata rows count like sidecars, without double-counting an image that
# has both.
ds_root, _ = dataset_roots
folder = ds_root / "meta-captioned"
folder.mkdir()
@ -969,10 +969,10 @@ def test_diffusion_dataset_upload_accumulates(client, dataset_roots):
def test_diffusion_dataset_upload_normalizes_windows_and_rejects_dotdot(client, dataset_roots):
ds_root, _ = dataset_roots
# A Windows client can send a backslash path in the multipart filename; POSIX Path.name does
# not split on backslash, so it must be folded to the true basename, or the stored name holds
# backslashes that _safe_dataset_image_path later rejects -- an image the labeling grid can
# list but never preview/caption/delete (an orphan).
# A Windows client can send a backslash path in the multipart filename, and POSIX Path.name does
# not split on backslash, so it must be folded to the true basename -- else the stored name holds
# backslashes that _safe_dataset_image_path later rejects, leaving an orphan the grid can list but
# never preview/caption/delete.
r = client.post(
"/api/train/diffusion/dataset",
data = {"name": "winset"},
@ -985,8 +985,8 @@ def test_diffusion_dataset_upload_normalizes_windows_and_rejects_dotdot(client,
assert any(rec["filename"] == "cat.png" for rec in recs)
assert client.get("/api/train/diffusion/dataset/winset/image/cat.png").status_code == 200
# A basename that still contains ".." (which _safe_dataset_image_path rejects) is refused at
# upload rather than persisted as an unmanageable entry.
# A basename that still contains ".." is refused at upload rather than persisted as an
# unmanageable entry.
r = client.post(
"/api/train/diffusion/dataset",
data = {"name": "winset"},
@ -996,11 +996,9 @@ def test_diffusion_dataset_upload_normalizes_windows_and_rejects_dotdot(client,
def test_diffusion_dataset_upload_rejects_case_insensitive_stem_clash(client, dataset_roots):
# Two images whose stems differ only by case (sample.png vs Sample.jpg) map to the SAME
# <stem>.txt caption sidecar on case-insensitive filesystems (Windows / default macOS), so
# keeping both would silently share -- and corrupt -- one caption during training. The clash
# check compares stems with casefold, so it must reject the pair (the comparison is pure
# string logic, so this fires regardless of the test host filesystem).
# Two images whose stems differ only by case map to the SAME <stem>.txt caption sidecar on
# case-insensitive filesystems, so keeping both would silently corrupt one caption. The clash
# check compares stems with casefold, so it fires regardless of the test host filesystem.
r = client.post(
"/api/train/diffusion/dataset",
data = {"name": "caseset"},
@ -1041,9 +1039,9 @@ def test_diffusion_dataset_upload_rejects_case_insensitive_stem_clash(client, da
def test_diffusion_dataset_upload_over_cap_keeps_existing_example(
client, dataset_roots, monkeypatch
):
# A re-upload that trips the size cap mid-write must not destroy an example already
# stored under the same name: the write goes to a sibling temp file and only atomically
# replaces the original on success, so the 413 leaves the prior good bytes intact.
# A re-upload that trips the size cap mid-write must not destroy an example already stored under
# the same name: the write goes to a sibling temp file and only atomically replaces the original
# on success, so the 413 leaves the prior good bytes intact.
import utils.upload_limits as ul
ds_root, _ = dataset_roots
@ -1066,9 +1064,8 @@ def test_diffusion_dataset_upload_over_cap_keeps_existing_example(
def test_diffusion_info_empty_sidecar_shadows_metadata_caption(client, dataset_roots):
# An empty (tombstone) .txt sidecar shadows a metadata row -- the trainer strips it and
# skips the image -- so the summary must not count it as captioned (which would report a
# dataset as captioned that the trainer would reject as having no captioned images).
# An empty (tombstone) .txt sidecar shadows a metadata row -- the trainer strips it and skips the
# image -- so the summary must not count it as captioned.
ds_root, _ = dataset_roots
folder = ds_root / "tombstoned"
folder.mkdir()
@ -1107,7 +1104,7 @@ def _png_bytes(width: int, height: int) -> bytes:
def test_diffusion_dataset_upload_rejects_oversized_image(client, dataset_roots):
# A decompression bomb: a small compressible PNG with huge dimensions passes the byte limit but
# would OOM the trainer on decode. It must 400 at upload (dimension check reads only the header).
# would OOM the trainer on decode, so it must 400 at upload (the check reads only the header).
pytest.importorskip("PIL")
big = _png_bytes(5000, 64) # > 4096 per side
r = client.post(
@ -1128,11 +1125,10 @@ def test_diffusion_dataset_upload_rejects_oversized_image(client, dataset_roots)
def test_diffusion_dataset_upload_maps_pillow_bomb_to_400(client, dataset_roots, monkeypatch):
# Past Pillow's hard bomb threshold (> 2 x Image.MAX_IMAGE_PIXELS) Image.open() itself raises
# DecompressionBombError, which derives straight from Exception -- not OSError/ValueError -- so
# the dimension guard's except clause missed it and the upload 500'd instead of returning the
# intended oversized-image 400. Shrink the limit (as Pillow's own tests do) so a small file
# crosses it, instead of building a 179 MP image in the test.
# Past Pillow's hard bomb threshold Image.open() itself raises DecompressionBombError, which
# derives straight from Exception, so the dimension guard's except clause missed it and the upload
# 500'd instead of the intended 400. Shrink the limit (as Pillow's own tests do) so a small file
# crosses it.
pytest.importorskip("PIL")
from PIL import Image
@ -1149,8 +1145,8 @@ def test_diffusion_dataset_upload_maps_pillow_bomb_to_400(client, dataset_roots,
def test_diffusion_info_tolerates_non_object_jsonl(client, dataset_roots):
# A metadata.jsonl line that is valid JSON but not an object ([]/null/string/number) or malformed
# must be skipped per-line, not 500 the info endpoint; a valid row in the same file still counts.
# A metadata.jsonl line that is valid JSON but not an object, or malformed, must be skipped
# per-line, not 500 the info endpoint; a valid row in the same file still counts.
ds_root, _ = dataset_roots
folder = ds_root / "weird-meta"
folder.mkdir()
@ -1169,8 +1165,8 @@ def test_diffusion_info_tolerates_non_object_jsonl(client, dataset_roots):
def test_diffusion_info_skips_null_metadata_captions(client, dataset_roots):
# A JSON null caption is "no caption": str(None) would store the literal "None" and
# both count as captioned and train on that text.
# A JSON null caption is "no caption": str(None) would store the literal "None" and both count as
# captioned and train on that text.
ds_root, _ = dataset_roots
folder = ds_root / "null-caption"
folder.mkdir()
@ -1196,7 +1192,7 @@ def test_diffusion_info_skips_null_metadata_captions(client, dataset_roots):
def test_diffusion_info_tolerates_invalid_utf8_jsonl(client, dataset_roots):
# Invalid UTF-8 in a metadata file must not 500 the info endpoint; the file is skipped, not decoded.
# Invalid UTF-8 in a metadata file must not 500 the info endpoint; the file is skipped.
ds_root, _ = dataset_roots
folder = ds_root / "bad-utf8-meta"
folder.mkdir()
@ -1209,8 +1205,8 @@ def test_diffusion_info_tolerates_invalid_utf8_jsonl(client, dataset_roots):
def test_diffusion_info_tolerates_invalid_utf8_sidecar(client, dataset_roots):
# Same for a per-image .txt sidecar: read_text raises UnicodeDecodeError, which is not an
# OSError, so an unguarded read 500s the info endpoint after the upload already committed.
# Same for a per-image .txt sidecar: read_text raises UnicodeDecodeError, which is not an OSError,
# so an unguarded read 500s the info endpoint after the upload already committed.
ds_root, _ = dataset_roots
folder = ds_root / "bad-utf8-sidecar"
folder.mkdir()
@ -1283,9 +1279,8 @@ def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_root
def test_free_gpu_for_diffusion_training_unloads_video(monkeypatch):
# A resident Video pipeline loads under the VIDEO arbiter owner, which the Images teardown
# does not free; starting diffusion training must unload it too or the trainer OOMs against
# the still-resident video model.
# A resident Video pipeline loads under the VIDEO arbiter owner, which the Images teardown does
# not free, so starting diffusion training must unload it too or the trainer OOMs.
import routes.training as tr
from core.inference import gpu_arbiter
@ -1325,10 +1320,9 @@ def test_free_gpu_for_diffusion_training_unloads_video(monkeypatch):
def test_keepwarm_tracks_image_video_generation_paths():
# The API-key training-start guard uses other_inference_request_count(), which only sees
# paths the keepwarm middleware tracks. Image/video generation must be tracked so a training
# start is refused (409) while one is in-flight rather than its unload cancelling it; the GET
# *-progress and */cancel variants must stay untracked.
# The API-key training-start guard uses other_inference_request_count(), which only sees paths the
# keepwarm middleware tracks. Image/video generation must be tracked so a training start is
# refused (409) while one is in flight; the GET *-progress and */cancel variants stay untracked.
from core.inference.llama_keepwarm import _is_inference_path
assert _is_inference_path("/api/inference/images/generate")
@ -1343,8 +1337,8 @@ def test_import_example_partial_failure_leaves_no_partial_dataset(
client, dataset_roots, monkeypatch
):
# A materialize that writes some images then fails must not leave a partial dataset: it stages
# into a discarded temp dir, so the target folder stays empty and a retry re-materializes
# instead of the image_count>0 idempotency check treating a truncated result as complete.
# into a discarded temp dir, so the target folder stays empty and a retry re-materializes instead
# of the idempotency check treating a truncated result as complete.
import routes.training as tr
ds_root, _ = dataset_roots
@ -1365,8 +1359,8 @@ def test_import_example_partial_failure_leaves_no_partial_dataset(
def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(
client, dataset_roots, monkeypatch
):
# All-or-nothing: a valid image ahead of the one that trips the size cap must NOT be
# left on disk (the 413 mid-batch rolls back every file written this request).
# All-or-nothing: a valid image ahead of the one that trips the size cap must NOT be left on disk
# (the 413 mid-batch rolls back every file written this request).
import utils.upload_limits as ul
monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100)
@ -1385,9 +1379,8 @@ def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(
def test_diffusion_dataset_upload_over_cap_preserves_existing_file(
client, dataset_roots, monkeypatch
):
# A failed batch that reuses an existing filename must NOT delete the user's
# pre-existing file (repeat uploads accumulate). Staging to a temp file keeps the
# original intact until the whole batch commits.
# A failed batch that reuses an existing filename must NOT delete the user's pre-existing file:
# staging to a temp file keeps the original intact until the whole batch commits.
import utils.upload_limits as ul
monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100)
@ -1407,8 +1400,8 @@ def test_diffusion_dataset_upload_over_cap_preserves_existing_file(
def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch):
# A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed,
# so a bad pick never unloads the user's working chat/Images model.
# A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed, so a bad pick
# never unloads the user's working chat/Images model.
import routes.training as tr
freed = []
@ -1423,9 +1416,8 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat
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
# A DiT precision the host cannot run must 400 BEFORE resident GPU workloads are freed, else the
# host tears down the user's model and the run 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
@ -1448,8 +1440,8 @@ def test_route_start_refuses_non_bf16_gpu_without_freeing_gpu(client, monkeypatc
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.
# 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
@ -1559,8 +1551,8 @@ def test_info_lists_trainable_families(client):
def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatch):
# A gated FLUX base with no valid token must 400 from the HEAD preflight BEFORE the GPU
# residents are freed, so a doomed start never evicts the user's loaded model.
# A gated FLUX base with no valid token must 400 from the HEAD preflight BEFORE the GPU residents
# are freed, so a doomed start never evicts the user's loaded model.
import urllib.error
import urllib.request
@ -1584,9 +1576,9 @@ def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatc
def test_route_start_refuses_a_dit_family_without_a_gpu_before_freeing(client, monkeypatch):
# nf4 is not a CPU fallback: the 4-bit base load needs CUDA/XPU/MPS. Without a pre-teardown
# gate the default Train pick on a GPU-less host unloaded the working Images pipeline, pulled
# the text encoders, and only then failed in the child.
# nf4 is not a CPU fallback: the 4-bit base load needs CUDA/XPU/MPS. Without a pre-teardown gate
# the default Train pick on a GPU-less host unloaded the working Images pipeline, pulled the text
# encoders, and only then failed in the child.
import torch
import routes.training as tr
@ -1606,7 +1598,7 @@ def test_route_start_refuses_a_dit_family_without_a_gpu_before_freeing(client, m
assert freed == []
assert client._fake.started_with is None
# SDXL trains fp32 on CPU (that is its documented fallback), so it is not gated.
# SDXL trains fp32 on CPU (its documented fallback), so it is not gated.
r2 = client.post("/api/train/diffusion/start", json = _BODY)
assert r2.status_code == 200, r2.text
@ -1618,8 +1610,8 @@ def test_start_ungated_base_preflight_is_noop(client, monkeypatch):
import routes.training as tr
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
# This asserts the gated-repo probe is a no-op, not device support: pin the precision
# preflight so a GPU-less test host does not 400 for an unrelated reason.
# This asserts the gated-repo probe is a no-op, not device support: pin the precision preflight so
# a GPU-less test host does not 400 for an unrelated reason.
monkeypatch.setattr(
"core.training.diffusion_train_common.training_precision_preflight_error",
lambda fam, prec: None,
@ -1736,8 +1728,8 @@ def test_runs_endpoints_list_and_detail(client, _isolated_runs_dir):
def test_list_diffusion_runs_skips_wrong_shape_records(_isolated_runs_dir):
# A valid-JSON file with the wrong shape (non-dict, or missing the required string
# job_id / status) must be skipped by list_diffusion_runs so it never reaches the route's
# A valid-JSON file with the wrong shape (non-dict, or missing the required string job_id /
# status) must be skipped by list_diffusion_runs so it never reaches the route's
# DiffusionTrainingRunSummary(**r) and takes down the whole Previous runs panel.
import json
@ -1760,9 +1752,9 @@ def test_list_diffusion_runs_skips_wrong_shape_records(_isolated_runs_dir):
def test_runs_route_tolerates_bad_field_record(client, _isolated_runs_dir):
# A record that passes the service's shape check but has a wrong-typed field (a
# non-numeric avg_loss) would raise pydantic ValidationError in the route; the route must
# catch it per record so one bad file never breaks the panel and the good runs still list.
# A record that passes the service's shape check but has a wrong-typed field would raise a
# pydantic ValidationError in the route, so the route must catch it per record and still list the
# good runs.
import json
good = {"job_id": "a" * 32, "status": "completed", "adapter": "good", "saved": True}
@ -1782,9 +1774,8 @@ def test_runs_route_tolerates_bad_field_record(client, _isolated_runs_dir):
def test_run_detail_route_non_object_record_is_404(client, _isolated_runs_dir):
# A valid-JSON but non-object record (a truncated / hand-edited [] file named with a real
# job id) makes DiffusionTrainingRunDetail(**rec) raise TypeError, not ValidationError; the
# detail route must shape-check like the list path and 404 instead of 500.
# A valid-JSON but non-object record makes DiffusionTrainingRunDetail(**rec) raise TypeError, not
# ValidationError, so the detail route must shape-check like the list path and 404 instead of 500.
import json
job_id = "a" * 32
@ -1795,8 +1786,8 @@ def test_run_detail_route_non_object_record_is_404(client, _isolated_runs_dir):
def test_request_model_rejects_non_finite_learning_rate():
# 1e309 parses as inf, which satisfies a gt-only bound; the route would then start
# AdamW with an infinite rate and save a destroyed adapter while progress looked fine.
# 1e309 parses as inf, which satisfies a gt-only bound; the route would then start AdamW with an
# infinite rate and save a destroyed adapter while progress looked fine.
from pydantic import ValidationError
from models.training import DiffusionTrainingStartRequest as R

38
studio/frontend/src/features/images/api.ts
View file

@ -256,16 +256,43 @@ export async function getDiffusionDownloadPlan(
);
}
/** The generate POST's response was lost in transit (proxy/tunnel timeout, or a dropped
* connection) rather than refused by the backend, so the generation is still running. */
export class GenerateResponseLostError extends Error {
constructor(message: string) {
super(message);
this.name = "GenerateResponseLostError";
}
}
// Gateway statuses that mean the origin never got to answer: 524 is Cloudflare's ~100s
// origin-response cap, which a native-CPU or high-step run routinely exceeds in secure mode.
const RESPONSE_LOST_STATUSES = new Set([408, 502, 503, 504, 522, 524]);
/** Generate synchronously: the response carries the finished images. A run longer than the
* proxy's response window throws GenerateResponseLostError with the generation still in
* flight, so the caller settles it via getGenerateProgress + the gallery instead of retrying
* (which would duplicate the work). */
export async function generateDiffusionImage(
body: DiffusionGenerateRequest,
): Promise<DiffusionGenerateResponse> {
return parseJson(
await authFetch("/api/inference/images/generate", {
let response: Response;
try {
response = await authFetch("/api/inference/images/generate", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(body),
}),
);
});
} catch (err) {
// fetch itself rejected: the connection dropped, not a backend refusal.
throw new GenerateResponseLostError(
err instanceof Error ? err.message : "Lost connection during image generation",
);
}
if (!response.ok && RESPONSE_LOST_STATUSES.has(response.status)) {
throw new GenerateResponseLostError(await readFastApiError(response));
}
return parseJson(response);
}
export async function unloadDiffusionModel(): Promise<DiffusionStatus> {
@ -392,6 +419,8 @@ export interface DiffusionTrainingStatus {
in_model_load: boolean;
output_dir: string | null;
lora_path: string | null;
// The second, EMA-averaged adapter, present only when the run enabled ema_decay.
ema_path?: string | null;
started_at: number | null;
updated_at: number | null;
// Where the trained adapter was mirrored into the Studio LoRA catalog, and the family /
@ -455,6 +484,7 @@ export interface DiffusionTrainingRunDetail extends DiffusionTrainingRunSummary
peak_memory_gb?: number | null;
num_images?: number | null;
lora_path?: string | null;
ema_path?: string | null;
config?: Record<string, unknown> | null;
metric_history?: DiffusionMetricHistory | null;
}

134
studio/frontend/src/features/images/images-page.tsx
View file

@ -81,11 +81,13 @@ import {
type ControlNetSpecInput,
type DiffusionControlNetInfo,
type DiffusionGenerateProgress,
type DiffusionGenerateResponse,
type DiffusionLoadProgress,
type DiffusionLoraInfo,
type DiffusionStatus,
type GalleryImage,
type LoraSpecInput,
GenerateResponseLostError,
deleteGalleryImage,
fetchGalleryObjectUrl,
generateDiffusionImage,
@ -381,6 +383,32 @@ function genStepLabel(p: DiffusionGenerateProgress): string {
return eta ? `${base} · ~${eta}` : base;
}
// Settling a generation whose POST response was lost (a tunnel timeout, say): the backend
// keeps denoising, so poll until it goes idle rather than failing the run.
const SETTLE_POLL_MS = 1000;
const SETTLE_MAX_MS = 6 * 60 * 60 * 1000; // hard cap; a native-CPU batch can run for hours
const SETTLE_MAX_FAILS = 5; // consecutive progress failures before calling the backend gone
/** Wait for the generation that outlived its POST to finish. Resolves once progress reports
* idle (the images are on disk by then: the backend keeps it active through persistence);
* throws if the backend stays unreachable, so a real outage still surfaces. */
async function settleLostGeneration(isCurrent: () => boolean): Promise<void> {
const start = Date.now();
let fails = 0;
while (Date.now() - start < SETTLE_MAX_MS) {
await new Promise((r) => setTimeout(r, SETTLE_POLL_MS));
if (!isCurrent()) return;
try {
const p = await getGenerateProgress();
fails = 0;
if (!p.active) return;
} catch {
fails += 1;
if (fails >= SETTLE_MAX_FAILS) throw new Error("Lost connection to the image server.");
}
}
}
// The chat tab's model-load toast styling, reused verbatim so the diffusion
// load toast is visually identical (persistent, progress bar, same chrome).
const LOAD_TOAST_CLASSNAMES = {
@ -2233,51 +2261,65 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
// The page truly unmounted mid-run (app close / chat-only eject): stop
// issuing more GPU generations. A plain tab switch keeps it mounted.
if (!isMounted.current) break;
const res = await generateDiffusionImage({
prompt: prompt.trim(),
// Only send a negative prompt when guidance uses it, so the recipe
// doesn't record one the model ignored.
negative_prompt: guidance > 0 ? negativePrompt.trim() || undefined : undefined,
width: w,
height: h,
steps,
guidance,
// Offset runs by the batch size: the native engine seeds image j of a
// run at seed+j, so a +1 run offset would regenerate the previous run's
// batch-mates. Unique per image on both engines, reproducible via recipes.
seed: baseSeed + i * batchSize,
batch_size: batchSize,
// Transform/Inpaint/Extend send the source image (+ mask for inpaint/extend) and
// a denoise strength, resolved above. The backend derives output size from the
// image, so width/height are advisory here.
init_image: condInit,
mask_image: condMask,
strength: condStrength,
upscale: condUpscale,
reference_images: condRefImages,
// Drop empty (no id yet) and zero-weight rows, and trim hand-typed repo ids, so the
// recipe records only adapters that applied. Gate on loraCapable: a restore can leave
// adapters in state while the loaded model doesn't support LoRA (picker hidden), and
// sending them would fail generation with no visible row to remove.
loras: (() => {
if (!loraCapable) return undefined;
const active = loras
.map((l) => ({ id: l.id.trim(), weight: l.weight }))
.filter((l) => l.id && l.weight > 0);
return active.length ? active : undefined;
})(),
// ControlNet: sent only when a model + control image are chosen; v1 conditions plain
// text-to-image only, so skip it for image-conditioned workflows.
controlnet:
controlnetCapable && controlnetId && controlImage && workflow === "create"
? {
id: controlnetId,
image: controlImage,
control_type: controlType,
strength: controlStrength,
}
: undefined,
});
let res: DiffusionGenerateResponse;
try {
res = await generateDiffusionImage({
prompt: prompt.trim(),
// Only send a negative prompt when guidance uses it, so the recipe
// doesn't record one the model ignored.
negative_prompt: guidance > 0 ? negativePrompt.trim() || undefined : undefined,
width: w,
height: h,
steps,
guidance,
// Offset runs by the batch size: the native engine seeds image j of a
// run at seed+j, so a +1 run offset would regenerate the previous run's
// batch-mates. Unique per image on both engines, reproducible via recipes.
seed: baseSeed + i * batchSize,
batch_size: batchSize,
// Transform/Inpaint/Extend send the source image (+ mask for inpaint/extend) and
// a denoise strength, resolved above. The backend derives output size from the
// image, so width/height are advisory here.
init_image: condInit,
mask_image: condMask,
strength: condStrength,
upscale: condUpscale,
reference_images: condRefImages,
// Drop empty (no id yet) and zero-weight rows, and trim hand-typed repo ids, so the
// recipe records only adapters that applied. Gate on loraCapable: a restore can leave
// adapters in state while the loaded model doesn't support LoRA (picker hidden), and
// sending them would fail generation with no visible row to remove.
loras: (() => {
if (!loraCapable) return undefined;
const active = loras
.map((l) => ({ id: l.id.trim(), weight: l.weight }))
.filter((l) => l.id && l.weight > 0);
return active.length ? active : undefined;
})(),
// ControlNet: sent only when a model + control image are chosen; v1 conditions plain
// text-to-image only, so skip it for image-conditioned workflows.
controlnet:
controlnetCapable && controlnetId && controlImage && workflow === "create"
? {
id: controlnetId,
image: controlImage,
control_type: controlType,
strength: controlStrength,
}
: undefined,
});
} catch (err) {
// The POST's response was lost while the backend kept generating (secure mode's
// tunnel caps the origin response near 100s, which a native-CPU or high-step run
// exceeds). Retrying would duplicate the work, so wait the run out and read its
// images off the gallery, which is where they are persisted anyway.
if (!(err instanceof GenerateResponseLostError)) throw err;
await settleLostGeneration(() => isMounted.current);
if (!isMounted.current) break;
await loadGallery();
setGenDone(i + 1);
continue;
}
if (!isMounted.current) break;
// Prepend this run's records (newest first) and load their blobs.
setImages((prev) => [...res.images, ...prev]);
@ -2303,7 +2345,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
setGenDone(null);
setGenStep(null);
}
}, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, loraCapable, controlnetCapable, controlnetId, controlImage, controlType, controlStrength, ensureSrc, refreshStatus]);
}, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, loraCapable, controlnetCapable, controlnetId, controlImage, controlType, controlStrength, ensureSrc, loadGallery, refreshStatus]);
// Keep the active workflow valid for the loaded model: an edit-only model (Qwen-Image-
// Edit) has no Create/Transform tabs, a base model has no Edit tab. Snap to the first

3
studio/frontend/src/features/images/train/diffusion-train-panel.tsx
View file

@ -1540,6 +1540,9 @@ export function DiffusionTrainPanel({
{status?.lora_path && (
<span className="mt-1 block break-all">Saved: {status.lora_path}</span>
)}
{status?.ema_path && (
<span className="mt-1 block break-all">EMA adapter: {status.ema_path}</span>
)}
</p>
<div className="mt-1 flex gap-2">
<Button type="button" size="sm" onClick={onDeployClick}>