diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 8e7cd4ea5a..3b33a7ce84 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -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", diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index de55a062e4..4d4f95541f 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.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 diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index e652b0068c..8b1ce29ec8 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -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): diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index b7a2b7ac45..3cb74d7f80 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index f49843ed3d..88ce4007dc 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -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) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 065fd4eb11..276dcad3eb 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -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: diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index c370777bab..09e3faa778 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -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: diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 577370468b..a6dd63f4bc 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -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]: diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index fe1e7ebc92..0838edf4d6 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -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( diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 50bb0f5d3b..54ff06d2ce 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -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, diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index cdf78995b3..aabfbf42a6 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -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: diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 3883b9847f..1d57322781 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1cc1c83ad4..32f7b48f58 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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")) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d80a8e3b06..554b3b85b7 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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} diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 3e43aab8a9..99c5e966f2 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 1fb6714cf7..57f9e5a24c 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -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() diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 62071d9aa6..befab46a6c 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index ced1b752ac..ab5eb71b22 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -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 == {} diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index 154b671215..451f594ba1 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -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. diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index d9078bf0e7..bd91eae9c8 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -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( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index cd271c2f05..9233a54def 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -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") diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index 4157d0e01f..fdc5aec863 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -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(): diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index daaf5223a8..8fee117b3a 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -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"): diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 336dfece69..f096af73c8 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -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,