Trim the verbose comments added by the fixes
This commit is contained in:
oobabooga
parent
3bb3734879
commit
756c715721
17 changed files with 75 additions and 170 deletions
|
|
@ -380,9 +380,8 @@ 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.
|
||||
# The native diffusion sibling (<custom root>.parent\stable-diffusion.cpp) is left
|
||||
# in place: sd.cpp writes no owner marker, so auto-deleting it could destroy a clone.
|
||||
}
|
||||
# Default install dir (always at %USERPROFILE%\.unsloth\studio when present).
|
||||
if ($defaultStudioHome) { _RemovePath $defaultStudioHome }
|
||||
|
|
|
|||
|
|
@ -217,11 +217,9 @@ _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()). 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.
|
||||
# studio like llama.cpp. Only the default location is removed; the env/custom-mode
|
||||
# sibling (<custom root>.parent/stable-diffusion.cpp) is left in place since sd.cpp
|
||||
# writes no owner marker and auto-deleting it could destroy a user-managed clone.
|
||||
_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
|
||||
|
|
|
|||
|
|
@ -404,18 +404,11 @@ class DiffusionBackend:
|
|||
"""True when ``load_pipeline`` may take the dense transformer-quant path, so
|
||||
the prefetch should also pull the base repo's ``transformer/`` shards.
|
||||
|
||||
Those shards are excluded from the prefetch by default (the GGUF supplies
|
||||
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. 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."""
|
||||
Those shards are excluded from the prefetch by default (the GGUF supplies the
|
||||
transformer), but ``_load_dense_quant_pipeline`` fetches them later under the
|
||||
load lock, where unload/cancellation cannot preempt the download. Checks only
|
||||
the dense-path gates knowable pre-download; skips the ``offload_policy`` gate
|
||||
since that needs the GGUF's on-disk size, not known until after this runs."""
|
||||
mode = normalize_transformer_quant(kwargs.get("transformer_quant"))
|
||||
if mode is None:
|
||||
return False
|
||||
|
|
@ -807,7 +800,7 @@ class DiffusionBackend:
|
|||
|
||||
@staticmethod
|
||||
def _hub_cache_repo_dir(repo_id: str) -> Path:
|
||||
"""The local HF hub cache dir for ``repo_id`` (``.../models--org--name``)."""
|
||||
"""Local HF hub cache dir for ``repo_id``."""
|
||||
from huggingface_hub import constants
|
||||
return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}"
|
||||
|
||||
|
|
@ -853,15 +846,10 @@ class DiffusionBackend:
|
|||
def _companion_cache_bytes(base: str) -> int:
|
||||
"""Resident companion (VAE + text-encoder) size for the memory plan.
|
||||
|
||||
Sums the cached VAE + text-encoder weights while EXCLUDING ``transformer/`` (the
|
||||
GGUF / single file supplies the transformer, so its ``transformer/`` shards are
|
||||
not resident here). This matters for the dense ``transformer_quant`` path: it
|
||||
prefetches the base repo's ``transformer/`` shards into the cache, and folding
|
||||
those multi-GB shards into the companion size would inflate the plan and wrongly
|
||||
force offload -- gating off the very quant path that fetched them. For a LOCAL
|
||||
diffusers base the blob cache is empty, so walk the on-disk weights; for a hub
|
||||
base, walk the snapshot (whose ``transformer/`` subfolder we can skip) instead of
|
||||
the flat, content-addressed ``blobs/`` dir, which carries no subfolder split."""
|
||||
Excludes ``transformer/`` (supplied by the GGUF/single file, not resident here) --
|
||||
otherwise the dense-quant prefetch's cached transformer shards would inflate this
|
||||
and wrongly force offload. Walks the snapshot dir, not the flat ``blobs/`` cache,
|
||||
since only the snapshot preserves the subfolder split needed to exclude it."""
|
||||
local = Path(base).expanduser()
|
||||
if local.is_dir():
|
||||
return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
|
||||
|
|
@ -1029,9 +1017,7 @@ class DiffusionBackend:
|
|||
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.
|
||||
# memory_mode forcing offload silently drops the requested quant; warn so it's diagnosable.
|
||||
logger.warning(
|
||||
"diffusion.transformer_quant: %s requested but memory_mode forces "
|
||||
"offload (%s); loading GGUF without dense quant",
|
||||
|
|
@ -1440,12 +1426,10 @@ class DiffusionBackend:
|
|||
companion_mib = None
|
||||
else:
|
||||
if kind == "single_file":
|
||||
# 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.
|
||||
# An fp8 transformer checkpoint loads via from_single_file with a bf16
|
||||
# compute dtype and no quantization_config, so diffusers upcasts it to
|
||||
# bf16 (~2x resident); detect it from the basename. Excludes the
|
||||
# single-file-is-pipeline (SDXL) case, which is already a bf16 pipeline.
|
||||
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
|
||||
)
|
||||
|
|
@ -1540,10 +1524,9 @@ 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.
|
||||
# Keep at most one ControlNet resident: evict the previous module + its
|
||||
# from_pipe wrapper before loading the new one, or swapping 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()
|
||||
|
|
@ -1718,10 +1701,8 @@ class DiffusionBackend:
|
|||
resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a
|
||||
transformer without the hook (uncached load) is a silent no-op."""
|
||||
transformer = getattr(pipe, "transformer", None)
|
||||
# A diffusers CacheMixin transformer clears FBCache via ``_reset_stateful_cache``
|
||||
# (which drives its HookRegistry.reset_stateful_hooks internally). The public
|
||||
# ``reset_stateful_hooks`` name lives only on the HookRegistry, not on the
|
||||
# transformer, so keep it only as a version fallback.
|
||||
# ``_reset_stateful_cache`` is the transformer-level entry point; the public
|
||||
# ``reset_stateful_hooks`` lives only on the HookRegistry, kept as a fallback.
|
||||
reset = getattr(transformer, "_reset_stateful_cache", None) or getattr(
|
||||
transformer, "reset_stateful_hooks", None
|
||||
)
|
||||
|
|
@ -1833,8 +1814,7 @@ class DiffusionBackend:
|
|||
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).
|
||||
# The edit family has no inpaint pipeline; a mask would be silently dropped.
|
||||
raise ValueError(
|
||||
f"{state.family.name} is an image-editing model and does not "
|
||||
"support masks (mask_image)."
|
||||
|
|
|
|||
|
|
@ -204,9 +204,8 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
|
|||
|
||||
# Union ControlNet mode indices. A single "union" model covers several control modes and
|
||||
# selects the active one via an integer ``control_mode`` argument; these are the standard
|
||||
# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" (an already-made
|
||||
# map) carries no intrinsic mode, so union_control_mode() defaults it to 0 (a union model
|
||||
# still requires a concrete mode).
|
||||
# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" carries no
|
||||
# intrinsic mode, so union_control_mode() defaults it to 0.
|
||||
_UNION_CONTROL_MODES: dict[str, int] = {
|
||||
"canny": 0,
|
||||
"tile": 1,
|
||||
|
|
@ -221,12 +220,9 @@ _UNION_CONTROL_MODES: dict[str, int] = {
|
|||
def union_control_mode(spec_id: str, control_type: str) -> Optional[int]:
|
||||
"""The integer ``control_mode`` for a union ControlNet, or None.
|
||||
|
||||
A union model REQUIRES a concrete ``control_mode`` (diffusers raises when it is None),
|
||||
so for a curated union entry always return an index: the mapped mode, or a default
|
||||
(0 / canny) for a type that carries no intrinsic mode such as 'passthrough' (an
|
||||
already-made control map, which is also the UI's default for these models). For a
|
||||
non-union entry return None so the caller omits the kwarg (it has a single fixed
|
||||
mode). Pure lookup, no network."""
|
||||
A union model requires a concrete mode (diffusers raises on None), so a curated union
|
||||
entry always gets an index, defaulting to 0 for types like 'passthrough' that carry
|
||||
none. A non-union entry returns None so the caller omits the kwarg."""
|
||||
entry = _catalog_by_id().get(spec_id)
|
||||
if entry is None or not entry.is_union:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -235,15 +235,10 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
|
|||
mps_available = False
|
||||
|
||||
if mps_available:
|
||||
# Relax the MPS memory watermark BEFORE the first MPS allocation (the bfloat16
|
||||
# probe just below). torch reads PYTORCH_MPS_HIGH_WATERMARK_RATIO exactly once,
|
||||
# when the MPS allocator first initializes, so setting it any later is a no-op.
|
||||
# The allocator otherwise caps a process at ~1.7x recommendedMaxWorkingSetSize,
|
||||
# and a model that fits in unified system RAM but exceeds that cap OOMs at
|
||||
# pipe.to("mps") (observed on an 8GB M1 mac mini: "MPS allocated 9.06 GiB, max
|
||||
# allowed 9.07 GiB"). CPU offload can't help on unified memory (it frees no
|
||||
# device bytes). Lifting the cap lets MPS spill into system RAM; a model larger
|
||||
# than RAM would fail either way. setdefault respects a user-provided override.
|
||||
# torch reads PYTORCH_MPS_HIGH_WATERMARK_RATIO once, at the first MPS allocation
|
||||
# (the bfloat16 probe below), so it must be relaxed before that. Otherwise the
|
||||
# allocator caps the process at ~1.7x recommendedMaxWorkingSetSize and can OOM a
|
||||
# model that would otherwise fit in unified RAM. setdefault respects an override.
|
||||
os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
|
||||
# Prefer bfloat16; otherwise fall back to float32, NEVER silent float16.
|
||||
# Modern diffusion transformers (Z-Image, FLUX.2, ...) produce activations
|
||||
|
|
|
|||
|
|
@ -267,11 +267,7 @@ def _pick_repo_weight_file(repo_id: str, hf_token: Optional[str]) -> str:
|
|||
|
||||
|
||||
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.
|
||||
"""
|
||||
"""Strip embedded http(s) URLs from a Hub error message before it hits a 400 body."""
|
||||
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()
|
||||
|
|
@ -285,20 +281,10 @@ def resolve_specs(
|
|||
) -> list[ResolvedLora]:
|
||||
"""Resolve request (id, weight) pairs, dropping zero-weight entries.
|
||||
|
||||
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."""
|
||||
Maps the named not-found/gated Hub errors (bad repo/revision/file/gating) to a 400
|
||||
and scrubs the URL from the message; deliberately does NOT catch the base
|
||||
HfHubHTTPError so a Hub 5xx stays a 500. A mid-download cancel also maps to a 409
|
||||
instead of a generic 500."""
|
||||
from huggingface_hub.errors import (
|
||||
EntryNotFoundError,
|
||||
GatedRepoError,
|
||||
|
|
|
|||
|
|
@ -241,18 +241,10 @@ def estimate_safetensors_dense_mib(
|
|||
) -> 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 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.
|
||||
"""
|
||||
Unlike a GGUF (dequantised on load, so a 4-bit file expands ~4x), a safetensors
|
||||
checkpoint usually loads near its on-disk size (None passes through unchanged).
|
||||
Exception: ``fp8_upcast`` -- an fp8 single-file transformer loads with no
|
||||
quantization_config, so diffusers upcasts it to bf16 (~2x on-disk resident)."""
|
||||
if storage_mib is None:
|
||||
return None
|
||||
if fp8_upcast:
|
||||
|
|
|
|||
|
|
@ -737,11 +737,8 @@ 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).
|
||||
# strength 0 (or None) disables ControlNet (documented on the request model, matches
|
||||
# the diffusers path), so it must be a no-op here too, not a hard 400.
|
||||
if controlnet is not None and controlnet[3] in (None, 0, 0.0):
|
||||
controlnet = None
|
||||
if controlnet is not None:
|
||||
|
|
|
|||
|
|
@ -127,10 +127,7 @@ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
|
|||
highest priority first: the cmake ``build/bin`` tree, then a Windows Release
|
||||
subdir, then the root itself, then the prebuilt archive's versioned subdir.
|
||||
|
||||
The prebuilt archive extracts into a top-level versioned dir
|
||||
(``sd-master-<tag>-bin-<host>/``) rather than flattening into ``root``, so without
|
||||
the ``root/*/`` glob a fresh prebuilt install is invisible here -- which silently
|
||||
demotes the persistent sd-server to one-shot mode and re-downloads on every start."""
|
||||
The prebuilt lands in its own versioned subdir rather than flattening into ``root``."""
|
||||
name = _binary_name(stem)
|
||||
cands = [
|
||||
root / "build" / "bin" / name,
|
||||
|
|
@ -138,9 +135,7 @@ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
|
|||
root / "bin" / name,
|
||||
root / name,
|
||||
]
|
||||
# Prebuilt archive layout: root/sd-master-<tag>-bin-<host>/<name> (+ its own bin/).
|
||||
# Newest install first (by mtime -- tag strings don't sort numerically, so a lexical
|
||||
# sort would rank build 99 above build 100).
|
||||
# Newest install first, by mtime -- tag strings don't sort numerically.
|
||||
try:
|
||||
subdirs = [p for p in root.iterdir() if p.is_dir()]
|
||||
subdirs.sort(key = lambda p: p.stat().st_mtime, reverse = True)
|
||||
|
|
|
|||
|
|
@ -63,10 +63,8 @@ def _select_lora_targets(
|
|||
) -> tuple[str, ...]:
|
||||
"""Pick the LoRA target modules for a DiT run.
|
||||
|
||||
``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."""
|
||||
An empty ``cfg_targets`` means "unset": use the family's ``spec.lora_targets``. Any
|
||||
explicit tuple is a deliberate override and still wins."""
|
||||
if not tuple(cfg_targets):
|
||||
return tuple(spec_targets)
|
||||
return tuple(cfg_targets)
|
||||
|
|
|
|||
|
|
@ -211,7 +211,7 @@ 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.
|
||||
# Empty (unset) config means use the family default.
|
||||
unet_targets = list(cfg.lora_target_modules) or list(DEFAULT_LORA_TARGETS)
|
||||
unet.add_adapter(
|
||||
LoraConfig(
|
||||
|
|
|
|||
|
|
@ -30,14 +30,11 @@ from core.inference.diffusion_families import (
|
|||
trainable_family_names,
|
||||
)
|
||||
|
||||
# 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 target modules: the SDXL U-Net attention projections. DiT trainers supply
|
||||
# their own wider set instead.
|
||||
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.
|
||||
# diffusers' SchedulerType names (diffusers.optimization.get_scheduler).
|
||||
_LR_SCHEDULERS: frozenset[str] = frozenset(
|
||||
{
|
||||
"linear",
|
||||
|
|
@ -50,10 +47,8 @@ _LR_SCHEDULERS: frozenset[str] = frozenset(
|
|||
}
|
||||
)
|
||||
|
||||
# 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.
|
||||
# DiT families whose fp32 RoPE/embedder overflow fp16, so they train in bf16 only. Must stay
|
||||
# in sync with the DiT trainer's own specs (kept separate to avoid an import cycle).
|
||||
_FORCE_BF16_FAMILIES: frozenset[str] = frozenset({"qwen-image", "z-image"})
|
||||
|
||||
_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
||||
|
|
@ -104,12 +99,8 @@ 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.
|
||||
# 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.
|
||||
# Exempt a local diffusers checkout that merely has "gguf" in its path, identified by its
|
||||
# ``model_index.json`` marker (same marker the loader uses), not a bare ``is_dir()``.
|
||||
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):
|
||||
|
|
@ -244,8 +235,7 @@ class DiffusionLoraConfig:
|
|||
lora_rank: int = 16
|
||||
lora_alpha: Optional[int] = None # defaults to lora_rank
|
||||
lora_dropout: float = 0.0
|
||||
# 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.
|
||||
# Empty = "unset": each trainer supplies its own family default.
|
||||
lora_target_modules: tuple[str, ...] = ()
|
||||
seed: int = 42
|
||||
mixed_precision: str = "bf16" # "bf16" | "fp16" | "no"
|
||||
|
|
@ -290,9 +280,7 @@ 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.
|
||||
# Refuse fp16 for a bf16-only DiT family up front, before evicting resident models.
|
||||
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 "
|
||||
|
|
@ -312,8 +300,7 @@ 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
|
||||
# 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.
|
||||
# Leave an unset (empty) target list empty so the trainer fills the family default.
|
||||
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.
|
||||
|
|
@ -441,9 +428,8 @@ def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Option
|
|||
alias = sanitize_alias(base)
|
||||
src_resolved = Path(lora_path).resolve()
|
||||
dest = loras_dir() / f"{alias}.safetensors"
|
||||
# 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.
|
||||
# A retrain with the same adapter name must not clobber a prior mirror: pick the next
|
||||
# free numeric suffix instead.
|
||||
if dest.exists() and dest.resolve() != src_resolved:
|
||||
n = 2
|
||||
while True:
|
||||
|
|
|
|||
|
|
@ -698,9 +698,8 @@ 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. 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.
|
||||
# sets them is not silently trained with defaults. Empty = "unset": the trainer fills in
|
||||
# its family default.
|
||||
lora_target_modules: List[str] = Field(
|
||||
default_factory = list,
|
||||
description = "Modules to attach LoRA to; empty = the trainer's family default",
|
||||
|
|
|
|||
|
|
@ -1389,10 +1389,8 @@ async def upload_diffusion_dataset(
|
|||
detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
|
||||
)
|
||||
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.
|
||||
# Roll back all files written this request on a mid-batch failure (size limit,
|
||||
# disk error, disconnect) so the dataset is never left partially populated.
|
||||
written: list[Path] = []
|
||||
committed = False
|
||||
try:
|
||||
|
|
@ -1884,8 +1882,7 @@ 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:
|
||||
# First writer wins over sorted manifests, so the plain manifest
|
||||
# (e.g. output_file.jsonl) is deterministic rather than OS-visit order.
|
||||
# First writer wins over sorted manifests, for deterministic results.
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -131,9 +131,7 @@ export function DiffusionTrainDialog({
|
|||
? Math.min(100, Math.round((status.step / status.total_steps) * 100))
|
||||
: 0;
|
||||
|
||||
// Notify the parent exactly once when a run finishes with a saved adapter, so it can
|
||||
// rescan the LoRA picker (a LoRA trained while a model is loaded is otherwise invisible
|
||||
// until a model swap re-runs the discovery effect).
|
||||
// Notify the parent exactly once per finished run so it rescans the LoRA picker.
|
||||
const [notifiedComplete, setNotifiedComplete] = useState(false);
|
||||
useEffect(() => {
|
||||
if (hasSavedAdapter && !notifiedComplete) {
|
||||
|
|
|
|||
|
|
@ -1237,9 +1237,7 @@ 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);
|
||||
// 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.
|
||||
// Negative prompt only applies when guidance>0; don't restore a hidden value.
|
||||
setNegativePrompt(image.guidance > 0 ? (image.negative_prompt ?? "") : "");
|
||||
setSteps(image.steps);
|
||||
setGuidance(image.guidance);
|
||||
|
|
@ -1275,9 +1273,7 @@ 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.
|
||||
// The control image isn't persisted, so clear any stale ControlNet selection.
|
||||
setControlnetId("");
|
||||
setControlImage(null);
|
||||
toast.success("Settings restored to inputs");
|
||||
|
|
@ -1528,9 +1524,7 @@ 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).
|
||||
// Optimistically clear the quant label, revert it if the load never starts.
|
||||
const prevQuant = quant;
|
||||
quantRevert.current = { prev: prevQuant };
|
||||
setQuant(null);
|
||||
|
|
@ -1600,9 +1594,7 @@ 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).
|
||||
// Optimistically clear the quant label, revert it if the load never starts.
|
||||
const prevQuant = quant;
|
||||
quantRevert.current = { prev: prevQuant };
|
||||
setQuant(null);
|
||||
|
|
|
|||
|
|
@ -318,11 +318,8 @@ export function DiffusionTrainPanel({
|
|||
// terminal "completed" status until the next start, so we can't rely on it clearing).
|
||||
const [dismissedJobId, setDismissedJobId] = useState<string | null>(null);
|
||||
const running = Boolean(status?.active) || status?.status === "running";
|
||||
// A stopped run still saves + catalog-publishes a real, deployable adapter (status
|
||||
// carries catalog_path), so treat "stopped with an adapter" as terminal-with-adapter
|
||||
// too -- otherwise the normal "stop once the loss looks good" flow leaves the trained
|
||||
// adapter with no Deploy button and no picker refresh. A save=False cancel has no
|
||||
// catalog_path, so it correctly still shows nothing.
|
||||
// A stopped run still saves + publishes a deployable adapter (catalog_path set), so
|
||||
// treat it as finished-with-adapter too; a save=False cancel has no catalog_path.
|
||||
const hasSavedAdapter =
|
||||
status?.status === "completed" ||
|
||||
(status?.status === "stopped" && Boolean(status?.catalog_path));
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue