Merge branch 'diffusion-train-precision' into diffusion-train-tab-2
This commit is contained in:
Daniel Han
commit
5979d68a03
6 changed files with 171 additions and 7 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1182,11 +1182,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:
|
||||
|
|
@ -1660,7 +1678,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
|
||||
|
|
@ -1745,7 +1767,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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ───────────────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -331,8 +331,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
|
||||
|
||||
|
||||
|
|
@ -521,6 +525,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")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue