diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index fda2a27a4b..0531b39aef 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -691,6 +691,10 @@ 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. + clear_gpu_cache() # 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 @@ -1971,7 +1975,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 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/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index d8cb21d409..78425833d6 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,16 @@ 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 ─────────────────────────────────────────────────────────── diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 50bb0f5d3b..d29a85a140 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) 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 54324dd544..a89a8d01c2 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1219,6 +1219,16 @@ 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. 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..c380e04244 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) &&