Merge remote-tracking branch 'origin/image-generation' into fix/imggen-review-bugs
# Conflicts: # studio/backend/core/inference/diffusion.py # studio/backend/core/training/diffusion_train_common.py # studio/backend/routes/training.py # studio/backend/tests/test_diffusion_dataset_api.py
This commit is contained in:
oobabooga
commit
9d8bf002f8
24 changed files with 633 additions and 52 deletions
|
|
@ -226,6 +226,22 @@
|
|||
"evidence": "FS: L528: history: list[Response] | None = None, sha256:f56272dccd651b2644aa41ef6e688e211462427aad07fef5150240ec7347446e\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):",
|
||||
"evidence_hash": "b32f79e58c938680d89efa74113eeba76c9fc5aedf5de18086f93bef274c4bda"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/_sandbox.py",
|
||||
"check": "C2 polling/beaconing loop detected",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38",
|
||||
"evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/_sandbox.py",
|
||||
"check": "Writes to /tmp and executes (staged dropper)",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"",
|
||||
"evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/hf_api.py",
|
||||
|
|
@ -242,6 +258,14 @@
|
|||
"evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272",
|
||||
"evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/hf_api.py",
|
||||
"check": "C2 polling/beaconing loop detected",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6",
|
||||
"evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/hf_api.py",
|
||||
|
|
@ -266,6 +290,14 @@
|
|||
"evidence": "L298: while True: sha256:6b8e5e569594caf7c4eca6137646dae471a7c3aae7294096cf876f30b5f90306",
|
||||
"evidence_hash": "c066cc27bce31ee7b6ce07411ee7a7d9ecfbf3aafc8848f6641fabfe522a7703"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/utils/_http.py",
|
||||
"check": "C2 polling/beaconing loop detected",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958",
|
||||
"evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e"
|
||||
},
|
||||
{
|
||||
"package": "ipython",
|
||||
"file": "IPython/core/debugger.py",
|
||||
|
|
|
|||
|
|
@ -59,18 +59,20 @@ from .diffusion_speed import (
|
|||
SPEED_OFF,
|
||||
apply_speed_optims,
|
||||
compile_eligible,
|
||||
normalize_speed_mode,
|
||||
resolve_speed_mode,
|
||||
restore_backend_flags,
|
||||
snapshot_backend_flags,
|
||||
)
|
||||
from .diffusion_attention import (
|
||||
apply_attention_backend,
|
||||
normalize_attention_backend,
|
||||
select_attention_backend,
|
||||
)
|
||||
from . import diffusion_compile_cache as compile_cache
|
||||
from . import diffusion_gguf_compile as gguf_compile
|
||||
from .diffusion_cache import apply_step_cache
|
||||
from .diffusion_precision import quantize_text_encoders
|
||||
from .diffusion_cache import apply_step_cache, normalize_transformer_cache
|
||||
from .diffusion_precision import normalize_te_quant, quantize_text_encoders
|
||||
from .diffusion_prequant import (
|
||||
load_prequantized_transformer,
|
||||
resolve_prequant_source,
|
||||
|
|
@ -436,11 +438,17 @@ class DiffusionBackend:
|
|||
base: str,
|
||||
base_files: list[str],
|
||||
hf_token: Optional[str],
|
||||
) -> None:
|
||||
) -> Optional[str]:
|
||||
"""Pre-download the GGUF + the given ``base_files`` into the HF cache,
|
||||
WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's
|
||||
from_single_file / from_pretrained hit the cache and the heavy download can
|
||||
be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``."""
|
||||
be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``.
|
||||
|
||||
Returns the base repo's local snapshot dir when the prefetched set includes
|
||||
the pipeline manifest, so from_pretrained can load from disk instead of
|
||||
re-sweeping the hub (its own sweep also pulls files the scoped list skips,
|
||||
e.g. the 24 GB packaged root singles in each FLUX.1 repo); None otherwise
|
||||
(estimate failure, config-only base, local repo) -> hub id as before."""
|
||||
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
|
||||
|
||||
# GGUF transformer (hub repos only; a local path is already on disk).
|
||||
|
|
@ -449,12 +457,16 @@ class DiffusionBackend:
|
|||
repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event
|
||||
)
|
||||
# Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
|
||||
snapshot_root: Optional[str] = None
|
||||
for rfilename in base_files:
|
||||
if self._cancel_event.is_set():
|
||||
raise RuntimeError("Cancelled")
|
||||
hf_hub_download_with_xet_fallback(
|
||||
local = hf_hub_download_with_xet_fallback(
|
||||
base, rfilename, hf_token, cancel_event = self._cancel_event
|
||||
)
|
||||
if rfilename == "model_index.json":
|
||||
snapshot_root = str(Path(local).parent)
|
||||
return snapshot_root
|
||||
|
||||
def validate_load_request(
|
||||
self,
|
||||
|
|
@ -671,7 +683,7 @@ class DiffusionBackend:
|
|||
self._loading.expected_bytes = expected
|
||||
# Download outside the lock so unload()/an eviction can preempt the
|
||||
# multi-GB pull; load_pipeline below then assembles from the cache.
|
||||
self._prefetch_files(
|
||||
kwargs["_base_local_dir"] = self._prefetch_files(
|
||||
kwargs["repo_id"],
|
||||
kwargs.get("gguf_filename"),
|
||||
base,
|
||||
|
|
@ -690,6 +702,15 @@ class DiffusionBackend:
|
|||
if self._load_token != token:
|
||||
return
|
||||
logger.error("diffusion.load_failed: %s", exc)
|
||||
# Free the debris of a failed construction (e.g. a load-time OOM): _state was
|
||||
# never committed, and the next load's _unload_locked early-returns on a None
|
||||
# state, so nothing else releases the reserved VRAM. Guarded: a sticky CUDA
|
||||
# error makes synchronize() raise, which would skip stamping the REAL error
|
||||
# below and leave the client polling forever.
|
||||
try:
|
||||
clear_gpu_cache()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
# Redact native paths: this error is surfaced verbatim via the
|
||||
# load-progress poll, and Studio can run as a shared server.
|
||||
from utils.native_path_leases import redact_native_paths
|
||||
|
|
@ -931,6 +952,7 @@ class DiffusionBackend:
|
|||
transformer_cache_threshold: Optional[float] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
_load_token: Optional[int] = None,
|
||||
_base_local_dir: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
# A blank / whitespace-only token must degrade to anonymous access, not be passed
|
||||
# as an explicit credential (from_single_file / from_pretrained / the Hub client
|
||||
|
|
@ -951,6 +973,14 @@ class DiffusionBackend:
|
|||
model_kind = model_kind,
|
||||
)
|
||||
kind = resolve_model_kind(gguf_filename, model_kind)
|
||||
# Validate every mode string that can raise NOW, before this load evicts the
|
||||
# previous pipeline below: their first in-line uses all sit past _unload_locked,
|
||||
# where a bad request would cost the user their working model.
|
||||
transformer_quant = normalize_transformer_quant(transformer_quant)
|
||||
normalize_speed_mode(speed_mode)
|
||||
normalize_attention_backend(attention_backend)
|
||||
normalize_transformer_cache(transformer_cache)
|
||||
normalize_te_quant(text_encoder_quant)
|
||||
# For a full pipeline the repo itself supplies every component, so it is its
|
||||
# own base; the single-file kinds resolve the companion base diffusers repo.
|
||||
base = (
|
||||
|
|
@ -1020,7 +1050,7 @@ class DiffusionBackend:
|
|||
# pipeline) do not have.
|
||||
dense_quant_requested = (
|
||||
kind == "gguf"
|
||||
and normalize_transformer_quant(transformer_quant) is not None
|
||||
and transformer_quant is not None # normalized above, pre-eviction
|
||||
and dense_transformer_supported(target)
|
||||
)
|
||||
# `plan` budgets the GGUF file, but this path materializes the base repo's
|
||||
|
|
@ -1076,6 +1106,7 @@ class DiffusionBackend:
|
|||
transformer_quant,
|
||||
transformer_quant_fast_accum,
|
||||
fam = fam,
|
||||
base_local_dir = _base_local_dir,
|
||||
prequant_path = transformer_prequant_path,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — fall back to the GGUF build
|
||||
|
|
@ -1090,7 +1121,12 @@ class DiffusionBackend:
|
|||
# clear_gpu_cache() could not otherwise reclaim that VRAM before the
|
||||
# GGUF build (the OOM-fallback path this cleanup exists for).
|
||||
del exc
|
||||
clear_gpu_cache()
|
||||
# Guarded: after an OOM/sticky CUDA error synchronize() can
|
||||
# raise, and this fallback path must still reach the GGUF build.
|
||||
try:
|
||||
clear_gpu_cache()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
elif dense_quant_requested:
|
||||
# Quant requested but the dense fast path needs a resident load that isn't
|
||||
# available here (an explicit memory_mode offload, or the dense transformer
|
||||
|
|
@ -1111,7 +1147,12 @@ class DiffusionBackend:
|
|||
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype}
|
||||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
|
||||
# The prefetched snapshot dir keeps from_pretrained off the hub:
|
||||
# its own snapshot sweep re-downloads files the scoped prefetch
|
||||
# skipped (root packaged singles, e.g. 24 GB per FLUX.1 repo).
|
||||
pipe = pipeline_cls.from_pretrained(
|
||||
_base_local_dir or repo_id, **pipe_kwargs
|
||||
)
|
||||
elif kind == "single_file" and fam.single_file_is_pipeline:
|
||||
# A single-file SDXL-style checkpoint is the WHOLE pipeline
|
||||
# (U-Net + VAE + both text encoders), not a transformer-only file,
|
||||
|
|
@ -1147,7 +1188,7 @@ class DiffusionBackend:
|
|||
pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
|
||||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
|
||||
pipe = pipeline_cls.from_pretrained(_base_local_dir or base, **pipe_kwargs)
|
||||
|
||||
# Resolve the effective speed mode: GGUF models default to the
|
||||
# near-lossless `default` profile (compile is ~2.2x and sits below
|
||||
|
|
@ -1379,6 +1420,7 @@ class DiffusionBackend:
|
|||
*,
|
||||
fam: Optional[DiffusionFamily] = None,
|
||||
prequant_path: Optional[str] = None,
|
||||
base_local_dir: Optional[str] = None,
|
||||
) -> tuple[Any, str]:
|
||||
"""Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
|
||||
|
||||
|
|
@ -1426,7 +1468,7 @@ class DiffusionBackend:
|
|||
)
|
||||
if transformer is not None:
|
||||
pipe = self._assemble_pipe(
|
||||
pipeline_cls, base, transformer, dtype, hf_token, device
|
||||
pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
|
||||
)
|
||||
return pipe, scheme
|
||||
|
||||
|
|
@ -1434,7 +1476,9 @@ class DiffusionBackend:
|
|||
transformer = transformer_cls.from_pretrained(
|
||||
base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
|
||||
)
|
||||
pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
|
||||
pipe = self._assemble_pipe(
|
||||
pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
|
||||
)
|
||||
scheme = quantize_transformer(
|
||||
pipe,
|
||||
target,
|
||||
|
|
@ -1455,13 +1499,14 @@ class DiffusionBackend:
|
|||
dtype: Any,
|
||||
hf_token: Optional[str],
|
||||
device: str,
|
||||
base_local_dir: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
|
||||
(a no-op for an already-placed pre-quantized transformer; it moves the companions)."""
|
||||
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
|
||||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
|
||||
pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs)
|
||||
pipe.to(device)
|
||||
return pipe
|
||||
|
||||
|
|
@ -1606,6 +1651,16 @@ class DiffusionBackend:
|
|||
if cn_model is None:
|
||||
if cancel.is_set():
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# resolve_controlnet accepts a bare owner/name repo without the non-GGUF base
|
||||
# trust gate, and from_pretrained below downloads and deserializes it. A
|
||||
# malicious pickle .bin would execute on load, so run the same Hub malware
|
||||
# preflight the chat/export loaders use before any remote ControlNet load. A
|
||||
# local dir the user picked has no Hub scan and is exempt (fail-open there).
|
||||
if not getattr(resolved_cn, "is_local", False):
|
||||
from utils.security import evaluate_file_security
|
||||
_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.
|
||||
|
|
@ -1656,7 +1711,14 @@ class DiffusionBackend:
|
|||
pipe = getattr(diffusers, pipe_cls_name).from_pipe(
|
||||
state.pipe, controlnet = cn_model, torch_dtype = None
|
||||
)
|
||||
self._cn_pipes[key] = pipe
|
||||
with self._lock:
|
||||
# Same race as the model cache above: an unload/superseding load may
|
||||
# have cleared _cn_pipes while from_pipe ran; caching now would pin a
|
||||
# pipeline built around the UNLOADED base and hand it to the next load.
|
||||
if cancel.is_set() or self._state is not state:
|
||||
del pipe
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
self._cn_pipes[key] = pipe
|
||||
return pipe
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -2102,7 +2164,8 @@ class DiffusionBackend:
|
|||
gen = _GenState(total_steps = steps)
|
||||
|
||||
def _on_step(pipe, step_index, timestep, callback_kwargs):
|
||||
now = time.time()
|
||||
# Monotonic: a wall-clock adjustment (NTP) mid-denoise would skew the ETA.
|
||||
now = time.monotonic()
|
||||
gen.step = step_index + 1
|
||||
if gen.first_step_at == 0.0:
|
||||
gen.first_step_at = now
|
||||
|
|
@ -2179,9 +2242,8 @@ class DiffusionBackend:
|
|||
self._cancel_event.set()
|
||||
with self._lock:
|
||||
# Abort an in-flight denoise too by setting ITS cancel event, so the step
|
||||
# callback stops it. unload does NOT take _generate_lock — it must return
|
||||
# promptly; the running generate keeps its own pipe reference, so freeing
|
||||
# _state here can't crash it, and its VRAM is reclaimed when it returns
|
||||
# callback stops it. The running generate keeps its own pipe reference, so
|
||||
# freeing _state here can't crash it; its VRAM is reclaimed when it exits
|
||||
# (within ~one step thanks to the cancel).
|
||||
if self._active_generate_cancel is not None:
|
||||
self._active_generate_cancel.set()
|
||||
|
|
@ -2190,6 +2252,14 @@ class DiffusionBackend:
|
|||
# committing) and drop the marker so the next load starts clean.
|
||||
self._load_token += 1
|
||||
self._loading = None
|
||||
# Wait for the signalled denoise to actually exit before reporting unloaded:
|
||||
# callers treat this return as "VRAM is free" (the GPU arbiter hands the GPU
|
||||
# to chat next; the training routes size their run against it), and the
|
||||
# denoise holds its pipe until the next step callback. generate() holds
|
||||
# _generate_lock for its full body, so a bare acquire is the exit barrier
|
||||
# (never while holding _lock -- generate takes _lock inside _generate_lock).
|
||||
with self._generate_lock:
|
||||
pass
|
||||
return self.status()
|
||||
|
||||
def _unload_locked(self) -> None:
|
||||
|
|
@ -2214,7 +2284,7 @@ class DiffusionBackend:
|
|||
uninstall_patches()
|
||||
uninstall_arch_patches()
|
||||
# NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload()
|
||||
# sets the cancel event but does not take _generate_lock, so a LoRA-backed denoise
|
||||
# only acquires _generate_lock AFTER this teardown, so a LoRA-backed denoise
|
||||
# can still be running on this same pipe for up to one more callback; mutating its
|
||||
# adapter layers now would race that in-flight generation. The whole pipe is dropped
|
||||
# just below (self._state = None; del state; clear_gpu_cache()), so the adapter
|
||||
|
|
|
|||
|
|
@ -129,6 +129,11 @@ def select_attention_backend(
|
|||
backend = _ALIASES[alias]
|
||||
if backend == "native":
|
||||
return None
|
||||
# Every explicit kernel here (cuDNN / flash* / sage) is CUDA+NVIDIA-only; on
|
||||
# ROCm / MPS / CPU diffusers accepts the name at set time and the first
|
||||
# generation crashes, so drop to the native default up front.
|
||||
if not _is_cuda_nvidia(target):
|
||||
return None
|
||||
# An arch-gated kernel (flash3/flash4) on a card that can't run it would set fine
|
||||
# then crash mid-generation, so drop it to the native default up front.
|
||||
if not _backend_arch_supported(backend):
|
||||
|
|
|
|||
|
|
@ -89,7 +89,10 @@ def apply_step_cache(
|
|||
_warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
|
||||
return None
|
||||
try:
|
||||
from diffusers import FirstBlockCacheConfig
|
||||
try:
|
||||
from diffusers import FirstBlockCacheConfig
|
||||
except ImportError: # older diffusers exports it only from diffusers.hooks
|
||||
from diffusers.hooks import FirstBlockCacheConfig
|
||||
|
||||
config = FirstBlockCacheConfig(threshold = thr)
|
||||
enable_cache(config)
|
||||
|
|
@ -101,6 +104,12 @@ def apply_step_cache(
|
|||
logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
|
||||
return mode
|
||||
except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached
|
||||
# enable_cache can fail after hooking some blocks; drop any partial hooks so
|
||||
# the reported-uncached model doesn't actually run half-cached.
|
||||
try:
|
||||
transformer.disable_cache()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
_warn(logger, mode, exc)
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -430,10 +430,19 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
|
|||
_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
|
||||
("z-image-turbo", 9, 0.0),
|
||||
("flux.1-schnell", 4, 0.0),
|
||||
# Kontext (editing) before the generic flux.1: ~28 steps, lower guidance (~2.5).
|
||||
("kontext", 28, 2.5),
|
||||
("flux.1", 28, 3.5),
|
||||
("flux.2-klein", 4, 0.0),
|
||||
# FLUX.2-dev is the full (non-distilled) model: more steps + real guidance.
|
||||
("flux.2-dev", 28, 4.0),
|
||||
("qwen-image", 20, 4.0),
|
||||
("z-image", 20, 4.0),
|
||||
# SDXL: Turbo is distilled (few steps, no CFG); base/full SDXL wants ~30 steps and
|
||||
# real CFG (~7). "sdxl-turbo" must precede the generic "sdxl" substring match.
|
||||
("sdxl-turbo", 3, 0.0),
|
||||
("stable-diffusion-xl", 30, 7.0),
|
||||
("sdxl", 30, 7.0),
|
||||
)
|
||||
# Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback.
|
||||
_GENERATION_DEFAULT_FALLBACK = (9, 0.0)
|
||||
|
|
|
|||
|
|
@ -507,6 +507,16 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
|
|||
import torch
|
||||
from diffusers.hooks import apply_group_offloading
|
||||
|
||||
# A dual-DiT pipeline (e.g. Ideogram 4's unconditional tower) carries a second
|
||||
# denoiser as large as the first; leaving it resident would defeat this tier
|
||||
# (the pair rarely fits where one alone did not). Stream every DiT and keep
|
||||
# only the genuinely smaller companions resident.
|
||||
streamed: dict[str, Any] = {"transformer": transformer}
|
||||
for extra in ("transformer_2", "unconditional_transformer"):
|
||||
module = getattr(pipe, extra, None)
|
||||
if isinstance(module, torch.nn.Module):
|
||||
streamed[extra] = module
|
||||
|
||||
onload = torch.device(device)
|
||||
use_stream = onload.type == "cuda" # overlap H2D copies with compute on CUDA
|
||||
gkwargs: dict[str, Any] = {
|
||||
|
|
@ -536,11 +546,12 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
|
|||
# load-time crash. The streamed transformer manages its own placement via the
|
||||
# offloading hooks applied next.
|
||||
for name, comp in getattr(pipe, "components", {}).items():
|
||||
if name == "transformer":
|
||||
if name in streamed:
|
||||
continue
|
||||
if isinstance(comp, torch.nn.Module):
|
||||
comp.to(onload)
|
||||
apply_group_offloading(transformer, **gkwargs)
|
||||
for module in streamed.values():
|
||||
apply_group_offloading(module, **gkwargs)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — fall back to whole-module offload
|
||||
if logger is not None:
|
||||
|
|
|
|||
|
|
@ -269,7 +269,12 @@ def build_sd_cpp_command(
|
|||
if params.seed is not None:
|
||||
cmd += ["--seed", str(int(params.seed))]
|
||||
if params.batch_count and params.batch_count != 1:
|
||||
cmd += ["--batch-count", str(int(params.batch_count))]
|
||||
# sd-cli names the extra batch images itself (output_2.png, ...) and the runner
|
||||
# collects only the literal --output path, so a CLI batch would silently drop
|
||||
# every image after the first. Batches go through the sdcpp server API instead.
|
||||
raise ValueError(
|
||||
"sd-cli runs are single-image; use the sdcpp server API for batch generation."
|
||||
)
|
||||
|
||||
cmd += ["--output", output_path]
|
||||
if threads is not None:
|
||||
|
|
|
|||
|
|
@ -241,6 +241,9 @@ class _SdLoading:
|
|||
|
||||
repo_id: str
|
||||
base_repo: str
|
||||
# Companion asset repos (VAE / text encoders) this load fetches, so the
|
||||
# delete-cached guard protects them for the whole download/finalize window.
|
||||
asset_repos: tuple[str, ...] = ()
|
||||
expected_bytes: int = 0
|
||||
downloaded_bytes: int = 0
|
||||
error: Optional[str] = None
|
||||
|
|
@ -407,7 +410,17 @@ class SdCppDiffusionBackend:
|
|||
self._load_token += 1
|
||||
token = self._load_token
|
||||
self._cancel_event.clear()
|
||||
self._loading = _SdLoading(repo_id = repo_id, base_repo = base)
|
||||
self._loading = _SdLoading(
|
||||
repo_id = repo_id,
|
||||
base_repo = base,
|
||||
asset_repos = tuple(
|
||||
dict.fromkeys(
|
||||
r
|
||||
for r, _f, kind in self._asset_specs(repo_id, gguf_filename, fam)
|
||||
if kind != "diffusion_model"
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
threading.Thread(
|
||||
target = self._run_load,
|
||||
|
|
@ -678,12 +691,14 @@ class SdCppDiffusionBackend:
|
|||
def loading_repo_ids(self) -> tuple[str, ...]:
|
||||
"""Repo ids an in-flight background load is downloading (empty when idle).
|
||||
Mirrors the diffusers backend so the delete-cached guard can query whichever
|
||||
engine is active without caring which one it got."""
|
||||
engine is active without caring which one it got. Includes the companion
|
||||
VAE / text-encoder repos: deleting one of those mid-load would remove files
|
||||
the committed SdCppModelFiles paths need."""
|
||||
with self._lock:
|
||||
loading = self._loading
|
||||
if loading is None or loading.error is not None:
|
||||
return ()
|
||||
return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
|
||||
return tuple(r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r)
|
||||
|
||||
# ── Generate ───────────────────────────────────────────────────────────
|
||||
|
||||
|
|
@ -1085,6 +1100,13 @@ class SdCppDiffusionBackend:
|
|||
state.server.stop()
|
||||
if pending is not None and pending is not (state.server if state else None):
|
||||
pending.stop()
|
||||
# Wait for a signalled one-shot generation to actually exit before reporting
|
||||
# unloaded: callers (the GPU arbiter, training cleanup) treat this return as
|
||||
# "the device is free", but a one-shot sd-cli child killed by the cancel above
|
||||
# unwinds under _generate_lock. A bare acquire is the exit barrier (never taken
|
||||
# while holding _lock; same pattern as DiffusionBackend.unload).
|
||||
with self._generate_lock:
|
||||
pass
|
||||
return self.status()
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
|
|
|
|||
|
|
@ -375,6 +375,9 @@ class SdCppEngine:
|
|||
def _prepare_out(output_path: str) -> Path:
|
||||
out = Path(output_path)
|
||||
out.parent.mkdir(parents = True, exist_ok = True)
|
||||
# Drop a stale file at the target so the post-run is_file() check proves THIS
|
||||
# run produced the image, not a leftover from an earlier run at the same path.
|
||||
out.unlink(missing_ok = True)
|
||||
return out
|
||||
|
||||
def _run(
|
||||
|
|
|
|||
|
|
@ -528,6 +528,13 @@ def run_dit_lora_training(
|
|||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
# The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is
|
||||
# unsupported for real runs but keeps import/unit tests architecture-agnostic).
|
||||
# Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the
|
||||
# run would otherwise die deep in model load with an opaque dtype error.
|
||||
if device == "cuda" and not torch.cuda.is_bf16_supported():
|
||||
raise ValueError(
|
||||
"This trainer requires a bfloat16-capable GPU (Ampere or newer); "
|
||||
"this CUDA device does not support bf16."
|
||||
)
|
||||
weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32
|
||||
use_lora_targets = _select_lora_targets(cfg.lora_target_modules, spec.lora_targets)
|
||||
|
||||
|
|
@ -585,6 +592,16 @@ def run_dit_lora_training(
|
|||
lora_params = [p for p in transformer.parameters() if p.requires_grad]
|
||||
|
||||
optimizer = _make_optimizer(lora_params, cfg.learning_rate)
|
||||
# One lr_sched.step() per optimizer update (cfg.train_steps total), matching the
|
||||
# SDXL trainer: counting micro-steps instead would stretch warmup past the run.
|
||||
from diffusers.optimization import get_scheduler
|
||||
|
||||
lr_sched = get_scheduler(
|
||||
cfg.lr_scheduler,
|
||||
optimizer = optimizer,
|
||||
num_warmup_steps = cfg.lr_warmup_steps,
|
||||
num_training_steps = cfg.train_steps,
|
||||
)
|
||||
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
||||
cfg.base_model, subfolder = "scheduler", token = cfg.hf_token
|
||||
)
|
||||
|
|
@ -597,10 +614,15 @@ def run_dit_lora_training(
|
|||
peak_gb = 0.0
|
||||
t_start = time.time()
|
||||
done = 0
|
||||
# Honor train_batch_size by folding it into the micro-step count: averaging the
|
||||
# gradient over batch * accum single-image passes is mathematically identical to
|
||||
# true batching with a mean loss, and keeps the QLoRA memory profile flat (one
|
||||
# image's activations at a time). Previously batch_size > 1 silently trained at 1.
|
||||
micro_steps = cfg.gradient_accumulation_steps * cfg.train_batch_size
|
||||
for opt_step in range(cfg.train_steps):
|
||||
optimizer.zero_grad(set_to_none = True)
|
||||
step_loss = 0.0
|
||||
for _ in range(cfg.gradient_accumulation_steps):
|
||||
for _ in range(micro_steps):
|
||||
i = rng.randrange(len(image_paths))
|
||||
px = (
|
||||
_load_pixel_tensor(
|
||||
|
|
@ -638,8 +660,8 @@ def run_dit_lora_training(
|
|||
)
|
||||
target = noise - latents
|
||||
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
|
||||
(loss / cfg.gradient_accumulation_steps).backward()
|
||||
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
|
||||
(loss / micro_steps).backward()
|
||||
step_loss += float(loss.detach()) / micro_steps
|
||||
|
||||
grad_norm: Optional[float] = None
|
||||
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
|
||||
|
|
@ -647,6 +669,7 @@ def run_dit_lora_training(
|
|||
# chart wants (spikes stay visible even when clipping flattens the update).
|
||||
grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
|
||||
optimizer.step()
|
||||
lr_sched.step()
|
||||
|
||||
running_loss += step_loss
|
||||
done = opt_step + 1
|
||||
|
|
@ -665,7 +688,7 @@ def run_dit_lora_training(
|
|||
total_steps = cfg.train_steps,
|
||||
loss = round(step_loss, 5),
|
||||
avg_loss = round(running_loss / done, 5),
|
||||
learning_rate = cfg.learning_rate,
|
||||
learning_rate = lr_sched.get_last_lr()[0],
|
||||
grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
|
||||
samples_per_second = sps,
|
||||
peak_memory_gb = peak_gb or None,
|
||||
|
|
|
|||
|
|
@ -291,6 +291,10 @@ class DiffusionLoraConfig:
|
|||
f"lr_scheduler must be one of {', '.join(sorted(_LR_SCHEDULERS))}; "
|
||||
f"got {self.lr_scheduler!r}"
|
||||
)
|
||||
# A zero/negative gamma would zero out (or invert) the min-SNR weight and
|
||||
# silently train on a degenerate loss; None is the documented disable.
|
||||
if self.snr_gamma is not None and float(self.snr_gamma) <= 0:
|
||||
raise ValueError("snr_gamma must be > 0, or null to disable min-SNR weighting")
|
||||
# learning_rate can arrive as a string ("1e-4") from the Studio config path, which
|
||||
# preserves it as a string after validation; coerce so AdamW receives a float.
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -704,10 +704,14 @@ class DiffusionTrainingStartRequest(BaseModel):
|
|||
default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"],
|
||||
description = "U-Net modules to attach LoRA to",
|
||||
)
|
||||
max_grad_norm: float = Field(1.0, gt = 0, description = "Gradient clipping max-norm")
|
||||
max_grad_norm: float = Field(
|
||||
1.0, ge = 0, description = "Gradient clipping max-norm; 0 disables clipping"
|
||||
)
|
||||
seed: int = Field(42)
|
||||
mixed_precision: Literal["bf16", "fp16", "no"] = Field("bf16")
|
||||
snr_gamma: Optional[float] = Field(5.0, description = "Min-SNR loss weighting; null disables")
|
||||
snr_gamma: Optional[float] = Field(
|
||||
5.0, gt = 0, description = "Min-SNR loss weighting; null disables"
|
||||
)
|
||||
gradient_checkpointing: bool = Field(True)
|
||||
lr_scheduler: Literal[
|
||||
"linear",
|
||||
|
|
|
|||
|
|
@ -12076,6 +12076,21 @@ async def openai_image_generations(
|
|||
# isn't loaded; the global handler turns this into the OpenAI envelope.
|
||||
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
|
||||
|
||||
# An edit-only model (Qwen-Image-Edit, FLUX Kontext) needs an input image this API
|
||||
# cannot supply; refuse up front with a 400 instead of letting the backend's
|
||||
# ValueError surface as a sanitized 500.
|
||||
workflows = status.get("workflows") or []
|
||||
if workflows and "txt2img" not in workflows:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = openai_error_body(
|
||||
"The loaded image model is edit-only (it requires an input image); "
|
||||
"load a text-to-image model to use this endpoint.",
|
||||
status = 400,
|
||||
param = "model",
|
||||
),
|
||||
)
|
||||
|
||||
# Fall back to the resolved base repo so a local-path load (whose repo_id is a
|
||||
# filesystem path) still gets the right per-model steps/guidance.
|
||||
steps, guidance = default_generation_params(status.get("repo_id"), status.get("base_repo"))
|
||||
|
|
|
|||
|
|
@ -1177,11 +1177,29 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
|
|||
|
||||
@router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
|
||||
async def start_diffusion_training(
|
||||
body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject)
|
||||
body: DiffusionTrainingStartRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
via_api_key: bool = Depends(authenticated_via_api_key),
|
||||
):
|
||||
"""Start an SDXL LoRA training job from an image + caption dataset."""
|
||||
from core.training.diffusion_training_service import get_diffusion_training_service
|
||||
|
||||
# When Studio is driven as an inference API (API-key auth), refuse to start training
|
||||
# while a request is in flight: _free_gpu_for_diffusion_training() below unloads the
|
||||
# chat backends to reclaim VRAM, which would kill the stream. Mirrors start_training so
|
||||
# a diffusion start cannot silently drop an active API inference request.
|
||||
if via_api_key is True:
|
||||
from core.inference.llama_keepwarm import other_inference_request_count
|
||||
if other_inference_request_count(current_request_counted = False) > 0:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
"Cannot start diffusion (Images) training over the API while an inference "
|
||||
"request is in progress. Wait for it to finish, or start training from the "
|
||||
"Studio UI."
|
||||
),
|
||||
)
|
||||
|
||||
# Interlock: refuse while an LLM training run holds the GPU (symmetric with the
|
||||
# diffusion check in start_training), so the two trainers never contend for VRAM.
|
||||
try:
|
||||
|
|
@ -1219,11 +1237,36 @@ async def start_diffusion_training(
|
|||
except ValueError as e:
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
|
||||
# Run the trainers' trust gate here too (both assert the same predicate before
|
||||
# from_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents
|
||||
# instead of tearing down the user's chat/Images model and failing in the child.
|
||||
from core.training.diffusion_train_common import _assert_trusted_base_model
|
||||
|
||||
try:
|
||||
_assert_trusted_base_model(config.get("base_model", ""))
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
|
||||
# Preflight access to a gated base repo with the user's token BEFORE freeing GPU
|
||||
# residents, so a missing/insufficient token fails fast (400) without tearing down the
|
||||
# user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
|
||||
_preflight_gated_base(config.get("base_model", ""), config.get("hf_token"))
|
||||
|
||||
# Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise
|
||||
# fails inside the spawned trainer AFTER the user's chat/Images model was
|
||||
# evicted. Same discovery the trainer runs, so the two cannot disagree.
|
||||
from core.training import diffusion_train_common as _dtc
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
_dtc.discover_image_caption_pairs,
|
||||
config["data_dir"],
|
||||
instance_prompt = config.get("instance_prompt") or None,
|
||||
caption_column = config.get("caption_column") or "text",
|
||||
)
|
||||
except (FileNotFoundError, ValueError) as e:
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
|
||||
# Free resident GPU workloads (export / Images pipeline / chat) before the trainer
|
||||
# loads its own pipeline.
|
||||
_free_gpu_for_diffusion_training()
|
||||
|
|
@ -1598,6 +1641,10 @@ async def get_diffusion_dataset_image(
|
|||
|
||||
thumbs_dir = folder / _THUMBS_DIRNAME
|
||||
thumbs_dir.mkdir(exist_ok = True)
|
||||
# Key on the full filename (stem + extension), not the stem: two images that
|
||||
# share a stem but differ by extension (sample.png / sample.jpg) would otherwise
|
||||
# collide on one cache file, and an mtime-newer cache built for the first would
|
||||
# be served for the second, showing the wrong image in the labeling grid.
|
||||
thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg"
|
||||
src_mtime = image_path.stat().st_mtime
|
||||
if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime:
|
||||
|
|
@ -1643,12 +1690,24 @@ async def set_diffusion_dataset_caption(
|
|||
sidecar = image_path.with_suffix(".txt")
|
||||
if caption:
|
||||
sidecar.write_text(caption, encoding = "utf-8")
|
||||
else:
|
||||
# Blank clears the sidecar; also drop a stale .caption so the image reads as
|
||||
# uncaptioned afterwards.
|
||||
sidecar.unlink(missing_ok = True)
|
||||
image_path.with_suffix(".caption").unlink(missing_ok = True)
|
||||
return _image_record(folder, image_path, _load_metadata_captions(folder))
|
||||
return _image_record(folder, image_path, _load_metadata_captions(folder))
|
||||
# Blank must actually clear. Unlinking alone would resurface this image's
|
||||
# metadata.jsonl / captions.jsonl caption (the fallback source), so when one
|
||||
# exists write an EMPTY sidecar instead: both the record reader and the
|
||||
# trainer's discovery treat an existing sidecar as authoritative even when
|
||||
# empty, which makes it a tombstone. No metadata caption -> plain cleanup.
|
||||
meta = _load_metadata_captions(folder)
|
||||
try:
|
||||
rel = image_path.relative_to(folder).as_posix()
|
||||
except ValueError:
|
||||
rel = image_path.name
|
||||
if image_path.name in meta or rel in meta:
|
||||
sidecar.write_text("", encoding = "utf-8")
|
||||
else:
|
||||
sidecar.unlink(missing_ok = True)
|
||||
image_path.with_suffix(".caption").unlink(missing_ok = True)
|
||||
return _image_record(folder, image_path, meta)
|
||||
|
||||
return await asyncio.to_thread(write)
|
||||
|
||||
|
|
@ -1671,6 +1730,9 @@ async def delete_diffusion_dataset_image(
|
|||
image_path.with_suffix(ext).unlink(missing_ok = True)
|
||||
thumbs_dir = folder / _THUMBS_DIRNAME
|
||||
if thumbs_dir.is_dir():
|
||||
# Thumbs are keyed on the full filename (stem + extension), so match that
|
||||
# here too; a stem-only glob would leave this image's thumbs behind and
|
||||
# could delete a same-stem sibling's (sample.png vs sample.jpg).
|
||||
for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
|
||||
t.unlink(missing_ok = True)
|
||||
return {"deleted": image_path.name}
|
||||
|
|
|
|||
|
|
@ -75,7 +75,7 @@ def test_auto_stays_native_off_nvidia(monkeypatch):
|
|||
|
||||
|
||||
def test_explicit_backend_honored_regardless_of_speed(monkeypatch):
|
||||
monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False)
|
||||
monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True)
|
||||
# Pin a high capability so the arch-gated flash4 isn't dropped by the runtime check.
|
||||
monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
|
||||
assert select_attention_backend(_target(), "sage", speed_active = False) == "sage"
|
||||
|
|
@ -83,6 +83,15 @@ def test_explicit_backend_honored_regardless_of_speed(monkeypatch):
|
|||
assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn"
|
||||
|
||||
|
||||
def test_explicit_backend_dropped_off_nvidia_cuda(monkeypatch):
|
||||
# Explicit cuDNN/flash/sage on ROCm / MPS / CPU passes diffusers' set-time check
|
||||
# and crashes at the first generation, so selection drops to the native default.
|
||||
monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False)
|
||||
monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
|
||||
for alias in ("sage", "flash", "flash4", "cudnn"):
|
||||
assert select_attention_backend(_target(device = "mps"), alias, speed_active = True) is None
|
||||
|
||||
|
||||
def test_explicit_native_returns_none():
|
||||
# native is the default -> nothing to set.
|
||||
assert select_attention_backend(_target(), "native", speed_active = True) is None
|
||||
|
|
|
|||
|
|
@ -1384,6 +1384,32 @@ def test_load_promotes_fp16_to_fp32_for_zimage_only(fake_runtime, monkeypatch, t
|
|||
assert q["dtype"] == "float16" # fp16-compatible family keeps fp16 on pre-Ampere
|
||||
|
||||
|
||||
def test_bad_mode_strings_fail_before_eviction(fake_runtime):
|
||||
# Every mode normalizer that can raise runs BEFORE the load evicts the previous
|
||||
# pipeline, so a bad request never costs the user their working model.
|
||||
backend = DiffusionBackend()
|
||||
fam = detect_family("unsloth/Z-Image-GGUF")
|
||||
backend._state = _LoadState(
|
||||
pipe = object(),
|
||||
family = fam,
|
||||
repo_id = "r",
|
||||
base_repo = "b",
|
||||
device = "cpu",
|
||||
dtype = "float32",
|
||||
cpu_offload = False,
|
||||
)
|
||||
for kwargs in (
|
||||
{"transformer_quant": "int7"},
|
||||
{"speed_mode": "warp"},
|
||||
{"attention_backend": "bogus"},
|
||||
{"transformer_cache": "bogus"},
|
||||
{"text_encoder_quant": "fp3"},
|
||||
):
|
||||
with pytest.raises(ValueError):
|
||||
backend.load_pipeline("unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs)
|
||||
assert backend._state is not None
|
||||
|
||||
|
||||
# Lock split + mid-denoise cancellation
|
||||
|
||||
|
||||
|
|
@ -1427,17 +1453,25 @@ def test_generate_lock_split_keeps_status_and_unload_responsive(fake_runtime):
|
|||
assert backend.status()["loaded"] is True
|
||||
assert backend.generate_progress()["active"] is True
|
||||
|
||||
# unload() must return promptly (it does not wait on _generate_lock) and signal
|
||||
# THIS in-flight generation's cancel event.
|
||||
cancel_ref = backend._active_generate_cancel
|
||||
assert cancel_ref is not None
|
||||
|
||||
# unload() signals THIS generation's cancel event, then waits for the denoise to
|
||||
# actually exit before returning: callers treat its return as "VRAM is free" (the
|
||||
# GPU arbiter hands the GPU to chat on it). Release the pipe once the cancel
|
||||
# lands, standing in for the step callback of a real pipeline.
|
||||
releaser = threading.Thread(target = lambda: (cancel_ref.wait(5), release.set()))
|
||||
releaser.start()
|
||||
backend.unload()
|
||||
assert backend._active_generate_cancel is not None
|
||||
assert backend._active_generate_cancel.is_set()
|
||||
releaser.join(5)
|
||||
assert cancel_ref.is_set()
|
||||
assert backend.status()["loaded"] is False
|
||||
|
||||
release.set()
|
||||
t.join(5)
|
||||
# The cancelled generation raised rather than returning a now-evicted image.
|
||||
# The cancelled generation raised rather than returning a now-evicted image, and
|
||||
# it had already exited (deregistering its cancel) before unload() returned.
|
||||
assert "exc" in out and "cancelled" in str(out["exc"]).lower()
|
||||
assert backend._active_generate_cancel is None
|
||||
|
||||
|
||||
def test_callback_cancellation_interrupts_denoise(fake_runtime):
|
||||
|
|
@ -2183,3 +2217,34 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
|
|||
backend.generate(prompt = "a sloth")
|
||||
backend.generate(prompt = "another sloth")
|
||||
assert resets == [True, True]
|
||||
|
||||
|
||||
def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch):
|
||||
# The prefetched pipeline manifest's directory is the local snapshot root; a
|
||||
# config-only base list (no manifest) returns None so the hub id stays in use.
|
||||
backend = DiffusionBackend()
|
||||
monkeypatch.setattr(
|
||||
"utils.hf_xet_fallback.hf_hub_download_with_xet_fallback",
|
||||
lambda repo, fn, tok, **k: f"/cache/snap/{fn}",
|
||||
)
|
||||
root = backend._prefetch_files(
|
||||
"base/repo", None, "base/repo", ["model_index.json", "vae/x.safetensors"], None
|
||||
)
|
||||
assert root == "/cache/snap"
|
||||
assert (
|
||||
backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) is None
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path):
|
||||
# With a prefetched snapshot, from_pretrained must receive the local dir --
|
||||
# its own hub sweep would re-download the root packaged singles the scoped
|
||||
# prefetch skips (24 GB per FLUX.1 repo).
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
"unsloth/Qwen-Image-2512-bnb-4bit",
|
||||
model_kind = "pipeline",
|
||||
_base_local_dir = str(tmp_path),
|
||||
)
|
||||
assert _FakePipeline.last["base"] == str(tmp_path)
|
||||
backend.unload()
|
||||
|
|
|
|||
|
|
@ -136,6 +136,29 @@ def test_incompatible_model_runs_uncached(monkeypatch):
|
|||
assert apply_step_cache(_pipe(t), mode = "fbcache") is None
|
||||
|
||||
|
||||
def test_enable_cache_failure_rolls_back_partial_hooks(monkeypatch):
|
||||
# enable_cache can raise after hooking some blocks; the reported-uncached model
|
||||
# must not actually run half-cached, so the failure path calls disable_cache.
|
||||
_stub_diffusers(monkeypatch)
|
||||
t = _MixinTransformer(fail = True)
|
||||
t.disabled = False
|
||||
t.disable_cache = lambda: setattr(t, "disabled", True)
|
||||
assert apply_step_cache(_pipe(t), mode = "fbcache") is None
|
||||
assert t.disabled is True
|
||||
|
||||
|
||||
def test_config_import_falls_back_to_hooks_module(monkeypatch):
|
||||
# Older diffusers exports FirstBlockCacheConfig only from diffusers.hooks.
|
||||
diffusers = types.ModuleType("diffusers") # no FirstBlockCacheConfig attribute
|
||||
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
|
||||
hooks = types.ModuleType("diffusers.hooks")
|
||||
hooks.FirstBlockCacheConfig = _Config
|
||||
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
|
||||
t = _MixinTransformer()
|
||||
assert apply_step_cache(_pipe(t), mode = "fbcache") == TC_FBCACHE
|
||||
assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD
|
||||
|
||||
|
||||
def test_missing_transformer_is_none(monkeypatch):
|
||||
_stub_diffusers(monkeypatch)
|
||||
pipe = types.SimpleNamespace(transformer = None)
|
||||
|
|
@ -143,7 +166,10 @@ def test_missing_transformer_is_none(monkeypatch):
|
|||
|
||||
|
||||
def test_diffusers_unavailable_runs_uncached(monkeypatch):
|
||||
# no diffusers import -> best-effort returns None, load proceeds uncached.
|
||||
# no diffusers import -> best-effort returns None, load proceeds uncached. Block the
|
||||
# hooks module too: the config import falls back to diffusers.hooks, which a REAL
|
||||
# earlier import in the test session may have left cached in sys.modules.
|
||||
monkeypatch.setitem(sys.modules, "diffusers", None)
|
||||
monkeypatch.setitem(sys.modules, "diffusers.hooks", None)
|
||||
t = _MixinTransformer()
|
||||
assert apply_step_cache(_pipe(t), mode = "fbcache") is None
|
||||
|
|
|
|||
|
|
@ -232,14 +232,28 @@ def _state():
|
|||
)
|
||||
|
||||
|
||||
def _allow_cn_security(monkeypatch):
|
||||
"""Stub the Hub malware preflight to allow the load (hermetic, no network)."""
|
||||
import utils.security
|
||||
monkeypatch.setattr(
|
||||
utils.security,
|
||||
"evaluate_file_security",
|
||||
lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = False, reason = ""),
|
||||
)
|
||||
|
||||
|
||||
def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
|
||||
import threading
|
||||
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
|
||||
_allow_cn_security(monkeypatch)
|
||||
b = DiffusionBackend()
|
||||
st = _state()
|
||||
# The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
|
||||
# from_pipe must not repopulate the cache), so mirror the loaded invariant.
|
||||
b._state = st
|
||||
resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
|
||||
p1 = b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel)
|
||||
|
|
@ -250,6 +264,69 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
|
|||
assert b._cn_models["flux-union-pro"] is p1.controlnet
|
||||
|
||||
|
||||
def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
|
||||
# A bare owner/name ControlNet is accepted by resolve_controlnet without the base
|
||||
# trust gate, so the load path must run the Hub malware preflight: a flagged remote
|
||||
# repo must raise BEFORE from_pretrained downloads/deserializes it.
|
||||
import threading
|
||||
|
||||
import utils.security
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
loaded = {"called": False}
|
||||
|
||||
class _TrapModel(_FakeCNModel):
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
path,
|
||||
torch_dtype = None,
|
||||
token = None,
|
||||
):
|
||||
loaded["called"] = True
|
||||
return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)
|
||||
|
||||
mod = _fake_diffusers()
|
||||
mod.FluxControlNetModel = _TrapModel
|
||||
monkeypatch.setitem(sys.modules, "diffusers", mod)
|
||||
monkeypatch.setattr(
|
||||
utils.security,
|
||||
"evaluate_file_security",
|
||||
lambda name, hf_token = None, **kw: types.SimpleNamespace(
|
||||
blocked = True, reason = "Hugging Face security scan flagged unsafe files: evil.bin"
|
||||
),
|
||||
)
|
||||
b = DiffusionBackend()
|
||||
st = _state()
|
||||
b._state = st
|
||||
resolved = dc.ResolvedControlNet("evil/cn", "evil/cn", is_local = False)
|
||||
with pytest.raises(ValueError, match = "security scan flagged"):
|
||||
b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert loaded["called"] is False
|
||||
|
||||
|
||||
def test_controlnet_pipe_skips_scan_for_local_dir(monkeypatch, tmp_path):
|
||||
# A local dir the user picked has no Hub scan; the preflight must not block it even
|
||||
# if the (unused) scan stub would say blocked.
|
||||
import threading
|
||||
|
||||
import utils.security
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
|
||||
monkeypatch.setattr(
|
||||
utils.security,
|
||||
"evaluate_file_security",
|
||||
lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = True, reason = "x"),
|
||||
)
|
||||
b = DiffusionBackend()
|
||||
st = _state()
|
||||
b._state = st
|
||||
resolved = dc.ResolvedControlNet("my-cn", str(tmp_path), is_local = True)
|
||||
p = b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert isinstance(p, _FakeCNPipe)
|
||||
|
||||
|
||||
def test_controlnet_pipe_rejects_family_without_classes():
|
||||
import threading
|
||||
|
||||
|
|
@ -260,3 +337,19 @@ def test_controlnet_pipe_rejects_family_without_classes():
|
|||
st.family.controlnet_pipeline_class = None
|
||||
with pytest.raises(ValueError, match = "not supported"):
|
||||
b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
|
||||
|
||||
|
||||
def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
|
||||
# An unload that lands while from_pipe is assembling must not let the wrapper
|
||||
# repopulate the cache around the torn-down base pipe.
|
||||
import threading
|
||||
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
|
||||
b = DiffusionBackend()
|
||||
st = _state() # never committed to b._state: the load is already gone
|
||||
resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
|
||||
with pytest.raises(RuntimeError, match = "cancelled"):
|
||||
b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert b._cn_pipes == {}
|
||||
|
|
|
|||
|
|
@ -177,10 +177,9 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
|
|||
folder.mkdir()
|
||||
_write_png(folder / "x.png")
|
||||
(folder / "x.txt").write_text("cap", encoding = "utf-8")
|
||||
# Generate a thumbnail so we can assert it is cleaned up too.
|
||||
# Generate a thumbnail so we can assert it is cleaned up too. Thumbs are keyed on
|
||||
# the full filename (stem + extension) to avoid same-stem collisions across formats.
|
||||
client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32")
|
||||
# Thumb cache key includes the extension (x.png_32.jpg), so png and jpg
|
||||
# siblings can't collide.
|
||||
assert list((folder / ".thumbs").glob("x.png_*.jpg"))
|
||||
|
||||
r = client.delete("/api/train/diffusion/dataset/d/image/x.png")
|
||||
|
|
@ -190,6 +189,19 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
|
|||
assert not list((folder / ".thumbs").glob("x.png_*.jpg"))
|
||||
|
||||
|
||||
def test_thumb_cache_key_distinguishes_same_stem_extensions(client, ds_root):
|
||||
# sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache
|
||||
# file, so the labeling grid never serves one image's thumbnail for the other.
|
||||
folder = ds_root / "d"
|
||||
folder.mkdir()
|
||||
Image.new("RGB", (8, 8), (10, 20, 30)).save(folder / "sample.png", format = "PNG")
|
||||
Image.new("RGB", (8, 8), (200, 210, 220)).save(folder / "sample.jpg", format = "JPEG")
|
||||
client.get("/api/train/diffusion/dataset/d/image/sample.png?thumb=32")
|
||||
client.get("/api/train/diffusion/dataset/d/image/sample.jpg?thumb=32")
|
||||
thumbs = sorted(p.name for p in (folder / ".thumbs").glob("*.jpg"))
|
||||
assert thumbs == ["sample.jpg_32.jpg", "sample.png_32.jpg"]
|
||||
|
||||
|
||||
# ── traversal / validation ───────────────────────────────────────────────────
|
||||
def test_dataset_name_traversal_rejected_over_http(client, ds_root):
|
||||
# A name that fails the folder-name validator returns 400, never touches disk.
|
||||
|
|
|
|||
|
|
@ -138,6 +138,16 @@ def test_config_rejects_zero_lora_alpha():
|
|||
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", lora_alpha = 0).normalized()
|
||||
|
||||
|
||||
def test_config_rejects_nonpositive_snr_gamma():
|
||||
# gamma <= 0 zeroes/inverts the min-SNR weight; None is the documented disable.
|
||||
with pytest.raises(ValueError, match = "snr_gamma"):
|
||||
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = 0).normalized()
|
||||
cfg = DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = None
|
||||
).normalized()
|
||||
assert cfg.snr_gamma is None
|
||||
|
||||
|
||||
def test_config_coerces_string_learning_rate():
|
||||
# The Studio config path preserves learning_rate as a string; normalize to float.
|
||||
cfg = DiffusionLoraConfig(
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ import pytest
|
|||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from auth.authentication import authenticated_via_api_key, get_current_subject
|
||||
from core.training.diffusion_training_service import DiffusionTrainingService
|
||||
from routes.training import router as training_router
|
||||
|
||||
|
|
@ -265,11 +265,22 @@ def client(monkeypatch):
|
|||
|
||||
monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False))
|
||||
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
|
||||
# The dataset preflight runs the trainer's discovery against _BODY's fake
|
||||
# data_dir; stub it here so wiring tests pass, and let the dedicated preflight
|
||||
# tests below re-point it at a real tmp dataset.
|
||||
monkeypatch.setattr(
|
||||
"core.training.diffusion_train_common.discover_image_caption_pairs",
|
||||
lambda data_dir, **kw: [("img.png", "caption")],
|
||||
)
|
||||
app = FastAPI()
|
||||
app.include_router(training_router, prefix = "/api/train")
|
||||
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
||||
# Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there,
|
||||
# so the wiring tests below behave as before. The guard test flips this override.
|
||||
app.dependency_overrides[authenticated_via_api_key] = lambda: False
|
||||
c = TestClient(app)
|
||||
c._fake = fake # type: ignore[attr-defined]
|
||||
c._app = app # type: ignore[attr-defined]
|
||||
return c
|
||||
|
||||
|
||||
|
|
@ -303,6 +314,20 @@ def test_route_start_forwards_extra_training_knobs(client):
|
|||
assert client._fake.started_with["lora_target_modules"] == ["to_q", "to_v"]
|
||||
|
||||
|
||||
def test_route_start_accepts_zero_max_grad_norm(client):
|
||||
# 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_);
|
||||
# the request model must not reject it.
|
||||
r = client.post("/api/train/diffusion/start", json = {**_BODY, "max_grad_norm": 0.0})
|
||||
assert r.status_code == 200, r.text
|
||||
assert client._fake.started_with["max_grad_norm"] == 0.0
|
||||
|
||||
|
||||
def test_route_start_rejects_nonpositive_snr_gamma(client):
|
||||
# gamma <= 0 zeroes/inverts the min-SNR loss weight; null is the disable value.
|
||||
r = client.post("/api/train/diffusion/start", json = {**_BODY, "snr_gamma": 0})
|
||||
assert r.status_code == 422
|
||||
|
||||
|
||||
def test_route_start_rejects_uncontained_paths(client):
|
||||
# An absolute path outside the Studio dataset roots is a 400, not silently accepted.
|
||||
r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"})
|
||||
|
|
@ -344,6 +369,37 @@ def test_route_start_conflict_maps_to_409(client):
|
|||
assert r.status_code == 409
|
||||
|
||||
|
||||
def test_route_start_over_api_with_inference_in_flight_is_409(client, monkeypatch):
|
||||
# An API-key client must not start diffusion training (which frees VRAM by unloading
|
||||
# chat) while an inference request is streaming; it should 409 instead of killing it.
|
||||
client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
|
||||
monkeypatch.setattr(
|
||||
"core.inference.llama_keepwarm.other_inference_request_count",
|
||||
lambda current_request_counted = False: 1,
|
||||
)
|
||||
freed = {"called": False}
|
||||
import routes.training as tr
|
||||
|
||||
monkeypatch.setattr(
|
||||
tr, "_free_gpu_for_diffusion_training", lambda: freed.__setitem__("called", True)
|
||||
)
|
||||
r = client.post("/api/train/diffusion/start", json = _BODY)
|
||||
assert r.status_code == 409
|
||||
# The guard must run BEFORE any GPU is freed, so the live inference stream survives.
|
||||
assert freed["called"] is False
|
||||
|
||||
|
||||
def test_route_start_over_api_without_inference_proceeds(client, monkeypatch):
|
||||
# Same API-key path but no inference in flight: the start proceeds normally.
|
||||
client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
|
||||
monkeypatch.setattr(
|
||||
"core.inference.llama_keepwarm.other_inference_request_count",
|
||||
lambda current_request_counted = False: 0,
|
||||
)
|
||||
r = client.post("/api/train/diffusion/start", json = _BODY)
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
def test_route_status_and_stop(client):
|
||||
client.post("/api/train/diffusion/start", json = _BODY)
|
||||
s = client.get("/api/train/diffusion/status")
|
||||
|
|
|
|||
|
|
@ -166,10 +166,18 @@ def test_build_appends_offload_and_extra_args_last():
|
|||
|
||||
def test_build_negative_prompt_and_batch():
|
||||
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
|
||||
params = SdCppGenParams(prompt = "x", negative_prompt = "blurry", batch_count = 3)
|
||||
params = SdCppGenParams(prompt = "x", negative_prompt = "blurry")
|
||||
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
|
||||
assert _pair(cmd, "--negative-prompt") == "blurry"
|
||||
assert _pair(cmd, "--batch-count") == "3"
|
||||
# A CLI batch would silently drop every image after the first (the runner only
|
||||
# collects the literal --output path), so the builder rejects it outright.
|
||||
with pytest.raises(ValueError, match = "single-image"):
|
||||
build_sd_cpp_command(
|
||||
"/bin/sd-cli",
|
||||
files,
|
||||
SdCppGenParams(prompt = "x", batch_count = 3),
|
||||
output_path = "/o.png",
|
||||
)
|
||||
|
||||
|
||||
def test_build_omits_unset_optional_params():
|
||||
|
|
|
|||
|
|
@ -323,6 +323,22 @@ def test_generate_raises_when_no_output_despite_success(tmp_path, monkeypatch):
|
|||
)
|
||||
|
||||
|
||||
def test_generate_does_not_return_stale_preexisting_output(tmp_path, monkeypatch):
|
||||
# A leftover file at the target path must not satisfy the post-run output check
|
||||
# when the run itself produced nothing: the target is cleared before the run.
|
||||
e = _engine(tmp_path)
|
||||
out = tmp_path / "img.png"
|
||||
out.write_bytes(b"stale")
|
||||
_patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out, write = False)
|
||||
with pytest.raises(RuntimeError, match = "no image"):
|
||||
e.generate(
|
||||
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
||||
SdCppGenParams(prompt = "x"),
|
||||
output_path = str(out),
|
||||
)
|
||||
assert not out.exists()
|
||||
|
||||
|
||||
def test_generate_raises_when_binary_missing():
|
||||
e = SdCppEngine(binary = None)
|
||||
with pytest.raises(RuntimeError, match = "not found"):
|
||||
|
|
|
|||
|
|
@ -1982,7 +1982,9 @@ export function HubModelPicker({
|
|||
);
|
||||
// Local ./models entries. Chat-only Studio runs GGUF (any host) and MLX (Mac
|
||||
// only), so raw checkpoints there are hidden (mirrors the cached non-GGUF
|
||||
// rule). An MLX build a Mac user dropped in ./models stays selectable.
|
||||
// rule). An MLX build a Mac user dropped in ./models stays selectable. A
|
||||
// task-scoped picker (Images) is exempt: the image backend loads local
|
||||
// diffusers/safetensors pipelines even on chat-only (no-GPU, native) hosts.
|
||||
const sortedLocalDir = useMemo(
|
||||
() =>
|
||||
sortLocalModels(
|
||||
|
|
@ -1990,6 +1992,7 @@ export function HubModelPicker({
|
|||
(m) =>
|
||||
passesTaskGate(m.task, m.model_id ?? m.id, task) &&
|
||||
(!chatOnly ||
|
||||
task != null ||
|
||||
localModelIsGguf(m) ||
|
||||
(isMac && localModelIsMlx(m))) &&
|
||||
localModelMatchesFormat(m, formatFilter) &&
|
||||
|
|
@ -2279,6 +2282,14 @@ export function HubModelPicker({
|
|||
}
|
||||
|
||||
if (section === "recommended") {
|
||||
// Curated safetensors rows render ABOVE the recommended rows (and call
|
||||
// getOptionProps), so their keys must lead here or they fall back to the
|
||||
// duplicate ...-option-missing id and drop out of arrow-key navigation.
|
||||
keys.push(
|
||||
...curatedSafetensorsRows.map((m) =>
|
||||
makeModelOptionKey("curated-safetensors", m.id),
|
||||
),
|
||||
);
|
||||
keys.push(
|
||||
...recommendedRows.map((r) => makeModelOptionKey("recommended", r.id)),
|
||||
);
|
||||
|
|
@ -2288,6 +2299,7 @@ export function HubModelPicker({
|
|||
}, [
|
||||
cachedReady,
|
||||
chatOnly,
|
||||
curatedSafetensorsRows,
|
||||
sortedCustomFolderModels,
|
||||
customFoldersCollapsed,
|
||||
downloadedCollapsed,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue