Fix diffusion training validation, dataset upload atomicity, and LoRA error mapping
This commit is contained in:
parent
26a81c5e17
commit
e5c63cdfff
21 changed files with 428 additions and 71 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ───────────────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue