Reject doomed local picks before the GPU handoff; load bare local safetensors

Several image/video/training preflights ran before the route acquires the GPU or
frees resident models, but let a doomed local pick through and only failed deep in
the background load, after the user's chat/Images/Video model was already evicted.

- Local base_repo / base_model: _is_trusted_diffusion_repo accepts any existing
  local path, but the base loads via from_pretrained (needs model_index.json). A
  local dir that is not a diffusers pipeline passed the trust gate, evicted the
  resident model, then failed. Add a shared _assert_local_base_is_pipeline check
  and call it in the image, video, and training preflights.
- Dataset images: discover_image_caption_pairs only checked filenames, so a
  corrupt or zero-byte upload passed the start-route preflight, freed the GPU, then
  crashed the spawned trainer in PIL. Add an opt-in verify_images decode probe
  (cheap PIL header check) that the start route enables; the trainers leave it off
  since they decode every image anyway.
- Local single-file safetensors: the On-Device scanner advertises a bare
  .safetensors directory (no model_index.json) as a text-to-image model, but the
  picker starts it as a pipeline with no filename, so every click 400s. Reinterpret
  such a pick as a single_file load of the sole checkpoint (resolve_local_single_file)
  so the advertised model is actually loadable.

Regression tests for each: local non-pipeline base (image/video/training), the
verify_images decode gate, and resolve_local_single_file.

studio/backend/core/inference/diffusion.py

