diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 50bc621417..956687c4d6 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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 diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 692a9d22a0..69447984d4 100644 --- a/studio/backend/core/inference/video.py +++ b/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: diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index e14022b73c..5d54e4717a 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/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]: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 96f94d0457..75c913e01c 100644 --- a/studio/backend/routes/inference.py +++ b/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. diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 8aedecbbba..d663c31034 100644 --- a/studio/backend/routes/training.py +++ b/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)) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 2a8b5922ca..6503bfa794 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 5b19e43212..ed6128c87b 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/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") diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index f2971451e9..17bee1c1ae 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/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"): diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 68a23b3ce9..a215275b88 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/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