Fix review findings on image generation: failed-load VRAM, API defaults, preflights

- Free reserved VRAM in the diffusion load worker's failure path: a load-time OOM
  never commits _state and the next load's _unload_locked early-returns, so nothing
  else reclaimed the half-built pipeline's memory
- Use a monotonic clock for the denoise ETA rate
- Sync _GENERATION_DEFAULTS with the UI table: kontext, flux.2-dev, sdxl-turbo and
  SDXL base rows so /v1/images/generations stops falling back to 9 steps / CFG 0
- 400 (not sanitized 500) when /v1/images/generations hits an edit-only model
- Fail fast on pre-Ampere CUDA in the DiT trainer instead of dying in model load
- Run the trainer trust gate in the diffusion training route before freeing GPU
  residents so an untrusted base cannot tear down loaded chat/Images models
- Protect native sd.cpp companion VAE/text-encoder repos from cache deletion while
  a load is downloading them
- Exempt the task-scoped Images picker from the chat-only GGUF/MLX format gate so
  local diffusers pipelines stay selectable on no-GPU hosts

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

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)

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 ───────────────────────────────────────────────────────────
 

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)
 

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"))

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.

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) &&