diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 007d6d349d..603e4ca637 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1537,6 +1537,17 @@ class DiffusionBackend: if cn_model is None: if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) + # resolve_controlnet accepts a bare owner/name repo without the non-GGUF base + # trust gate, and from_pretrained below downloads and deserializes it. A + # malicious pickle .bin would execute on load, so run the same Hub malware + # preflight the chat/export loaders use before any remote ControlNet load. A + # local dir the user picked has no Hub scan and is exempt (fail-open there). + if not getattr(resolved_cn, "is_local", False): + from utils.security import evaluate_file_security + + _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None) + if _cn_fs.blocked: + raise ValueError(_cn_fs.reason) import torch # state.dtype is the display string saved at load ("bfloat16"), NOT a diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6666af1016..c579e1ca24 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1178,11 +1178,29 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None: @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse) async def start_diffusion_training( - body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject) + body: DiffusionTrainingStartRequest, + current_subject: str = Depends(get_current_subject), + via_api_key: bool = Depends(authenticated_via_api_key), ): """Start an SDXL LoRA training job from an image + caption dataset.""" from core.training.diffusion_training_service import get_diffusion_training_service + # When Studio is driven as an inference API (API-key auth), refuse to start training + # while a request is in flight: _free_gpu_for_diffusion_training() below unloads the + # chat backends to reclaim VRAM, which would kill the stream. Mirrors start_training so + # a diffusion start cannot silently drop an active API inference request. + if via_api_key is True: + from core.inference.llama_keepwarm import other_inference_request_count + if other_inference_request_count(current_request_counted = False) > 0: + raise HTTPException( + status_code = 409, + detail = ( + "Cannot start diffusion (Images) training over the API while an inference " + "request is in progress. Wait for it to finish, or start training from the " + "Studio UI." + ), + ) + # Interlock: refuse while an LLM training run holds the GPU (symmetric with the # diffusion check in start_training), so the two trainers never contend for VRAM. try: @@ -1613,7 +1631,11 @@ async def get_diffusion_dataset_image( thumbs_dir = folder / _THUMBS_DIRNAME thumbs_dir.mkdir(exist_ok = True) - thumb_path = thumbs_dir / f"{image_path.stem}_{size}.jpg" + # Key on the full filename (stem + extension), not the stem: two images that + # share a stem but differ by extension (sample.png / sample.jpg) would otherwise + # collide on one cache file, and an mtime-newer cache built for the first would + # be served for the second, showing the wrong image in the labeling grid. + thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg" src_mtime = image_path.stat().st_mtime if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime: return thumb_path @@ -1698,7 +1720,10 @@ async def delete_diffusion_dataset_image( image_path.with_suffix(ext).unlink(missing_ok = True) thumbs_dir = folder / _THUMBS_DIRNAME if thumbs_dir.is_dir(): - for t in thumbs_dir.glob(f"{image_path.stem}_*.jpg"): + # Thumbs are keyed on the full filename (stem + extension), so match that + # here too; a stem-only glob would leave this image's thumbs behind and + # could delete a same-stem sibling's (sample.png vs sample.jpg). + for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"): t.unlink(missing_ok = True) return {"deleted": image_path.name} diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index 50bf7159b0..b7ae037197 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -231,12 +231,24 @@ def _state(): ) +def _allow_cn_security(monkeypatch): + """Stub the Hub malware preflight to allow the load (hermetic, no network).""" + import utils.security + + monkeypatch.setattr( + utils.security, + "evaluate_file_security", + lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = False, reason = ""), + ) + + def test_controlnet_pipe_loads_once_and_caches(monkeypatch): import threading from core.inference.diffusion import DiffusionBackend monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers()) + _allow_cn_security(monkeypatch) b = DiffusionBackend() st = _state() # The pipe cache only commits while ``st`` is the CURRENT load (an unload racing @@ -252,6 +264,64 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch): assert b._cn_models["flux-union-pro"] is p1.controlnet +def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch): + # A bare owner/name ControlNet is accepted by resolve_controlnet without the base + # trust gate, so the load path must run the Hub malware preflight: a flagged remote + # repo must raise BEFORE from_pretrained downloads/deserializes it. + import threading + + import utils.security + from core.inference.diffusion import DiffusionBackend + + loaded = {"called": False} + + class _TrapModel(_FakeCNModel): + @classmethod + def from_pretrained(cls, path, torch_dtype = None, token = None): + loaded["called"] = True + return super().from_pretrained(path, torch_dtype = torch_dtype, token = token) + + mod = _fake_diffusers() + mod.FluxControlNetModel = _TrapModel + monkeypatch.setitem(sys.modules, "diffusers", mod) + monkeypatch.setattr( + utils.security, + "evaluate_file_security", + lambda name, hf_token = None, **kw: types.SimpleNamespace( + blocked = True, reason = "Hugging Face security scan flagged unsafe files: evil.bin" + ), + ) + b = DiffusionBackend() + st = _state() + b._state = st + resolved = dc.ResolvedControlNet("evil/cn", "evil/cn", is_local = False) + with pytest.raises(ValueError, match = "security scan flagged"): + b._controlnet_pipe(st, resolved, threading.Event()) + assert loaded["called"] is False + + +def test_controlnet_pipe_skips_scan_for_local_dir(monkeypatch, tmp_path): + # A local dir the user picked has no Hub scan; the preflight must not block it even + # if the (unused) scan stub would say blocked. + import threading + + import utils.security + from core.inference.diffusion import DiffusionBackend + + monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers()) + monkeypatch.setattr( + utils.security, + "evaluate_file_security", + lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = True, reason = "x"), + ) + b = DiffusionBackend() + st = _state() + b._state = st + resolved = dc.ResolvedControlNet("my-cn", str(tmp_path), is_local = True) + p = b._controlnet_pipe(st, resolved, threading.Event()) + assert isinstance(p, _FakeCNPipe) + + def test_controlnet_pipe_rejects_family_without_classes(): import threading diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index b10dd04f4c..451f594ba1 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -177,15 +177,29 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root): folder.mkdir() _write_png(folder / "x.png") (folder / "x.txt").write_text("cap", encoding = "utf-8") - # Generate a thumbnail so we can assert it is cleaned up too. + # Generate a thumbnail so we can assert it is cleaned up too. Thumbs are keyed on + # the full filename (stem + extension) to avoid same-stem collisions across formats. client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32") - assert list((folder / ".thumbs").glob("x_*.jpg")) + assert list((folder / ".thumbs").glob("x.png_*.jpg")) r = client.delete("/api/train/diffusion/dataset/d/image/x.png") assert r.status_code == 200, r.text assert not (folder / "x.png").exists() assert not (folder / "x.txt").exists() - assert not list((folder / ".thumbs").glob("x_*.jpg")) + assert not list((folder / ".thumbs").glob("x.png_*.jpg")) + + +def test_thumb_cache_key_distinguishes_same_stem_extensions(client, ds_root): + # sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache + # file, so the labeling grid never serves one image's thumbnail for the other. + folder = ds_root / "d" + folder.mkdir() + Image.new("RGB", (8, 8), (10, 20, 30)).save(folder / "sample.png", format = "PNG") + Image.new("RGB", (8, 8), (200, 210, 220)).save(folder / "sample.jpg", format = "JPEG") + client.get("/api/train/diffusion/dataset/d/image/sample.png?thumb=32") + client.get("/api/train/diffusion/dataset/d/image/sample.jpg?thumb=32") + thumbs = sorted(p.name for p in (folder / ".thumbs").glob("*.jpg")) + assert thumbs == ["sample.jpg_32.jpg", "sample.png_32.jpg"] # ── traversal / validation ─────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index a9f67a942e..a03803f1e6 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -20,7 +20,7 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from auth.authentication import get_current_subject +from auth.authentication import authenticated_via_api_key, get_current_subject from core.training.diffusion_training_service import DiffusionTrainingService from routes.training import router as training_router @@ -316,8 +316,12 @@ def client(monkeypatch): app = FastAPI() app.include_router(training_router, prefix = "/api/train") app.dependency_overrides[get_current_subject] = lambda: "test-user" + # Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there, + # so the wiring tests below behave as before. The guard test flips this override. + app.dependency_overrides[authenticated_via_api_key] = lambda: False c = TestClient(app) c._fake = fake # type: ignore[attr-defined] + c._app = app # type: ignore[attr-defined] return c @@ -406,6 +410,37 @@ def test_route_start_conflict_maps_to_409(client): assert r.status_code == 409 +def test_route_start_over_api_with_inference_in_flight_is_409(client, monkeypatch): + # An API-key client must not start diffusion training (which frees VRAM by unloading + # chat) while an inference request is streaming; it should 409 instead of killing it. + client._app.dependency_overrides[authenticated_via_api_key] = lambda: True + monkeypatch.setattr( + "core.inference.llama_keepwarm.other_inference_request_count", + lambda current_request_counted = False: 1, + ) + freed = {"called": False} + import routes.training as tr + + monkeypatch.setattr( + tr, "_free_gpu_for_diffusion_training", lambda: freed.__setitem__("called", True) + ) + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 409 + # The guard must run BEFORE any GPU is freed, so the live inference stream survives. + assert freed["called"] is False + + +def test_route_start_over_api_without_inference_proceeds(client, monkeypatch): + # Same API-key path but no inference in flight: the start proceeds normally. + client._app.dependency_overrides[authenticated_via_api_key] = lambda: True + monkeypatch.setattr( + "core.inference.llama_keepwarm.other_inference_request_count", + lambda current_request_counted = False: 0, + ) + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 200 + + def test_route_status_and_stop(client): client.post("/api/train/diffusion/start", json = _BODY) s = client.get("/api/train/diffusion/status") diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index c380e04244..f096af73c8 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -2282,6 +2282,14 @@ export function HubModelPicker({ } if (section === "recommended") { + // Curated safetensors rows render ABOVE the recommended rows (and call + // getOptionProps), so their keys must lead here or they fall back to the + // duplicate ...-option-missing id and drop out of arrow-key navigation. + keys.push( + ...curatedSafetensorsRows.map((m) => + makeModelOptionKey("curated-safetensors", m.id), + ), + ); keys.push( ...recommendedRows.map((r) => makeModelOptionKey("recommended", r.id)), ); @@ -2291,6 +2299,7 @@ export function HubModelPicker({ }, [ cachedReady, chatOnly, + curatedSafetensorsRows, sortedCustomFolderModels, customFoldersCollapsed, downloadedCollapsed,