Fix diffusion training validation, dataset upload atomicity, and LoRA error mapping

This commit is contained in:
oobabooga 2026-07-04 03:17:12 -03:00
commit e5c63cdfff
21 changed files with 428 additions and 71 deletions

View file

@ -24,7 +24,7 @@ import numpy as np
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_t211_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_t211_images"
# ----------------------------------------------------------------------------- diag

View file

@ -27,7 +27,7 @@ import numpy as np
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/sparse_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "sparse_images"
def _psnr(a, b):

View file

@ -315,7 +315,7 @@ function Uninstall-UnslothStudio {
# 2. A loaded module under a target root (orphaned mp-fork python holding a
# venv DLL). Scoped to names that load our DLLs to keep the scan fast.
try {
$cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli -ErrorAction SilentlyContinue
$cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli, sd-cli, sd-server -ErrorAction SilentlyContinue
foreach ($proc in $cands) {
$hit = $false
try {
@ -366,7 +366,7 @@ function Uninstall-UnslothStudio {
_StopStudioProcesses -KnownRoots $knownRoots
# Also stop anything holding a handle on the exact paths we delete (llama-server,
# the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused.
_StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode))
_StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultSdCpp, $defaultCache, $defaultNode))
# ── Remove custom-root install trees ──
_Step "Removing data and install directories..."
@ -380,6 +380,9 @@ function Uninstall-UnslothStudio {
continue
}
_RemovePath $r
# The native diffusion sibling (<custom root>.parent\stable-diffusion.cpp) is
# intentionally NOT removed: sd.cpp writes no owner marker and sits in the user's
# own parent dir, so auto-deleting it could destroy a user-managed clone.
}
# Default install dir (always at %USERPROFILE%\.unsloth\studio when present).
if ($defaultStudioHome) { _RemovePath $defaultStudioHome }

View file

@ -217,8 +217,11 @@ _remove_path "$HOME/.unsloth/studio"
# when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept.
_remove_path "$HOME/.unsloth/llama.cpp"
# Default-mode native diffusion (stable-diffusion.cpp / sd-cli) build, a sibling of
# studio like llama.cpp (install_sd_cpp_prebuilt.default_install_dir()). No-op in
# env/custom mode and when absent. A user-set UNSLOTH_SD_CPP_PATH is kept.
# studio like llama.cpp (install_sd_cpp_prebuilt.default_install_dir()). Only the default
# location is removed. In env/custom mode the install is <custom root>.parent/
# stable-diffusion.cpp, which is intentionally left in place: sd.cpp writes no owner
# marker and sits in the user's own parent dir, so auto-deleting it could destroy a
# user-managed stable-diffusion.cpp clone. A user-set UNSLOTH_SD_CPP_PATH is kept.
_remove_path "$HOME/.unsloth/stable-diffusion.cpp"
_remove_path "$HOME/.unsloth/.cache"
# Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in

View file

