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

Refactor diffusion LoRA training into a family-aware platform

Browse source 
Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:

- New core/training/diffusion_train_common.py holds the config + validation, dataset
  discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
  registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
  re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
  from the base model (or an explicit model_family) via the diffusion family registry,
  and a known-but-not-yet-trainable family is refused with a clear message. Unknown
  custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
  DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
  trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
  picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
  and Qwen-Image repos (safetensors-only, no remote code).
This commit is contained in:
Daniel Han 2026-07-02 14:55:09 +00:00
commit 7f0a9ebd2f
5 changed files with 500 additions and 300 deletions
Download patch file Download diff file Expand all files Collapse all files

7
studio/backend/core/inference/diffusion.py
View file

@ -196,6 +196,13 @@ _TRUSTED_NON_GGUF_REPOS = frozenset(
{
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/sdxl-turbo",
# Official vendor, safetensors-only base repos allowlisted as LoRA TRAINING bases
# (diffusion training loads the full pipeline from these). Same rule as above: no
# pickled weights, no remote code, exact-match lowercased. FLUX.1-dev is gated on
# the Hub (needs the user's token); the other two are open.
"black-forest-labs/flux.1-dev",
"tongyi-mai/z-image-turbo",
"qwen/qwen-image",
}
)

20
studio/backend/core/inference/diffusion_families.py
View file

@ -118,6 +118,14 @@ class DiffusionFamily:
# leaves sd-cli's defaults (correct for the distilled flux/z-image families).
sd_cpp_sampling_method: Optional[str] = None
sd_cpp_flow_shift: Optional[float] = None
# True when Studio can TRAIN a LoRA on this family (a trainer is registered for it in
# core.training). Loadable-for-inference is the default; training is opt-in per family
# because each architecture needs its own training loop. The diffusion training start
# path resolves the base model's family and refuses a non-trainable one up front.
trainable: bool = False
# Recommended base repos to train FROM, most-preferred first (e.g. a QLoRA-friendly
# prequant repo, then a bf16 repo). Surfaced by the Train UI as the base-model choices.
train_base_repos: tuple[str, ...] = field(default_factory = tuple)
# Keyed by architecture, not per model variant: a checkpoint's specific base repo
@ -286,9 +294,21 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
inpaint_pipeline_class = "StableDiffusionXLInpaintPipeline",
controlnet_pipeline_class = "StableDiffusionXLControlNetPipeline",
controlnet_model_class = "ControlNetModel",
# SDXL is the one family with a shipped LoRA trainer today (the U-Net trainer).
# DiT families (flux.1 / qwen-image / z-image) become trainable in a follow-up.
trainable = True,
train_base_repos = (
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/sdxl-turbo",
),
),
)
def trainable_family_names() -> tuple[str, ...]:
"""Names of families Studio can train a LoRA on, in registry order."""
return tuple(fam.name for fam in _FAMILIES if fam.trainable)
# Editing / inpaint checkpoints share an arch keyword but need a different
# pipeline and an input image, which this text-to-image backend doesn't drive.
# "layered" rejects Qwen-Image-Layered: its transformer sets additional_t_cond=True

33
studio/backend/core/inference/diffusion_lora.py
View file

