diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 2184a955a6..ee8222d2f7 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -141,6 +141,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "FluxTransformer2DModel", base_repo = "black-forest-labs/FLUX.1-schnell", aliases = ("flux1", "flux-1"), + # LoRA training targets the guidance-distilled FLUX.1-dev via the DiT trainer + # (QLoRA nf4). The dev repo is gated on the Hub, so a user HF token is required. + trainable = True, + train_base_repos = ("black-forest-labs/FLUX.1-dev",), img2img_pipeline_class = "FluxImg2ImgPipeline", inpaint_pipeline_class = "FluxInpaintPipeline", controlnet_pipeline_class = "FluxControlNetPipeline", @@ -238,6 +242,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( base_repo = "Qwen/Qwen-Image", cfg_kwarg = "true_cfg_scale", aliases = ("qwen_image", "qwenimage"), + # LoRA training via the DiT trainer, defaulting to the prequant nf4 repo (QLoRA). + trainable = True, + train_base_repos = ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "Qwen/Qwen-Image"), img2img_pipeline_class = "QwenImageImg2ImgPipeline", inpaint_pipeline_class = "QwenImageInpaintPipeline", controlnet_pipeline_class = "QwenImageControlNetPipeline", @@ -262,6 +269,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "ZImageTransformer2DModel", base_repo = "Tongyi-MAI/Z-Image-Turbo", aliases = ("zimage", "z_image"), + # LoRA training via the DiT trainer (bf16 only). Defaults to the prequant nf4 repo + # for QLoRA; the bf16 Tongyi-MAI base is the alternative. + trainable = True, + train_base_repos = ("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", "Tongyi-MAI/Z-Image-Turbo"), img2img_pipeline_class = "ZImageImg2ImgPipeline", inpaint_pipeline_class = "ZImageInpaintPipeline", # Z-Image's MLP down-projections peak near 9e5, which overflows float16. @@ -309,7 +320,6 @@ 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/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py new file mode 100644 index 0000000000..82b65ee5ad --- /dev/null +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -0,0 +1,706 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Flow-matching LoRA training for the DiT image families (FLUX.1-dev, Qwen-Image, Z-Image). + +These are rectified-flow transformers, not the SDXL U-Net, so they share only the plumbing +in ``diffusion_train_common`` (config, dataset discovery, events, stop, publishing). The +training math here is flow matching: sample a sigma with the logit-normal density used by +the diffusers dreambooth scripts, form ``noisy = (1 - sigma) * latents + sigma * noise``, +predict the velocity, and regress it onto ``target = noise - latents``. + +The per-family differences (latent normalisation + packing, the transformer forward +signature, and the LoRA save entrypoint) live in small ``_FamilySpec`` objects; the loop +itself is family-agnostic. Verified against diffusers 0.38.0. + +Memory: the text encoder(s) are the largest module (T5-XXL ~9 GB for FLUX, Qwen2.5-VL ~7 GB +for Qwen-Image, Qwen3 for Z-Image), so captions are encoded ONCE up front and the encoders +are freed before the loop. The transformer trains as a QLoRA (nf4) adapter by default with +gradient checkpointing and 8-bit AdamW, so only the (small) LoRA params + optimizer state +and the frozen 4-bit base sit in VRAM during the loop. +""" + +from __future__ import annotations + +import gc +import random +import time +from contextlib import nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Optional + +from core.training.diffusion_train_common import ( + DEFAULT_LORA_FILENAME, + DiffusionLoraConfig, + EventCb, + StopCb, + _assert_trusted_base_model, + _emit, + _publish_to_lora_catalog, + discover_image_caption_pairs, +) + +# 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. +_FLUX_TARGETS = ( + "to_q", + "to_k", + "to_v", + "to_out.0", + "add_q_proj", + "add_k_proj", + "add_v_proj", + "to_add_out", +) +_QWEN_TARGETS = _FLUX_TARGETS +_ZIMAGE_TARGETS = ("to_q", "to_k", "to_v", "to_out.0") + + +@dataclass +class _FamilySpec: + """Everything the shared loop needs that differs by family.""" + + family: str + lora_targets: tuple[str, ...] + # bf16 only (Z-Image overflows fp16 and its RoPE/embedder run in fp32). + force_bf16: bool + # Builds (pipe, transformer, vae) with the transformer loaded as a trainable nf4 QLoRA + # when qlora=True. Returns the pipeline (for save_lora_weights + encode_prompt), the + # transformer to attach LoRA to, and the VAE (kept resident for latent encoding). + load: Callable[..., tuple[Any, Any, 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] + # One transformer forward: (transformer, noisy, timesteps, sigmas, embeds_batch, cfg, + # device, weight_dtype) -> model_pred aligned with target = noise - latents. + forward: Callable[..., Any] + # Save the LoRA in diffusers format via the family pipeline's save_lora_weights. + save: Callable[..., None] + + +# ── shared flow-matching helpers ────────────────────────────────────────────── +def _get_sigmas(scheduler, timesteps, device, dtype, n_dim): + """Gather per-sample sigmas for ``timesteps`` and broadcast to ``n_dim`` (matches the + diffusers dreambooth get_sigmas helper).""" + import torch + + sigmas = scheduler.sigmas.to(device = device, dtype = dtype) + schedule_timesteps = scheduler.timesteps.to(device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + sigma = sigmas[step_indices].flatten() + while sigma.ndim < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + +def _sample_timesteps(scheduler, batch_size, device): + """Logit-normal density timestep sampling (weighting_scheme='logit_normal'), returning + (timesteps, indices) into the scheduler's schedule.""" + import torch + from diffusers.training_utils import compute_density_for_timestep_sampling + + u = compute_density_for_timestep_sampling( + weighting_scheme = "logit_normal", + batch_size = batch_size, + logit_mean = 0.0, + logit_std = 1.0, + mode_scale = 1.29, + ) + num_train = scheduler.config.num_train_timesteps + indices = (u * num_train).long().clamp(0, num_train - 1) + timesteps = scheduler.timesteps.to(device)[indices].to(device) + return timesteps + + +def _encoders_to_device(pipe, device) -> None: + """Move the pipeline's (non-quantized) text encoders to ``device`` before encoding. + + A QLoRA FLUX load places the nf4 transformer on GPU but leaves the text encoders on + CPU, so encode_prompt would mix devices. Best-effort per encoder: a 4-bit encoder that + is already placed raises on .to() and is left as-is.""" + for attr in ("text_encoder", "text_encoder_2", "text_encoder_3"): + enc = getattr(pipe, attr, None) + if enc is None: + continue + try: + enc.to(device) + except (ValueError, RuntimeError, NotImplementedError): + pass # already-placed 4-bit encoder / non-movable module + + +def _bnb_4bit_config(): + from diffusers import BitsAndBytesConfig as DiffusersBnb + import torch + return DiffusersBnb( + load_in_4bit = True, + bnb_4bit_quant_type = "nf4", + bnb_4bit_compute_dtype = torch.bfloat16, + ) + + +def _repo_is_prequantized(base_model: str) -> bool: + """Heuristic: a repo whose name marks a bitsandbytes 4-bit build already ships a + quantized transformer, so we load it as-is rather than re-quantizing on the fly. A + dense (bf16) base instead gets on-the-fly nf4 quantization for QLoRA.""" + name = str(base_model or "").lower() + return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name + + +def _load_quantized_transformer(transformer_cls, cfg): + """Load ``cfg.base_model``'s transformer subfolder as a trainable nf4 QLoRA module.""" + import torch + return transformer_cls.from_pretrained( + cfg.base_model, + subfolder = "transformer", + quantization_config = _bnb_4bit_config(), + torch_dtype = torch.bfloat16, + token = cfg.hf_token, + ) + + +# ── FLUX.1-dev ──────────────────────────────────────────────────────────────── +def _flux_load(cfg, device, weight_dtype, qlora): + import torch + from diffusers import FluxPipeline, FluxTransformer2DModel + + if qlora: + transformer = FluxTransformer2DModel.from_pretrained( + cfg.base_model, + subfolder = "transformer", + quantization_config = _bnb_4bit_config(), + torch_dtype = torch.bfloat16, + token = cfg.hf_token, + ) + pipe = FluxPipeline.from_pretrained( + cfg.base_model, + transformer = transformer, + torch_dtype = torch.bfloat16, + token = cfg.hf_token, + ) + else: + pipe = FluxPipeline.from_pretrained( + cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token + ) + transformer = pipe.transformer + pipe.vae.to(device, dtype = torch.float32) + return pipe, transformer, pipe.vae + + +def _flux_encode_prompts(pipe, captions, device): + import torch + + _encoders_to_device(pipe, device) + out = [] + with torch.no_grad(): + for cap in captions: + pe, pooled, text_ids = pipe.encode_prompt( + prompt = cap, + prompt_2 = cap, + device = device, + num_images_per_prompt = 1, + max_sequence_length = 512, + ) + out.append((pe.cpu(), pooled.cpu(), text_ids.cpu())) + return out + + +def _flux_encode_latents(vae, pixel_values): + import torch + + with torch.no_grad(): + lat = vae.encode(pixel_values.to(torch.float32)).latent_dist.sample() + lat = (lat - vae.config.shift_factor) * vae.config.scaling_factor + return lat + + +def _flux_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): + import torch + from diffusers import FluxPipeline + + pe, pooled, text_ids = embeds_batch + bsz, c, h, w = noisy.shape + packed = FluxPipeline._pack_latents(noisy, bsz, c, h, w) + # Position ids drive RoPE and are indices, not activations -- keep them float32 (the + # dtype diffusers' own pipeline builds) regardless of the bf16 training dtype. + img_ids = FluxPipeline._prepare_latent_image_ids(bsz, h // 2, w // 2, device, torch.float32) + guidance = torch.full((bsz,), 1.0, device = device, dtype = torch.float32) + model_pred = transformer( + hidden_states = packed, + timestep = timesteps / 1000, + guidance = guidance, + pooled_projections = pooled.to(weight_dtype), + encoder_hidden_states = pe.to(weight_dtype), + txt_ids = text_ids.to(torch.float32), + img_ids = img_ids, + return_dict = False, + )[0] + return FluxPipeline._unpack_latents(model_pred, h * 8, w * 8, 8) + + +def _flux_save(pipe_cls, out_dir, transformer_lora_layers): + from diffusers import FluxPipeline + FluxPipeline.save_lora_weights( + save_directory = out_dir, + transformer_lora_layers = transformer_lora_layers, + weight_name = DEFAULT_LORA_FILENAME, + ) + + +# ── Qwen-Image ──────────────────────────────────────────────────────────────── +def _qwen_load(cfg, device, weight_dtype, qlora): + import torch + from diffusers import QwenImagePipeline, QwenImageTransformer2DModel + + # The prequant default (unsloth/Qwen-Image-2512-unsloth-bnb-4bit) ships the transformer + # 4-bit, so from_pretrained loads it trainable as-is. A dense (bf16) base -- the 20B + # Qwen/Qwen-Image -- is quantized to nf4 on the fly so QLoRA still fits. + kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token} + if qlora and not _repo_is_prequantized(cfg.base_model): + kwargs["transformer"] = _load_quantized_transformer(QwenImageTransformer2DModel, cfg) + pipe = QwenImagePipeline.from_pretrained(cfg.base_model, **kwargs) + pipe.vae.to(device, dtype = torch.float32) + return pipe, pipe.transformer, pipe.vae + + +def _qwen_encode_prompts(pipe, captions, device): + import torch + + _encoders_to_device(pipe, device) + out = [] + with torch.no_grad(): + for cap in captions: + pe, mask = pipe.encode_prompt( + prompt = cap, + device = device, + num_images_per_prompt = 1, + max_sequence_length = 1024, + ) + out.append((pe.cpu(), mask.cpu() if mask is not None else None)) + return out + + +def _qwen_encode_latents(vae, pixel_values): + import torch + + # AutoencoderKLQwenImage is a 3D (video) VAE: add a temporal dim, encode, drop it back + # into a [B,16,1,H,W] latent normalised by the per-channel latents_mean / latents_std. + px = pixel_values.to(torch.float32).unsqueeze(2) # [B,3,1,H,W] + with torch.no_grad(): + lat = vae.encode(px).latent_dist.sample() # [B,16,1,h,w] + z = vae.config.z_dim + mean = torch.tensor(vae.config.latents_mean, device = lat.device, dtype = lat.dtype) + std = torch.tensor(vae.config.latents_std, device = lat.device, dtype = lat.dtype) + mean = mean.view(1, z, 1, 1, 1) + std = std.view(1, z, 1, 1, 1) + return (lat - mean) / std + + +def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): + import torch + from diffusers import QwenImagePipeline + + pe, mask = embeds_batch + bsz, c, f, h, w = noisy.shape + packed = QwenImagePipeline._pack_latents(noisy, bsz, c, h, w) + # Each batch entry is a LIST of one (frame, h/2, w/2) tuple: the transformer indexes + # sample[0] / sample[1:] per entry (transformer_qwenimage.py), so a flat list breaks it. + img_shapes = [[(1, h // 2, w // 2)]] * bsz + pred = transformer( + hidden_states = packed, + encoder_hidden_states = pe.to(weight_dtype), + encoder_hidden_states_mask = mask.to(device) if mask is not None else None, + timestep = timesteps / 1000, + img_shapes = img_shapes, + return_dict = False, + )[0] + return QwenImagePipeline._unpack_latents(pred, h * 8, w * 8, 8) + + +def _qwen_save(pipe_cls, out_dir, transformer_lora_layers): + from diffusers import QwenImagePipeline + QwenImagePipeline.save_lora_weights( + save_directory = out_dir, + transformer_lora_layers = transformer_lora_layers, + weight_name = DEFAULT_LORA_FILENAME, + ) + + +# ── Z-Image ─────────────────────────────────────────────────────────────────── +def _zimage_load(cfg, device, weight_dtype, qlora): + import torch + from diffusers import ZImagePipeline, ZImageTransformer2DModel + + # Prequant default loads 4-bit as-is; the dense bf16 Tongyi-MAI base is quantized to nf4 + # on the fly. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows). + kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token} + if qlora and not _repo_is_prequantized(cfg.base_model): + kwargs["transformer"] = _load_quantized_transformer(ZImageTransformer2DModel, cfg) + pipe = ZImagePipeline.from_pretrained(cfg.base_model, **kwargs) + pipe.vae.to(device, dtype = torch.float32) + return pipe, pipe.transformer, pipe.vae + + +def _zimage_encode_prompts(pipe, captions, device): + import torch + + _encoders_to_device(pipe, device) + out = [] + with torch.no_grad(): + for cap in captions: + pe, _neg = pipe.encode_prompt( + prompt = cap, + device = device, + do_classifier_free_guidance = False, + max_sequence_length = 512, + ) + # pe is a list of one variable-length [seq, 2560] tensor per prompt. + emb = pe[0] if isinstance(pe, (list, tuple)) else pe + out.append((emb.cpu(),)) + return out + + +def _zimage_encode_latents(vae, pixel_values): + import torch + with torch.no_grad(): + lat = vae.encode(pixel_values.to(torch.float32)).latent_dist.mode() + return (lat - vae.config.shift_factor) * vae.config.scaling_factor + + +def _zimage_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype): + import torch + + (emb,) = embeds_batch + # List I/O: one [C,1,H,W] latent + one [seq,2560] caption per sample. The timestep + # convention is REVERSED ((1000 - t) / 1000) and the prediction is NEGATED. + x_list = list(noisy.unsqueeze(2).unbind(dim = 0)) + cap_list = [emb.to(device = device, dtype = weight_dtype)] + t_norm = (1000 - timesteps) / 1000 + out = transformer(x_list, t_norm, cap_list, return_dict = False)[0] + return -torch.stack(out, dim = 0).squeeze(2) + + +def _zimage_save(pipe_cls, out_dir, transformer_lora_layers): + from diffusers import ZImagePipeline + ZImagePipeline.save_lora_weights( + save_directory = out_dir, + transformer_lora_layers = transformer_lora_layers, + weight_name = DEFAULT_LORA_FILENAME, + ) + + +_SPECS: dict[str, _FamilySpec] = { + "flux.1": _FamilySpec( + family = "flux.1", + lora_targets = _FLUX_TARGETS, + force_bf16 = False, + load = _flux_load, + encode_prompts = _flux_encode_prompts, + encode_latents = _flux_encode_latents, + forward = _flux_forward, + save = _flux_save, + ), + "qwen-image": _FamilySpec( + family = "qwen-image", + lora_targets = _QWEN_TARGETS, + force_bf16 = True, + load = _qwen_load, + encode_prompts = _qwen_encode_prompts, + encode_latents = _qwen_encode_latents, + forward = _qwen_forward, + save = _qwen_save, + ), + "z-image": _FamilySpec( + family = "z-image", + lora_targets = _ZIMAGE_TARGETS, + force_bf16 = True, + load = _zimage_load, + encode_prompts = _zimage_encode_prompts, + encode_latents = _zimage_encode_latents, + forward = _zimage_forward, + save = _zimage_save, + ), +} + + +# 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. +_GATED_TRAIN_REPOS = frozenset({"black-forest-labs/flux.1-dev"}) + + +def _assert_gated_access(base_model: str, hf_token: Optional[str]) -> None: + """Raise a clear error before loading a gated base without a token.""" + name = str(base_model or "").strip().lower() + if name in _GATED_TRAIN_REPOS and not (hf_token and str(hf_token).strip()): + raise ValueError( + f"'{base_model}' is a gated Hugging Face repo. Accept its license on the Hub " + f"and add your HF token in Studio settings before training from it." + ) + + +def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng): + """Load an image -> a normalised [3,H,W] tensor in [-1,1]. Same geometry as the SDXL + loader but without the SDXL time-ids (DiT families don't use them).""" + import numpy as np + import torch + from PIL import Image, ImageOps + + img = ImageOps.exif_transpose(Image.open(path)).convert("RGB") + w0, h0 = img.size + scale = resolution / min(w0, h0) + rw, rh = max(resolution, round(w0 * scale)), max(resolution, round(h0 * scale)) + img = img.resize((rw, rh), Image.LANCZOS) + if center_crop: + left, top = (rw - resolution) // 2, (rh - resolution) // 2 + else: + left = rng.randint(0, max(0, rw - resolution)) + top = rng.randint(0, max(0, rh - resolution)) + img = img.crop((left, top, left + resolution, top + resolution)) + if random_flip and rng.random() < 0.5: + img = img.transpose(Image.FLIP_LEFT_RIGHT) + arr = np.asarray(img, dtype = np.float32) / 255.0 + return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0 + + +def run_dit_lora_training( + config: DiffusionLoraConfig, + *, + on_event: Optional[EventCb] = None, + should_stop: Optional[StopCb] = None, +) -> str: + """Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it.""" + import torch + import torch.nn.functional as F + from diffusers import FlowMatchEulerDiscreteScheduler + from diffusers.training_utils import cast_training_params + from peft import LoraConfig + from peft.utils import get_peft_model_state_dict + + cfg = config.normalized() + spec = _SPECS.get(cfg.resolved_family) + if spec is None: + raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}") + + rng = random.Random(cfg.seed) + torch.manual_seed(cfg.seed) + + save_on_stop = True + + def _check_stop() -> bool: + nonlocal save_on_stop + if should_stop is None: + return False + sig = should_stop() + if not sig: + return False + if isinstance(sig, dict) and sig.get("save") is False: + save_on_stop = False + return True + + # 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. + if cfg.mixed_precision == "fp16" and spec.force_bf16: + raise ValueError( + f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / " + f"embedder internals. Set mixed precision to bf16." + ) + + 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). + weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32 + use_lora_targets = tuple(cfg.lora_target_modules) or spec.lora_targets + + _assert_trusted_base_model(cfg.base_model) + _assert_gated_access(cfg.base_model, cfg.hf_token) + pairs = discover_image_caption_pairs( + cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column + ) + _emit(on_event, "model_load_started", num_images = len(pairs)) + if _check_stop(): + out_dir = Path(cfg.output_dir).expanduser() + _emit( + on_event, "complete", output_dir = str(out_dir), lora_path = None, stopped = True, steps_run = 0 + ) + return str(out_dir) + + # QLoRA by default for the big DiTs (nf4 transformer). The prequant Qwen/Z-Image repos + # are already 4-bit; FLUX quantizes its transformer on the fly. + pipe, transformer, vae = spec.load(cfg, device, weight_dtype, qlora = True) + + # Precompute all caption embeddings, then free the (large) text encoder(s): captions are + # constant and the encoders are frozen, so this is exact and the biggest memory win. + image_paths = [p for p, _ in pairs] + captions = [c for _, c in pairs] + uniq = sorted(set(captions)) + encoded = spec.encode_prompts(pipe, uniq, device) + caption_embeds = {cap: emb for cap, emb in zip(uniq, encoded)} + _free_text_encoders(pipe) + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + + # Freeze the base; attach the trainable LoRA to the transformer. + transformer.requires_grad_(False) + transformer.add_adapter( + LoraConfig( + r = cfg.lora_rank, + lora_alpha = cfg.lora_alpha, + lora_dropout = cfg.lora_dropout, + init_lora_weights = "gaussian", + target_modules = list(use_lora_targets), + ) + ) + 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). + import functools + import torch.utils.checkpoint as _ckpt + transformer.enable_gradient_checkpointing( + gradient_checkpointing_func = functools.partial(_ckpt.checkpoint, use_reentrant = False) + ) + cast_training_params(transformer, dtype = torch.float32) + lora_params = [p for p in transformer.parameters() if p.requires_grad] + + optimizer = _make_optimizer(lora_params, cfg.learning_rate) + scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + cfg.base_model, subfolder = "scheduler", token = cfg.hf_token + ) + + _emit(on_event, "model_load_completed") + + transformer.train() + stopped = False + 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 + for _ in range(cfg.gradient_accumulation_steps): + i = rng.randrange(len(image_paths)) + px = ( + _load_pixel_tensor( + image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng + ) + .unsqueeze(0) + .to(device) + ) + latents = spec.encode_latents(vae, px).to(weight_dtype) + + noise = torch.randn_like(latents) + timesteps = _sample_timesteps(scheduler, latents.shape[0], device) + sigmas = _get_sigmas(scheduler, timesteps, device, weight_dtype, latents.ndim) + noisy = (1.0 - sigmas) * latents + sigmas * noise + + emb = caption_embeds[captions[i]] + emb_dev = tuple( + t.to(device = device, dtype = weight_dtype) + if (t is not None and t.is_floating_point()) + else (t.to(device) if t is not None else None) + for t in emb + ) + # 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. + autocast = ( + torch.autocast(device_type = "cuda", dtype = torch.bfloat16) + if device == "cuda" + else nullcontext() + ) + with autocast: + model_pred = spec.forward( + transformer, noisy, timesteps, sigmas, emb_dev, cfg, device, weight_dtype + ) + target = noise - latents + loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean") + (loss / cfg.gradient_accumulation_steps).backward() + step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps + + if cfg.max_grad_norm and cfg.max_grad_norm > 0: + torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + optimizer.step() + + running_loss += step_loss + done = opt_step + 1 + if done % cfg.log_every == 0 or done == cfg.train_steps: + if device == "cuda": + peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) + sps = round( + (done * cfg.train_batch_size * cfg.gradient_accumulation_steps) + / max(time.time() - t_start, 1e-6), + 3, + ) + _emit( + on_event, + "progress", + step = done, + total_steps = cfg.train_steps, + loss = round(step_loss, 5), + avg_loss = round(running_loss / done, 5), + learning_rate = cfg.learning_rate, + samples_per_second = sps, + peak_memory_gb = peak_gb or None, + ) + if _check_stop(): + stopped = True + break + + out_dir = Path(cfg.output_dir).expanduser() + lora_path: Optional[str] = None + catalog_path: Optional[str] = None + if not (stopped and not save_on_stop): + out_dir.mkdir(parents = True, exist_ok = True) + layers = get_peft_model_state_dict(transformer) + spec.save(pipe, str(out_dir), layers) + lora_path = str(out_dir / DEFAULT_LORA_FILENAME) + catalog_path = _publish_to_lora_catalog(lora_path, cfg) + _emit( + on_event, + "complete", + 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 _make_optimizer(params, lr): + """8-bit AdamW (bitsandbytes) when available -- half the optimizer state, no accuracy + regression for LoRA -- else the torch AdamW fallback.""" + import torch + try: + import bitsandbytes as bnb + return bnb.optim.AdamW8bit(params, lr = lr) + except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW + return torch.optim.AdamW(params, lr = lr) + + +def _free_text_encoders(pipe) -> None: + """Drop every text-encoder / tokenizer the pipeline holds, so the (large) encoders do + not sit in VRAM during training. The embeddings are already precomputed.""" + for attr in ( + "text_encoder", + "text_encoder_2", + "text_encoder_3", + "tokenizer", + "tokenizer_2", + "tokenizer_3", + ): + if getattr(pipe, attr, None) is not None: + try: + setattr(pipe, attr, None) + except Exception: # noqa: BLE001 + pass diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 795bf191d2..deef4d03ce 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -25,6 +25,7 @@ Design: from __future__ import annotations import argparse +import gc import os import random import time @@ -226,7 +227,7 @@ def run_diffusion_lora_training( cast_training_params(unet, dtype = torch.float32) lora_params = [p for p in unet.parameters() if p.requires_grad] - optimizer = torch.optim.AdamW(lora_params, lr = cfg.learning_rate) + optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate) # The scheduler advances once per optimizer update: lr_sched.step() runs a single # time per outer opt_step (after the accumulation inner loop), for cfg.train_steps # total. Count warmup/decay in those optimizer steps -- multiplying by the @@ -241,6 +242,24 @@ def run_diffusion_lora_training( vae_scale = vae.config.scaling_factor prediction_type = noise_scheduler.config.prediction_type + # Precompute text embeddings once per unique caption, then free the CLIP text encoders. + # SDXL re-encoded captions every step (pure waste: captions are constant) and kept both + # text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no + # torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and + # lighter. The env toggle exists purely so the accuracy guard can A/B the two paths. + precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true") + caption_embeds: dict[str, tuple] = {} + if precompute: + for cap in sorted({c for _, c in pairs}): + pe, pooled_c = _encode_sdxl_prompts([cap], tokenizers, text_encoders, device) + caption_embeds[cap] = (pe.cpu(), pooled_c.cpu()) + for te in text_encoders: + te.to("cpu") + text_encoders = [] + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + _emit(on_event, "model_load_completed") def _next_batch() -> tuple[list[str], list[str]]: @@ -281,9 +300,13 @@ def run_diffusion_lora_training( ).long() noisy = noise_scheduler.add_noise(latents, noise, timesteps) - prompt_embeds, pooled = _encode_sdxl_prompts( - captions, tokenizers, text_encoders, device - ) + if precompute: + prompt_embeds = torch.cat([caption_embeds[c][0] for c in captions]).to(device) + pooled = torch.cat([caption_embeds[c][1] for c in captions]).to(device) + else: + prompt_embeds, pooled = _encode_sdxl_prompts( + captions, tokenizers, text_encoders, device + ) prompt_embeds = prompt_embeds.to(dtype = weight_dtype) pooled = pooled.to(dtype = weight_dtype) added = {"text_embeds": pooled, "time_ids": batch_time_ids} @@ -381,6 +404,21 @@ def run_diffusion_lora_training( return str(out_dir) +def _make_lora_optimizer(params: list, lr: float) -> Any: + """8-bit AdamW (bitsandbytes) by default -- half the optimizer state, no meaningful + quality cost for LoRA -- falling back to fp32 AdamW when unavailable or when + UNSLOTH_DIFFUSION_FP32_OPTIM is set (used by the accuracy guard).""" + import torch + + if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") not in ("1", "true"): + try: + import bitsandbytes as bnb + return bnb.optim.AdamW8bit(params, lr = lr) + except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW + pass + return torch.optim.AdamW(params, lr = lr) + + 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``. Dispatches to diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 3594a88724..f89b05ff80 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -80,9 +80,13 @@ 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() - if name.endswith(".gguf"): + # 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. + if name.endswith(".gguf") or "gguf" in name: raise ValueError( - f"'{base_model}' is a GGUF checkpoint, which can't be trained. {_trainable_hint()}" + f"'{base_model}' is a GGUF checkpoint/repo, which can't be a training base " + f"(training needs the full diffusers model). {_trainable_hint()}" ) if model_family and str(model_family).strip(): key = str(model_family).strip().lower() @@ -91,7 +95,9 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = 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()}") + raise ValueError( + f"'{fam.name}' models can't be trained yet. {_trainable_hint()}" + ) return fam.name fam = detect_family_for_pick(base_model) @@ -126,9 +132,72 @@ def get_trainer(family: str) -> Callable[..., str]: if key == "sdxl": from core.training.diffusion_lora_trainer import run_diffusion_lora_training return run_diffusion_lora_training + if key in ("flux.1", "qwen-image", "z-image"): + from core.training.diffusion_dit_trainer import run_dit_lora_training + return run_dit_lora_training 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. +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}, + "qwen-image": {"lora_rank": 16, "learning_rate": 5e-5, "resolution": 512}, + "z-image": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 768}, +} + + +def train_defaults(family: str) -> dict[str, Any]: + """Recommended starting hyperparameters for ``family`` (empty if unknown).""" + return dict(FAMILY_TRAIN_DEFAULTS.get((family or "").strip().lower(), {})) + + +# Display labels + a short VRAM/access note per trainable family, surfaced by the Train UI +# so users pick a base with realistic expectations. Kept next to the defaults they pair with. +_FAMILY_LABELS = { + "sdxl": "SDXL", + "flux.1": "FLUX.1-dev", + "qwen-image": "Qwen-Image", + "z-image": "Z-Image", +} +_FAMILY_VRAM_NOTES = { + "sdxl": "Trains on ~12 GB+ (bf16 LoRA). The lightest, fastest option.", + "flux.1": ( + "12B model, QLoRA (nf4) by default (~16 GB+). Gated on Hugging Face: accept the " + "FLUX.1-dev license and add your HF token before training." + ), + "qwen-image": "20B model, QLoRA (nf4) by default (~24 GB+). The heaviest option.", + "z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.", +} + + +def family_train_infos() -> list[dict[str, Any]]: + """Describe every trainable family for the Train UI: name, label, the default + allowed + base repos, the recommended starting hyperparameters, and a VRAM/access note. Built from + the family registry so it stays in sync with what the trainers actually support.""" + from core.inference.diffusion_families import detect_family + + infos: list[dict[str, Any]] = [] + for name in trainable_family_names(): + fam = detect_family("", override = name) + if fam is None: + continue + repos = list(fam.train_base_repos) or [fam.base_repo] + infos.append( + { + "name": name, + "label": _FAMILY_LABELS.get(name, name), + "default_base": repos[0], + "base_repos": repos, + "defaults": train_defaults(name), + "vram_note": _FAMILY_VRAM_NOTES.get(name, ""), + } + ) + return infos + + @dataclass class DiffusionLoraConfig: """Everything a diffusion LoRA training run needs. Only ``base_model`` / diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 8385bdebb2..b36661cfef 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -772,15 +772,25 @@ class DiffusionDatasetSummary(BaseModel): caption_count: int -class DiffusionTrainingInfoResponse(BaseModel): - """Where diffusion training reads/writes on this Studio, plus usable datasets. +class DiffusionTrainableFamily(BaseModel): + """A base-model family the diffusion trainer supports, with UI-facing metadata.""" - Lets the UI show real on-disk locations and offer existing dataset folders, - instead of asking users to know the Studio home layout.""" + name: str + label: str + default_base: str + base_repos: List[str] = Field(default_factory = list) + defaults: dict = Field(default_factory = dict) + vram_note: str = "" + + +class DiffusionTrainingInfoResponse(BaseModel): + """Where diffusion training reads/writes on this Studio, plus usable datasets and the + trainable model families (so the UI can offer a base picker with realistic guidance).""" datasets_root: str outputs_root: str datasets: List[DiffusionDatasetSummary] + families: List[DiffusionTrainableFamily] = Field(default_factory = list) class DiffusionDatasetUploadResponse(BaseModel): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c8741785fa..0eb7e00561 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -61,6 +61,7 @@ from models.training import ( DiffusionDatasetSummary, DiffusionDatasetUploadResponse, DiffusionMetricHistory, + DiffusionTrainableFamily, DiffusionTrainingInfoResponse, DiffusionTrainingStartRequest, DiffusionTrainingStartResponse, @@ -1125,6 +1126,44 @@ def _free_gpu_for_diffusion_training() -> None: logger.warning("Could not free chat models for diffusion training: %s", e) +def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None: + """HEAD a remote base repo's model_index.json with the caller's token; raise HTTP 400 on + 401/403 (gated / unauthorized) with an actionable message. Best-effort: a local path, + a non-repo string, or a network hiccup passes through so the trainer can surface any real + load error itself. Runs before GPU teardown so a doomed start never evicts a loaded model.""" + import urllib.error + import urllib.request + + repo = (base_model or "").strip() + # Only remote 'org/name' repos are gated; skip local paths and single-file names. + if ( + not repo + or repo.count("/") != 1 + or repo.startswith((".", "/", "~")) + or repo.endswith(".gguf") + ): + return + url = f"https://huggingface.co/{repo}/resolve/main/model_index.json" + headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {} + req = urllib.request.Request(url, method = "HEAD", headers = headers) + try: + urllib.request.urlopen(req, timeout = 5) + except urllib.error.HTTPError as e: + if e.code in (401, 403): + raise HTTPException( + status_code = 400, + detail = ( + f"Access to '{repo}' is gated or unauthorized. Accept the model's license " + f"on its Hugging Face page and add your HF token in Studio settings, then " + f"try again." + ), + ) + # 404 (e.g. a repo without a root model_index.json) and other codes are not an + # access problem -- let the trainer surface any genuine load error. + except Exception: # noqa: BLE001 -- network/DNS hiccup must not block a start + return + + @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse) async def start_diffusion_training( body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject) @@ -1169,8 +1208,13 @@ async def start_diffusion_training( except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) + # Preflight access to a gated base repo with the user's token BEFORE freeing GPU + # residents, so a missing/insufficient token fails fast (400) without tearing down the + # user's loaded chat/Images model, and never surfaces as a confusing mid-load 401. + _preflight_gated_base(config.get("base_model", ""), config.get("hf_token")) + # Free resident GPU workloads (export / Images pipeline / chat) before the trainer - # loads its own SDXL pipeline. + # loads its own pipeline. _free_gpu_for_diffusion_training() service = get_diffusion_training_service() @@ -1259,8 +1303,14 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub continue if summary.image_count > 0: found.append(summary) + from core.training.diffusion_train_common import family_train_infos + + families = [DiffusionTrainableFamily(**info) for info in family_train_infos()] return DiffusionTrainingInfoResponse( - datasets_root = str(root), outputs_root = str(outputs_root()), datasets = found + datasets_root = str(root), + outputs_root = str(outputs_root()), + datasets = found, + families = families, ) return await asyncio.to_thread(scan) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py new file mode 100644 index 0000000000..ba54480003 --- /dev/null +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the flow-matching DiT LoRA trainer (FLUX.1 / Qwen-Image / Z-Image). + +CPU-only: cover family resolution, the per-family spec table, the QLoRA prequant +heuristic, the bf16-only guard, and the gated-repo name check. The full training loop is +exercised by the live GPU smokes, not here.""" + +from __future__ import annotations + +import pytest + +from core.training.diffusion_dit_trainer import ( + _GATED_TRAIN_REPOS, + _SPECS, + _assert_gated_access, + _repo_is_prequantized, + run_dit_lora_training, +) +from core.training.diffusion_train_common import DiffusionLoraConfig, family_train_infos + + +def test_specs_cover_the_three_dit_families(): + assert set(_SPECS) == {"flux.1", "qwen-image", "z-image"} + # FLUX / Qwen share the added-kv attention target set; Z-Image is single-stream. + assert "add_q_proj" in _SPECS["flux.1"].lora_targets + assert "add_q_proj" in _SPECS["qwen-image"].lora_targets + assert "add_q_proj" not in _SPECS["z-image"].lora_targets + # Z-Image and Qwen are bf16-only. + assert _SPECS["z-image"].force_bf16 is True + assert _SPECS["qwen-image"].force_bf16 is True + + +@pytest.mark.parametrize( + "repo, expected", + [ + ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", True), + ("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", True), + ("some/model-int4", True), + ("black-forest-labs/FLUX.1-dev", False), + ("Tongyi-MAI/Z-Image-Turbo", False), + ], +) +def test_prequant_heuristic(repo, expected): + assert _repo_is_prequantized(repo) is expected + + +def test_zimage_rejects_fp16_before_loading(): + # bf16-only families must refuse an explicit fp16 request up front (no model load). + cfg = DiffusionLoraConfig( + base_model = "Tongyi-MAI/Z-Image-Turbo", + data_dir = "does-not-exist", + output_dir = "o", + mixed_precision = "fp16", + ) + with pytest.raises(ValueError, match = "bf16"): + run_dit_lora_training(cfg) + + +def test_gated_access_requires_token(): + assert "black-forest-labs/flux.1-dev" in _GATED_TRAIN_REPOS + # No token -> clear, actionable error before any download. + with pytest.raises(ValueError, match = "gated"): + _assert_gated_access("black-forest-labs/FLUX.1-dev", None) + with pytest.raises(ValueError, match = "gated"): + _assert_gated_access("black-forest-labs/FLUX.1-dev", " ") + # With a token, or for a non-gated repo, it is a no-op. + _assert_gated_access("black-forest-labs/FLUX.1-dev", "hf_realtoken") + _assert_gated_access("Tongyi-MAI/Z-Image-Turbo", None) + + +def test_family_train_infos_lists_dit_families(): + infos = {i["name"]: i for i in family_train_infos()} + for fam in ("sdxl", "flux.1", "qwen-image", "z-image"): + assert fam in infos, f"{fam} missing from family_train_infos" + assert infos[fam]["default_base"] + assert infos[fam]["base_repos"] + assert "resolution" in infos[fam]["defaults"] + # FLUX default base is the gated dev repo; its note flags the license requirement. + assert infos["flux.1"]["default_base"] == "black-forest-labs/FLUX.1-dev" + assert "gated" in infos["flux.1"]["vram_note"].lower() + # Z-Image defaults to the prequant nf4 repo for QLoRA. + assert "4bit" in infos["z-image"]["default_base"].lower() diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index dadd5559fd..15eb70f3a2 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -195,22 +195,31 @@ def test_config_rejects_nonpositive_learning_rate(): ).normalized() -def test_config_rejects_known_non_sdxl_base_models(): - # Known DiT families and GGUF checkpoints must fail at normalise time (an instant - # 400 via the API) instead of minutes later inside from_pretrained. +def test_config_rejects_untrainable_base_models(): + # GGUF checkpoints and families without a trainer (Kontext editing, SD3) must fail at + # normalise time (an instant 400 via the API), not minutes later inside from_pretrained. for bad in ( "unsloth/FLUX.1-dev-GGUF", - "black-forest-labs/FLUX.1-schnell", - "unsloth/Qwen-Image-2512-unsloth-bnb-4bit", - "Tongyi-MAI/Z-Image-Turbo", + "z-image-turbo-Q4_K_M.gguf", "stabilityai/stable-diffusion-3-medium", "unsloth/FLUX.1-Kontext-dev", - "z-image-turbo-Q4_K_M.gguf", ): - with pytest.raises(ValueError, match = "SDXL"): + with pytest.raises(ValueError): DiffusionLoraConfig(base_model = bad, data_dir = "d", output_dir = "o").normalized() +def test_config_resolves_dit_families(): + # FLUX.1 / Qwen-Image / Z-Image bases now resolve to their DiT trainer families. + for base, fam in ( + ("black-forest-labs/FLUX.1-dev", "flux.1"), + ("black-forest-labs/FLUX.1-schnell", "flux.1"), + ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "qwen-image"), + ("Tongyi-MAI/Z-Image-Turbo", "z-image"), + ): + cfg = DiffusionLoraConfig(base_model = base, data_dir = "d", output_dir = "o").normalized() + assert cfg.resolved_family == fam + + def test_config_accepts_sdxl_and_unknown_base_models(): # SDXL names and unclassifiable custom names/paths must pass the guard (a wrong # custom pick still fails cleanly in from_pretrained). @@ -227,14 +236,23 @@ def test_config_accepts_sdxl_and_unknown_base_models(): # ── trainer registry + family resolution + metadata sidecar (PR A platform) ── def test_get_trainer_resolves_sdxl(): from core.training.diffusion_lora_trainer import get_trainer, run_diffusion_lora_training + assert get_trainer("sdxl") is run_diffusion_lora_training assert get_trainer("SDXL") is run_diffusion_lora_training # case-insensitive def test_get_trainer_unknown_family_raises(): from core.training.diffusion_lora_trainer import get_trainer + with pytest.raises(ValueError, match = "No trainer"): - get_trainer("flux.1") # not registered until the DiT trainers ship + get_trainer("flux.2-dev") # a real family with no registered trainer + + +def test_get_trainer_resolves_dit_families(): + from core.training.diffusion_dit_trainer import run_dit_lora_training + from core.training.diffusion_lora_trainer import get_trainer + for fam in ("flux.1", "qwen-image", "z-image"): + assert get_trainer(fam) is run_dit_lora_training def test_normalized_sets_resolved_family(): @@ -242,9 +260,7 @@ def test_normalized_sets_resolved_family(): base_model = "stabilityai/stable-diffusion-xl-base-1.0", data_dir = "d", output_dir = "o" ).normalized() assert cfg.resolved_family == "sdxl" - cfg2 = DiffusionLoraConfig( - base_model = "my-custom-thing", data_dir = "d", output_dir = "o" - ).normalized() + cfg2 = DiffusionLoraConfig(base_model = "my-custom-thing", data_dir = "d", output_dir = "o").normalized() assert cfg2.resolved_family == "sdxl" # unknown -> default SDXL trainer @@ -254,9 +270,16 @@ def test_explicit_model_family_validated(): # A bogus explicit family is rejected up front. with pytest.raises(ValueError, match = "Unknown model_family"): C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "not-a-family").normalized() - # A known-but-not-yet-trainable family is rejected with the SDXL hint. - with pytest.raises(ValueError, match = "SDXL"): - C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1").normalized() + # A known-but-not-trainable family (Kontext editing) is rejected with a helpful hint. + with pytest.raises(ValueError): + C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1-kontext").normalized() + # A DiT family that IS trainable resolves to itself. + assert ( + C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1") + .normalized() + .resolved_family + == "flux.1" + ) # SDXL explicit passes. assert ( C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "sdxl") diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 0a273a5a95..6931418cde 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -484,15 +484,8 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat def test_apply_event_records_metric_history_and_perf(): svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) svc._apply_event( - { - "type": "progress", - "step": 1, - "total_steps": 10, - "loss": 0.5, - "learning_rate": 1e-4, - "samples_per_second": 3.2, - "peak_memory_gb": 7.1, - } + {"type": "progress", "step": 1, "total_steps": 10, "loss": 0.5, + "learning_rate": 1e-4, "samples_per_second": 3.2, "peak_memory_gb": 7.1} ) svc._apply_event( {"type": "progress", "step": 2, "total_steps": 10, "loss": 0.4, "learning_rate": 9e-5} @@ -538,14 +531,8 @@ def test_metric_history_decimates_at_cap(): def test_complete_event_records_family_and_catalog(): svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) svc._apply_event( - { - "type": "complete", - "output_dir": "/o", - "lora_path": "/o/w.safetensors", - "catalog_path": "/loras/w.safetensors", - "family": "sdxl", - "base_model": "b", - } + {"type": "complete", "output_dir": "/o", "lora_path": "/o/w.safetensors", + "catalog_path": "/loras/w.safetensors", "family": "sdxl", "base_model": "b"} ) st = svc.status() assert st["status"] == "completed" @@ -557,12 +544,8 @@ def test_complete_event_records_family_and_catalog(): def test_status_route_nests_metric_history(client): # The status route folds the service's flat arrays into a nested metric_history object. client._fake.status_extra = { - "metric_steps": [1, 2], - "metric_loss": [0.5, 0.4], - "metric_lr": [1e-4, 9e-5], - "family": "sdxl", - "samples_per_second": 2.0, - "peak_memory_gb": 6.0, + "metric_steps": [1, 2], "metric_loss": [0.5, 0.4], "metric_lr": [1e-4, 9e-5], + "family": "sdxl", "samples_per_second": 2.0, "peak_memory_gb": 6.0, } r = client.get("/api/train/diffusion/status") assert r.status_code == 200, r.text @@ -571,3 +554,55 @@ def test_status_route_nests_metric_history(client): assert body["metric_history"]["loss"] == [0.5, 0.4] assert body["family"] == "sdxl" assert body["samples_per_second"] == 2.0 + + +# ── /diffusion/info families + gated-repo preflight (PR B) ────────────────────── +def test_info_lists_trainable_families(client): + r = client.get("/api/train/diffusion/info") + assert r.status_code == 200, r.text + families = {f["name"]: f for f in r.json()["families"]} + for fam in ("sdxl", "flux.1", "qwen-image", "z-image"): + assert fam in families + assert families["flux.1"]["default_base"] == "black-forest-labs/FLUX.1-dev" + assert families["z-image"]["defaults"]["resolution"] == 768 + + +def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatch): + # A gated FLUX base with no valid token must 400 from the HEAD preflight BEFORE the GPU + # residents are freed, so a doomed start never evicts the user's loaded model. + import urllib.error + import urllib.request + + import routes.training as tr + + freed: list[int] = [] + monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1)) + + def _fake_urlopen(req, timeout = None): + raise urllib.error.HTTPError(req.full_url, 403, "Forbidden", {}, None) + + monkeypatch.setattr(urllib.request, "urlopen", _fake_urlopen) + r = client.post( + "/api/train/diffusion/start", + json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"}, + ) + assert r.status_code == 400 + assert "gated" in r.json()["detail"].lower() + assert freed == [] + assert client._fake.started_with is None + + +def test_start_ungated_base_preflight_is_noop(client, monkeypatch): + # A reachable base (HEAD 200) proceeds to start normally. + import urllib.request + + import routes.training as tr + + monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None) + monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: object()) + r = client.post( + "/api/train/diffusion/start", + json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"}, + ) + assert r.status_code == 200, r.text + assert client._fake.started_with["base_model"] == "black-forest-labs/FLUX.1-dev"