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

Tighten comments in the diffusion training core and API models

Browse source
This commit is contained in:
Daniel Han 2026-07-12 12:06:44 +00:00
commit ecae46cbfb
6 changed files with 247 additions and 269 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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"

178
studio/backend/core/training/diffusion_train_common.py
View file

@ -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):

14
studio/backend/core/training/training.py
View file

@ -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")

38
studio/backend/main.py
View file

@ -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
)

84
studio/backend/models/inference.py
View file

@ -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), "

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

@ -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