@ -15,6 +15,7 @@ directory / the HF hub before anything is loaded.
from __future__ import annotations
import json
import os
import re
import threading
@ -118,6 +119,11 @@ def _scan_local() -> list[LoraCatalogEntry]:
except OSError:
size = 0
entry_id = p.name if stem_counts.get(p.stem, 0) > 1 else p.stem
# A ``<stem>.json`` sidecar (written by the diffusion trainer on publish) records the
# adapter's family + default weight, so a trained adapter is family-gated in the
# picker instead of showing as "unknown" for every model. Best-effort: a missing or
# malformed sidecar just leaves the defaults (families = (), weight_default = 1.0).
families, weight_default = _read_lora_sidecar(p)
entries.append(
LoraCatalogEntry(
id = entry_id,
@ -126,11 +132,38 @@ def _scan_local() -> list[LoraCatalogEntry]:
fmt = "gguf" if ext == ".gguf" else "safetensors",
local_path = str(p),
size_bytes = size,
families = families,
weight_default = weight_default,
)
)
return entries
def _read_lora_sidecar(weight_path: Path) -> tuple[tuple[str, ...], float]:
"""Read the ``<stem>.json`` metadata sidecar next to a local adapter, returning
``(families, weight_default)``. Returns ``((), 1.0)`` when the sidecar is absent or
unreadable, so discovery never fails on a bad file."""
sidecar = weight_path.with_suffix(".json")
try:
data = json.loads(sidecar.read_text(encoding = "utf-8"))
except (OSError, ValueError):
return (), 1.0
if not isinstance(data, dict):
return (), 1.0
raw_family = data.get("family")
raw_families = data.get("families")
names: list[str] = []
if isinstance(raw_families, (list, tuple)):
names = [str(f).strip().lower() for f in raw_families if str(f).strip()]
elif isinstance(raw_family, str) and raw_family.strip():
names = [raw_family.strip().lower()]
weight_default = 1.0
raw_weight = data.get("weight_default")
if isinstance(raw_weight, (int, float)) and raw_weight > 0:
weight_default = float(raw_weight)
return tuple(names), weight_default
def list_loras(*, family: Optional[str] = None) -> list[LoraCatalogEntry]:
"""Return the merged catalog (curated + local), optionally family-filtered.

346
studio/backend/core/training/diffusion_lora_trainer.py
View file

@ -8,207 +8,47 @@ and exports it as a diffusers-format ``.safetensors`` that the Studio diffusion
(and any diffusers pipeline via ``load_lora_weights``) can load.
Design:
- The pure helpers (dataset discovery, config normalisation, SDXL add-time-ids) have no
torch/diffusers dependency at call time and are unit-tested without a GPU.
- ``run_diffusion_lora_training`` is the training loop. It reports progress through an
- Family-agnostic building blocks (dataset discovery, config normalisation + validation,
event emission, the stop protocol, adapter publishing, and the family/trainer registry)
live in ``diffusion_train_common`` and are shared with the DiT trainers. They are
re-exported here so existing import paths keep working.
- ``run_diffusion_lora_training`` is the SDXL training loop. It reports progress through an
``on_event`` callback whose payloads match the training worker's event protocol
(``{"type": ..., "ts": ...}``) so it can be spawned as a subprocess and streamed to the
UI, and it polls a ``should_stop`` callback so a stop request ends it cleanly (with a
partial save).
- ``run_diffusion_training_process`` is the thin mp.Queue adapter, and ``main`` is a CLI.
Only the SDXL (U-Net) architecture is trained here; DiT families (FLUX / Qwen-Image /
Z-Image) are a follow-up. The exported adapter is loaded by the existing LoRA path.
- ``run_diffusion_training_process`` is the thin mp.Queue adapter; it dispatches to the
trainer registered for the resolved family (SDXL here, DiT families in a follow-up).
``main`` is a CLI.
"""
from __future__ import annotations
import argparse
import json
import math
import os
import random
import re
import time
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Callable, Optional
from typing import Any, Optional
# Default LoRA target modules: the U-Net attention projections. These are the standard
# SDXL LoRA targets (the diffusers/kohya convention) and keep the adapter small.
DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
_CAPTION_EXTS = (".txt", ".caption")
# diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it.
DEFAULT_LORA_FILENAME = "pytorch_lora_weights.safetensors"
# Families Studio can LOAD but not train (DiT architectures). A base-model name that
# clearly belongs to one is refused in normalized(), so a wrong pick fails at start
# (an instant HTTP 400 through the API) instead of minutes later inside
# StableDiffusionXLPipeline.from_pretrained. Single tokens match on word boundaries;
# hyphenated markers match as substrings of the hyphen-condensed name.
_NON_SDXL_TOKENS = frozenset({"flux", "sd3", "kontext", "pixart", "sana", "lumina", "cogview"})
_NON_SDXL_PHRASES = ("qwen-image", "z-image", "stable-diffusion-3", "hunyuan-dit")
_ONLY_SDXL_HINT = (
"Only SDXL bases can be trained right now (e.g. stabilityai/stable-diffusion-xl-base-1.0 "
"or stabilityai/sdxl-turbo). Other families can load LoRAs but not train them yet."
# Shared, family-agnostic building blocks. Re-exported so callers/tests that import them
# from this module (its historical home) keep working unchanged.
from core.training.diffusion_train_common import ( # noqa: F401
DEFAULT_LORA_FILENAME,
DEFAULT_LORA_TARGETS,
EventCb,
StopCb,
DiffusionLoraConfig,
_assert_trusted_base_model,
_coerce_gradient_checkpointing,
_config_from_dict,
_CONFIG_ALIASES,
_emit,
_publish_to_lora_catalog,
discover_image_caption_pairs,
get_trainer,
)
EventCb = Callable[[dict[str, Any]], None]
# Returns a falsy value to keep training, or a truthy stop signal: bare True, or a dict
# that may carry ``save=False`` to cancel without saving a partial adapter.
StopCb = Callable[[], Any]
@dataclass
class DiffusionLoraConfig:
"""Everything a diffusion LoRA training run needs. Only ``base_model`` /
``data_dir`` / ``output_dir`` are required; the rest have sensible defaults."""
base_model: str
data_dir: str
output_dir: str
# Dreambooth-style caption applied to any image without its own caption. Required if
# the dataset has no captions.jsonl / sidecar files.
instance_prompt: Optional[str] = None
resolution: int = 1024
train_steps: int = 500
learning_rate: float = 1e-4
train_batch_size: int = 1
gradient_accumulation_steps: int = 1
lora_rank: int = 16
lora_alpha: Optional[int] = None # defaults to lora_rank
lora_dropout: float = 0.0
lora_target_modules: tuple[str, ...] = DEFAULT_LORA_TARGETS
seed: int = 42
mixed_precision: str = "bf16" # "bf16" | "fp16" | "no"
snr_gamma: Optional[float] = 5.0 # min-SNR loss weighting; None disables
gradient_checkpointing: bool = True
max_grad_norm: float = 1.0
lr_scheduler: str = "constant"
lr_warmup_steps: int = 0
center_crop: bool = False
random_flip: bool = True
caption_column: str = "text" # column in metadata.jsonl
adapter_name: str = "default"
hf_token: Optional[str] = None
# How often to emit a progress event (in optimizer steps).
log_every: int = 1
def normalized(self) -> "DiffusionLoraConfig":
"""Return a copy with derived/validated fields filled in. Raises ValueError on a
request that cannot train (bad numbers, or no caption source).
Also coerces values that arrive as strings/blanks through the Studio config path
(``learning_rate`` is preserved as a string there; ``hf_token`` defaults to "")."""
assert_trainable_base_model(self.base_model)
if self.train_steps < 1:
raise ValueError("train_steps must be >= 1")
if self.train_batch_size < 1:
raise ValueError("train_batch_size must be >= 1")
if self.gradient_accumulation_steps < 1:
raise ValueError("gradient_accumulation_steps must be >= 1")
if self.lora_rank < 1:
raise ValueError("lora_rank must be >= 1")
if self.lora_alpha is not None and self.lora_alpha < 1:
raise ValueError(
"lora_alpha must be >= 1 (a zero/negative alpha scales the adapter to nothing)"
)
if self.resolution < 64 or self.resolution % 8 != 0:
raise ValueError("resolution must be a multiple of 8 and >= 64")
if self.mixed_precision not in ("bf16", "fp16", "no"):
raise ValueError("mixed_precision must be one of bf16 / fp16 / no")
# learning_rate can arrive as a string ("1e-4") from the Studio config path, which
# preserves it as a string after validation; coerce so AdamW receives a float.
try:
learning_rate = float(self.learning_rate)
except (TypeError, ValueError) as exc:
raise ValueError(f"learning_rate must be a number, got {self.learning_rate!r}") from exc
if learning_rate <= 0:
raise ValueError("learning_rate must be > 0")
alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank
targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS
# A blank Hub token (the Studio default when none is configured) must load
# anonymously, not as an explicit empty credential.
token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token
return replace(
self,
learning_rate = learning_rate,
lora_alpha = alpha,
lora_target_modules = targets,
max_grad_norm = float(self.max_grad_norm),
hf_token = token or None,
)
def discover_image_caption_pairs(
data_dir: str | os.PathLike[str],
*,
instance_prompt: Optional[str] = None,
caption_column: str = "text",
) -> list[tuple[str, str]]:
"""Resolve ``(image_path, caption)`` pairs from a dataset directory.
Caption sources, in priority order per image:
1. a ``metadata.jsonl`` / ``captions.jsonl`` row keyed by ``file_name`` (or ``image``)
carrying the caption in ``caption_column`` (default ``text``),
2. a per-image sidecar ``<stem>.txt`` / ``<stem>.caption``,
3. ``instance_prompt`` (dreambooth) for any remaining image.
Images with no caption from any source are skipped. Pure filesystem + JSON, so it is
unit-testable without torch. Raises FileNotFoundError for a missing dir and ValueError
when nothing is captionable.
"""
root = Path(data_dir).expanduser()
if not root.is_dir():
raise FileNotFoundError(f"data_dir is not a directory: {data_dir}")
images = sorted(p for p in root.iterdir() if p.is_file() and p.suffix.lower() in _IMAGE_EXTS)
# 1. metadata.jsonl / captions.jsonl (either name accepted).
meta_caption: dict[str, str] = {}
for meta_name in ("metadata.jsonl", "captions.jsonl"):
meta_path = root / meta_name
if not meta_path.is_file():
continue
for line in meta_path.read_text(encoding = "utf-8").splitlines():
line = line.strip()
if not line:
continue
try:
row = json.loads(line)
except json.JSONDecodeError:
continue
key = row.get("file_name") or row.get("image") or row.get("file")
cap = row.get(caption_column)
if isinstance(key, str) and isinstance(cap, str) and cap.strip():
meta_caption[Path(key).name] = cap.strip()
break # first present metadata file wins
pairs: list[tuple[str, str]] = []
for img in images:
caption = meta_caption.get(img.name)
if caption is None:
for ext in _CAPTION_EXTS:
sidecar = img.with_suffix(ext)
if sidecar.is_file():
text = sidecar.read_text(encoding = "utf-8").strip()
if text:
caption = text
break
if caption is None and instance_prompt and instance_prompt.strip():
caption = instance_prompt.strip()
if caption:
pairs.append((str(img), caption))
if not pairs:
raise ValueError(
f"No captioned images found under {data_dir}. Provide captions via a "
f"metadata.jsonl, per-image .txt sidecars, or an instance_prompt."
)
return pairs
def compute_sdxl_add_time_ids(resolution: int) -> tuple[int, int, int, int, int, int]:
"""SDXL micro-conditioning ``add_time_ids`` for a square ``resolution`` train crop:
@ -218,11 +58,6 @@ def compute_sdxl_add_time_ids(resolution: int) -> tuple[int, int, int, int, int,
return (resolution, resolution, 0, 0, resolution, resolution)
def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None:
if on_event is not None:
on_event({"type": type_, "ts": time.time(), **kw})
def _load_image_tensor(
path: str, resolution: int, center_crop: bool, random_flip: bool, rng: random.Random
) -> tuple[Any, tuple[int, int, int, int, int, int]]:
@ -290,45 +125,6 @@ def _encode_sdxl_prompts(
return prompt_embeds, pooled
def assert_trainable_base_model(base_model: str) -> None:
"""Refuse base models that are recognisably not SDXL, before anything is downloaded.
Purely name-based: a GGUF filename or a known DiT-family name (FLUX / Qwen-Image /
Z-Image / SD3 / ...) can never train on the SDXL U-Net trainer, so failing here turns
a confusing mid-run crash into an immediate, actionable error. Names this cannot
classify pass through; from_pretrained still fails cleanly on a genuinely wrong pick."""
name = str(base_model or "").strip().lower()
if name.endswith(".gguf"):
raise ValueError(
f"'{base_model}' is a GGUF checkpoint, which can't be trained. {_ONLY_SDXL_HINT}"
)
condensed = re.sub(r"[^a-z0-9]+", "-", name)
hit = next(
(p for p in _NON_SDXL_PHRASES if p in condensed),
None,
) or next(
(t for t in condensed.split("-") if t in _NON_SDXL_TOKENS),
None,
)
if hit:
raise ValueError(
f"'{base_model}' looks like a {hit} model, which isn't trainable. {_ONLY_SDXL_HINT}"
)
def _assert_trusted_base_model(base_model: str) -> None:
"""Gate the training base model the same way the inference backend gates non-GGUF loads:
a local path or a trusted repo (``unsloth/*`` or an allowlisted official base). This runs
BEFORE ``from_pretrained`` so an untrusted remote repo (which could ship pickle weights)
is never fetched or deserialised."""
from core.inference.diffusion import _is_trusted_diffusion_repo
if not _is_trusted_diffusion_repo(base_model):
raise ValueError(
f"Refusing to train from untrusted base model '{base_model}'. Use a local path or "
f"a trusted repo (an unsloth/* repo or an official SDXL base)."
)
def run_diffusion_lora_training(
config: DiffusionLoraConfig,
*,
@ -456,6 +252,9 @@ def run_diffusion_lora_training(
stopped = False
micro = 0
running_loss = 0.0
peak_gb = 0.0
t_start = time.time()
done = 0
for opt_step in range(cfg.train_steps):
optimizer.zero_grad(set_to_none = True)
step_loss = 0.0
@ -527,6 +326,13 @@ def run_diffusion_lora_training(
# ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so
# these progress events are directly consumable by the existing training
# status/SSE machinery when the diffusion trainer is wired into the worker.
if device == "cuda":
peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
samples_per_second = round(
(done * cfg.train_batch_size * cfg.gradient_accumulation_steps)
/ max(time.time() - t_start, 1e-6),
3,
)
_emit(
on_event,
"progress",
@ -535,6 +341,8 @@ def run_diffusion_lora_training(
loss = round(step_loss, 5),
avg_loss = round(running_loss / done, 5),
learning_rate = lr_sched.get_last_lr()[0],
samples_per_second = samples_per_second,
peak_memory_gb = peak_gb or None,
)
if _check_stop():
@ -565,38 +373,19 @@ def run_diffusion_lora_training(
output_dir = str(out_dir),
lora_path = lora_path,
catalog_path = catalog_path,
family = cfg.resolved_family,
base_model = cfg.base_model,
stopped = stopped,
steps_run = done if cfg.train_steps else 0,
)
return str(out_dir)
def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Optional[str]:
"""Best-effort copy of the trained adapter into the Studio diffusion LoRA directory so
the Images LoRA picker (which scans only files directly under ``loras/diffusion``) finds
it without the user moving files. Returns the published path, or None on any failure."""
try:
import shutil
from core.inference.diffusion_lora import loras_dir, sanitize_alias
base = (
cfg.adapter_name
if cfg.adapter_name and cfg.adapter_name != "default"
else Path(cfg.output_dir).name
)
dest = loras_dir() / f"{sanitize_alias(base)}.safetensors"
if Path(lora_path).resolve() != dest.resolve():
shutil.copy2(lora_path, dest)
return str(dest)
except Exception: # noqa: BLE001 -- the catalog mirror is best-effort, never fatal
return None
def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
"""mp.Queue subprocess adapter: run a training job, translating the ``on_event``
callback to ``event_queue`` and a ``stop_queue`` poll to ``should_stop``. Any
unexpected exception is reported as an ``error`` event rather than crashing silently."""
callback to ``event_queue`` and a ``stop_queue`` poll to ``should_stop``. Dispatches to
the trainer registered for the resolved family. Any unexpected exception is reported as
an ``error`` event rather than crashing silently."""
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
def on_event(ev: dict) -> None:
@ -616,8 +405,11 @@ def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config:
return got if saw else False
try:
cfg = _config_from_dict(config)
run_diffusion_lora_training(cfg, on_event = on_event, should_stop = should_stop)
# normalized() resolves + validates the family; dispatch through the registry so a
# DiT family runs its own trainer while SDXL keeps this module's loop.
cfg = _config_from_dict(config).normalized()
trainer = get_trainer(cfg.resolved_family)
trainer(cfg, on_event = on_event, should_stop = should_stop)
except Exception as exc: # noqa: BLE001 -- surfaced to the parent as an error event
# Emit both keys: the diffusion service reads ``message``, but the generic Studio
# training worker reads ``error``; carrying both keeps the real failure visible on
@ -627,50 +419,6 @@ def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config:
)
# 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).
_CONFIG_ALIASES = {
"model_name": "base_model",
"max_steps": "train_steps",
"batch_size": "train_batch_size",
"lora_r": "lora_rank",
"lr_scheduler_type": "lr_scheduler",
"random_seed": "seed",
"lr": "learning_rate",
}
def _coerce_gradient_checkpointing(value: Any) -> bool:
"""Studio sends gradient_checkpointing as a string ("none" / "true" / "unsloth"); the
disable words are False, anything else truthy is True. A real bool passes through."""
if isinstance(value, str):
return value.strip().lower() not in ("", "none", "false", "0", "no", "off")
return bool(value)
def _config_from_dict(config: dict) -> DiffusionLoraConfig:
"""Build a DiffusionLoraConfig from a plain dict. Unknown keys are ignored so a richer
request payload (UI form) does not break construction; a small set of generic Studio
training keys are aliased onto the diffusion field names, and string flags are coerced."""
valid = DiffusionLoraConfig.__dataclass_fields__.keys()
kwargs: dict[str, Any] = {}
# Aliases first (lowest priority); a canonical key present in the payload overrides.
for src, dst in _CONFIG_ALIASES.items():
if src in config and config[src] is not None and dst in valid:
kwargs[dst] = config[src]
for k, v in config.items():
if k in valid:
kwargs[k] = v
if kwargs.get("lora_target_modules"):
kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"])
if "gradient_checkpointing" in kwargs:
kwargs["gradient_checkpointing"] = _coerce_gradient_checkpointing(
kwargs["gradient_checkpointing"]
)
return DiffusionLoraConfig(**kwargs)
def main(argv: Optional[list[str]] = None) -> int:
p = argparse.ArgumentParser(description = "Train an SDXL LoRA and export it.")
p.add_argument("--base-model", required = True, help = "HF repo or local path to an SDXL pipeline")

392
studio/backend/core/training/diffusion_train_common.py Normal file
View file

@ -0,0 +1,392 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared, family-agnostic building blocks for diffusion LoRA training.
Everything here is independent of a specific model architecture, so the SDXL U-Net
trainer (``diffusion_lora_trainer``) and the flow-matching DiT trainers share it:
dataset discovery, the request config + validation, image loading, event emission,
the stop protocol, adapter publishing, and the family/trainer registry.
The pure helpers have no torch/diffusers import at call time and are unit-tested
without a GPU. ``run_diffusion_lora_training`` and its per-family siblings own the
actual training loop; this module only routes a request to the right one.
"""
from __future__ import annotations
import json
import os
import re
import time
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Callable, Optional
from core.inference.diffusion_families import (
detect_family,
detect_family_for_pick,
supported_family_names,
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_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
_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.
_NON_TRAINABLE_RESIDUAL_TOKENS = frozenset({"sd3", "pixart", "sana", "lumina", "cogview"})
_NON_TRAINABLE_RESIDUAL_PHRASES = ("stable-diffusion-3", "hunyuan-dit")
EventCb = Callable[[dict[str, Any]], None]
# Returns a falsy value to keep training, or a truthy stop signal: bare True, or a dict
# that may carry ``save=False`` to cancel without saving a partial adapter.
StopCb = Callable[[], Any]
def _trainable_hint() -> str:
"""A user-facing hint listing the families Studio can train today. Always names SDXL
explicitly so the message is actionable even as more families become trainable."""
names = ", ".join(trainable_family_names()) or "sdxl"
return (
f"Trainable families right now: {names} "
f"(for example the SDXL base stabilityai/stable-diffusion-xl-base-1.0). "
f"Other families can load LoRAs but not train them yet."
)
def resolve_trainable_family(base_model: str, model_family: Optional[str] = None) -> str:
"""Resolve the trainer family for a base model, or raise ValueError with a clear reason.
Positive resolution, before anything is downloaded:
- a ``.gguf`` name can never be a training base -> reject;
- an explicit ``model_family`` must name a registry family, and that family must be
trainable;
- otherwise the family is detected from the base-model name; a KNOWN but non-trainable
family (e.g. a DiT family before its trainer ships) is rejected;
- a name that resolves to no registry family but matches a known non-trainable
architecture (SD3 / PixArt / ...) is rejected;
- an unclassifiable custom name/path falls through to the SDXL trainer (backwards
compatible: a genuinely wrong pick still fails cleanly later in from_pretrained).
"""
name = str(base_model or "").strip().lower()
if name.endswith(".gguf"):
raise ValueError(
f"'{base_model}' is a GGUF checkpoint, which can't be trained. {_trainable_hint()}"
)
if model_family and str(model_family).strip():
key = str(model_family).strip().lower()
fam = detect_family("", override = key)
if fam is None:
known = ", ".join(supported_family_names())
raise ValueError(f"Unknown model_family {model_family!r}. Known families: {known}.")
if not fam.trainable:
raise ValueError(
f"'{fam.name}' models can't be trained yet. {_trainable_hint()}"
)
return fam.name
fam = detect_family_for_pick(base_model)
if fam is not None:
if not fam.trainable:
raise ValueError(
f"'{base_model}' looks like a {fam.name} model, which isn't trainable yet. "
f"{_trainable_hint()}"
)
return fam.name
condensed = re.sub(r"[^a-z0-9]+", "-", name)
hit = next(
(p for p in _NON_TRAINABLE_RESIDUAL_PHRASES if p in condensed),
None,
) or next(
(t for t in condensed.split("-") if t in _NON_TRAINABLE_RESIDUAL_TOKENS),
None,
)
if hit:
raise ValueError(
f"'{base_model}' looks like a {hit} model, which isn't trainable. {_trainable_hint()}"
)
# Unknown custom name / local path: default to the SDXL trainer (unchanged behaviour).
return "sdxl"
def get_trainer(family: str) -> Callable[..., str]:
"""Return the training entrypoint for ``family``. Imports the trainer module lazily so
this shared module stays free of the heavy trainer imports (and any import cycle)."""
key = (family or "sdxl").strip().lower()
if key == "sdxl":
from core.training.diffusion_lora_trainer import run_diffusion_lora_training
return run_diffusion_lora_training
raise ValueError(f"No trainer is registered for family {family!r}.")
@dataclass
class DiffusionLoraConfig:
"""Everything a diffusion LoRA training run needs. Only ``base_model`` /
``data_dir`` / ``output_dir`` are required; the rest have sensible defaults."""
base_model: str
data_dir: str
output_dir: str
# Dreambooth-style caption applied to any image without its own caption. Required if
# the dataset has no captions.jsonl / sidecar files.
instance_prompt: Optional[str] = None
resolution: int = 1024
train_steps: int = 500
learning_rate: float = 1e-4
train_batch_size: int = 1
gradient_accumulation_steps: int = 1
lora_rank: int = 16
lora_alpha: Optional[int] = None # defaults to lora_rank
lora_dropout: float = 0.0
lora_target_modules: tuple[str, ...] = DEFAULT_LORA_TARGETS
seed: int = 42
mixed_precision: str = "bf16" # "bf16" | "fp16" | "no"
snr_gamma: Optional[float] = 5.0 # min-SNR loss weighting; None disables
gradient_checkpointing: bool = True
max_grad_norm: float = 1.0
lr_scheduler: str = "constant"
lr_warmup_steps: int = 0
center_crop: bool = False
random_flip: bool = True
caption_column: str = "text" # column in metadata.jsonl
adapter_name: str = "default"
hf_token: Optional[str] = None
# 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.
model_family: Optional[str] = None
resolved_family: str = "sdxl"
def normalized(self) -> "DiffusionLoraConfig":
"""Return a copy with derived/validated fields filled in. Raises ValueError on a
request that cannot train (bad numbers, or an untrainable base model).
Also coerces values that arrive as strings/blanks through the Studio config path
(``learning_rate`` is preserved as a string there; ``hf_token`` defaults to "")."""
resolved_family = resolve_trainable_family(self.base_model, self.model_family)
if self.train_steps < 1:
raise ValueError("train_steps must be >= 1")
if self.train_batch_size < 1:
raise ValueError("train_batch_size must be >= 1")
if self.gradient_accumulation_steps < 1:
raise ValueError("gradient_accumulation_steps must be >= 1")
if self.lora_rank < 1:
raise ValueError("lora_rank must be >= 1")
if self.lora_alpha is not None and self.lora_alpha < 1:
raise ValueError(
"lora_alpha must be >= 1 (a zero/negative alpha scales the adapter to nothing)"
)
if self.resolution < 64 or self.resolution % 8 != 0:
raise ValueError("resolution must be a multiple of 8 and >= 64")
if self.mixed_precision not in ("bf16", "fp16", "no"):
raise ValueError("mixed_precision must be one of bf16 / fp16 / no")
# learning_rate can arrive as a string ("1e-4") from the Studio config path, which
# preserves it as a string after validation; coerce so AdamW receives a float.
try:
learning_rate = float(self.learning_rate)
except (TypeError, ValueError) as exc:
raise ValueError(f"learning_rate must be a number, got {self.learning_rate!r}") from exc
if learning_rate <= 0:
raise ValueError("learning_rate must be > 0")
alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank
targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS
# A blank Hub token (the Studio default when none is configured) must load
# anonymously, not as an explicit empty credential.
token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token
return replace(
self,
learning_rate = learning_rate,
lora_alpha = alpha,
lora_target_modules = targets,
max_grad_norm = float(self.max_grad_norm),
hf_token = token or None,
resolved_family = resolved_family,
)
def discover_image_caption_pairs(
data_dir: str | os.PathLike[str],
*,
instance_prompt: Optional[str] = None,
caption_column: str = "text",
) -> list[tuple[str, str]]:
"""Resolve ``(image_path, caption)`` pairs from a dataset directory.
Caption sources, in priority order per image:
1. a ``metadata.jsonl`` / ``captions.jsonl`` row keyed by ``file_name`` (or ``image``)
carrying the caption in ``caption_column`` (default ``text``),
2. a per-image sidecar ``<stem>.txt`` / ``<stem>.caption``,
3. ``instance_prompt`` (dreambooth) for any remaining image.
Images with no caption from any source are skipped. Pure filesystem + JSON, so it is
unit-testable without torch. Raises FileNotFoundError for a missing dir and ValueError
when nothing is captionable.
"""
root = Path(data_dir).expanduser()
if not root.is_dir():
raise FileNotFoundError(f"data_dir is not a directory: {data_dir}")
images = sorted(p for p in root.iterdir() if p.is_file() and p.suffix.lower() in _IMAGE_EXTS)
# 1. metadata.jsonl / captions.jsonl (either name accepted).
meta_caption: dict[str, str] = {}
for meta_name in ("metadata.jsonl", "captions.jsonl"):
meta_path = root / meta_name
if not meta_path.is_file():
continue
for line in meta_path.read_text(encoding = "utf-8").splitlines():
line = line.strip()
if not line:
continue
try:
row = json.loads(line)
except json.JSONDecodeError:
continue
key = row.get("file_name") or row.get("image") or row.get("file")
if key and caption_column in row:
meta_caption[str(key)] = str(row[caption_column])
pairs: list[tuple[str, str]] = []
for img in images:
caption: Optional[str] = None
# 1. metadata row keyed by file name (basename or the name as written).
caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root)))
# 2. per-image sidecar caption file.
if caption is None:
for ext in _CAPTION_EXTS:
sidecar = img.with_suffix(ext)
if sidecar.is_file():
caption = sidecar.read_text(encoding = "utf-8").strip()
break
# 3. dreambooth instance prompt.
if caption is None and instance_prompt:
caption = instance_prompt
if caption:
pairs.append((str(img), caption))
if not pairs:
raise ValueError(
"No captioned images found. Provide a metadata.jsonl / captions.jsonl, per-image "
".txt captions, or an instance prompt."
)
return pairs
def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None:
if on_event is not None:
on_event({"type": type_, "ts": time.time(), **kw})
def _assert_trusted_base_model(base_model: str) -> None:
"""Gate the training base model the same way the inference backend gates non-GGUF loads:
a local path or a trusted repo (``unsloth/*`` or an allowlisted official base). This runs
BEFORE ``from_pretrained`` so an untrusted remote repo (which could ship pickle weights)
is never fetched or deserialised."""
from core.inference.diffusion import _is_trusted_diffusion_repo
if not _is_trusted_diffusion_repo(base_model):
raise ValueError(
f"Refusing to train from untrusted base model '{base_model}'. Use a local path or "
f"a trusted repo (an unsloth/* repo or an official base)."
)
def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Optional[str]:
"""Best-effort copy of the trained adapter into the Studio diffusion LoRA directory so
the Images LoRA picker (which scans only files directly under ``loras/diffusion``) finds
it without the user moving files. Also writes a ``<alias>.json`` metadata sidecar so the
picker can family-gate the adapter (family, base model, trigger prompt, ...). Returns the
published path, or None on any failure."""
try:
import shutil
from core.inference.diffusion_lora import loras_dir, sanitize_alias
base = (
cfg.adapter_name
if cfg.adapter_name and cfg.adapter_name != "default"
else Path(cfg.output_dir).name
)
alias = sanitize_alias(base)
dest = loras_dir() / f"{alias}.safetensors"
if Path(lora_path).resolve() != dest.resolve():
shutil.copy2(lora_path, dest)
_write_lora_sidecar(dest.with_suffix(".json"), cfg)
return str(dest)
except Exception: # noqa: BLE001 -- the catalog mirror is best-effort, never fatal
return None
def _write_lora_sidecar(sidecar_path: Path, cfg: DiffusionLoraConfig) -> None:
"""Write the adapter metadata sidecar read back by diffusion_lora._scan_local. Best
effort: a failure here must not fail publishing, so callers wrap it."""
meta = {
"family": cfg.resolved_family,
"families": [cfg.resolved_family],
"base_model": cfg.base_model,
"lora_rank": cfg.lora_rank,
"lora_alpha": cfg.lora_alpha,
"steps": cfg.train_steps,
"resolution": cfg.resolution,
"trigger_prompt": cfg.instance_prompt,
"created_at": time.time(),
"source": "studio-trained",
}
sidecar_path.write_text(json.dumps(meta, indent = 2), encoding = "utf-8")
# 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).
_CONFIG_ALIASES = {
"model_name": "base_model",
"max_steps": "train_steps",
"batch_size": "train_batch_size",
"lora_r": "lora_rank",
"lr_scheduler_type": "lr_scheduler",
"random_seed": "seed",
"lr": "learning_rate",
}
def _coerce_gradient_checkpointing(value: Any) -> bool:
"""Studio sends gradient_checkpointing as a string ("none" / "true" / "unsloth"); the
disable words are False, anything else truthy is True. A real bool passes through."""
if isinstance(value, str):
return value.strip().lower() not in ("", "none", "false", "0", "no", "off")
return bool(value)
def _config_from_dict(config: dict) -> DiffusionLoraConfig:
"""Build a DiffusionLoraConfig from a plain dict. Unknown keys are ignored so a richer
request payload (UI form) does not break construction; a small set of generic Studio
training keys are aliased onto the diffusion field names, and string flags are coerced."""
valid = DiffusionLoraConfig.__dataclass_fields__.keys()
kwargs: dict[str, Any] = {}
# Aliases first (lowest priority); a canonical key present in the payload overrides.
for src, dst in _CONFIG_ALIASES.items():
if src in config and config[src] is not None and dst in valid:
kwargs[dst] = config[src]
for k, v in config.items():
if k in valid:
kwargs[k] = v
if kwargs.get("lora_target_modules"):
kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"])
if "gradient_checkpointing" in kwargs:
kwargs["gradient_checkpointing"] = _coerce_gradient_checkpointing(
kwargs["gradient_checkpointing"]
)
return DiffusionLoraConfig(**kwargs)