diff --git a/scripts/nvfp4_t211_probe.py b/scripts/nvfp4_t211_probe.py index e69e01b2de..de697165ea 100644 --- a/scripts/nvfp4_t211_probe.py +++ b/scripts/nvfp4_t211_probe.py @@ -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 diff --git a/scripts/sparse_accum_probe.py b/scripts/sparse_accum_probe.py index 4f91f905b6..74006334f3 100644 --- a/scripts/sparse_accum_probe.py +++ b/scripts/sparse_accum_probe.py @@ -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): diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 1887143bf7..742a07c221 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -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 { diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 956687c4d6..c04907f9c5 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -19,11 +19,12 @@ bar. GPU-handoff policy lives in the arbiter the routes call, not here. from __future__ import annotations import inspect +import json import threading import time from dataclasses import dataclass from pathlib import Path -from typing import Any, Optional +from typing import Any, Callable, Optional from loggers import get_logger from utils.hardware import clear_gpu_cache @@ -1010,10 +1011,14 @@ class DiffusionBackend: return total, base_files @staticmethod - def _cache_bytes(repo_id: str) -> int: + def _hub_cache_repo_dir(repo_id: str) -> Path: + """Local HF hub cache dir for ``repo_id``.""" from huggingface_hub import constants + return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" - blobs = Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" / "blobs" + @staticmethod + def _cache_bytes(repo_id: str) -> int: + blobs = DiffusionBackend._hub_cache_repo_dir(repo_id) / "blobs" total = 0 try: for entry in blobs.iterdir(): @@ -1049,19 +1054,72 @@ class DiffusionBackend: continue return total + @staticmethod + def _max_over_cached_revs(base: str, fn: Callable[[Path], int]) -> int: + """Apply ``fn`` to a LOCAL diffusers dir, or to the fullest cached hub snapshot + revision (the active one is the fullest), returning that count. 0 when nothing is + cached. Multiple revisions may be cached, so take the max.""" + local = Path(base).expanduser() + if local.is_dir(): + return fn(local) + snapshots = DiffusionBackend._hub_cache_repo_dir(base) / "snapshots" + if not snapshots.is_dir(): + return 0 + return max((fn(rev) for rev in snapshots.iterdir() if rev.is_dir()), default = 0) + @staticmethod def _companion_cache_bytes(base: str) -> int: """Resident companion (VAE + text-encoder) size for the memory plan. - For a hub base repo this is the cached blob total (``_cache_bytes``). For a - LOCAL diffusers base directory the blob cache is empty, so sum the on-disk - component weights instead, excluding ``transformer/`` (the GGUF supplies the - transformer). Without this a local base folds its multi-GB VAE / text-encoder - weights to zero and auto planning can pick a resident placement that OOMs.""" - local = Path(base).expanduser() - if local.is_dir(): - return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True) - return DiffusionBackend._cache_bytes(base) + 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.""" + return DiffusionBackend._max_over_cached_revs( + base, lambda d: DiffusionBackend._local_dir_weight_bytes(d, exclude_transformer = True) + ) + + @staticmethod + def _safetensors_param_count(path: Path) -> int: + """Total tensor elements in a safetensors file, read from its JSON header without + touching the tensor data. 0 on any read/parse failure.""" + try: + with open(path, "rb") as fh: + header_len = int.from_bytes(fh.read(8), "little") + header = json.loads(fh.read(header_len)) + total = 0 + for name, meta in header.items(): + if name == "__metadata__" or not isinstance(meta, dict): + continue + numel = 1 + for dim in meta.get("shape", []): + numel *= dim + total += numel + return total + except Exception: # noqa: BLE001 — best-effort estimate; a corrupt/crafted shard + # (bad header length, non-dict header, odd shape) must degrade to 0 so the caller + # gates on the plain plan, never crash the load. + return 0 + + @staticmethod + def _dense_transformer_resident_bytes(base: str) -> int: + """Resident bf16 size of the base repo's dense ``transformer/`` for the dense-quant + preflight. That fast path loads the transformer at the compute dtype (bf16, 2 + bytes/param) before quantizing, so budget num_params * 2 -- NOT the on-disk bytes, + which for an F32 base (e.g. Z-Image) are ~2x the resident size. Read from the + safetensors shard headers. Returns 0 when no ``transformer/*.safetensors`` shards + are present (an uncached base, or a .bin-only transformer); the caller then gates + the fast path on the plain plan.""" + + def _params(d: Path) -> int: + tdir = d / "transformer" + if not tdir.is_dir(): + return 0 + return sum( + DiffusionBackend._safetensors_param_count(s) for s in tdir.glob("*.safetensors") + ) + + return DiffusionBackend._max_over_cached_revs(base, _params) * 2 # bf16: 2 bytes/param # ── Synchronous load / generate / unload ─────────────────────────────── @@ -1181,9 +1239,8 @@ class DiffusionBackend: pipeline_cls = getattr(diffusers, fam.pipeline_class) # Decide placement up front (the weights are still on CPU, so free VRAM is - # the real budget) -- this also doubles as the dense-quant preflight: the - # dense bf16 transformer must fit resident, so the fast path is offered only - # when the plan is `none`. + # the real budget). This plan budgets the GGUF file and places the plain + # load; the dense-quant fast path is preflighted separately below. plan = self._plan_memory( target, single_file_path, @@ -1227,48 +1284,96 @@ class DiffusionBackend: pipe = None transformer_quant_engaged = None quant_plan = None + # The GGUF-size `plan` can mis-budget the dense-quant fast path two ways, so + # preflight the real footprint BEFORE evicting the current pipeline. Both + # branches need the base repo + a resolved scheme, so gate on the dense-path + # preconditions first. + dense_declined = False if ( kind == "gguf" and normalize_transformer_quant(transformer_quant) is not None and dense_transformer_supported(target) - and plan.offload_policy != OFFLOAD_NONE ): - # The GGUF-size plan picked offload, but the dense-quant artifact has - # a DIFFERENT footprint: int8/fp8 weights are ~half the bf16 bytes, and - # a pre-quantized checkpoint never materialises dense bf16 at all. Ask - # the auto-policy for the candidate's estimate and re-plan against it: - # a resident quantised build beats an offloaded GGUF on speed AND - # quality, so it must be attempted before settling for offload. - candidate = resolve_dense_quant_candidate( - fam = fam, - target = target, - requested = transformer_quant, - base_repo = base, - prequant_path = transformer_prequant_path, - logger = logger, - ) - if candidate is not None: - replanned = self._plan_memory( - target, - single_file_path, - base, - fam, - memory_mode, - cpu_offload, - kind = kind, - repo_id = repo_id, - transformer_resident_override_mib = (candidate.transient_transformer_mib), - # The dense path prefetches the base transformer/ shards into the - # cache _companion_cache_bytes reads; pass the auto-policy's own - # companion estimate so the re-plan does not double-count them. - companion_override_mib = candidate.companions_mib, + if plan.offload_policy != OFFLOAD_NONE: + # The GGUF-size plan picked offload, but the dense-quant artifact has + # a DIFFERENT footprint: int8/fp8 weights are ~half the bf16 bytes, and + # a pre-quantized checkpoint never materialises dense bf16 at all. Ask + # the auto-policy for the candidate's estimate and re-plan against it: + # a resident quantised build beats an offloaded GGUF on speed AND + # quality, so it must be attempted before settling for offload. + candidate = resolve_dense_quant_candidate( + fam = fam, + target = target, + requested = transformer_quant, + base_repo = base, + prequant_path = transformer_prequant_path, + logger = logger, ) - if replanned.offload_policy == OFFLOAD_NONE: - quant_plan = replanned + if candidate is not None: + replanned = self._plan_memory( + target, + single_file_path, + base, + fam, + memory_mode, + cpu_offload, + kind = kind, + repo_id = repo_id, + transformer_resident_override_mib = ( + candidate.transient_transformer_mib + ), + # The dense path prefetches the base transformer/ shards into the + # cache _companion_cache_bytes reads; pass the auto-policy's own + # companion estimate so the re-plan does not double-count them. + companion_override_mib = candidate.companions_mib, + ) + if replanned.offload_policy == OFFLOAD_NONE: + quant_plan = replanned + else: + # The GGUF fits resident, but this path first materialises the base + # repo's dense bf16 transformer -- bigger than the quantised GGUF -- so + # re-check the fit against THAT. A card that fits the GGUF but not the + # dense transformer must skip the fast path up front, not evict the + # current pipeline then OOM in finalization. A prequant checkpoint loads + # a small quantised file (no dense bf16), so skip the re-check there + # (mirrors the prefetch guard). _dense_transformer_resident_bytes reads + # the on-disk shard headers, so it also covers families the size table + # (resolve_dense_quant_candidate) does not list; it returns 0 when the + # shards are absent, in which case the fast path keeps today's behaviour. + scheme = select_transformer_quant_scheme( + target, + transformer_quant, # normalized above + family = getattr(fam, "name", None), + ) + prequant = ( + resolve_prequant_source( + fam, scheme, path_override = transformer_prequant_path + ) + if scheme is not None + else None + ) + if prequant is None: + dense_mib = int( + self._dense_transformer_resident_bytes(base) // (1024 * 1024) + ) + if dense_mib > 0: + dense_plan = self._plan_memory( + target, + single_file_path, + base, + fam, + memory_mode, + cpu_offload, + kind = kind, + repo_id = repo_id, + transformer_resident_override_mib = dense_mib, + ) + dense_declined = dense_plan.offload_policy != OFFLOAD_NONE if ( kind == "gguf" and normalize_transformer_quant(transformer_quant) is not None and dense_transformer_supported(target) + and not dense_declined and (plan.offload_policy == OFFLOAD_NONE or quant_plan is not None) ): try: @@ -1837,7 +1942,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. ``transformer_resident_override_mib`` replaces the file-size transformer estimate when the loader is planning for a DIFFERENT artifact than the file on disk (the @@ -1900,9 +2006,15 @@ class DiffusionBackend: # file-size derivation; companions below stay measured from the cache. transformer_resident = transformer_resident_override_mib elif kind == "single_file": - # Safetensors single-file: no dequant expansion (it carries its dtype). + # 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 + ) 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)) @@ -2011,6 +2123,13 @@ class DiffusionBackend: _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None) if _cn_fs.blocked: raise ValueError(_cn_fs.reason) + # 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() + clear_gpu_cache() import torch # state.dtype is the display string saved at load ("bfloat16"), NOT a @@ -2415,6 +2534,12 @@ 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 mask would be silently dropped. + 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: diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index b318b040c8..a4e482121a 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -204,8 +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 it maps to nothing (the caller omits control_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, @@ -220,13 +220,24 @@ _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. - Returns a mode only for a curated *union* catalog entry AND a control type that maps to a - known index; otherwise None so the caller omits the kwarg (a non-union ControlNet has a - single fixed mode, and 'passthrough' does not name one). Pure lookup, no network.""" + A union model requires a concrete mode (diffusers raises on None). A known mode maps to its + index; ``passthrough`` (or an empty type) carries no intrinsic mode and defaults to 0 (the + canny head). An unknown/typo'd type (e.g. 'detph') raises ValueError so the route rejects it + with a 400 instead of silently running the canny head against a map meant for another mode, + which would produce wrong conditioning. 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 - return _UNION_CONTROL_MODES.get((control_type or "").strip().lower()) + ct = (control_type or "").strip().lower() + if ct in _UNION_CONTROL_MODES: + return _UNION_CONTROL_MODES[ct] + if ct in ("", "passthrough"): + return 0 # already-preprocessed map with no intrinsic mode; canny is the default head + raise ValueError( + f"Unknown control type {control_type!r} for a union ControlNet. Use one of: " + f"{', '.join(sorted(_UNION_CONTROL_MODES))}, or passthrough." + ) def preprocess_control(image: Any, control_type: str) -> Any: diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py index 2d2cc11343..f0dadcf312 100644 --- a/studio/backend/core/inference/diffusion_device.py +++ b/studio/backend/core/inference/diffusion_device.py @@ -14,6 +14,7 @@ dtype choice and the capability flags the backend keys optimisation paths off. from __future__ import annotations +import os from dataclasses import dataclass from typing import Any, Optional @@ -234,6 +235,11 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget: mps_available = False if mps_available: + # 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 # far outside float16's finite range (~6.5e4) -- Z-Image's MLP @@ -254,7 +260,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", diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index a971ffb8b3..5d886cf1bd 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -293,6 +293,13 @@ 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.""" + 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]], *, @@ -301,20 +308,31 @@ 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.""" + 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, + 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 diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 913574d2b6..cc672a3741 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -236,14 +236,19 @@ 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). - """ + 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: + return storage_mib * 2 return storage_mib diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 6f0d7cae96..d8f9e89429 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -768,6 +768,10 @@ 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, 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: raise ValueError( "ControlNet is not yet supported on the native sd.cpp engine; run on a GPU " diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index c6582b0cb8..0838edf4d6 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -125,7 +125,9 @@ def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]: """``stem`` locations under a stable-diffusion.cpp checkout/install ``root``, highest priority first: the cmake ``build/bin`` tree, then a Windows Release - subdir, then the root itself.""" + subdir, then the root itself, then the prebuilt archive's versioned subdir. + + The prebuilt lands in its own versioned subdir rather than flattening into ``root``.""" name = _binary_name(stem) cands = [ root / "build" / "bin" / name, @@ -133,6 +135,15 @@ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]: root / "bin" / name, root / name, ] + # 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) + for sub in subdirs: + cands.append(sub / name) + cands.append(sub / "bin" / name) + except OSError: + pass return cands diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 5d54e4717a..693fe964cd 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -37,6 +37,28 @@ from core.inference.diffusion_families import ( # set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback. DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0") +# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). piecewise_constant is +# intentionally excluded: it is the only scheduler that needs a `step_rules` string, which the +# trainers never pass (get_scheduler is called with only warmup/training steps, and there is no +# config field for it). Accepting it would pass normalized(), free the resident GPU workloads, +# then crash in the trainer subprocess (get_piecewise_constant_schedule does step_rules.split(",") +# on None) -- the exact evict-then-fail the up-front validation exists to prevent. The remaining +# six all run with only warmup/training steps. +_LR_SCHEDULERS: frozenset[str] = frozenset( + { + "linear", + "cosine", + "cosine_with_restarts", + "polynomial", + "constant", + "constant_with_warmup", + } +) + +# 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"} _CAPTION_EXTS = (".txt", ".caption") # diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it. @@ -85,7 +107,11 @@ 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: + # 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): 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()}" @@ -471,6 +497,17 @@ 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") + # 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 " + 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}" + ) if not 1 <= int(self.cache_variants) <= 16: raise ValueError("cache_variants must be between 1 and 16") compile_transformer = str(self.compile_transformer or "auto").strip().lower() @@ -870,8 +907,19 @@ 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: pick the next + # free numeric suffix instead. + 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) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index f8120c42ee..0fa43ae5f6 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -722,7 +722,14 @@ class DiffusionTrainingStartRequest(BaseModel): 5.0, gt = 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", + ] = Field("constant") lr_warmup_steps: int = Field(0, ge = 0) center_crop: bool = Field(False) random_flip: bool = Field(True) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 00ddae684d..b35ac90692 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1581,6 +1581,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: # Normalise to a safe basename. Path.name does not split on a backslash on POSIX, so a # Windows client that sends a backslash path in the multipart filename would otherwise be @@ -1597,40 +1600,61 @@ async def upload_diffusion_dataset( status_code = 400, detail = f"Unsupported file '{f.filename}'. Allowed: {exts}", ) - dest = folder / filename # Reject a second IMAGE that shares this one's stem but differs by extension (sample.png # vs sample.jpg): both resolve to the same .txt caption sidecar (the kohya/diffusers # convention the reader, editor, and delete paths all use), so keeping both would silently - # make them share -- and corrupt -- one caption during training. Scan the whole folder, not - # just this batch, because uploads accumulate (earlier-in-batch files are already on disk). - # Re-uploading the exact same name (same stem AND extension) stays allowed as an overwrite; - # caption/text files are exempt (attaching sample.txt to sample.png is the intended flow). + # share -- and corrupt -- one caption during training. Check both files already on disk + # (uploads accumulate) and earlier images validated in THIS batch (nothing is on disk yet + # in this up-front pass). Re-uploading the exact same name (same stem AND extension) stays + # an overwrite; caption/text files are exempt (sample.txt for sample.png is intended). if ext in _DIFFUSION_DATASET_IMAGE_EXTS: stem = Path(filename).stem - for existing in folder.iterdir(): - if ( - existing.is_file() - and existing.name != filename - and existing.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS - and existing.stem == stem - ): - raise HTTPException( - status_code = 400, - detail = ( - f"Duplicate image name '{stem}'. '{existing.name}' is already in this " - f"dataset; two images sharing a name would share one '{stem}.txt' " - f"caption. Rename one before uploading." - ), - ) - # Stream into a sibling temp file and only atomically promote it once the whole file - # is written and within the limit. A mid-stream 413 (or any abort) then removes the - # TEMP file, never dest, so re-uploading a batch that trips the limit can no longer - # truncate/delete an example that was already stored under the same name. - fd, tmp_name = tempfile.mkstemp(dir = folder, prefix = ".upload-", suffix = ext) - tmp = Path(tmp_name) - complete = False - try: - with os.fdopen(fd, "wb") as out: + clash = next( + ( + p.name + for p in folder.iterdir() + if p.is_file() + and p.name != filename + and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS + and p.stem == stem + ), + None, + ) + if clash is None: + clash = next( + ( + n + for n in names + if n != filename + and Path(n).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS + and Path(n).stem == stem + ), + None, + ) + if clash is not None: + raise HTTPException( + status_code = 400, + detail = ( + f"Duplicate image name '{stem}'. '{clash}' is already in this " + f"dataset; two images sharing a name would share one '{stem}.txt' " + f"caption. Rename one before uploading." + ), + ) + names.append(filename) + # Stage each file to a temp name and only move it into place once the whole batch is + # written, so a mid-batch failure (size limit, disk error, disconnect) leaves the + # dataset untouched -- including any pre-existing file that shares a name, which a + # direct write would have truncated (repeat uploads into the same name accumulate). + staged: list[tuple[Path, Path]] = [] # (temp, final) + committed = False + try: + for f, filename in zip(files, names): + dest = folder / filename + # A filename-independent temp name so a long (but valid, <= NAME_MAX) filename + # can't overflow NAME_MAX once the staging suffix is added. + tmp = folder / f".upload-{_uuid.uuid4().hex}.part" + staged.append((tmp, dest)) + with open(tmp, "wb") as out: while chunk := await f.read(1024 * 1024): total_bytes += len(chunk) if total_bytes > limit_bytes: @@ -1643,15 +1667,17 @@ async def upload_diffusion_dataset( ), ) out.write(chunk) - os.replace(tmp, dest) - complete = True - finally: - if not complete: + uploaded += 1 + for tmp, dest in staged: + tmp.replace(dest) # atomic on the same filesystem + committed = True + finally: + if not committed: + for tmp, _ in staged: try: tmp.unlink(missing_ok = True) except OSError: pass - uploaded += 1 summary = _diffusion_dataset_summary(folder) return DiffusionDatasetUploadResponse( @@ -2123,7 +2149,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: @@ -2134,7 +2160,8 @@ 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, 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( p diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 6503bfa794..b6e8ee6570 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2396,6 +2396,112 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo assert _FakeTransformer.last["path"] # GGUF path used +def test_dense_quant_skipped_when_dense_transformer_does_not_fit( + fake_runtime, tmp_path, monkeypatch +): + # The GGUF fits resident (plan `none`), but the DENSE bf16 transformer the fast path + # materializes does not. The fast path must be skipped up front (preflighted against + # the dense transformer, not the GGUF), and GGUF loads RESIDENT -- not evicted, OOMed + # in finalization, then offloaded. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + # A scheme resolves and there is no prequant, so the dense bf16 is materialized and the + # dense-fit re-check runs against a large (won't-fit) dense transformer. + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" + ) + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) + monkeypatch.setattr( + DiffusionBackend, + "_dense_transformer_resident_bytes", + staticmethod(lambda base: 40 * 1024**3), + ) + orig_plan = DiffusionBackend._plan_memory + + def plan_wrap( + self, + *a, + transformer_resident_override_mib = None, + **k, + ): + # GGUF budget fits (real plan -> none); the dense-transformer preflight does not. + if transformer_resident_override_mib is not None: + return types.SimpleNamespace(offload_policy = "model") + return orig_plan(self, *a, **k) + + monkeypatch.setattr(DiffusionBackend, "_plan_memory", plan_wrap) + + @classmethod + def _fp_fail(cls, *a, **k): + pytest.fail("dense transformer must not load when it won't fit resident") + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False) + (tmp_path / "m.gguf").write_bytes(b"x") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + ) + assert status["transformer_quant"] is None # dense quant skipped + assert status["offload_policy"] == "none" # GGUF loaded resident, not offloaded + assert _FakeTransformer.last["path"] # GGUF path used + + +def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypatch): + # With a prequant checkpoint, the fast path loads the small quantized file, not the + # dense bf16 -- so the dense-transformer re-check must NOT run and must NOT decline the + # fast path, even when the base's dense shards happen to be cached and large. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" + ) + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: "prequant/path") + # Large dense shards cached: if the re-check ran, it would wrongly decline the fast path. + monkeypatch.setattr( + DiffusionBackend, + "_dense_transformer_resident_bytes", + staticmethod(lambda base: 999 * 1024**3), + ) + dense_refit_ran = [] + orig_plan = DiffusionBackend._plan_memory + + def spy_plan( + self, + *a, + transformer_resident_override_mib = None, + **k, + ): + if transformer_resident_override_mib is not None: + dense_refit_ran.append(True) + return orig_plan(self, *a, **k) + + monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan) + attempted = [] + + def fake_dense_load(self, *a, **k): + attempted.append(True) + return None, None # fall through to GGUF; we only assert the path was reached + + monkeypatch.setattr(DiffusionBackend, "_load_dense_quant_pipeline", fake_dense_load) + (tmp_path / "m.gguf").write_bytes(b"x") + backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + ) + assert dense_refit_ran == [] # prequant -> dense re-check skipped + assert attempted == [True] # fast path still attempted (with the prequant) + + def test_transformer_quant_unsupported_scheme_skips_dense_download( fake_runtime, tmp_path, monkeypatch ): diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index bcb7393457..d5282299d4 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -56,15 +56,31 @@ def test_resolve_controlnet_enforces_family_match(): def test_union_control_mode_maps_only_union_entries(): - # Union entries map a known control type to its integer mode; passthrough / unknown - # types and non-union ids return None so the caller omits control_mode. + # Union entries map a known control type to its integer mode; a union model always + # needs a concrete mode, so an unmapped type (passthrough) defaults to 0. A non-union + # id returns None so the caller omits control_mode. assert dc.union_control_mode("flux-union-pro", "canny") == 0 assert dc.union_control_mode("flux-union-pro", "depth") == 2 assert dc.union_control_mode("flux-union-pro", "pose") == 4 - assert dc.union_control_mode("flux-union-pro", "passthrough") is None + assert dc.union_control_mode("flux-union-pro", "passthrough") == 0 assert dc.union_control_mode("some/bare-repo", "canny") is None +def test_union_control_mode_rejects_unknown_type(): + # An unknown / typo'd control type (e.g. 'detph') must NOT silently fall back to the canny + # head (0): preprocess_control passes non-canny maps through unchanged, so mode 0 would + # condition a map meant for another mode as canny -- silently wrong. Only passthrough (or an + # empty type) defaults to 0; anything else raises so the route returns a 400. + with pytest.raises(ValueError, match = "Unknown control type"): + dc.union_control_mode("flux-union-pro", "detph") + with pytest.raises(ValueError, match = "Unknown control type"): + dc.union_control_mode("flux-union-pro", "scribble") + # passthrough and empty still default to 0 (the intended no-intrinsic-mode case); a non-union + # entry is unaffected (returns None, never raises). + assert dc.union_control_mode("flux-union-pro", "") == 0 + assert dc.union_control_mode("some/bare-repo", "detph") is None + + def test_resolve_controlnet_local(tmp_path, monkeypatch): d = tmp_path / "controlnets" d.mkdir() diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py index a60140c7fc..6b1a980e7f 100644 --- a/studio/backend/tests/test_diffusion_lora.py +++ b/studio/backend/tests/test_diffusion_lora.py @@ -190,6 +190,29 @@ 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): + # response is optional in huggingface_hub 0.x but required in 1.x; both only + # read .headers / .request, so a stub keeps the test working on either. + raise RepositoryNotFoundError( + "404 Client Error. Repository Not Found for url: " + "https://huggingface.co/api/models/nope/nope (Request ID: abc)", + response = types.SimpleNamespace(headers = {}, request = None), + ) + + 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. diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index 17bee1c1ae..dd065b155e 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -240,6 +240,49 @@ def test_config_from_dict_threads_num_epochs(): assert cfg.num_epochs == 12 +def test_normalized_rejects_piecewise_constant(): + # piecewise_constant needs a step_rules string the trainers never supply, so get_scheduler() + # would crash in the trainer subprocess AFTER the resident GPU workloads are freed. It must be + # rejected up front (a clean ValueError -> 400), not accepted like the other schedulers. + with pytest.raises(ValueError, match = "lr_scheduler"): + DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "piecewise_constant" + ).normalized() + + +def test_normalized_accepts_supported_schedulers(): + # Every scheduler in the allow-list runs with only warmup/training steps (no extra required arg). + for sched in ( + "linear", + "cosine", + "cosine_with_restarts", + "polynomial", + "constant", + "constant_with_warmup", + ): + cfg = DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = sched + ).normalized() + assert cfg.lr_scheduler == sched + + +def test_api_scheduler_enum_never_advertises_a_rejected_scheduler(): + # The request-model enum must not offer a scheduler that normalized() rejects: a client that + # picks it straight from the schema would get a 400. Every option the API advertises must be in + # the validation allow-list (this guards against the enum and allow-list drifting apart again, + # e.g. piecewise_constant left in one but removed from the other). + import typing + + from core.training.diffusion_train_common import _LR_SCHEDULERS + from models.training import DiffusionTrainingStartRequest + + api_options = set( + typing.get_args(DiffusionTrainingStartRequest.model_fields["lr_scheduler"].annotation) + ) + assert api_options and api_options <= _LR_SCHEDULERS, api_options - _LR_SCHEDULERS + assert "piecewise_constant" not in api_options + + def test_compute_sdxl_add_time_ids(): assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024) @@ -471,3 +514,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") diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index c0cd6b129b..4fe30a8d58 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -1015,6 +1015,50 @@ def test_import_example_partial_failure_leaves_no_partial_dataset( assert not any(p.name.startswith(".dreambooth-dog.import-") for p in ds_root.iterdir()) +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_diffusion_dataset_upload_over_cap_preserves_existing_file( + client, dataset_roots, monkeypatch +): + # A failed batch that reuses an existing filename must NOT delete the user's + # pre-existing file (repeat uploads accumulate). Staging to a temp file keeps the + # original intact until the whole batch commits. + import utils.upload_limits as ul + + monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100) + ds_root, _ = dataset_roots + folder = ds_root / "keep" + folder.mkdir(parents = True) + (folder / "existing.png").write_bytes(b"ORIGINAL") # from an earlier upload + files = [ + ("files", ("existing.png", b"NEW", "image/png")), # re-upload, small + ("files", ("big.png", b"y" * 200, "image/png")), # trips the cap + ] + r = client.post("/api/train/diffusion/dataset", data = {"name": "keep"}, files = files) + assert r.status_code == 413, r.text + assert (folder / "existing.png").read_bytes() == b"ORIGINAL" # untouched + assert not (folder / "big.png").exists() + assert not list(folder.glob(".*.part")) # no leftover temp files + + 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. diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index be2fc22a4b..68caeb58dd 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -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. diff --git a/studio/frontend/src/features/images/diffusion-train-dialog.tsx b/studio/frontend/src/features/images/diffusion-train-dialog.tsx index 25b260a03e..17d6dc8193 100644 --- a/studio/frontend/src/features/images/diffusion-train-dialog.tsx +++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx @@ -120,24 +120,27 @@ export function DiffusionTrainDialog({ }, [open, poll]); const active = Boolean(status?.active) || status?.status === "running"; - const completed = status?.status === "completed"; + // 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)); + const completed = hasSavedAdapter; const pct = status && status.total_steps > 0 ? Math.min(100, Math.round((status.step / status.total_steps) * 100)) : 0; - // Notify the parent exactly once when a run reaches "completed", 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 (status?.status === "completed" && !notifiedComplete) { + if (hasSavedAdapter && !notifiedComplete) { setNotifiedComplete(true); onTrainingComplete?.(); } else if (status?.status === "running" && notifiedComplete) { setNotifiedComplete(false); // arm again for the next run } - }, [status?.status, notifiedComplete, onTrainingComplete]); + }, [hasSavedAdapter, status?.status, notifiedComplete, onTrainingComplete]); const selectedDataset = dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 41b82cdc68..fb51b0ffcf 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1277,7 +1277,8 @@ 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 ?? ""); + // 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); setSeed(String(image.seed)); @@ -1312,6 +1313,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight }); } setLoras(restoredLoras); + // The control image isn't persisted, so clear any stale ControlNet selection. + setControlnetId(""); + setControlImage(null); toast.success("Settings restored to inputs"); }, []); @@ -1649,11 +1653,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) { toast.error("Only unsloth or on-device image models can be loaded here"); return; } + // Optimistically clear the quant label, revert it if the load never starts. + 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], );