@ -408,9 +408,14 @@ class DiffusionBackend:
the transformer), but ``_load_dense_quant_pipeline`` fetches them with
``from_pretrained(subfolder = "transformer")`` under the load lock during
"finalizing", after the previous pipeline was already evicted, where
unload/cancellation cannot preempt the download. Mirrors the dense-path
gates in ``load_pipeline``: quant requested and supported for this device,
and no pre-quantized checkpoint that would shortcut the dense build."""
unload/cancellation cannot preempt the download. Checks the dense-path gates
in ``load_pipeline`` that are knowable pre-download: quant requested and
supported for this device, and no pre-quantized checkpoint that would shortcut
the dense build. It deliberately does NOT mirror the ``plan.offload_policy ==
OFFLOAD_NONE`` gate: the memory plan needs the GGUF's on-disk size, which
isn't known until the GGUF is cached (after this prefetch runs). So the
transformer/ shards can be prefetched for a load that the plan then routes to
offload -- they stay cached for a later resident load rather than being wasted."""
mode = normalize_transformer_quant(kwargs.get("transformer_quant"))
if mode is None:
return False
@ -1014,6 +1019,21 @@ class DiffusionBackend:
# GGUF build (the OOM-fallback path this cleanup exists for).
del exc
clear_gpu_cache()
elif (
kind == "gguf"
and normalize_transformer_quant(transformer_quant) is not None
and dense_transformer_supported(target)
and plan.offload_policy != OFFLOAD_NONE
):
# The dense fast path needs the transformer resident, so a memory_mode
# (balanced / low_vram) that forces offload silently drops the requested
# quant. Warn so the disengage is diagnosable rather than a null status.
logger.warning(
"diffusion.transformer_quant: %s requested but memory_mode forces "
"offload (%s); loading GGUF without dense quant",
normalize_transformer_quant(transformer_quant),
plan.offload_policy,
)
if pipe is None:
if kind == "pipeline":
@ -1390,7 +1410,8 @@ class DiffusionBackend:
The size estimate is per-kind: diffusers keeps GGUF weights packed (per-matmul
transient dequant), so a GGUF loads near its on-disk size; a safetensors
single-file loads near its on-disk size (it carries its dtype); and a full
single-file loads near its on-disk size (it carries its dtype), except an fp8
transformer file that gets upcast to bf16 on load (~2x resident); and a full
pipeline is one cached download (transformer + companions), already compressed."""
device_memory = snapshot_device_memory(target)
if kind == "pipeline":
@ -1408,9 +1429,19 @@ class DiffusionBackend:
companion_mib = None
else:
if kind == "single_file":
# Safetensors single-file: no dequant expansion (it carries its dtype).
# Safetensors single-file. A dense bf16 file loads near its on-disk size,
# but a transformer-only fp8 checkpoint is loaded via from_single_file with
# a bf16 compute dtype and NO quantization_config, so diffusers upcasts it
# fp8 -> bf16 (~2x resident). Detect fp8 from the basename and budget the
# expansion. The single-file-is-pipeline (SDXL) path is a full bf16 pipeline
# checkpoint, not this fp8 transformer path, so it stays at on-disk size.
fp8_upcast = not getattr(fam, "single_file_is_pipeline", False) and (
"fp8" in Path(single_file_path).name.lower()
if single_file_path
else False
)
transformer_resident = estimate_safetensors_dense_mib(
file_size_mib(single_file_path)
file_size_mib(single_file_path), fp8_upcast = fp8_upcast
)
else:
transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path))
@ -1500,6 +1531,14 @@ class DiffusionBackend:
if cn_model is None:
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# A single generation uses exactly one ControlNet, so keep at most one resident:
# on a miss for a new id, drop the previously-cached module + its from_pipe wrapper
# (both dicts, kept consistent) and free the VRAM before loading the new one, or
# swapping distinct ControlNets within a base-model load accumulates until OOM.
if self._cn_models or self._cn_pipes:
self._cn_models.clear()
self._cn_pipes.clear()
clear_gpu_cache()
import torch
# state.dtype is the display string saved at load ("bfloat16"), NOT a
@ -1784,6 +1823,13 @@ class DiffusionBackend:
raise ValueError(
f"{state.family.name} is an image-editing model: provide an input image."
)
if mask_image is not None:
# The edit family has no inpaint pipeline; a supplied mask would be
# silently dropped (this branch wins over the inpaint branch below).
raise ValueError(
f"{state.family.name} is an image-editing model and does not "
"support masks (mask_image)."
)
workflow = "edit"
init_pil = _decode_b64_image(init_image, mode = "RGB")
elif mask_image is not None and init_image is not None:

View file

@ -265,7 +265,9 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
def _cpu_target(torch: Any, dtype: Any = None) -> DiffusionDeviceTarget:
if dtype is None:
# torch is None on the no-torch CPU fallback; leave dtype=None then (matching the
# no-torch DiffusionDeviceTarget elsewhere) rather than crashing on torch.float32.
if dtype is None and torch is not None:
dtype = torch.float32
return DiffusionDeviceTarget(
device = "cpu",

View file

@ -266,6 +266,17 @@ def _pick_repo_weight_file(repo_id: str, hf_token: Optional[str]) -> str:
raise FileNotFoundError(f"no .safetensors/.gguf LoRA file found in '{repo_id}'")
def _scrub_hub_url(msg: str) -> str:
"""Strip embedded http(s) URLs from a Hub error message before it hits a 400 body.
huggingface_hub errors interpolate the request URL (and a request id) into their
message; a raw endpoint URL is noise in a client-facing 400, so drop it.
"""
cleaned = re.sub(r"https?://\S+", "", msg)
# Collapse the whitespace / stray separators the URL removal leaves behind.
return re.sub(r"\s{2,}", " ", cleaned).strip()
def resolve_specs(
specs: list[tuple[str, float]],
*,
@ -274,20 +285,41 @@ def resolve_specs(
) -> list[ResolvedLora]:
"""Resolve request (id, weight) pairs, dropping zero-weight entries.
A stale / unknown id raises FileNotFoundError inside resolve_one; convert it to
ValueError so the route (which maps only ValueError to a 400) reports bad client
input instead of a generic 500. A Hub download can also raise
``RuntimeError("Cancelled")`` when the user unloads / starts a superseding load
mid-download; convert that to the diffusion cancellation sentinel so the route
maps it to a 409 instead of a generic server error toast."""
A stale / unknown id raises FileNotFoundError inside resolve_one; a mistyped Hub
repo id makes the Hub resolution raise a huggingface_hub client error (a missing
repo -> RepositoryNotFoundError, a bad revision -> RevisionNotFoundError, a missing
weight file -> EntryNotFoundError, a gated model -> GatedRepoError). Convert those
NAMED not-found/gated errors to ValueError so the route (which maps only ValueError
to a 400) reports bad client input instead of a generic 500 -- Hub error messages
embed the request URL, so scrub it out before it reaches the 400 body. Catch them by
name rather than their common HfHubHTTPError base on purpose: a Hub-side 5xx / 429
(an outage, not bad input) is a bare HfHubHTTPError and must stay a 500. A Hub
download can also raise ``RuntimeError("Cancelled")`` when the user unloads / starts a
superseding load mid-download; convert that to the diffusion cancellation sentinel so
the route maps it to a 409 instead of a generic server error toast. A non-cancellation
RuntimeError (e.g. a stalled download, disk full) stays a 500 -- it is not bad
client input."""
from huggingface_hub.errors import (
EntryNotFoundError,
GatedRepoError,
RepositoryNotFoundError,
RevisionNotFoundError,
)
out: list[ResolvedLora] = []
try:
for spec_id, weight in specs:
if weight == 0:
continue
out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event))
except FileNotFoundError as exc:
raise ValueError(str(exc)) from exc
except (
FileNotFoundError,
RepositoryNotFoundError,
RevisionNotFoundError,
EntryNotFoundError,
GatedRepoError,
) as exc:
raise ValueError(_scrub_hub_url(str(exc))) from exc
except RuntimeError as exc:
if str(exc) == "Cancelled":
raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc

