From fa2cb600ee365c564617cf55d1a8acea373e887b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 15:25:43 +0000 Subject: [PATCH 1/4] Add flow-matching DiT LoRA trainers (FLUX.1-dev, Qwen-Image, Z-Image) Extends diffusion LoRA training beyond SDXL to the three popular DiT families via a single shared flow-matching loop parameterised by small per-family specs (loading, prompt/latent encoding, transformer forward, save). Verified against diffusers 0.38.0: - FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the user's HF token). - Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std, img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16 base). - Z-Image: list I/O with the reversed timestep convention and a negated prediction, bf16 only. The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now route these families to the DiT trainer; the SDXL blocklist guard is replaced by a positive family resolution that also rejects GGUF repos (inference-only) and still-unsupported families. Per-family defaults + labels + VRAM notes are exposed via family_train_infos for the Train UI. Memory: caption embeddings are precomputed once and the text encoders freed before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit) and 8-bit AdamW are on by default. --- .../core/inference/diffusion_families.py | 11 + .../core/training/diffusion_dit_trainer.py | 650 ++++++++++++++++++ .../core/training/diffusion_train_common.py | 71 +- 3 files changed, 730 insertions(+), 2 deletions(-) create mode 100644 studio/backend/core/training/diffusion_dit_trainer.py diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index ac2140d819..7cff2c3fec 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. 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..efaa4d9408 --- /dev/null +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -0,0 +1,650 @@ +# 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 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 + ) + 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_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 74a33331f1..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() @@ -128,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`` / From b28793418d50951fbcc9d6bfd844e85856109c36 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 15:25:55 +0000 Subject: [PATCH 2/4] Speed up + shrink SDXL LoRA training (precompute text embeds, 8-bit AdamW) SDXL re-encoded every caption with both CLIP text encoders on every step (pure waste, since captions are constant) and kept the encoders resident. Precompute each unique caption's embeddings once, then free the text encoders before the loop: numerically identical (embeddings are deterministic and this consumes no torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32 fallback, halving optimizer state with no meaningful LoRA quality cost. Env toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard A/B the paths. --- .../core/training/diffusion_lora_trainer.py | 47 +++++++++++++++++-- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 795bf191d2..59da0b6288 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,22 @@ 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 From afd93591d6f372b0ff508eb8f946726e0eb9b538 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 15:26:05 +0000 Subject: [PATCH 3/4] Expose trainable families in /diffusion/info and preflight gated bases The training info endpoint now returns the trainable model families (name, label, default + allowed base repos, recommended defaults, and a VRAM/access note) so the Train UI can offer a base picker with realistic guidance. The start route preflights a gated base repo (HEAD model_index.json with the user's token) BEFORE freeing resident GPU workloads, so a missing FLUX.1-dev license/token fails fast with an actionable 400 instead of evicting the loaded model and then hitting a confusing mid-load 401. --- studio/backend/models/training.py | 18 +++++++++--- studio/backend/routes/training.py | 49 +++++++++++++++++++++++++++++-- 2 files changed, 61 insertions(+), 6 deletions(-) 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..9e32253090 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,39 @@ 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 +1203,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 +1298,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) From 60268a77b27352fb367abe63f1405f86dbb4b8b5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 15:26:05 +0000 Subject: [PATCH 4/4] Tests for DiT trainers, family resolution, info families, gated preflight Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image are trainable (and GGUF repos are rejected as inference-only), the families list in /diffusion/info, and the gated-base 400 preflight that leaves the GPU untouched. --- .../tests/test_diffusion_dit_trainer.py | 84 +++++++++++++++++++ .../tests/test_diffusion_lora_trainer.py | 48 ++++++++--- .../backend/tests/test_diffusion_training.py | 52 ++++++++++++ 3 files changed, 172 insertions(+), 12 deletions(-) create mode 100644 studio/backend/tests/test_diffusion_dit_trainer.py 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 e229712087..8db45c588e 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). @@ -236,7 +245,15 @@ 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(): @@ -254,9 +271,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 039c5940b9..6931418cde 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -554,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"