diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 78c2e242b6..0defec3fd1 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -306,6 +306,9 @@ class DiffusionBackend: transformer_cache_threshold: Optional[float] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" + # A blank token (the Studio default when none is configured) must mean + # "anonymous", not an explicit empty credential the Hub rejects with 401. + hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, family_override = family_override ) @@ -416,6 +419,18 @@ class DiffusionBackend: fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0 return _progress("downloading", downloaded, expected, fraction) + def loading_repo_ids(self) -> tuple[str, ...]: + """Repo ids an in-flight background load is downloading (empty when idle). + + The delete-cached guard needs this: during a load ``status()["loaded"]`` is + still False, but deleting the target repo (or its companion base) would yank + blobs and snapshot files from under the download/assembly.""" + 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) + @staticmethod def _estimate_download_bytes( repo_id: str, gguf_filename: Optional[str], base_repo: str, hf_token: Optional[str] @@ -483,7 +498,10 @@ class DiffusionBackend: _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad - # family fails with ValueError even in a no-diffusers runtime. + # family fails with ValueError even in a no-diffusers runtime. Sanitize the + # token here too (direct callers bypass begin_load): a blank string must + # load anonymously, not 401 as an explicit empty credential. + hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, family_override = family_override ) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 88661f2be1..cdd2b1c543 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -446,6 +446,16 @@ class SdCppDiffusionBackend: fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0 return _progress("downloading", downloaded, expected, fraction) + 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.""" + 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) + # ── Generate ─────────────────────────────────────────────────────────── def generate( diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 2c1701f6ea..e2e57a4648 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1806,8 +1806,11 @@ class DiffusionGenerateRequest(BaseModel): ) steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps") guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") + # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript + # rounds integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then + # generate a different image. Random seeds are already masked to this range. seed: Optional[int] = Field( - None, ge = 0, le = 2**64 - 1, description = "Seed for reproducibility (random if omitted)" + None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)" ) batch_size: int = Field( 1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)" diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 321f87072f..da23205f07 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -27,6 +27,10 @@ class CachedModelRepo(BaseModel): repo_id: str size_bytes: int last_modified: Optional[float] = None + # "text-to-image" for cached diffusers image repos; response_model would silently + # drop the value the handler sets, letting image-only repos pass the chat picker's + # task gate. + task: Optional[str] = None class CachedModelsResponse(BaseModel): @@ -3320,8 +3324,13 @@ async def delete_cached_model( # delete guard is otherwise chat-only, so its GGUF could be removed from # under a live pipeline. Repo-level match, like the chat guards above. try: - from core.inference.diffusion import get_diffusion_backend - diffusion_status = get_diffusion_backend().status() + # The ACTIVE engine (diffusers or native sd_cpp): on a native selection the + # diffusers singleton reports unloaded while sd-cli still generates from the + # cached GGUF, so checking it alone would let the files be deleted mid-use. + from core.inference.diffusion_engine_router import get_active_diffusion_engine + + engine = get_active_diffusion_engine() + diffusion_status = engine.status() if diffusion_status.get("loaded") and diffusion_status.get("repo_id"): loaded_id = str(diffusion_status["repo_id"]).lower() if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()): @@ -3329,6 +3338,17 @@ async def delete_cached_model( status_code = 400, detail = "Unload the model before deleting", ) + # Also refuse while a background image load is DOWNLOADING this repo (or its + # companion base): status().loaded is still False in that window, but deleting + # would remove blobs from under the in-flight download/assembly. + loading_ids = getattr(engine, "loading_repo_ids", tuple)() + for lid in loading_ids: + lid = str(lid).lower() + if lid == repo_id.lower() or lid.startswith(repo_id.lower()): + raise HTTPException( + status_code = 400, + detail = "An Images model load is using this repo; wait for it to finish", + ) except HTTPException: raise except Exception: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d5ab4ae29a..bcb0c344e0 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -372,9 +372,14 @@ async def start_training( # release the arbiter so it doesn't think the gone pipeline owns # the GPU. Must precede the chat block, which early-returns. from core.inference import gpu_arbiter - from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion_engine_router import ( + get_active_diffusion_engine, + ) - diffusion = get_diffusion_backend() + # The ACTIVE engine, not the diffusers singleton: on a native + # (sd_cpp) selection the diffusers backend reports unloaded while + # the native engine still holds model state / a live generation. + diffusion = get_active_diffusion_engine() if diffusion.is_loaded: logger.info( "Unloading diffusion (Images) model to free GPU memory for training" diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index bee562f4bd..73e2b6c5d8 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -1203,10 +1203,7 @@ export function AppSidebar() { icon={PaintBrush02Icon} label={t("shell.navigation.images")} active={pathname === "/images" || pathname.startsWith("/images/")} - disabled={chatOnly} - tooltip={trainExportDisabledHint} onClick={() => { - if (chatOnly) return; navigate({ to: "/images" }); closeMobileIfOpen(); }} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index d2421ef8a1..14174985de 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -695,7 +695,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { }, [refreshStatus, dismissLoadToast, pollLoadProgress]); const handleLoad = useCallback( - async (repoId: string, ggufFilename: string) => { + // Resolves true when the background load STARTED (callers may revert + // optimistic picker state on false); poll outcomes are handled internally. + async (repoId: string, ggufFilename: string): Promise => { // Cancel any prior poll loop so two can't run at once. if (pollTimer.current) clearTimeout(pollTimer.current); setBusy("loading"); @@ -717,9 +719,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) { toast.error(err instanceof Error ? err.message : "Failed to start load"); setBusy(null); void refreshStatus(); - return; + return false; } void pollLoadProgress(); + return true; }, [pollLoadProgress, refreshStatus, dismissLoadToast], ); @@ -733,11 +736,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // busy, while the backend rejects the second load with a 409. if (busy !== null) return; if (meta.ggufVariant && meta.ggufFilename) { + // Optimistic for instant picker feedback, but revert if the load fails to + // START (400/409/network): the selector must not advertise a quant that + // is not the loaded one. Poll-phase failures re-sync via refreshStatus. + const prevQuant = quant; setQuant(meta.ggufVariant); const d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); - void handleLoad(id, meta.ggufFilename); + void handleLoad(id, meta.ggufFilename).then((started) => { + if (!started) setQuant(prevQuant); + }); return; } // A direct single-file local .gguf pick has no variant/filename (custom folder / @@ -755,7 +764,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { void handleLoad(dir, filename); } }, - [busy, handleLoad], + [busy, handleLoad, quant], ); const handleUnload = useCallback(async () => { @@ -841,7 +850,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) { height: h, steps, guidance, - seed: baseSeed + i, + // Offset runs by the batch size: the native engine seeds image j of a + // run at seed+j, so a +1 run offset would regenerate the previous run's + // batch-mates. Unique per image on both engines, reproducible via recipes. + seed: baseSeed + i * batchSize, batch_size: batchSize, }); if (!isMounted.current) break;