diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 79b8a50b11..9210249c6a 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/") @@ -1362,6 +1371,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..5aa8c9a71e 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -468,6 +468,38 @@ 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