diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 19f1aff2a8..4e5a2efec6 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -273,24 +273,33 @@ class DiffusionBackend: return self._repo_id def status(self) -> dict[str, Any]: + # Take _lock so the snapshot cannot observe a torn state where + # _pipe was already swapped but _family/_repo_id haven't been + # updated yet (or vice versa). Frontend polling at 1 Hz would + # otherwise render impossible "loaded but no repo_id" states. # Only echo the GGUF basename; full absolute path leaks the # local HF cache layout (and the system username on default # POSIX layouts) to any authenticated Studio session. - gguf_basename = Path(self._gguf_path).name if self._gguf_path else None - return { - "is_loaded": self.is_loaded, - "is_loading": self._loading, - "repo_id": self._repo_id, - "family": self._family.name if self._family else None, - "pipeline_class": self._family.pipeline_class if self._family else None, - "base_repo": self._base_repo, - "gguf_filename": gguf_basename, - "device": self._device, - "dtype": self._dtype, - "loaded_at": self._loaded_at, - "last_error": self._last_error, - "supported_families": supported_families(), - } + with self._lock: + gguf_basename = ( + Path(self._gguf_path).name if self._gguf_path else None + ) + return { + "is_loaded": self._pipe is not None, + "is_loading": self._loading, + "repo_id": self._repo_id, + "family": self._family.name if self._family else None, + "pipeline_class": ( + self._family.pipeline_class if self._family else None + ), + "base_repo": self._base_repo, + "gguf_filename": gguf_basename, + "device": self._device, + "dtype": self._dtype, + "loaded_at": self._loaded_at, + "last_error": self._last_error, + "supported_families": supported_families(), + } def _pick_device_and_dtype(self) -> tuple[str, "Any"]: """Pick (device, dtype) for the current host. @@ -506,6 +515,14 @@ class DiffusionBackend: return self.status() except Exception as exc: + # Scrub hf_token and pipe_kwargs from frame locals BEFORE + # logger.exception() captures them. Rich tracebacks and + # some structlog formatters render frame locals, which + # would otherwise echo the raw hf_... token into logs + # and any error reporting sink the user has wired up. + hf_token = None # noqa: F841 + pipe_kwargs = None # noqa: F841 + single_file_kwargs = None # noqa: F841 with self._lock: self._last_error = str(exc) logger.exception("Diffusion load failed for %s", repo_id) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 9ea113e488..13e2a4e83f 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2632,6 +2632,29 @@ async def delete_cached_model( except Exception: pass + # Also refuse to delete the cache underlying a loaded diffusion + # pipeline. The diffusion backend mmap's the GGUF + base repo + # weights and continues to read from the cache long after load, + # so deleting them out from under it would corrupt generation. + try: + from core.inference.diffusion import get_diffusion_backend + + diff_backend = get_diffusion_backend() + diff_status = diff_backend.status() + if diff_status.get("is_loaded"): + diff_repo = (diff_status.get("repo_id") or "").lower() + diff_base = (diff_status.get("base_repo") or "").lower() + needle = repo_id.lower() + if diff_repo.startswith(needle) or diff_base.startswith(needle): + raise HTTPException( + status_code = 400, + detail = "Unload the diffusion image model before deleting", + ) + except HTTPException: + raise + except Exception: + pass + try: cache_scans = _all_hf_cache_scans() diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 1cc795345c..05eb5c3a89 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -200,10 +200,28 @@ export function ImagesPage() { } setBusy("generating"); try { - const parsedSeed = seed.trim() ? Number(seed.trim()) : undefined; - if (parsedSeed !== undefined && !Number.isFinite(parsedSeed)) { - toast.error("Seed must be a number"); - return; + // Reject non-integer or out-of-safe-integer-range seeds rather + // than silently rounding via Number(). The backend takes an int + // and a precision loss here would yield a different image than + // the seed the user typed. + const seedStr = seed.trim(); + let parsedSeed: number | undefined; + if (seedStr) { + if (!/^-?\d+$/.test(seedStr)) { + toast.error("Seed must be an integer"); + return; + } + const candidate = Number(seedStr); + if ( + !Number.isFinite(candidate) || + !Number.isSafeInteger(candidate) + ) { + toast.error( + "Seed must fit in a JavaScript safe integer (<= 2^53 - 1)", + ); + return; + } + parsedSeed = candidate; } const out = await generateDiffusionImage({ prompt,