From 94d74e8bbe896118960794d74baac457aa5c5241 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:13:04 +0000 Subject: [PATCH 1/2] Diffusion: guard trust check against OSError and validate conditioning inputs - _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid characters (or a bare owner/name id) can't raise OSError; treat any failure as not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request still raises the clear FileNotFoundError for a genuinely missing local pick. - generate(): reject mask_image / upscale / reference_images supplied without an input image, and reject reference_images on a family that does not support reference conditioning, instead of silently degrading to txt2img / img2img. --- studio/backend/core/inference/diffusion.py | 31 +++++++++++++++++-- .../backend/tests/test_diffusion_backend.py | 30 ++++++++++++++++++ 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index af8b14df0d..a221021358 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -191,9 +191,18 @@ def _is_trusted_diffusion_repo(repo_id: str) -> bool: on an arbitrary repo, which fetches and deserialises third-party weights. So the non-GGUF paths are gated to the ``unsloth/*`` org (the curated safetensors models) and to local paths the user explicitly pointed at (already on their disk). The GGUF path - is unchanged and stays open to any repo, as before.""" - if Path(repo_id).expanduser().exists(): - return True + is unchanged and stays open to any repo, as before. + + A bare ``owner/name`` HF id is never a real filesystem path, and an id with invalid + characters makes ``Path.exists()`` raise OSError; treat any such failure as "not a + local path" so the trust decision falls through to the unsloth/ check (the loader's + validate_load_request raises the clear FileNotFoundError for a genuinely missing + local pick).""" + try: + if Path(repo_id).expanduser().exists(): + return True + except OSError: + pass return repo_id.strip().lower().startswith("unsloth/") @@ -1273,6 +1282,22 @@ class DiffusionBackend: pipe = state.pipe init_pil = mask_pil = None ref_extra: list = [] + # Validate parameter dependencies up front: mask / upscale / reference all + # need an input image, and reference conditioning needs a family that + # supports it. Without these guards an unsupported combination would be + # silently ignored and quietly fall back to txt2img / img2img. + if init_image is None: + if mask_image is not None: + raise ValueError("mask_image requires an input image (init_image).") + if upscale is not None and upscale > 1.0: + raise ValueError("upscale requires an input image (init_image).") + if reference_images: + raise ValueError("reference_images require an input image (init_image).") + if reference_images and not getattr(state.family, "reference", False): + raise ValueError( + f"Reference images are not supported for the '{state.family.name}' " + "model family." + ) if getattr(state.family, "edit", False): # Instruction editing: the loaded pipe is the edit pipeline. It always # needs an input image; the prompt is the edit instruction. No mask, no diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 6e88968316..947d9e7ae3 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -468,6 +468,36 @@ def test_generate_img2img_unsupported_family_raises(fake_runtime, tmp_path, monk backend.generate(prompt = "x", steps = 4, init_image = _tiny_png_b64()) +def test_generate_rejects_conditioning_without_init_image(fake_runtime, tmp_path): + """mask / upscale / reference all need an input image; without one they must raise a + clear ValueError rather than silently degrading to txt2img.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + with pytest.raises(ValueError, match = "mask_image requires"): + backend.generate(prompt = "x", steps = 4, mask_image = _mask_b64(64)) + with pytest.raises(ValueError, match = "upscale requires"): + backend.generate(prompt = "x", steps = 4, upscale = 2.0) + with pytest.raises(ValueError, match = "reference_images require"): + backend.generate(prompt = "x", steps = 4, reference_images = [_tiny_png_b64()]) + + +def test_generate_rejects_reference_on_unsupported_family(fake_runtime, tmp_path): + """A non-reference family rejects reference_images instead of silently dropping them.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + with pytest.raises(ValueError, match = "Reference images are not supported"): + backend.generate( + prompt = "x", steps = 4, init_image = _tiny_png_b64(), + reference_images = [_tiny_png_b64()], + ) + + def test_generate_upscale_enlarges_and_low_strength(fake_runtime, tmp_path): """An init_image + upscale factor routes generate() through the family's img2img pipeline (hires fix): the source is enlarged to size*factor (rounded to /16) before the From dcf6cf511742e21d747d9ea7b21bc6e1ea787479 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 01:14:19 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 947d9e7ae3..5aa8c9a71e 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -493,7 +493,9 @@ def test_generate_rejects_reference_on_unsupported_family(fake_runtime, tmp_path ) with pytest.raises(ValueError, match = "Reference images are not supported"): backend.generate( - prompt = "x", steps = 4, init_image = _tiny_png_b64(), + prompt = "x", + steps = 4, + init_image = _tiny_png_b64(), reference_images = [_tiny_png_b64()], )