From ecae46cbfb61fd51a761df4560656f8a505f3e1a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 12:06:44 +0000 Subject: [PATCH] Tighten comments in the diffusion training core and API models --- .../core/training/diffusion_dit_trainer.py | 177 +++++++++-------- .../core/training/diffusion_train_common.py | 178 +++++++++--------- studio/backend/core/training/training.py | 14 +- studio/backend/main.py | 38 ++-- studio/backend/models/inference.py | 84 ++++----- studio/backend/models/training.py | 25 ++- 6 files changed, 247 insertions(+), 269 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index d78ddf9d6b..ba220bd514 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -59,8 +59,8 @@ from core.training.diffusion_train_common import ( ) # Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks -# also carry added-kv projections; Z-Image is single-stream. Kept here (not in the generic -# DEFAULT_LORA_TARGETS) because they are architecture-specific. +# 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", @@ -73,9 +73,9 @@ _FLUX_TARGETS = ( ) _QWEN_TARGETS = _FLUX_TARGETS _ZIMAGE_TARGETS = ("to_q", "to_k", "to_v", "to_out.0") -# The Krea 2 authors' recommended default target set (their DreamBooth reference script): -# attention + SwiGLU + the text-fusion projector + the 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", @@ -120,18 +120,18 @@ class _FamilySpec: # Approximate dense-bf16 transformer weight size, used by base_precision="auto" to # decide which mode fits the free VRAM (with headroom for activations + optimizer). dense_bf16_gb: float - # Phased load, so the (multi-GB) transformer never has to coexist with the text - # encoders + VAE: ``load_conditioners`` builds the pipeline WITHOUT its transformer - # (encode_prompt + VAE only) and returns (pipe, vae); ``load_transformer`` loads the - # transformer alone (as a trainable nf4 QLoRA when qlora=True) once the conditioning - # modules are freed. Roughly halves peak VRAM for the big DiTs. + # 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: Callable[..., tuple[Any, Any]] load_transformer: Callable[..., Any] # Encode a list of captions -> a per-caption tuple of CPU tensors (the family's embeds). encode_prompts: Callable[..., list[tuple]] # Encode a pixel tensor [B,3,H,W] in [-1,1] -> latents (family-normalised, on device). encode_latents: Callable[..., Any] - # Encode a pixel tensor -> (A, B) affine posterior parameters so a per-step sample is + # 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_latent_stats: Callable[..., tuple] @@ -260,9 +260,9 @@ 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 (missing + # get_apply_tensor_subclass). return transformer_cls.from_pretrained( cfg.base_model, subfolder = "transformer", @@ -332,10 +332,10 @@ def _mx_module_filter(mod, fqn: str) -> bool: return False 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 tensor - # whose linear override computes input @ weight_t and drops the bias entirely, so an mxfp8'd - # FROZEN base linear would silently lose its bias and change the output the LoRA regresses - # against (verified on Blackwell: the bias term is fully dropped). Keep biased linears in bf16. + # 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. if getattr(mod, "bias", None) is not None: return False return mod.in_features % 32 == 0 and mod.out_features % 32 == 0 @@ -424,12 +424,11 @@ 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 as if it were int8 -- the memory saving - # silently gone and a likely OOM. The auto pick and /info already gate on a FUNCTIONAL - # torchao; apply the same gate to the explicit request so it fails fast with a clear - # message. fp8 keeps its own graceful fallback (_apply_fp8_training), so this is int8-only. + # 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. if mode == "int8" and not has_functional_torchao(): raise ValueError( "base_precision='int8' needs a functional torchao install; this host's " @@ -437,9 +436,9 @@ def _resolve_base_precision(cfg, spec, device) -> str: "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-transformer load. /info only advertises mxfp8 on - # sm100+ (train_precision_modes), so re-check it here to fail fast for a stale or direct - # client on an older CUDA GPU instead of crashing mid-run. + # 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. if mode == "mxfp8" and device == "cuda": try: import torch @@ -460,11 +459,10 @@ def _resolve_base_precision(cfg, spec, device) -> str: free_gb = None capability = None has_fp8 = False - # int8 quantization has no runtime fallback, so gate the auto pick on a FUNCTIONAL - # torchao: a plain find_spec("torchao") is satisfied by the Windows-ROCm import stub, - # whose quantize_ is a no-op that would leave the transformer dense while compile is - # disabled as if it were int8. has_functional_torchao imports the exact symbols - # _int8_quantize_base uses and rejects the stub. + # 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. has_torchao = has_functional_torchao() if device == "cuda": try: @@ -543,9 +541,9 @@ def _flux_collate( return (pe, pooled, text_ids) -# Per-run cache of the step-invariant FLUX conditioning tensors (RoPE image ids + the -# guidance vector): their shapes are fixed once resolution/batch are, so rebuilding them -# every step is pure allocator churn. Cleared at run start (subprocess-local anyway). +# 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). _FLUX_STATIC: dict[tuple, tuple] = {} @@ -601,8 +599,7 @@ 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. + # 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) @@ -794,10 +791,9 @@ 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 cannot - # parse, so the conditioning pipeline is assembled per-component (with the tokenizer - # and rope compat) instead of from_pretrained(transformer = None); see - # core/inference/diffusion_krea2.py for the exact compat story. + # 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. import torch from core.inference.diffusion_krea2 import load_krea2_pipeline @@ -822,10 +818,10 @@ def _krea2_encode_prompts(pipe, captions, device): out = [] 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 (the padding sits mid-template, BEFORE the assistant suffix, - # matching how the model was sampled at training time). + # 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). pe, mask = pipe.encode_prompt( prompt = cap, device = device, @@ -862,8 +858,8 @@ def _krea2_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, dev 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 exactly like the reference DreamBooth script (patch_size = 2). + # _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) @@ -951,9 +947,9 @@ _SPECS: dict[str, _FamilySpec] = { } -# HF repos that gate access behind a license acceptance: training needs a token whose -# account has 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 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. _GATED_TRAIN_REPOS = frozenset({"black-forest-labs/flux.1-dev"}) @@ -1067,10 +1063,10 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev a, b = _hold(a), _hold(b) if not forced and not gated: # Size-gate the automatic cache off the first REAL encoded variant, before - # building the rest: packed 16-channel DiT latents x variants x images of two - # fp32 tensors can exhaust host/pinned RAM. Over budget we bail with the VAE - # still resident so the loop encodes latents per step instead. ``b`` is None - # for a deterministic-latent family, so only ``a`` contributes bytes there. + # 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() @@ -1132,10 +1128,10 @@ def _should_compile( 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") @@ -1174,18 +1170,17 @@ def _maybe_compile_transformer( dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None) if dynamo_cfg is not None: # Heterogeneous-block DiTs (Z-Image: ~11 distinct block shapes) exceed dynamo's - # default recompile limit of 8; bump it the same way the inference speed layer - # does (diffusers' documented regional-compile fix). + # default recompile limit of 8; bump it like the inference speed layer does + # (diffusers' documented regional-compile fix). for attr in ("recompile_limit", "cache_size_limit"): if hasattr(dynamo_cfg, attr): setattr(dynamo_cfg, attr, max(getattr(dynamo_cfg, attr) or 0, 64)) if hasattr(dynamo_cfg, "suppress_errors"): dynamo_cfg.suppress_errors = True - # dynamic=True matches the inference speed layer's proven default: on torch 2.10 / - # B200 the dynamic=False specialisation fused a gemm_and_bias epilogue that failed - # with CUBLAS_STATUS_EXECUTION_FAILED then an illegal memory access on the FLUX - # training graph. fullgraph only on a dense base: bnb 4-bit layers graph-break by - # design. + # dynamic=True matches the inference speed layer's proven default: on torch 2.10 / B200 + # the dynamic=False specialisation fused a gemm_and_bias epilogue that failed with + # CUBLAS_STATUS_EXECUTION_FAILED then an illegal memory access on the FLUX training graph. + # fullgraph only on a dense base: bnb 4-bit layers graph-break by design. fn(fullgraph = not base_is_bnb, dynamic = True) return True except Exception as exc: # noqa: BLE001 -- optimisation only, never fatal @@ -1205,10 +1200,10 @@ 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). A caller that - # explicitly asks for fp16 on a bf16-only family is refused rather than 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 / " @@ -1235,12 +1230,12 @@ 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, which is - # unsupported for real runs but keeps import/unit tests architecture-agnostic). - # Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the run - # would otherwise die deep in model load with an opaque dtype error. Gate on NATIVE bf16 - # (capability major >= 8) -- is_bf16_supported() counts pre-Ampere emulation as supported, - # which is what this guard exists to reject; shared with the /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 + # (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. if device == "cuda" and not native_bf16_supported(): raise ValueError( "This trainer requires a bfloat16-capable GPU (Ampere or newer); " @@ -1253,9 +1248,9 @@ 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 that the dataset size is known, and - # rebind cfg so every downstream read (scheduler length, the loop range, progress - # total_steps, steps_run) sees the same resolved value. + # 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. 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(): @@ -1353,8 +1348,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto 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). + # default; bf16 / int8 / fp8 / mxfp8 are the dense speed modes; "auto" picks from free VRAM + # measured before the load). base_precision = _resolve_base_precision(cfg, spec, device) transformer = spec.load_transformer(cfg, device, weight_dtype, base_precision) base_is_bnb = base_precision == "nf4" @@ -1371,10 +1366,10 @@ 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 do not 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, which happens with a frozen 4-bit base). import functools import torch.utils.checkpoint as _ckpt transformer.enable_gradient_checkpointing( @@ -1406,8 +1401,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto cfg.base_model, subfolder = "scheduler", token = cfg.hf_token ) # The LR schedule advances once per optimizer update, so warmup/decay are counted in - # optimizer steps (matching the SDXL trainer; multiplying by the accumulation factor - # would stretch warmup past the run and never reach the decay). + # optimizer steps (matching the SDXL trainer; multiplying by the accumulation factor would + # stretch warmup past the run and never reach the decay). lr_sched = get_scheduler( cfg.lr_scheduler, optimizer = optimizer, @@ -1420,8 +1415,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto transformer.train() 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 + # 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. index_sampler = PermutationBatchSampler(n_images, rng) stopped = False @@ -1431,9 +1426,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto t_steady = None done = 0 # bf16 autocast around the forward + loss, matching the diffusers dreambooth scripts' - # accelerator.autocast: it reconciles the fp32 LoRA params with the bnb 4-bit base - # matmuls in one compute dtype. Without it the 4-bit backward on FLUX dies with an - # illegal-address / CUBLAS failure. + # 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. autocast = ( torch.autocast(device_type = "cuda", dtype = torch.bfloat16) if device == "cuda" diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index f67a6fd0f6..89efb7802b 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -32,18 +32,17 @@ from core.inference.diffusion_families import ( trainable_family_names, ) -# Default LoRA target modules: the attention projections common to the SDXL U-Net and the -# DiT transformers (the diffusers/kohya convention). A family whose trainer wants a wider -# set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback. +# Default LoRA target modules: the attention projections common to the SDXL U-Net and the DiT +# transformers (the diffusers/kohya convention). A family wanting a wider set overrides this in +# its own defaults; kept here so DiffusionLoraConfig has a sane fallback. DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0") # diffusers' SchedulerType names (diffusers.optimization.get_scheduler). piecewise_constant is -# intentionally excluded: it is the only scheduler that needs a `step_rules` string, which the -# trainers never pass (get_scheduler is called with only warmup/training steps, and there is no -# config field for it). Accepting it would pass normalized(), free the resident GPU workloads, -# then crash in the trainer subprocess (get_piecewise_constant_schedule does step_rules.split(",") -# on None) -- the exact evict-then-fail the up-front validation exists to prevent. The remaining -# six all run with only warmup/training steps. +# excluded: it is the only scheduler needing a `step_rules` string, which the trainers never pass +# (and there is no config field for it). Accepting it would pass normalized(), free the resident +# GPU workloads, then crash in the child (get_piecewise_constant_schedule does +# step_rules.split(",") on None) -- the evict-then-fail this validation prevents. The other six +# run with only warmup/training steps. _LR_SCHEDULERS: frozenset[str] = frozenset( { "linear", @@ -64,11 +63,10 @@ _CAPTION_EXTS = (".txt", ".caption") # diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it. DEFAULT_LORA_FILENAME = "pytorch_lora_weights.safetensors" -# Architectures Studio can neither train nor even load, so they are not in the family -# registry, but are recognisable by name. Rejecting them by name turns a confusing -# mid-run crash into an immediate, clear error. Families that ARE in the registry -# (flux / qwen-image / z-image / kontext) are handled by the positive registry check in -# ``resolve_trainable_family`` instead, so they are intentionally absent here. +# Architectures Studio can neither train nor load, so they are not in the family registry but are +# recognisable by name. Rejecting them by name turns a confusing mid-run crash into a clear error. +# Registry families (flux / qwen-image / z-image / kontext) are handled by the positive check in +# ``resolve_trainable_family``, so they are intentionally absent here. _NON_TRAINABLE_RESIDUAL_TOKENS = frozenset({"sd3", "pixart", "sana", "lumina", "cogview"}) _NON_TRAINABLE_RESIDUAL_PHRASES = ("stable-diffusion-3", "hunyuan-dit") @@ -104,11 +102,10 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None compatible: a genuinely wrong pick still fails cleanly later in from_pretrained). """ name = str(base_model or "").strip().lower() - # GGUF weights (a ``.gguf`` file or a ``*-GGUF`` repo) are inference-only: training needs - # the full diffusers pipeline (transformer + VAE + text encoders), which a GGUF repo does - # not provide. Reject by name even when the family itself is trainable. - # Exempt a local diffusers checkout that merely has "gguf" in its path, identified by its - # ``model_index.json`` marker (same marker the loader uses), not a bare ``is_dir()``. + # GGUF weights (a ``.gguf`` file or ``*-GGUF`` repo) are inference-only: training needs the + # full diffusers pipeline (transformer + VAE + text encoders), which a GGUF repo lacks. Reject + # by name even when the family is trainable. Exempt a local diffusers checkout that merely has + # "gguf" in its path, identified by its ``model_index.json`` marker, not a bare ``is_dir()``. local = Path(base_model).expanduser() if base_model else None is_local_diffusers = bool(local and (local / "model_index.json").is_file()) if name.endswith(".gguf") or ("gguf" in name and not is_local_diffusers): @@ -243,9 +240,9 @@ def get_trainer(family: str) -> Callable[..., str]: raise ValueError(f"No trainer is registered for family {family!r}.") -# Per-family training defaults surfaced by the Train UI. Distilled/turbo bases and the big -# DiTs want different rank / learning rate / resolution; these are starting points, not -# hard limits. Families absent here fall back to the DiffusionLoraConfig defaults. +# Per-family training defaults surfaced by the Train UI. Distilled/turbo bases and the big DiTs +# want different rank / learning rate / resolution; these are starting points, not hard limits. +# Families absent here fall back to the DiffusionLoraConfig defaults. FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = { "sdxl": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 1024}, "flux.1": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512}, @@ -285,10 +282,10 @@ _FAMILY_VRAM_NOTES = { ), } -# The flow-matching DiT families (run by diffusion_dit_trainer). They expose the -# base_precision / compile levers and require bf16 compute on CUDA; SDXL is absent because -# it uses its own mixed_precision path. Kept as a set so the UI gate, the bf16 preflight, -# and any future dispatch stay in sync. +# The flow-matching DiT families (run by diffusion_dit_trainer). They expose the base_precision / +# compile levers and require bf16 compute on CUDA; SDXL is absent (it uses its own +# mixed_precision path). A set so the UI gate, the bf16 preflight, and any future dispatch stay +# in sync. _DIT_TRAIN_FAMILIES = frozenset({"flux.1", "qwen-image", "z-image", "krea-2"}) @@ -351,8 +348,8 @@ def training_precision_preflight_error(resolved_family: str, base_precision: str if fam in _DIT_TRAIN_FAMILIES and mode in ("bf16", "int8", "fp8", "mxfp8"): # The DiT trainer's dense precisions all require CUDA (_resolve_base_precision rejects # bf16/int8/fp8/mxfp8 on device != "cuda"). bf16_unsupported_reason exempts a CPU-only host - # (the fp32 fallback for import/unit tests), so without this a dense request on a GPU-less - # host would pass the preflight, evict resident workloads, then raise only in the child. + # (the fp32 fallback for tests), so without this a dense request on a GPU-less host would + # pass the preflight, evict residents, then raise only in the child. try: import torch has_cuda = torch.cuda.is_available() @@ -369,9 +366,9 @@ def training_precision_preflight_error(resolved_family: str, base_precision: str "missing or the non-functional Windows-ROCm stub. Use '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-transformer load. Re-check here (mirroring - # _resolve_base_precision) so a stale or direct client on an older CUDA GPU fails fast - # before eviction instead of crashing mid-run. + # first training step, AFTER a full dense load. Re-check here (mirroring + # _resolve_base_precision) so a stale/direct client on an older GPU fails fast before + # eviction instead of crashing mid-run. if mode == "mxfp8": try: import torch @@ -400,16 +397,15 @@ def family_train_infos() -> list[dict[str, Any]]: if fam is None: continue repos = list(fam.train_base_repos) or [fam.base_repo] - # base_precision applies to the DiT trainer only; SDXL keeps its mixed_precision - # lever, so the UI hides the precision selector for it. compile applies everywhere: - # the SDXL trainer regionally compiles the U-Net's transformer blocks too. + # base_precision applies to the DiT trainer only; SDXL keeps its mixed_precision lever, so + # the UI hides the precision selector for it. compile applies everywhere: the SDXL trainer + # regionally compiles the U-Net's transformer blocks too. is_dit = name in _DIT_TRAIN_FAMILIES - # On a non-bf16 CUDA GPU the start route's preflight rejects EVERY DiT family (even nf4, - # since the DiT trainer requires bf16 unconditionally on CUDA), so advertise no precision - # for it -- otherwise /info offers an nf4 DiT option that always 400s. Otherwise drop any - # scheme this family's DiT corrupts (fp8 on Qwen-Image: activation outliers exceed fp8's - # range; the inference path denies the same set), so the UI never offers a mode - # normalized() would then reject. + # On a non-bf16 CUDA GPU the start preflight rejects EVERY DiT family (even nf4, since the + # DiT trainer requires bf16 on CUDA), so advertise no precision -- else /info offers an nf4 + # DiT option that always 400s. Otherwise drop any scheme this family's DiT corrupts (fp8 on + # Qwen-Image: activation outliers exceed fp8's range; the inference path denies the same + # set), so the UI never offers a mode normalized() would reject. dit_block = bf16_unsupported_reason(name) if is_dit else None if not is_dit or dit_block: fam_modes: list[str] = [] @@ -470,9 +466,9 @@ class DiffusionLoraConfig: caption_column: str = "text" # column in metadata.jsonl adapter_name: str = "default" hf_token: Optional[str] = None - # Precompute the VAE latents once (freeing the VAE for the whole run) instead of - # re-encoding every step. ``cache_variants`` crop/flip draws are frozen per image; - # the per-step VAE sampling noise itself is preserved (see the DiT trainer docstring). + # Precompute the VAE latents once (freeing the VAE for the run) instead of re-encoding every + # step. ``cache_variants`` crop/flip draws are frozen per image; the per-step VAE sampling + # noise itself is preserved (see the DiT trainer docstring). cache_latents: bool = True cache_variants: int = 4 # Regional torch.compile of the transformer blocks: "off" | "on" | "auto" (auto turns @@ -481,17 +477,16 @@ class DiffusionLoraConfig: # TF32 matmuls + high fp32 matmul precision + cudnn autotuning for the run. Near-lossless; # disable for strict bit-reproducibility A/Bs. enable_tf32: bool = True - # DiT base transformer precision: "nf4" (bitsandbytes QLoRA, the memory floor and the - # default), "bf16" (dense, fastest eager, compile-friendly), "int8" (torchao - # weight-only, half of bf16), "fp8" (torchao float8 training compute on the frozen - # linears, Ada/Hopper/Blackwell + compile), or "auto" (pick by free VRAM + GPU class). - # Non-nf4 modes need a dense base repo (not a prequant bnb-4bit one). SDXL ignores it. + # DiT base transformer precision: "nf4" (bitsandbytes QLoRA, the memory floor and default), + # "bf16" (dense, fastest eager, compile-friendly), "int8" (torchao weight-only, half of bf16), + # "fp8" (torchao float8 training on the frozen linears, Ada/Hopper/Blackwell + compile), or + # "auto" (by free VRAM + GPU class). Non-nf4 modes need a dense base repo. SDXL ignores it. base_precision: str = "nf4" # How often to emit a progress event (in optimizer steps). log_every: int = 1 - # Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from - # base_model. ``resolved_family`` is filled by normalized() with the trainer family - # that will actually run, so the process adapter can dispatch through the registry. + # Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from base_model. + # ``resolved_family`` is filled by normalized() with the trainer family that will run, so the + # process adapter can dispatch through the registry. model_family: Optional[str] = None resolved_family: str = "sdxl" @@ -539,10 +534,10 @@ class DiffusionLoraConfig: base_precision = str(self.base_precision or "nf4").strip().lower() if base_precision not in ("nf4", "bf16", "int8", "fp8", "mxfp8", "auto"): raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto") - # base_precision is a DiT-only lever (the transformer load precision); SDXL uses its - # own mixed_precision path and ignores base_precision entirely, so the dense-mode - # gates (prequant base / non-bf16 compute) apply only to the DiT families. The - # mode-name validity check above still runs for every family. + # base_precision is a DiT-only lever (transformer load precision); SDXL uses its own + # mixed_precision path and ignores it, so the dense-mode gates (prequant base / non-bf16 + # compute) apply only to the DiT families. The mode-name check above still runs for every + # family. if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8", "mxfp8"): if repo_is_prequantized(self.base_model): raise ValueError( @@ -556,10 +551,9 @@ class DiffusionLoraConfig: f"mixed_precision to bf16." ) # Some DiT families are corrupted by fp8's activation range: outliers exceed even - # per-row fp8's dynamic range, so the frozen linears' float8 training compute - # learns against a garbage forward pass. The inference path already denies these - # schemes; mirror that deny here so the run fails fast instead of silently - # producing a broken adapter. int8 (per-token) is unaffected and stays allowed. + # per-row fp8's range, so the frozen linears' float8 compute learns against a garbage + # forward pass. The inference path already denies these schemes; mirror it here so the + # run fails fast instead of producing a broken adapter. int8 (per-token) is unaffected. from core.inference.diffusion_transformer_quant import _family_denied if _family_denied(resolved_family, base_precision): @@ -640,9 +634,9 @@ class PermutationBatchSampler: self._pos = 0 def next_batch(self, k: int) -> list[int]: - # k may exceed n (batch larger than the dataset): the permutation is refilled across - # as many cycles as needed so the caller always gets exactly k indices and the batch - # never shrinks, matching the old sampler's fixed batch shape. + # k may exceed n (batch larger than the dataset): the permutation is refilled across as + # many cycles as needed so the caller always gets exactly k indices and the batch never + # shrinks, matching the old sampler's fixed batch shape. out: list[int] = [] while len(out) < k: if self._pos >= len(self._order): @@ -726,11 +720,10 @@ def discover_image_caption_pairs( caption = instance_prompt if caption: if verify_images: - # Reject a corrupt / zero-byte / truncated image now via a cheap PIL header - # probe (verify() does not decode the full pixels): otherwise it passes this - # filename-only discovery, the start route frees the resident GPU models, and - # the spawned trainer only then crashes in Image.open -- the eviction this - # preflight exists to prevent. + # Reject a corrupt/zero-byte/truncated image now via a cheap PIL header probe + # (verify() doesn't decode full pixels): otherwise it passes this filename-only + # discovery, the start route frees the resident GPU models, and the trainer only + # then crashes in Image.open -- the eviction this preflight prevents. try: from PIL import Image with Image.open(img) as _probe: @@ -780,18 +773,18 @@ def _plan_cache_variants( return plan -# Host-memory budget for the AUTOMATIC latent cache. The cache holds two fp32 posterior -# tensors (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA -# host. At 1024px an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times -# that, so a few thousand images x cache_variants can exhaust host or pinned RAM with no -# fallback. Over this budget the default falls back to per-step VAE encoding. A fixed -# constant (rather than a psutil RAM fraction) keeps the gate dependency-free and identical -# across hosts; it is deliberately conservative, well under a typical training host's RAM. +# Host-memory budget for the AUTOMATIC latent cache. The cache holds two fp32 posterior tensors +# (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA host. At +# 1024px an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times that, so a few +# thousand images x cache_variants can exhaust host or pinned RAM. Over budget the default falls +# back to per-step VAE encoding. A fixed constant (not a psutil RAM fraction) keeps the gate +# dependency-free and identical across hosts; deliberately conservative, well under a typical +# host's RAM. _LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3 # 4 GiB -# Returned by the cache builders when the estimated cache exceeds the budget: the caller -# keeps the VAE resident and encodes each step's latents in-loop. A distinct sentinel from -# ``None`` (which means a stop was requested mid-build) so the two are not conflated. +# Returned by the cache builders when the estimate exceeds budget: the caller keeps the VAE +# resident and encodes each step's latents in-loop. A distinct sentinel from ``None`` (a stop +# requested mid-build) so the two are not conflated. LATENT_CACHE_OVER_BUDGET: Any = object() @@ -851,10 +844,10 @@ def _apply_perf_flags( torch.set_float32_matmul_precision("highest") if cudnn_benchmark: torch.backends.cudnn.benchmark = True - # The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes - # on torch 2.10 + cu130 (B200): mha_graph.execute fails, then poisons the context - # into illegal memory accesses. Flash / mem-efficient SDPA are mathematically - # equivalent, so pin those for the run (restored on exit). + # The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes on + # torch 2.10 + cu130 (B200): mha_graph.execute fails, then poisons the context into + # illegal memory accesses. Flash / mem-efficient SDPA are equivalent, so pin those for the + # run (restored on exit). cuda_backends = getattr(torch.backends, "cuda", None) if cuda_backends is not None and hasattr(cuda_backends, "enable_cudnn_sdp"): try: @@ -907,9 +900,9 @@ def _assert_trusted_base_model(base_model: str) -> None: f"a trusted repo (an unsloth/* repo or an official base)." ) # An existing LOCAL base is loaded as a full pipeline (from_pretrained(base_model)) by the - # spawned trainer, which needs a model_index.json. Any existing path is "trusted" above, so - # reject a non-pipeline local dir here -- before /diffusion/start frees the resident GPU - # models -- rather than have the child fail after the teardown. + # trainer, which needs a model_index.json. Any existing path is "trusted" above, so reject a + # non-pipeline local dir here -- before /diffusion/start frees the GPU models -- rather than + # have the child fail after teardown. _assert_local_base_is_pipeline(base_model) @@ -969,13 +962,13 @@ def _write_lora_sidecar(sidecar_path: Path, cfg: DiffusionLoraConfig) -> None: # Aliases from the generic Studio training payload onto DiffusionLoraConfig fields, so the -# diffusion trainer can also be driven by the shared training request shape (not only its -# own request model whose keys already match). +# diffusion trainer can also be driven by the shared training request shape (not only its own +# request model, whose keys already match). _CONFIG_ALIASES = { "model_name": "base_model", "max_steps": "train_steps", - # The generic payload's num_epochs already matches the diffusion field name, but list it - # so the epochs override is threaded through the shared-payload path as explicitly as + # The generic payload's num_epochs already matches the diffusion field name, but list it so + # the epochs override is threaded through the shared-payload path as explicitly as # max_steps -> train_steps is. "num_epochs": "num_epochs", "batch_size": "train_batch_size", @@ -1016,11 +1009,10 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig: for k, v in config.items(): if k in valid: kwargs[k] = v - # Epoch-mode payloads from the generic Studio UI carry max_steps: 0 as the "use epochs" - # sentinel, which the max_steps -> train_steps alias copies as train_steps: 0. Since - # normalized() rejects train_steps < 1 before resolve_train_steps() can apply num_epochs, - # drop a falsy/0 train_steps when num_epochs > 0 so the dataclass default stands in until - # epoch resolution replaces it. + # Epoch-mode payloads from the generic UI carry max_steps: 0 as the "use epochs" sentinel, + # which the max_steps -> train_steps alias copies as train_steps: 0. Since normalized() + # rejects train_steps < 1 before resolve_train_steps() applies num_epochs, drop a falsy/0 + # train_steps when num_epochs > 0 so the dataclass default stands in until epoch resolution. try: _num_epochs = int(kwargs.get("num_epochs") or 0) except (TypeError, ValueError): diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 0481bb90ba..a13fd99d1e 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -739,9 +739,8 @@ class TrainingBackend: # True while a pump thread should be running; cleared on intended exits. # Left True after an abnormal death so _ensure_pump_alive spots a crash. self._pump_running: bool = False - # True from the start_training() guard passing until its spawn attempt - # finishes; blocks a second concurrent start (routes call start_training - # from a worker thread, so overlapping requests are possible). + # True from the start_training() guard passing until its spawn finishes; blocks a + # second concurrent start (routes call it from a worker thread, so starts can overlap). self._start_in_progress: bool = False self._lock = threading.Lock() @@ -804,11 +803,10 @@ class TrainingBackend: still letting auto-selection place training against the freed memory. Hook failures never block the start. """ - # Compare-and-set start guard: the route runs this whole method on a worker - # thread (asyncio.to_thread), so two overlapping /train/start requests can - # reach it concurrently. Without the flag both would pass the alive-check - # below (the proc is only assigned at the end) and double-spawn. Mirrors the - # diffusion training service's reserve(). + # Compare-and-set start guard: the route runs this method on a worker thread, so two + # overlapping /train/start requests can reach it concurrently. Without the flag both + # would pass the alive-check below (the proc is only assigned at the end) and + # double-spawn. Mirrors the diffusion training service's reserve(). with self._lock: if self._start_in_progress: logger.warning("Training start already in progress") diff --git a/studio/backend/main.py b/studio/backend/main.py index 92d9753753..4ec83c1b70 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -727,10 +727,10 @@ from utils.upload_limits import ( # noqa: E402 ) _BODY_PROTECTED_PREFIXES = ( - # Blanket-protect the whole OpenAI-compatible /v1 surface, like /api/inference below: every - # /v1 POST route (chat/completions, completions, images/generations, audio, embeddings, - # responses, messages, ...) buffers a JSON body and none is a multipart-upload passthrough, - # so a single prefix caps them all -- an enumerated list silently left new routes uncapped. + # Blanket-protect the whole /v1 surface, like /api/inference below: every /v1 POST route + # (chat/completions, images/generations, audio, embeddings, responses, ...) buffers a JSON + # body and none is a multipart passthrough, so one prefix caps them all -- an enumerated + # list would silently leave new routes uncapped. "/v1", "/p/", "/api/inference", @@ -746,14 +746,13 @@ _DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload" _DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = ( "/api/data-recipe/seed/upload-unstructured-file" ) -# The diffusion dataset upload route (POST /api/train/diffusion/dataset) is a multipart -# image upload under the protected /api/train prefix. Like /api/datasets/upload it enforces -# its own get_upload_limit_bytes() cap, so it must bypass the default body cap here or the -# middleware would 413 near-limit batches (and ignore a raised max_upload_size_mb) before the -# handler runs. Matched as an EXACT path (not a prefix): its JSON sub-routes -# (PUT .../{name}/caption/{filename}, POST .../import-example) live under the same prefix but -# must keep the normal small-JSON cap, or a large caption/import body would be buffered and -# parsed up to the far larger upload limit before the route-level length checks run. +# The diffusion dataset upload route (POST /api/train/diffusion/dataset) is a multipart image +# upload under the protected /api/train prefix. Like /api/datasets/upload it enforces its own +# get_upload_limit_bytes() cap, so it must bypass the default body cap here or the middleware +# would 413 near-limit batches (ignoring a raised max_upload_size_mb) before the handler runs. +# Matched as an EXACT path (not a prefix): its JSON sub-routes (.../caption/{filename}, +# .../import-example) share the prefix but must keep the small-JSON cap, or a large +# caption/import body would be buffered up to the far larger upload limit. _DIFFUSION_DATASET_UPLOAD_PATH = "/api/train/diffusion/dataset" _BODY_UPLOAD_PASSTHROUGH_PREFIXES = ( _DATASET_UPLOAD_PASSTHROUGH_PREFIX, @@ -767,11 +766,10 @@ _BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS = (_DIFFUSION_DATASET_UPLOAD_PATH,) def _get_upload_passthrough_request_max_bytes(path: str) -> int: if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX): return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES) - # The trailing-slash variant (/api/train/diffusion/dataset/) reaches this middleware - # BEFORE the router's redirect_slashes 307, so it must resolve to the same upload cap - # as the canonical path or a large upload 413s on the default /api/train body cap. - # Stripping slashes cannot promote a JSON sub-route: those all keep extra path - # components after normalization and still miss the exact match. + # The trailing-slash variant reaches this middleware BEFORE the router's redirect_slashes + # 307, so it must resolve to the same upload cap as the canonical path or a large upload + # 413s on the default /api/train cap. Stripping slashes can't promote a JSON sub-route: + # those keep extra path components after normalization and still miss the exact match. if ( path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX) or path.rstrip("/") == _DIFFUSION_DATASET_UPLOAD_PATH @@ -836,9 +834,9 @@ class MaxBodyMiddleware: self.upload_passthrough_exact_paths = upload_passthrough_exact_paths def _is_upload_passthrough(self, path: str) -> bool: - # Exact paths also match their trailing-slash variant: the middleware runs - # before the router's redirect_slashes 307, and a JSON sub-route can never - # normalize down to the exact path (it keeps extra components). + # Exact paths also match their trailing-slash variant: the middleware runs before the + # router's redirect_slashes 307, and a JSON sub-route can never normalize to the exact + # path (it keeps extra components). return path.rstrip("/") in self.upload_passthrough_exact_paths or any( path.startswith(p) for p in self.upload_passthrough_prefixes ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3fb47886c7..a6800087ec 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1894,9 +1894,9 @@ class DiffusionLoadRequest(BaseModel): @field_validator("attention_backend", mode = "before") @classmethod def _normalize_attention_backend(cls, value): - # The dispatcher accepts case/whitespace variants ("CuDNN", " sage "), but the - # Literal above is validated before any normaliser runs, so fold a string to its - # canonical lower/stripped form here -- otherwise valid casing gets a 422. + # The dispatcher accepts case/whitespace variants ("CuDNN", " sage "), but the Literal + # above is validated before any normaliser runs, so fold a string to its canonical + # lower/stripped form here -- otherwise valid casing gets a 422. return value.strip().lower() if isinstance(value, str) else value @@ -1953,8 +1953,8 @@ class ControlNetSpec(BaseModel): @model_validator(mode = "after") def _check_guidance_range(self) -> "ControlNetSpec": - # An inverted range (start > end) means "act over no steps"; reject it as a clean - # 422 instead of letting the diffusers pipeline raise a 500 deep in the denoise. + # An inverted range (start > end) means "act over no steps"; reject it as a clean 422 + # instead of letting the diffusers pipeline 500 deep in the denoise. if self.guidance_start > self.guidance_end: raise ValueError("guidance_start must be <= guidance_end") return self @@ -1973,21 +1973,19 @@ class DiffusionGenerateRequest(BaseModel): ) steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps") guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") - # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript - # rounds integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then - # generate a different image. Random seeds are already masked to this range. + # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds + # integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then generate a + # different image. Random seeds are already masked to this range. seed: Optional[int] = Field( None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)" ) batch_size: int = Field( 1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)" ) - # Image-conditioned workflows (base64 or data-URL). An init_image alone runs img2img; - # init_image + mask_image runs inpaint. Both require a model family with the matching - # pipeline (img2img/inpaint) or the load is rejected with a clear message. - # Cap each base64 image string so a single request can't buffer a multi-GB payload (the - # decoded dimensions are bounded separately in the backend). ~32 MiB comfortably fits a - # full 4096px image yet rejects abuse. + # Image-conditioned workflows (base64 or data-URL): init_image alone runs img2img, + # init_image + mask_image runs inpaint. Both require a family with the matching pipeline or + # the load is rejected. Cap each base64 string so one request can't buffer a multi-GB payload + # (decoded dimensions are bounded separately); ~32 MiB fits a full 4096px image yet rejects abuse. init_image: Optional[str] = Field( None, max_length = 32 * 1024 * 1024, @@ -2038,9 +2036,9 @@ class DiffusionGenerateRequest(BaseModel): @classmethod def _unique_lora_ids(cls, value: Optional[list[LoraSpec]]) -> Optional[list[LoraSpec]]: # Both apply paths break alias collisions by suffixing the adapter name/file, so a - # repeated id would load the SAME adapter as several distinct adapters and stack - # its effect past the per-adapter weight bound. The UI already prevents duplicates; - # reject them for API clients too so each adapter takes effect at most once. + # repeated id would load the SAME adapter several times and stack its effect past the + # per-adapter weight bound. The UI already blocks duplicates; reject them for API clients + # too so each adapter takes effect at most once. if value: seen: set[str] = set() for spec in value: @@ -2054,8 +2052,8 @@ class DiffusionGenerateRequest(BaseModel): @field_validator("reference_images") @classmethod def _bounded_reference_items(cls, value: Optional[list[str]]) -> Optional[list[str]]: - # Each reference is a base64 image; bound its length like init_image/mask_image so a - # request carrying several references can't buffer a multi-GB payload. + # Each reference is a base64 image; bound its length like init_image/mask_image so + # several references can't buffer a multi-GB payload. if value is not None: for item in value: if len(item) > 32 * 1024 * 1024: @@ -2065,9 +2063,8 @@ class DiffusionGenerateRequest(BaseModel): @field_validator("width", "height") @classmethod def _multiple_of_16(cls, value: int) -> int: - # Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x - # patch). Non-multiples crash deep in the pipeline, so reject them here - # for a clean 422 instead of a cryptic 500. + # Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x patch). + # Non-multiples crash deep in the pipeline, so reject them here for a clean 422. if value % 16 != 0: raise ValueError("must be a multiple of 16") return value @@ -2213,10 +2210,10 @@ class DiffusionStatusResponse(BaseModel): "picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False " "for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.", ) - # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Present - # only on backends that record it; null when nothing is loaded or on older backends. The - # frontend renders an "Auto: X" badge next to each control whose source == "auto". Declared - # explicitly so pydantic's default extra='ignore' does not silently drop the resolved record. + # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Present only + # on backends that record it; null when nothing is loaded or on older backends. The frontend + # renders an "Auto: X" badge next to each control whose source == "auto". Declared explicitly + # so pydantic's extra='ignore' doesn't drop the resolved record. resolved: Optional[Dict[str, DiffusionResolvedControl]] = Field( None, description = "Per-control resolved value + provenance (source auto|explicit + reason), " @@ -2253,12 +2250,11 @@ class DiffusionInferenceInfoResponse(BaseModel): # ── OpenAI-compatible images API (POST /v1/images/generations) ── # -# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf -# OpenAI clients work unchanged. The loaded image GGUF stands in for the model; -# GPT-image-only knobs (quality, style, background, output_format, ...) are -# accepted and ignored, exactly as dall-e-2 ignores them. The size string is -# parsed and `stream` is rejected in the route, where the diffusion backend is -# in reach; everything Pydantic can check declaratively lives here. +# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf clients work +# unchanged. The loaded image GGUF stands in for the model; GPT-image-only knobs (quality, +# style, background, output_format, ...) are accepted and ignored, like dall-e-2. The size +# string is parsed and `stream` is rejected in the route (where the diffusion backend is in +# reach); everything Pydantic can check declaratively lives here. class ImageGenerationRequest(BaseModel): @@ -2280,8 +2276,8 @@ class ImageGenerationRequest(BaseModel): "url", description = "Return each image as a URL or a base64-encoded PNG." ) user: Optional[str] = Field(None, description = "End-user identifier (accepted, unused).") - # gpt-image-only; declared so we can reject it with a clear error instead of - # silently returning JSON to a client that asked for an SSE stream. + # gpt-image-only; declared so we can reject it clearly instead of returning JSON to a + # client that asked for an SSE stream. stream: Optional[bool] = Field( None, description = "Streaming image generation is not supported; omit or set false." ) @@ -2289,8 +2285,8 @@ class ImageGenerationRequest(BaseModel): @field_validator("n", "size", "response_format", mode = "before") @classmethod def _null_means_default(cls, value, info): - # OpenAI marks these nullable WITH a default, so an explicit null means - # "use the default" — coalesce it instead of 400-ing a spec-valid body. + # OpenAI marks these nullable WITH a default, so an explicit null means "use the + # default" -- coalesce it instead of 400-ing a spec-valid body. if value is None: return cls.model_fields[info.field_name].default return value @@ -2435,8 +2431,8 @@ class VideoGenerateRequest(BaseModel): negative_prompt: Optional[str] = Field( None, description = "What to avoid (if the model supports it)" ) - # Width/height/num_frames/fps default per loaded family (the backend snaps them to - # the family's required multiples/lattice), so they are optional here. + # Width/height/num_frames/fps default per loaded family (the backend snaps them to its + # required multiples/lattice), so they are optional here. width: Optional[int] = Field( None, ge = 32, le = 2048, description = "Frame width in pixels (family multiple)" ) @@ -2467,9 +2463,9 @@ class VideoGenerateRequest(BaseModel): "pipeline default it to the main guidance. Ignored by single-DiT families (their pipeline " "signature has no second guidance kwarg).", ) - # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript - # rounds integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then - # generate a different clip. Random seeds are already masked to this range. + # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds + # integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then generate a + # different clip. Random seeds are already masked to this range. seed: Optional[int] = Field( None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)" ) @@ -2617,9 +2613,9 @@ class VideoStatusResponse(BaseModel): defaults: Optional[VideoGenerationDefaults] = Field( None, description = "Per-family generation defaults + shape constraints; null when unloaded" ) - # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Same - # shape as the diffusion status uses; null when nothing is loaded. The frontend renders - # an "Auto: X" badge next to each control whose source == "auto". + # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Same shape + # as the diffusion status; null when nothing is loaded. The frontend renders an "Auto: X" + # badge next to each control whose source == "auto". resolved: Optional[Dict[str, DiffusionResolvedControl]] = Field( None, description = "Per-control resolved value + provenance (source auto|explicit + reason), " diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 0fa43ae5f6..d695b1cdb6 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -706,9 +706,9 @@ class DiffusionTrainingStartRequest(BaseModel): lora_rank: int = Field(16, ge = 1, le = 320) lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank") lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0) - # Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that - # sets them is not silently trained with defaults. Default the target list 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", @@ -796,15 +796,15 @@ 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 (the training health - # signal the UI charts alongside the loss). + # Total pre-clip gradient norm from the last optimizer step (health signal the UI charts + # alongside the loss). 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 - # Where the trained adapter was mirrored into the Studio LoRA catalog, and what family - # / base it was trained from -- lets the UI deploy the adapter onto the right base. + # 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. catalog_path: Optional[str] = None family: Optional[str] = None base_model: Optional[str] = None @@ -872,15 +872,14 @@ class DiffusionTrainableFamily(BaseModel): base_repos: List[str] = Field(default_factory = list) defaults: dict = Field(default_factory = dict) vram_note: str = "" - # base_precision modes this machine supports for the family (empty = the family has no - # precision selector, e.g. SDXL), plus the recommended pick and whether regional - # torch.compile applies. Defaults keep older backends' payloads valid. + # 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 - # When set, deploying a LoRA trained on this family previews it on this repo instead of - # the training base (Krea trains on Raw but runs adapters on Turbo). Null for families - # that deploy on the base they trained on. + # When set, a LoRA trained on this family previews on this repo instead of the training base + # (Krea trains on Raw but runs adapters on Turbo). Null when it deploys on the training base. deploy_base: Optional[str] = None