From 756c7157218ea2fb90f61b893923fe5fc3bb2790 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Sat, 4 Jul 2026 12:54:47 -0300 Subject: [PATCH] Trim the verbose comments added by the fixes --- scripts/uninstall.ps1 | 5 +- scripts/uninstall.sh | 8 +-- studio/backend/core/inference/diffusion.py | 62 +++++++------------ .../core/inference/diffusion_controlnet.py | 14 ++--- .../core/inference/diffusion_device.py | 13 ++-- .../backend/core/inference/diffusion_lora.py | 24 ++----- .../core/inference/diffusion_memory.py | 16 ++--- .../backend/core/inference/sd_cpp_backend.py | 7 +-- .../backend/core/inference/sd_cpp_engine.py | 9 +-- .../core/training/diffusion_dit_trainer.py | 6 +- .../core/training/diffusion_lora_trainer.py | 2 +- .../core/training/diffusion_train_common.py | 38 ++++-------- studio/backend/models/training.py | 5 +- studio/backend/routes/training.py | 9 +-- .../images/diffusion-train-dialog.tsx | 4 +- .../src/features/images/images-page.tsx | 16 ++--- .../images/train/diffusion-train-panel.tsx | 7 +-- 17 files changed, 75 insertions(+), 170 deletions(-) diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 2de89637b6..0f2aca669b 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -380,9 +380,8 @@ function Uninstall-UnslothStudio { continue } _RemovePath $r - # The native diffusion sibling (.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 (.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 } diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index 36d974844d..320eb29c07 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -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 .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 (.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 diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index fa3e168c48..cdac37c2b1 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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)." diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index 9f42350c73..170b48bd0f 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py index 1c5dff9580..f0dadcf312 100644 --- a/studio/backend/core/inference/diffusion_device.py +++ b/studio/backend/core/inference/diffusion_device.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index bd3e2d3784..cee8d51242 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -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, diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 090e295a2e..065fd4eb11 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -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: diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 6c54af5e87..577370468b 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -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: diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index af896007b5..fe1e7ebc92 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -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--bin-/``) 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--bin-/ (+ 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) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 80e40abcaa..226b011320 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -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) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index b96dc1fa07..3424df7dc3 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -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( diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 0040650f0c..e4a83b8c50 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -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 (-2, -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: diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 522d5d3521..224caad35d 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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", diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index a82491c04f..134fdbf8c2 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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( diff --git a/studio/frontend/src/features/images/diffusion-train-dialog.tsx b/studio/frontend/src/features/images/diffusion-train-dialog.tsx index 016087bd8d..17d6dc8193 100644 --- a/studio/frontend/src/features/images/diffusion-train-dialog.tsx +++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx @@ -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) { diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index b518ad2cb9..bb4f95b8e3 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -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); diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 04fdeded5c..7111307907 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -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(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));