View file

@ -236,14 +236,27 @@ def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]:
return int(storage_mib * 1.05) # small margin for allocator + bf16 norms/biases
def estimate_safetensors_dense_mib(storage_mib: Optional[int]) -> Optional[int]:
def estimate_safetensors_dense_mib(
storage_mib: Optional[int], *, fp8_upcast: bool = False
) -> Optional[int]:
"""Resident size of a safetensors checkpoint, in MiB.
Unlike a GGUF (which is dequantised to bf16/fp16 on load, so a 4-bit file
expands ~4x), a safetensors checkpoint loads near its on-disk size: a dense
bf16 file is already bf16, and a bnb-4bit / fp8 file stays compressed in VRAM.
So the on-disk size is the estimate, returned unchanged (None passes through).
expands ~4x), a safetensors checkpoint usually loads near its on-disk size: a
dense bf16 file is already bf16, and a bnb-4bit file stays compressed in VRAM
(it carries its own quantization_config). So the on-disk size is the estimate,
returned unchanged (None passes through).
The exception is ``fp8_upcast``: the fp8 single-file transformer path loads via
``from_single_file`` with a bf16 compute dtype and NO quantization_config, so
diffusers upcasts the fp8 weights (1 byte/param) to bf16 (2 bytes/param) --
roughly 2x the on-disk bytes resident. Budget that, or the plan under-reserves
and OOMs.
"""
if storage_mib is None:
return None
if fp8_upcast:
return storage_mib * 2
return storage_mib

View file

@ -737,6 +737,13 @@ class SdCppDiffusionBackend:
"img2img / inpaint / reference / upscale are not yet supported on the native "
"sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows."
)
# strength 0 (or None) disables ControlNet -- documented on the request model, and
# the diffusers path treats it as plain txt2img -- so a strength-0 spec must be a
# no-op here too, not a hard 400. Only a genuinely active (strength > 0) ControlNet
# is rejected. Strength is element 3 of the tuple
# (id, image, type, strength, guidance_start, guidance_end).
if controlnet is not None and controlnet[3] in (None, 0, 0.0):
controlnet = None
if controlnet is not None:
raise ValueError(
"ControlNet is not yet supported on the native sd.cpp engine; run on a GPU "

View file

@ -32,7 +32,6 @@ from typing import Any, Callable, Optional
from core.training.diffusion_train_common import (
DEFAULT_LORA_FILENAME,
DEFAULT_LORA_TARGETS,
DiffusionLoraConfig,
EventCb,
StopCb,
@ -64,11 +63,11 @@ def _select_lora_targets(
) -> tuple[str, ...]:
"""Pick the LoRA target modules for a DiT run.
``normalized()`` always fills ``lora_target_modules`` with the generic
``DEFAULT_LORA_TARGETS`` when a caller does not set it, so that value means "unset"
here: prefer the family's ``spec.lora_targets`` (which add the DiT-specific
projections). Any OTHER explicit tuple is a deliberate override and still wins."""
if tuple(cfg_targets) == DEFAULT_LORA_TARGETS:
``normalized()`` leaves ``lora_target_modules`` empty when a caller does not set it, so
an empty tuple means "unset" here: use the family's ``spec.lora_targets`` (which add the
DiT-specific joint-attention projections). Any explicit tuple is a deliberate override
and still wins."""
if not tuple(cfg_targets):
return tuple(spec_targets)
return tuple(cfg_targets)

View file

@ -211,13 +211,15 @@ def run_diffusion_lora_training(
for m in (unet, *text_encoders):
m.to(device, dtype = weight_dtype)
# An empty (unset) config means "use the family default": the SDXL attention projections.
unet_targets = list(cfg.lora_target_modules) or list(DEFAULT_LORA_TARGETS)
unet.add_adapter(
LoraConfig(
r = cfg.lora_rank,
lora_alpha = cfg.lora_alpha,
lora_dropout = cfg.lora_dropout,
init_lora_weights = "gaussian",
target_modules = list(cfg.lora_target_modules),
target_modules = unet_targets,
)
)
if cfg.gradient_checkpointing:

View file

@ -30,11 +30,32 @@ from core.inference.diffusion_families import (
trainable_family_names,
)
# Default LoRA target modules: the attention projections common to the SDXL U-Net and the
# DiT transformers (the diffusers/kohya convention). A family whose trainer wants a wider
# set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback.
# Default LoRA target modules: the attention projections of the SDXL U-Net (the
# diffusers/kohya convention). Used by the SDXL trainer as its fallback when the config
# leaves ``lora_target_modules`` empty; the DiT trainers supply their own wider set. Kept
# here so the SDXL trainer has a named default even for an empty (unset) config.
DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). Validated in
# normalized() so a typo fails fast at request time, not minutes later in the subprocess.
_LR_SCHEDULERS: frozenset[str] = frozenset(
{
"linear",
"cosine",
"cosine_with_restarts",
"polynomial",
"constant",
"constant_with_warmup",
"piecewise_constant",
}
)
# DiT families that overflow fp16 (their RoPE / embedder run in fp32), so they train in bf16
# only. Encoded here -- keyed by resolved family -- so normalized() can reject an fp16 request
# before spawn without importing the DiT trainer's _SPECS (which would create an import
# cycle). The DiT trainer keeps a matching guard as defense in depth.
_FORCE_BF16_FAMILIES: frozenset[str] = frozenset({"qwen-image", "z-image"})
_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
_CAPTION_EXTS = (".txt", ".caption")
# diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it.
@ -83,7 +104,15 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
# 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:
# A ``.gguf`` file always rejects. The broad ``"gguf" in name`` catch (for ``*-GGUF``
# repos) must NOT reject a real local diffusers directory that merely has "gguf" in its
# path, so it is skipped for a local diffusers checkout -- identified by its
# ``model_index.json`` marker (the same marker the loader uses), NOT a bare ``is_dir()``:
# a GGUF-only folder must still reject here and fail fast, rather than pass and fail late
# in the subprocess after the resident chat/Images models were already evicted.
local = Path(base_model).expanduser() if base_model else None
is_local_diffusers = bool(local and (local / "model_index.json").is_file())
if name.endswith(".gguf") or ("gguf" in name and not is_local_diffusers):
raise ValueError(
f"'{base_model}' is a GGUF checkpoint/repo, which can't be a training base "
f"(training needs the full diffusers model). {_trainable_hint()}"
@ -215,7 +244,9 @@ class DiffusionLoraConfig:
lora_rank: int = 16
lora_alpha: Optional[int] = None # defaults to lora_rank
lora_dropout: float = 0.0
lora_target_modules: tuple[str, ...] = DEFAULT_LORA_TARGETS
# Empty = "unset": each trainer supplies its family default (SDXL DEFAULT_LORA_TARGETS,
# or the DiT family's wider joint-attention set). A non-empty tuple is an explicit override.
lora_target_modules: tuple[str, ...] = ()
seed: int = 42
mixed_precision: str = "bf16" # "bf16" | "fp16" | "no"
snr_gamma: Optional[float] = 5.0 # min-SNR loss weighting; None disables
@ -259,6 +290,19 @@ class DiffusionLoraConfig:
raise ValueError("resolution must be a multiple of 8 and >= 64")
if self.mixed_precision not in ("bf16", "fp16", "no"):
raise ValueError("mixed_precision must be one of bf16 / fp16 / no")
# A bf16-only DiT family (Qwen-Image / Z-Image) must refuse fp16 up front rather than
# accepting the request, evicting resident models, and only then failing in the
# subprocess. The DiT trainer keeps a matching guard as defense in depth.
if self.mixed_precision == "fp16" and resolved_family in _FORCE_BF16_FAMILIES:
raise ValueError(
f"'{resolved_family}' LoRA training requires bf16: fp16 overflows its fp32 "
f"RoPE / embedder internals. Set mixed precision to bf16."
)
if str(self.lr_scheduler) not in _LR_SCHEDULERS:
raise ValueError(
f"lr_scheduler must be one of {', '.join(sorted(_LR_SCHEDULERS))}; "
f"got {self.lr_scheduler!r}"
)
# learning_rate can arrive as a string ("1e-4") from the Studio config path, which
# preserves it as a string after validation; coerce so AdamW receives a float.
try:
@ -268,7 +312,9 @@ class DiffusionLoraConfig:
if learning_rate <= 0:
raise ValueError("learning_rate must be > 0")
alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank
targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS
# Leave an unset (empty) target list empty: the trainer fills the family default
# (SDXL DEFAULT_LORA_TARGETS, or the DiT family's wider set) so the family spec wins.
targets = tuple(self.lora_target_modules)
# A blank Hub token (the Studio default when none is configured) must load
# anonymously, not as an explicit empty credential.
token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token
@ -390,8 +436,20 @@ def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Option
else Path(cfg.output_dir).name
)
alias = sanitize_alias(base)
src_resolved = Path(lora_path).resolve()
dest = loras_dir() / f"{alias}.safetensors"
if Path(lora_path).resolve() != dest.resolve():
# A retrain with the same adapter name must not clobber a prior mirror: if the
# destination already exists and is a different file, pick the next free numeric
# suffix (<alias>-2, <alias>-3, ...) for both the weights and their .json sidecar.
if dest.exists() and dest.resolve() != src_resolved:
n = 2
while True:
candidate = loras_dir() / f"{alias}-{n}.safetensors"
if not candidate.exists() or candidate.resolve() == src_resolved:
dest = candidate
break
n += 1
if src_resolved != dest.resolve():
shutil.copy2(lora_path, dest)
_write_lora_sidecar(dest.with_suffix(".json"), cfg)
return str(dest)

View file

@ -698,18 +698,27 @@ class DiffusionTrainingStartRequest(BaseModel):
lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank")
lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0)
# Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that
# sets them is not silently trained with defaults. Default the target list to the SDXL
# attention projections (the trainer's DEFAULT_LORA_TARGETS) so it is never None.
# sets them is not silently trained with defaults. Empty (the default) means "unset": each
# trainer supplies its own family targets -- the SDXL DEFAULT_LORA_TARGETS for SDXL, the
# wider joint-attention set for the DiT families. A non-empty list is an explicit override.
lora_target_modules: List[str] = Field(
default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"],
description = "U-Net modules to attach LoRA to",
default_factory = list,
description = "Modules to attach LoRA to; empty = the trainer's family default",
)
max_grad_norm: float = Field(1.0, gt = 0, description = "Gradient clipping max-norm")
seed: int = Field(42)
mixed_precision: Literal["bf16", "fp16", "no"] = Field("bf16")
snr_gamma: Optional[float] = Field(5.0, description = "Min-SNR loss weighting; null disables")
gradient_checkpointing: bool = Field(True)
lr_scheduler: str = Field("constant")
lr_scheduler: Literal[
"linear",
"cosine",
"cosine_with_restarts",
"polynomial",
"constant",
"constant_with_warmup",
"piecewise_constant",
] = Field("constant")
lr_warmup_steps: int = Field(0, ge = 0)
center_crop: bool = Field(False)
random_flip: bool = Field(True)

View file

@ -1375,6 +1375,9 @@ async def upload_diffusion_dataset(
total_bytes = 0
uploaded = 0
allowed = _DIFFUSION_DATASET_IMAGE_EXTS | _DIFFUSION_DATASET_TEXT_EXTS
# Validate every filename up front so a valid image ahead of a bad one is not left
# written on disk when the 400 fires -- make the upload all-or-nothing.
names: list[str] = []
for f in files:
filename = Path(f.filename or "").name.strip().replace("\x00", "")
ext = Path(filename).suffix.lower()
@ -1384,9 +1387,17 @@ async def upload_diffusion_dataset(
status_code = 400,
detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
)
dest = folder / filename
complete = False
try:
names.append(filename)
# Roll back every file written this request if the batch does not fully commit, so a
# mid-batch 413 (or a disk error / client disconnect) leaves the dataset unchanged
# rather than partially populated -- the upload is all-or-nothing, not just for the
# extension check above but for the size limit too.
written: list[Path] = []
committed = False
try:
for f, filename in zip(files, names):
dest = folder / filename
written.append(dest)
with open(dest, "wb") as out:
while chunk := await f.read(1024 * 1024):
total_bytes += len(chunk)
@ -1400,14 +1411,15 @@ async def upload_diffusion_dataset(
),
)
out.write(chunk)
complete = True
finally:
if not complete:
uploaded += 1
committed = True
finally:
if not committed:
for p in written:
try:
dest.unlink(missing_ok = True)
p.unlink(missing_ok = True)
except OSError:
pass
uploaded += 1
summary = _diffusion_dataset_summary(folder)
return DiffusionDatasetUploadResponse(
@ -1572,7 +1584,7 @@ async def get_diffusion_dataset_image(
thumbs_dir = folder / _THUMBS_DIRNAME
thumbs_dir.mkdir(exist_ok = True)
thumb_path = thumbs_dir / f"{image_path.stem}_{size}.jpg"
thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg"
src_mtime = image_path.stat().st_mtime
if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime:
return thumb_path
@ -1645,7 +1657,7 @@ async def delete_diffusion_dataset_image(
image_path.with_suffix(ext).unlink(missing_ok = True)
thumbs_dir = folder / _THUMBS_DIRNAME
if thumbs_dir.is_dir():
for t in thumbs_dir.glob(f"{image_path.stem}_*.jpg"):
for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
t.unlink(missing_ok = True)
return {"deleted": image_path.name}
@ -1852,7 +1864,7 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
)
# Map basename -> caption from every jsonl carrying file_name + caption column.
captions: dict[str, str] = {}
for jf in snap.rglob("*.jsonl"):
for jf in sorted(snap.rglob("*.jsonl")):
for line in jf.read_text(encoding = "utf-8").splitlines():
line = line.strip()
if not line:
@ -1863,7 +1875,9 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
continue
fn = row.get("file_name") or row.get("image") or row.get("file")
if fn and caption_col in row:
captions[Path(str(fn)).name] = str(row[caption_col])
# First writer wins over sorted manifests, so the plain manifest
# (e.g. output_file.jsonl) is deterministic rather than OS-visit order.
captions.setdefault(Path(str(fn)).name, str(row[caption_col]))
# Copy images (those with a caption first, so a cap keeps captioned pairs).
images = sorted(
p

View file

@ -179,13 +179,15 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
(folder / "x.txt").write_text("cap", encoding = "utf-8")
# Generate a thumbnail so we can assert it is cleaned up too.
client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32")
assert list((folder / ".thumbs").glob("x_*.jpg"))
# Thumb cache key includes the extension (x.png_32.jpg), so png and jpg
# siblings can't collide.
assert list((folder / ".thumbs").glob("x.png_*.jpg"))
r = client.delete("/api/train/diffusion/dataset/d/image/x.png")
assert r.status_code == 200, r.text
assert not (folder / "x.png").exists()
assert not (folder / "x.txt").exists()
assert not list((folder / ".thumbs").glob("x_*.jpg"))
assert not list((folder / ".thumbs").glob("x.png_*.jpg"))
# ── traversal / validation ───────────────────────────────────────────────────

View file

@ -41,23 +41,25 @@ def test_specs_cover_the_three_dit_families():
def test_select_lora_targets_uses_family_default_for_generic_config():
# normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a
# caller doesn't set it, so that value must resolve to the family's targets (which add
# the DiT-specific projections), not stay stuck on the generic SDXL list.
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == _FLUX_TARGETS
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _QWEN_TARGETS) == _QWEN_TARGETS
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
# normalized() leaves lora_target_modules empty when a caller doesn't set it, so an empty
# tuple must resolve to the family's targets (which add the DiT-specific joint-attention
# projections), not the generic SDXL list.
assert _select_lora_targets((), _FLUX_TARGETS) == _FLUX_TARGETS
assert _select_lora_targets((), _QWEN_TARGETS) == _QWEN_TARGETS
assert _select_lora_targets((), _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
# The generic SDXL default is NOT treated as unset here: an explicit list wins.
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == DEFAULT_LORA_TARGETS
def test_select_lora_targets_explicit_override_wins():
# Any OTHER explicit tuple is a deliberate override and must win over the family spec.
# Any explicit tuple is a deliberate override and must win over the family spec.
override = ("to_q", "to_k")
assert _select_lora_targets(override, _FLUX_TARGETS) == override
# The default request path (config carrying the generic default) reaches the spec.
# The default request path (config leaving targets unset) reaches the spec.
cfg = DiffusionLoraConfig(
base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o"
).normalized()
assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
assert cfg.lora_target_modules == ()
assert (
_select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets)
== _FLUX_TARGETS

View file

@ -189,6 +189,26 @@ def test_resolve_specs_maps_unknown_id_to_valueerror(tmp_path, monkeypatch):
dl.resolve_specs([("nope", 1.0)])
def test_resolve_specs_maps_hub_error_to_valueerror(tmp_path, monkeypatch):
# A mistyped Hub repo id makes the Hub resolution raise a huggingface_hub client error
# (RepositoryNotFoundError, an HfHubHTTPError). resolve_specs must surface it as
# ValueError so the route returns 400, not a generic 500. The Hub message embeds the
# request URL, which must be scrubbed out of the client-facing 400.
from huggingface_hub.errors import RepositoryNotFoundError
def _boom(spec_id, weight, **kw):
raise RepositoryNotFoundError(
"404 Client Error. Repository Not Found for url: "
"https://huggingface.co/api/models/nope/nope (Request ID: abc)"
)
monkeypatch.setattr(dl, "resolve_one", _boom)
with pytest.raises(ValueError) as ei:
dl.resolve_specs([("nope/nope", 1.0)])
assert "http" not in str(ei.value) # request URL scrubbed
assert "Repository Not Found" in str(ei.value)
def test_scan_local_disambiguates_identical_stems(tmp_path, monkeypatch):
# foo.safetensors and foo.gguf must get distinct ids so each is addressable; a
# unique stem keeps its clean stem id.

View file

@ -94,7 +94,8 @@ def test_discover_missing_dir_raises(tmp_path):
def test_config_normalized_defaults():
cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o").normalized()
assert cfg.lora_alpha == cfg.lora_rank # alpha defaults to rank
assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
# Targets stay empty (unset) after normalize; each trainer fills its own family default.
assert cfg.lora_target_modules == ()
@pytest.mark.parametrize(
@ -335,3 +336,85 @@ def test_publish_writes_metadata_sidecar(tmp_path, monkeypatch):
assert meta["lora_rank"] == 8
assert meta["trigger_prompt"] == "a photo in sks style"
assert meta["source"] == "studio-trained"
def test_publish_does_not_clobber_same_name_adapter(tmp_path, monkeypatch):
# A retrain with the same adapter name must not overwrite a prior mirror: the second
# publish lands under a numeric suffix (my-style -> my-style-2), sidecar alongside it.
from pathlib import Path
from core.inference import diffusion_lora
from core.training.diffusion_lora_trainer import _publish_to_lora_catalog
loras = tmp_path / "loras"
loras.mkdir()
monkeypatch.setattr(diffusion_lora, "loras_dir", lambda: loras)
def _publish(payload: bytes) -> str:
src = tmp_path / "run" / "pytorch_lora_weights.safetensors"
src.parent.mkdir(parents = True, exist_ok = True)
src.write_bytes(payload)
cfg = DiffusionLoraConfig(
base_model = "stabilityai/sdxl-turbo",
data_dir = "d",
output_dir = str(tmp_path / "run"),
adapter_name = "my-style",
).normalized()
return _publish_to_lora_catalog(str(src), cfg)
first = _publish(b"adapter-v1")
second = _publish(b"adapter-v2")
assert Path(first).name == "my-style.safetensors"
assert Path(second).name == "my-style-2.safetensors"
# The first mirror is intact (not clobbered) and the second is the new content.
assert Path(first).read_bytes() == b"adapter-v1"
assert Path(second).read_bytes() == b"adapter-v2"
assert Path(second).with_suffix(".json").is_file()
def test_config_rejects_bad_lr_scheduler():
# A typo'd scheduler ('constnat') must fail at normalize time, not later in the subprocess.
with pytest.raises(ValueError, match = "lr_scheduler"):
DiffusionLoraConfig(
base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "constnat"
).normalized()
# A valid diffusers scheduler passes.
cfg = DiffusionLoraConfig(
base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "cosine"
).normalized()
assert cfg.lr_scheduler == "cosine"
def test_config_rejects_fp16_on_bf16_only_family():
# qwen-image / z-image are bf16-only: an fp16 request must be rejected before spawn,
# in normalized(), not only by the subprocess-side guard.
for base in ("Tongyi-MAI/Z-Image-Turbo", "unsloth/Qwen-Image-2512-unsloth-bnb-4bit"):
with pytest.raises(ValueError, match = "bf16"):
DiffusionLoraConfig(
base_model = base, data_dir = "d", output_dir = "o", mixed_precision = "fp16"
).normalized()
# FLUX (not force-bf16) still accepts fp16.
cfg = DiffusionLoraConfig(
base_model = "black-forest-labs/FLUX.1-dev",
data_dir = "d",
output_dir = "o",
mixed_precision = "fp16",
).normalized()
assert cfg.mixed_precision == "fp16"
def test_gguf_substring_does_not_reject_local_diffusers_dir(tmp_path):
# A local diffusers directory whose path merely contains 'gguf' is a valid training base
# (it carries model_index.json, not GGUF weights); the broad substring must not reject it.
from core.training.diffusion_train_common import resolve_trainable_family
local = tmp_path / "my-gguf-experiments" / "sdxl-finetune"
local.mkdir(parents = True)
(local / "model_index.json").write_text("{}", encoding = "utf-8")
assert resolve_trainable_family(str(local)) == "sdxl"
# A real .gguf file still rejects even inside such a dir.
with pytest.raises(ValueError, match = "GGUF"):
resolve_trainable_family(str(local / "weights.gguf"))
# A *-GGUF repo id (not a local dir) still rejects.
with pytest.raises(ValueError, match = "GGUF"):
resolve_trainable_family("unsloth/FLUX.1-dev-GGUF")

View file

@ -492,6 +492,24 @@ def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_root
assert "Unsupported file" in r.json()["detail"]
def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(client, dataset_roots, monkeypatch):
# All-or-nothing: a valid image ahead of the one that trips the size cap must NOT be
# left on disk (the 413 mid-batch rolls back every file written this request).
import utils.upload_limits as ul
monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100)
ds_root, _ = dataset_roots
files = [
("files", ("small.png", b"x" * 50, "image/png")),
("files", ("big.png", b"y" * 200, "image/png")),
]
r = client.post("/api/train/diffusion/dataset", data = {"name": "rollback"}, files = files)
assert r.status_code == 413, r.text
folder = ds_root / "rollback"
assert not (folder / "small.png").exists() # the earlier valid file was rolled back
assert not (folder / "big.png").exists()
def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch):
# A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed,
# so a bad pick never unloads the user's working chat/Images model.

View file

@ -771,6 +771,22 @@ def test_generate_rejects_controlnet_on_native_engine():
b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0))
@pytest.mark.parametrize("cn_strength", [0, 0.0, None])
def test_generate_treats_zero_strength_controlnet_as_disabled(cn_strength):
# strength 0 (or None) disables ControlNet -- the diffusers path treats it as plain
# txt2img and the request model documents it -- so a strength-0 spec must succeed on the
# native engine too, not 400. Only a genuinely active (strength > 0) ControlNet is rejected.
eng = _FakeEngine()
b = _loaded_backend(engine = eng)
out = b.generate(
prompt = "x",
steps = 4,
seed = 1,
controlnet = ("id", "img", "canny", cn_strength, 0.0, 1.0),
)
assert len(out["images"]) == 1
def test_generate_rejects_image_conditioned_on_native_engine():
# img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call
# with an init image on the native engine gets a clean ValueError, not a silent txt2img.

View file

@ -1237,7 +1237,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
// Load an image's recipe back into the form inputs.
const restoreSettings = useCallback((image: GalleryImage) => {
setPrompt(image.prompt);
setNegativePrompt(image.negative_prompt ?? "");
// The Negative-prompt field only renders (and submits) when guidance>0, so a
// guidance=0 recipe must not restore a hidden negative prompt that would
// resurface if guidance is later raised. Mirror the submit-path gating.
setNegativePrompt(image.guidance > 0 ? (image.negative_prompt ?? "") : "");
setSteps(image.steps);
setGuidance(image.guidance);
setSeed(String(image.seed));
@ -1272,6 +1275,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight });
}
setLoras(restoredLoras);
// The recipe carries no control image (it isn't persisted), so a faithful
// restore can't reproduce a ControlNet run -- clear any stale form selection
// rather than leaking it into the restored recipe, mirroring the LoRA clear.
setControlnetId("");
setControlImage(null);
toast.success("Settings restored to inputs");
}, []);
@ -1520,11 +1528,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
// Curated non-GGUF model: load as a full pipeline or single-file safetensors.
const spec = SAFETENSORS_MODELS[id];
if (spec) {
// Optimistically drop the quant label, but revert if the load never starts
// so a failed load doesn't leave a stale 'GGUF · variant' label (mirrors the
// GGUF branches below; the poll owns the after-start revert via quantRevert).
const prevQuant = quant;
quantRevert.current = { prev: prevQuant };
setQuant(null);
const d = defaultsFor(id);
setSteps(d.steps);
setGuidance(d.guidance);
void handleLoad(id, { kind: spec.kind, filename: spec.filename });
void handleLoad(id, { kind: spec.kind, filename: spec.filename }).then((started) => {
if (!started) {
setQuant(prevQuant);
quantRevert.current = null;
}
});
return;
}
// GGUF quant pick from the variant expander. Optimistic for instant picker
@ -1582,11 +1600,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
toast.error("Only unsloth or on-device image models can be loaded here");
return;
}
// Optimistically drop the quant label, but revert if the load never starts so
// a failed load doesn't leave a stale 'GGUF · variant' label (mirrors the GGUF
// branches above; the poll owns the after-start revert via quantRevert).
const prevQuant = quant;
quantRevert.current = { prev: prevQuant };
setQuant(null);
const d = defaultsFor(id);
setSteps(d.steps);
setGuidance(d.guidance);
void handleLoad(id, { kind: "pipeline" });
void handleLoad(id, { kind: "pipeline" }).then((started) => {
if (!started) {
setQuant(prevQuant);
quantRevert.current = null;
}
});
},
[busy, handleLoad, quant],
);