Studio: gate untrusted companion base_repo on image loads, track image/video generation for the training guard, and body-cap the whole /v1 surface

- Image load now trust-gates a client-supplied base_repo. validate_load_request
  rejects a base_repo that is not an unsloth/* repo, an allowlisted official base, or a
  local path, mirroring the repo_id gate and the video loader. The route passes
  base_repo into that pre-eviction validation, so an authenticated client can no longer
  keep model_path on a trusted GGUF while pointing base_repo at an arbitrary remote repo
  that the server would download and deserialize (a from_pretrained pickle/config path),
  and no resident model is evicted for the rejected load.
- The keepwarm middleware now tracks the image and video generation routes
  (/images/generate, /images/generations, /video/generate), so
  other_inference_request_count() sees an in-flight generation and an API-key training
  start is refused (409) before its unload would cancel that generation. endswith keeps
  the GET *-progress and */cancel variants untracked.
- The OpenAI-compatible /v1 surface is now blanket body-capped like /api/inference,
  instead of only /v1/chat/completions and /v1/completions. Every /v1 POST route
  (images/generations, audio, embeddings, responses, messages, ...) buffers a JSON body
  and none is a multipart-upload passthrough, so an unbounded ImageGenerationRequest
  prompt on /v1/images/generations can no longer be buffered outside the request limit.

Adds regression tests: the base_repo trust gate at both the backend (untrusted remote
raises, local passes) and the route (untrusted base_repo returns 400 with no load), the
keepwarm tracking of the image/video generation paths (and not the progress/cancel
variants), and the /v1 surface being body-protected.

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,

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
 )
 
 

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",

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

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.

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"}

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
 ):

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,