diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 97b86e2120..0bf1a1efbd 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -577,6 +577,7 @@ class DiffusionBackend: gguf_filename: Optional[str] = None, family_override: Optional[str] = None, model_kind: Optional[str] = None, + base_repo: Optional[str] = None, ) -> DiffusionFamily: """Cheap, network-free validation shared by the route (before it evicts the chat model) and the load paths, so an unloadable pick fails BEFORE the GPU @@ -628,6 +629,16 @@ class DiffusionBackend: f"Non-GGUF diffusion loads are restricted to unsloth/* repos (or a local " f"path); got '{repo_id}'. Pass a gguf_filename to load a GGUF instead." ) + # A companion base repo also loads via from_pretrained (its diffusers pipeline is + # assembled around the GGUF/single-file transformer), so it must clear the same trust + # bar as a non-GGUF repo_id -- otherwise a trusted GGUF model_path could smuggle in an + # arbitrary remote base that gets downloaded and deserialised. Gate it here (before the + # route evicts the resident model), mirroring the video loader's base_repo check. + if base_repo and base_repo.strip() and not _is_trusted_diffusion_repo(base_repo): + raise ValueError( + f"base_repo is restricted to unsloth/* repos (or a local path); got " + f"'{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 @@ -714,6 +725,10 @@ class DiffusionBackend: # A blank token (the Studio default when none is configured) must mean # "anonymous", not an explicit empty credential the Hub rejects with 401. hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None + # base_repo is gated at the /images/load route's pre-eviction validate_load_request + # (the client entry point); the re-validation here is a redundant cheap-fail guard for + # the resolved repo/family, so it does not re-gate base_repo (which internal callers pass + # through already-validated). fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, @@ -1029,6 +1044,9 @@ class DiffusionBackend: # token here too (direct callers bypass begin_load): a blank string must # load anonymously, not 401 as an explicit empty credential. hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None + # base_repo is gated at the route before eviction (validate_load_request there); this + # direct-load re-validation only cheap-fails the resolved repo/family, so it does not + # re-gate an already-validated base_repo. fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py index 4ce663c3ce..211b7e273f 100644 --- a/studio/backend/core/inference/llama_keepwarm.py +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -61,6 +61,13 @@ _INFERENCE_SUFFIXES = ( "/responses", "/generate/stream", # Studio's own streaming route on the same llama-server "/audio/generate", # direct GGUF TTS; can outlive the idle TTL + # Image/video generation holds a multi-GB diffusion/video pipeline for the whole request. + # Tracking them here lets other_inference_request_count() see an in-flight generation, so an + # API-key training start is refused (409) before its unload cancels the generation. endswith + # so the GET *-progress and */cancel variants are not matched. + "/images/generate", # /api/inference/images/generate + "/images/generations", # /v1/images/generations (+ /api/inference/images/generations) + "/video/generate", # /api/inference/video/generate ) diff --git a/studio/backend/main.py b/studio/backend/main.py index 6297bf21d7..c947c5df5d 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -727,8 +727,11 @@ from utils.upload_limits import ( # noqa: E402 ) _BODY_PROTECTED_PREFIXES = ( - "/v1/chat/completions", - "/v1/completions", + # Blanket-protect the whole OpenAI-compatible /v1 surface, like /api/inference below: every + # /v1 POST route (chat/completions, completions, images/generations, audio, embeddings, + # responses, messages, ...) buffers a JSON body and none is a multipart-upload passthrough, + # so a single prefix caps them all -- an enumerated list silently left new routes uncapped. + "/v1", "/p/", "/api/inference", "/api/data-recipe", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ed66b46113..505617dc64 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11948,6 +11948,7 @@ async def load_diffusion_model( gguf_filename = request.gguf_filename, family_override = request.family_override, model_kind = kind, + base_repo = request.base_repo, ) # Refuse while training is running: a multi-GB diffusion pipeline would # compete with the training subprocess for VRAM. The chat path does the diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 0ccc7a5d23..2c64e43ebe 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1210,6 +1210,29 @@ def test_load_sdxl_rejects_untrusted_repo(fake_runtime): backend.load_pipeline("randomorg/my-sdxl-merge", family_override = "sdxl") +def test_validate_gates_untrusted_base_repo(fake_runtime, tmp_path): + # A companion base_repo also loads via from_pretrained, so a trusted GGUF model_path must + # not smuggle in an arbitrary remote base: base_repo clears the same trust bar as a non-GGUF + # repo id (mirrors the video loader), and the check runs before any GPU handoff. + backend = DiffusionBackend() + with pytest.raises(ValueError, match = "base_repo"): + backend.validate_load_request( + "unsloth/Qwen-Image-2512-GGUF", + gguf_filename = "x.gguf", + model_kind = "gguf", + base_repo = "evil/companions", + ) + # A local base_repo (already on disk) still passes the gate. + (tmp_path / "model_index.json").write_text("{}") + fam = backend.validate_load_request( + "unsloth/Qwen-Image-2512-GGUF", + gguf_filename = "x.gguf", + model_kind = "gguf", + base_repo = str(tmp_path), + ) + assert fam is not None + + def test_detect_family_rejects_layered(): # Qwen-Image-Layered needs a dedicated pipeline (additional_t_cond); it must be # rejected so it fails fast at load instead of crashing at the first denoise step. diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 6a4b53961b..c9dca95e8e 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -36,6 +36,7 @@ class _FakeBackend: gguf_filename = None, family_override = None, model_kind = None, + base_repo = None, ): # Mirror the real backend's cheap validation so the route's # validate-before-evict ordering is exercised. @@ -50,6 +51,12 @@ class _FakeBackend: raise ValueError( f"Non-GGUF diffusion loads are restricted to unsloth/* repos; got '{model_path}'." ) + # A client-supplied base_repo clears the same trust bar (mirrors the real backend's + # gate), so the route's validate-before-evict rejects an untrusted companion base. + if base_repo and base_repo.strip() and not base_repo.lower().startswith("unsloth/"): + raise ValueError( + f"base_repo is restricted to unsloth/* repos (or a local path); got '{base_repo}'." + ) fam = detect_family(model_path, family_override) if fam is None: raise ValueError(f"Could not infer a diffusion family for '{model_path}'.") @@ -192,7 +199,7 @@ def test_load_generate_status_unload_roundtrip(client): json = { "model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "z-image-turbo-Q4_K_S.gguf", - "base_repo": "base/repo", + "base_repo": "unsloth/Z-Image-base", }, ) assert loaded.status_code == 200 @@ -221,6 +228,23 @@ def test_load_generate_status_unload_roundtrip(client): assert client.get("/api/inference/images/status").json()["loaded"] is False +def test_load_rejects_untrusted_base_repo(client): + # A trusted GGUF model_path paired with an untrusted remote base_repo is rejected at the + # route (validate runs before the GPU handoff), so an authenticated client cannot make the + # server fetch and deserialize an arbitrary companion repo, and no model is loaded/evicted. + r = client.post( + "/api/inference/images/load", + json = { + "model_path": "unsloth/Z-Image-Turbo-GGUF", + "gguf_filename": "z-image-turbo-Q4_K_S.gguf", + "base_repo": "evil/companions", + }, + ) + assert r.status_code == 400 + assert "base_repo" in r.json()["detail"] + assert client.get("/api/inference/images/status").json()["loaded"] is False + + def test_generate_batch_size_persists_each_image(client): client.post( "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index a35f6ee1a3..d5317c2b74 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -816,6 +816,21 @@ def test_free_gpu_for_diffusion_training_unloads_video(monkeypatch): assert gpu_arbiter.VIDEO in released +def test_keepwarm_tracks_image_video_generation_paths(): + # The API-key training-start guard uses other_inference_request_count(), which only sees + # paths the keepwarm middleware tracks. Image/video generation must be tracked so a training + # start is refused (409) while one is in-flight rather than its unload cancelling it; the GET + # *-progress and */cancel variants must stay untracked. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/images/generate") + assert _is_inference_path("/v1/images/generations") + assert _is_inference_path("/api/inference/video/generate") + assert not _is_inference_path("/api/inference/images/generate-progress") + assert not _is_inference_path("/api/inference/video/generate-progress") + assert not _is_inference_path("/api/inference/video/generate/cancel") + + def test_import_example_partial_failure_leaves_no_partial_dataset( client, dataset_roots, monkeypatch ): diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 21e8cf1058..14deb9a4c1 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -183,6 +183,23 @@ class TestMaxBodyMiddleware: assert cap == upload_request_limit_bytes() # DB-aware cap + multipart overhead assert cap > default_request_body_limit_bytes() # not the plain default body cap + def test_v1_surface_is_body_protected(self, main_module): + # /images/generations is mounted at both /api/inference and /v1; the /v1 alias (and every + # other /v1 POST route) must be body-capped via the /v1 blanket prefix, or an unbounded + # ImageGenerationRequest.prompt buffers outside the Studio request limit. Also confirms + # the blanket did not drop protection for the original /v1 chat/completions route. + for path in ( + "/v1/images/generations", + "/v1/audio/generate", + "/v1/embeddings", + "/v1/responses", + "/v1/messages", + "/v1/chat/completions", + ): + assert any( + path.startswith(p) for p in main_module._BODY_PROTECTED_PREFIXES + ), path + def test_upload_passthrough_rejects_declared_body_over_dedicated_cap(self, main_module): app = _make_protected_app( 128,