diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index ed40e9c3fe..4aac8d7144 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -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 - # /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( diff --git a/studio/backend/core/training/diffusion_train_extras.py b/studio/backend/core/training/diffusion_train_extras.py index a25e45198d..8381d377d5 100644 --- a/studio/backend/core/training/diffusion_train_extras.py +++ b/studio/backend/core/training/diffusion_train_extras.py @@ -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 diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index f55636f75a..b69e30cfb5 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -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") diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 162c5e727c..da325aacec 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_train_extras.py b/studio/backend/tests/test_diffusion_train_extras.py index d5b2740314..550d63783d 100644 --- a/studio/backend/tests/test_diffusion_train_extras.py +++ b/studio/backend/tests/test_diffusion_train_extras.py @@ -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) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 28f13cd133..b5f20f51bc 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -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 - # .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 .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 diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index f360427e68..bf31b685e7 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -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 { - 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 { @@ -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 | null; metric_history?: DiffusionMetricHistory | null; } diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index e73eec453d..4325d3d123 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -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 { + 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 diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index e68277c253..bbd4a3cc7f 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -1540,6 +1540,9 @@ export function DiffusionTrainPanel({ {status?.lora_path && ( Saved: {status.lora_path} )} + {status?.ema_path && ( + EMA adapter: {status.ema_path} + )}