Merge remote-tracking branch 'origin/diffusion-train-dit' into diffusion-train-datasets
This commit is contained in:
Daniel Han
commit
8547fd14f7
9 changed files with 1006 additions and 24 deletions
|
|
@ -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.
|
||||
|
|
|
|||
650
studio/backend/core/training/diffusion_dit_trainer.py
Normal file
650
studio/backend/core/training/diffusion_dit_trainer.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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`` /
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ from models.training import (
|
|||
DiffusionDatasetSummary,
|
||||
DiffusionDatasetUploadResponse,
|
||||
DiffusionMetricHistory,
|
||||
DiffusionTrainableFamily,
|
||||
DiffusionTrainingInfoResponse,
|
||||
DiffusionTrainingStartRequest,
|
||||
DiffusionTrainingStartResponse,
|
||||
|
|
@ -1132,6 +1133,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)
|
||||
|
|
@ -1176,8 +1210,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()
|
||||
|
|
@ -1266,8 +1305,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)
|
||||
|
|
|
|||
84
studio/backend/tests/test_diffusion_dit_trainer.py
Normal file
84
studio/backend/tests/test_diffusion_dit_trainer.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue