diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 53d83ac4a4..28eafe6829 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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", } ) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 760f00f0ea..ac2140d819 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index 7db07dcd60..cad6d3eae7 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -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 ``.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 ``.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. diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index b213626b54..795bf191d2 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -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 ``.txt`` / ``.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") diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py new file mode 100644 index 0000000000..74a33331f1 --- /dev/null +++ b/studio/backend/core/training/diffusion_train_common.py @@ -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 ``.txt`` / ``.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 ``.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)