@@ -147,6 +147,30 @@ def resolve_model_kind(gguf_filename: Optional[str], model_kind: Optional[str] =
     return "single_file"
 
 
+def resolve_local_single_file(model_path: str) -> Optional[str]:
+    """The sole single-file checkpoint basename in a local ``model_path`` directory that is NOT a
+    diffusers pipeline (no ``model_index.json``) and holds exactly one ``.safetensors`` file, else
+    None.
+
+    The On-Device scanner advertises a bare single-file safetensors directory as a text-to-image
+    model (it matches a known family by name), but the local picker starts it as a ``pipeline``
+    with no filename, so a pipeline load 400s on the missing ``model_index.json`` and the
+    advertised model is unusable. The images load route uses this to reinterpret such a pick as a
+    ``single_file`` load of the sole checkpoint. A real pipeline dir (has ``model_index.json``) or
+    an ambiguous one (0 or more than 1 ``.safetensors``, e.g. a sharded pipeline) returns None and
+    loads unchanged. Never raises."""
+    try:
+        root = Path(model_path).expanduser()
+        if not root.is_dir() or (root / "model_index.json").is_file():
+            return None
+        checkpoints = [
+            p.name for p in root.iterdir() if p.is_file() and p.suffix.lower() == ".safetensors"
+        ]
+    except OSError:
+        return None
+    return checkpoints[0] if len(checkpoints) == 1 else None
+
+
 def _decode_b64_image(data: str, *, mode: str = "RGB") -> Any:
     """Decode a base64 (optionally ``data:`` URL) image string to a PIL image.
 
@@ -276,6 +300,32 @@ def _is_trusted_diffusion_repo(repo_id: str) -> bool:
     return rid.startswith("unsloth/") or rid in _TRUSTED_NON_GGUF_REPOS
 
 
+def _assert_local_base_is_pipeline(base_repo: str) -> None:
+    """A companion ``base_repo`` fed to ``from_pretrained(base)`` (or ``config=base``) must be a
+    diffusers PIPELINE directory (has ``model_index.json``). ``_is_trusted_diffusion_repo`` accepts
+    ANY existing local path, so without this a local base that is not a pipeline dir would pass the
+    preflight, let the route evict the resident GPU model, then fail deep in the background load --
+    the eviction this validation exists to prevent. A non-existent local base is already rejected
+    by the trust check (it is neither an existing path nor an unsloth/*/allowlisted repo); a bare
+    remote id is left for the loader to resolve. Shared by the image, video, and training preflights
+    so their local-base shape check stays in sync. Never evicts; raises ValueError on a bad local
+    base."""
+    base = (base_repo or "").strip()
+    if not base:
+        return
+    try:
+        root = Path(base).expanduser()
+        exists = root.exists()
+    except OSError:
+        return  # invalid path characters -> a remote id, not a local path
+    if not exists:
+        return
+    if not root.is_dir() or not (root / "model_index.json").is_file():
+        raise ValueError(
+            f"Local base_repo is not a diffusers pipeline directory (no model_index.json): {base}"
+        )
+
+
 @dataclass(frozen = True)
 class _LoadState:
     """Everything about the currently-loaded pipeline, swapped as one unit."""
@@ -639,6 +689,11 @@ class DiffusionBackend:
                 f"base_repo is restricted to unsloth/* repos (or a local path); got "
                 f"'{base_repo}'."
             )
+        # An existing LOCAL base_repo is loaded as a full pipeline (from_pretrained(base) /
+        # config=base), which needs a model_index.json. Any existing path passes the trust check
+        # above, so reject a non-pipeline local base here -- before the route evicts the resident
+        # model -- rather than deep in the background load. Mirrors the repo_id check below.
+        _assert_local_base_is_pipeline(base_repo)
         # Reject a bad LOCAL pick now (the same checks the load would hit later), so
         # the route never evicts a working chat model for a request that can't load.
         # A path-shaped repo_id (absolute / ~ / ./ / ..) is meant to be on disk, so a

studio/backend/core/inference/video.py

@@ -426,6 +426,14 @@ class VideoBackend:
                 f"base_repo is limited to unsloth/* repos, the official family base "
                 f"repos, and local paths; '{base_repo}' is neither."
             )
+        # An existing LOCAL base_repo loads as a full pipeline (from_pretrained(base) / config=base),
+        # which needs a model_index.json. The pipeline-kind shape check below covers only repo_id,
+        # and an explicit base_repo is only meaningful for gguf/single_file kinds, so a non-pipeline
+        # local base would otherwise pass here and fail deep in the background load AFTER the route
+        # evicted the resident model. Shared helper, so image/video/training stay in sync.
+        from core.inference.diffusion import _assert_local_base_is_pipeline
+
+        _assert_local_base_is_pipeline(base_repo)
         if kind in ("gguf", "single_file") and not gguf_filename:
             raise ValueError("A gguf/single_file load needs the checkpoint filename.")
         if kind in ("gguf", "single_file") and fam.is_moe:

studio/backend/core/training/diffusion_train_common.py

@@ -598,6 +598,7 @@ def discover_image_caption_pairs(
     *,
     instance_prompt: Optional[str] = None,
     caption_column: str = "text",
+    verify_images: bool = False,
 ) -> list[tuple[str, str]]:
     """Resolve ``(image_path, caption)`` pairs from a dataset directory.
 
@@ -614,6 +615,12 @@ def discover_image_caption_pairs(
     Images with no caption from any source are skipped. Pure filesystem + JSON, so it is
     unit-testable without torch. Raises FileNotFoundError for a missing dir and ValueError
     when nothing is captionable.
+
+    ``verify_images`` (opt-in) additionally runs a cheap PIL header probe on each captioned
+    image and raises ValueError on a corrupt/zero-byte/truncated file. The start route enables
+    it so a bad upload is rejected BEFORE the resident GPU models are freed, instead of crashing
+    the spawned trainer after teardown; the trainers leave it off (they decode every image
+    anyway, so a second probe pass would be redundant).
     """
     root = Path(data_dir).expanduser()
     if not root.is_dir():
@@ -658,6 +665,21 @@ def discover_image_caption_pairs(
         if caption is None and instance_prompt:
             caption = instance_prompt
         if caption:
+            if verify_images:
+                # Reject a corrupt / zero-byte / truncated image now via a cheap PIL header
+                # probe (verify() does not decode the full pixels): otherwise it passes this
+                # filename-only discovery, the start route frees the resident GPU models, and
+                # the spawned trainer only then crashes in Image.open -- the eviction this
+                # preflight exists to prevent.
+                try:
+                    from PIL import Image
+                    with Image.open(img) as _probe:
+                        _probe.verify()
+                except Exception as e:  # noqa: BLE001 -- corrupt/zero-byte/truncated file
+                    raise ValueError(
+                        f"Image cannot be decoded: {img.name} ({e}). Remove or replace the "
+                        f"corrupt or zero-byte file before training."
+                    ) from e
             pairs.append((str(img), caption))
 
     if not pairs:
@@ -817,12 +839,18 @@ def _assert_trusted_base_model(base_model: str) -> None:
     a local path or a trusted repo (``unsloth/*`` or an allowlisted official base). This runs
     BEFORE ``from_pretrained`` so an untrusted remote repo (which could ship pickle weights)
     is never fetched or deserialised."""
-    from core.inference.diffusion import _is_trusted_diffusion_repo
+    from core.inference.diffusion import _assert_local_base_is_pipeline, _is_trusted_diffusion_repo
+
     if not _is_trusted_diffusion_repo(base_model):
         raise ValueError(
             f"Refusing to train from untrusted base model '{base_model}'. Use a local path or "
             f"a trusted repo (an unsloth/* repo or an official base)."
         )
+    # An existing LOCAL base is loaded as a full pipeline (from_pretrained(base_model)) by the
+    # spawned trainer, which needs a model_index.json. Any existing path is "trusted" above, so
+    # reject a non-pipeline local dir here -- before /diffusion/start frees the resident GPU
+    # models -- rather than have the child fail after the teardown.
+    _assert_local_base_is_pipeline(base_model)
 
 
 def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Optional[str]:

studio/backend/routes/inference.py

@@ -12433,7 +12433,11 @@ def _guard_diffusion_load_against_training() -> None:
 async def load_diffusion_model(
     request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
 ):
-    from core.inference.diffusion import get_diffusion_backend, resolve_model_kind
+    from core.inference.diffusion import (
+        get_diffusion_backend,
+        resolve_local_single_file,
+        resolve_model_kind,
+    )
     from core.inference.diffusion_device import resolve_diffusion_device_target
     from core.inference.diffusion_engine_router import (
         annotate_status,
@@ -12447,6 +12451,17 @@ async def load_diffusion_model(
         # Resolve the load kind once (gguf / single_file / pipeline) so validation,
         # engine selection, and the load all agree. A bad explicit kind raises here -> 400.
         kind = resolve_model_kind(request.gguf_filename, request.model_kind)
+        # A local On-Device pick can be a bare single-file .safetensors directory (no
+        # model_index.json): the scanner advertises it as a text-to-image model, but the local
+        # picker starts it as a pipeline with no filename, so a pipeline load would 400 on the
+        # missing model_index.json and the advertised model is unusable. If the directory holds
+        # exactly one checkpoint, reinterpret the pick as a single_file load of it (the only
+        # loadable shape for that dir), so validation, engine selection, and the load all agree.
+        if kind == "pipeline" and not request.gguf_filename:
+            sole = await asyncio.to_thread(resolve_local_single_file, request.model_path)
+            if sole is not None:
+                request.gguf_filename = sole
+                kind = resolve_model_kind(sole)
         # Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family,
         # missing local GGUF, a non-unsloth non-GGUF repo) must not evict a working chat
         # model and then 400. The validated family also drives engine selection below.

studio/backend/routes/training.py

@@ -1314,6 +1314,10 @@ async def start_diffusion_training(
             config["data_dir"],
             instance_prompt = config.get("instance_prompt") or None,
             caption_column = config.get("caption_column") or "text",
+            # Decode-probe every image now (cheap PIL header check) so a corrupt / zero-byte
+            # upload is rejected with a 400 BEFORE _free_gpu_for_diffusion_training() tears down
+            # the user's resident models, instead of crashing the spawned trainer post-eviction.
+            verify_images = True,
         )
     except (FileNotFoundError, ValueError) as e:
         raise HTTPException(status_code = 400, detail = str(e))

studio/backend/tests/test_diffusion_backend.py

@@ -1223,7 +1223,20 @@ def test_validate_gates_untrusted_base_repo(fake_runtime, tmp_path):
             model_kind = "gguf",
             base_repo = "evil/companions",
         )
-    # A local base_repo (already on disk) still passes the gate.
+    # A local base_repo dir that is NOT a diffusers pipeline (no model_index.json) is rejected
+    # HERE, before the GPU handoff: it passes the any-existing-path trust check but the base loads
+    # via from_pretrained (needs model_index.json), so it would otherwise evict the resident model
+    # and only then fail in the background load.
+    bad_base = tmp_path / "bare-base"
+    bad_base.mkdir()
+    with pytest.raises(ValueError, match = "model_index.json"):
+        backend.validate_load_request(
+            "unsloth/Qwen-Image-2512-GGUF",
+            gguf_filename = "x.gguf",
+            model_kind = "gguf",
+            base_repo = str(bad_base),
+        )
+    # A local base_repo that IS a real pipeline dir (model_index.json) passes the gate.
     (tmp_path / "model_index.json").write_text("{}")
     fam = backend.validate_load_request(
         "unsloth/Qwen-Image-2512-GGUF",
@@ -1234,6 +1247,32 @@ def test_validate_gates_untrusted_base_repo(fake_runtime, tmp_path):
     assert fam is not None
 
 
+def test_resolve_local_single_file(tmp_path):
+    # A bare single-file safetensors directory (no model_index.json) resolves to that checkpoint's
+    # basename, so the images load route can reinterpret an On-Device "pipeline" pick as a
+    # single_file load instead of 400ing on the missing model_index.json.
+    from core.inference.diffusion import resolve_local_single_file
+
+    d = tmp_path / "solo"
+    d.mkdir()
+    (d / "model.safetensors").write_bytes(b"w")
+    assert resolve_local_single_file(str(d)) == "model.safetensors"
+
+    # A real diffusers pipeline dir (has model_index.json) loads as a pipeline unchanged -> None.
+    (d / "model_index.json").write_text("{}")
+    assert resolve_local_single_file(str(d)) is None
+
+    # Ambiguous (two checkpoints, e.g. a sharded pipeline) or empty dirs -> None (unchanged load).
+    d2 = tmp_path / "shards"
+    d2.mkdir()
+    (d2 / "a.safetensors").write_bytes(b"w")
+    (d2 / "b.safetensors").write_bytes(b"w")
+    assert resolve_local_single_file(str(d2)) is None
+    assert resolve_local_single_file(str(tmp_path / "empty-nonexistent")) is None
+    # A remote repo id (not a local dir) -> None.
+    assert resolve_local_single_file("unsloth/Qwen-Image-2512-GGUF") is None
+
+
 def test_resolve_base_repo_drops_untrusted_card_tag(monkeypatch):
     # When no base_repo is passed, the base is resolved from the GGUF repo's base_model card
     # tag -- attacker-controlled metadata on any remote repo -- and then loaded via

studio/backend/tests/test_diffusion_base_precision.py

@@ -563,3 +563,21 @@ def test_request_model_base_precision():
         }
     )
     assert cfg.base_precision == "bf16"
+
+
+def test_assert_trusted_base_model_rejects_local_non_pipeline(tmp_path):
+    # A local base_model dir that is NOT a diffusers pipeline (no model_index.json) is "trusted"
+    # (any existing path passes the trust check), but the spawned trainer loads it via
+    # from_pretrained, so it must be rejected in the /diffusion/start preflight BEFORE
+    # _free_gpu_for_diffusion_training tears down the resident models -- not fail the child after
+    # the eviction.
+    bad = tmp_path / "bare-base"
+    bad.mkdir()
+    with pytest.raises(ValueError, match = "model_index.json"):
+        common._assert_trusted_base_model(str(bad))
+    # A real local pipeline dir (model_index.json) is accepted.
+    (bad / "model_index.json").write_text("{}")
+    common._assert_trusted_base_model(str(bad))  # no raise
+    # An untrusted remote base is still rejected by the trust gate.
+    with pytest.raises(ValueError, match = "untrusted"):
+        common._assert_trusted_base_model("evil/base")

studio/backend/tests/test_diffusion_lora_trainer.py

@@ -81,6 +81,35 @@ def test_discover_custom_caption_column(tmp_path):
     assert discover_image_caption_pairs(tmp_path, caption_column = "caption")[0][1] == "col"
 
 
+def test_discover_verify_images_rejects_undecodable(tmp_path):
+    # verify_images (opt-in, enabled by the start route) rejects a corrupt / zero-byte image with
+    # a clear ValueError -> 400 BEFORE the route frees the resident GPU models, instead of letting
+    # the spawned trainer crash in PIL after the teardown. The trainers leave it off (default),
+    # since they decode every image anyway.
+    from PIL import Image
+
+    good = tmp_path / "good.png"
+    Image.new("RGB", (8, 8), "white").save(good)
+    (tmp_path / "good.txt").write_text("ok", encoding = "utf-8")
+    # A zero-byte file with an image extension + a caption passes filename-only discovery.
+    bad = tmp_path / "bad.png"
+    bad.write_bytes(b"")
+    (tmp_path / "bad.txt").write_text("broken", encoding = "utf-8")
+
+    # Default (verify off): the bad file is accepted (filename-only), matching trainer behavior.
+    pairs = dict(discover_image_caption_pairs(tmp_path))
+    assert str(bad) in pairs and str(good) in pairs
+
+    # verify_images on: the undecodable file raises a clear ValueError.
+    with pytest.raises(ValueError, match = "cannot be decoded"):
+        discover_image_caption_pairs(tmp_path, verify_images = True)
+
+    # A dataset of only valid images passes the verify.
+    bad.unlink()
+    (tmp_path / "bad.txt").unlink()
+    assert discover_image_caption_pairs(tmp_path, verify_images = True) == [(str(good), "ok")]
+
+
 def test_discover_empty_raises(tmp_path):
     _touch(tmp_path / "x.png")  # no captions anywhere, no instance prompt
     with pytest.raises(ValueError, match = "No captioned images"):

studio/backend/tests/test_video_backend.py

@@ -478,6 +478,32 @@ def test_validate_rejects_local_pipeline_without_model_index(tmp_path):
     assert fam.name == "ltx-2"
 
 
+def test_validate_rejects_local_base_repo_without_model_index(tmp_path):
+    backend = VideoBackend()
+    # A local base_repo dir that is NOT a diffusers pipeline (no model_index.json) passes the
+    # any-existing-path trust check, but the base loads via from_pretrained (needs model_index),
+    # so reject it HERE before the route hands the GPU to VIDEO -- the pipeline-kind shape check
+    # covers only repo_id, and an explicit base_repo is only meaningful for a gguf/single_file load.
+    bad_base = tmp_path / "bare-base"
+    bad_base.mkdir()
+    with pytest.raises(ValueError, match = "model_index.json"):
+        backend.validate_load_request(
+            "unsloth/LTX-2.3-GGUF",
+            gguf_filename = "x.gguf",
+            model_kind = "gguf",
+            base_repo = str(bad_base),
+        )
+    # A local base_repo that IS a real pipeline dir passes the gate.
+    (bad_base / "model_index.json").write_text("{}")
+    fam = backend.validate_load_request(
+        "unsloth/LTX-2.3-GGUF",
+        gguf_filename = "x.gguf",
+        model_kind = "gguf",
+        base_repo = str(bad_base),
+    )
+    assert fam.name == "ltx-2"
+
+
 def test_validate_rejects_gguf_repo_as_pipeline():
     backend = VideoBackend()
     # A -GGUF repo with no quant filename resolves to the pipeline kind and would