From 0a2a423b08ac6448142bb178bcbb0ef905747be6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 03:21:52 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion.py | 22 +++++--- studio/backend/core/inference/gpu_arbiter.py | 1 - .../backend/core/inference/image_gallery.py | 2 +- studio/backend/models/inference.py | 16 +++--- studio/backend/routes/inference.py | 52 +++++++++---------- studio/backend/routes/models.py | 26 +++++++--- .../backend/tests/test_diffusion_backend.py | 32 +++++++++--- studio/backend/tests/test_diffusion_routes.py | 47 ++++++++++++----- 8 files changed, 129 insertions(+), 69 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index d2fcbb2058..5ca18831e2 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -134,7 +134,9 @@ class DiffusionBackend: ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" if not gguf_filename: - raise ValueError("gguf_filename is required: this backend loads single-file GGUF checkpoints only.") + raise ValueError( + "gguf_filename is required: this backend loads single-file GGUF checkpoints only." + ) fam = detect_family(repo_id, family_override) if fam is None: raise ValueError( @@ -173,7 +175,9 @@ class DiffusionBackend: # calls) so begin_load returns instantly; the bar shows raw bytes until # the total lands. This is the only writer of _loading's fields here. fam = detect_family(kwargs["repo_id"], kwargs.get("family_override")) - base = _resolve_base_repo(kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token")) + base = _resolve_base_repo( + kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token") + ) kwargs["base_repo"] = base loading = self._loading if loading is not None: @@ -270,7 +274,9 @@ class DiffusionBackend: import diffusers if not gguf_filename: - raise ValueError("gguf_filename is required: this backend loads single-file GGUF checkpoints only.") + raise ValueError( + "gguf_filename is required: this backend loads single-file GGUF checkpoints only." + ) fam = detect_family(repo_id, family_override) if fam is None: raise ValueError( @@ -343,7 +349,6 @@ class DiffusionBackend: batch_size: int = 1, ) -> dict[str, Any]: import torch - with self._lock: state = self._state if state is None: @@ -406,7 +411,13 @@ class DiffusionBackend: """Live per-step progress for an in-flight generation (lock-free read).""" gen = self._gen if gen is None or gen.total_steps <= 0: - return {"active": False, "step": 0, "total_steps": 0, "fraction": 0.0, "eta_seconds": None} + return { + "active": False, + "step": 0, + "total_steps": 0, + "fraction": 0.0, + "eta_seconds": None, + } return { "active": True, "step": gen.step, @@ -474,7 +485,6 @@ def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]: return None try: from huggingface_hub import HfApi - meta = HfApi().model_info(repo_id, token = hf_token).cardData or {} except Exception: # noqa: BLE001 — best-effort; fall back to the family default return None diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index 1b52cf775c..70a2e2f65b 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -52,7 +52,6 @@ def _evict_chat() -> None: def _evict_diffusion() -> None: from core.inference.diffusion import get_diffusion_backend - get_diffusion_backend().unload() diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index 38c0ff9264..66016f807a 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -143,7 +143,7 @@ def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, except OSError: return [] paths.sort(key = _mtime, reverse = True) - window = paths[offset:] if limit is None else paths[offset:offset + limit] + window = paths[offset:] if limit is None else paths[offset : offset + limit] records = [] for path in window: meta = _read_meta(path) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index a492cc390c..da5bedc098 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1690,7 +1690,9 @@ class DiffusionLoadRequest(BaseModel): """Request to load a local diffusion (text-to-image) checkpoint.""" model_path: str = Field(..., description = "Diffusion GGUF repo id or local path") - gguf_filename: str = Field(..., description = "The chosen single-file GGUF quant inside model_path") + gguf_filename: str = Field( + ..., description = "The chosen single-file GGUF quant inside model_path" + ) base_repo: Optional[str] = Field( None, description = "Companion diffusers repo for VAE/text-encoders (default: family base)" ) @@ -1705,9 +1707,13 @@ class DiffusionGenerateRequest(BaseModel): """Request to generate one image from the loaded diffusion model.""" prompt: str = Field(..., min_length = 1, description = "Text prompt") - negative_prompt: Optional[str] = Field(None, description = "What to avoid (if the model supports it)") + negative_prompt: Optional[str] = Field( + None, description = "What to avoid (if the model supports it)" + ) width: int = Field(1024, ge = 256, le = 2048, description = "Image width in pixels (multiple of 16)") - height: int = Field(1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 16)") + height: int = Field( + 1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 16)" + ) steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps") guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") seed: Optional[int] = Field( @@ -1748,9 +1754,7 @@ class GalleryImage(BaseModel): class DiffusionGenerateResponse(BaseModel): """The persisted gallery records for one generation call (a batch).""" - images: list[GalleryImage] = Field( - ..., description = "Saved records, one per image in the batch" - ) + images: list[GalleryImage] = Field(..., description = "Saved records, one per image in the batch") class GalleryListResponse(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 83ef7aad69..91da38109e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2524,6 +2524,7 @@ async def load_model( # ownership is asserted without depending on chat/diffusion exclusivity # holding. No-op when diffusion isn't loaded. from core.inference.gpu_arbiter import acquire_for, CHAT + await asyncio.to_thread(acquire_for, CHAT) # ── Already-loaded check: skip reload if the exact model is active ── @@ -10054,8 +10055,7 @@ async def _openai_passthrough_non_streaming( @studio_router.post("/images/load", response_model = DiffusionStatusResponse) async def load_diffusion_model( - request: DiffusionLoadRequest, - current_subject: str = Depends(get_current_subject), + request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject) ): from core.inference.diffusion import get_diffusion_backend from core.inference.gpu_arbiter import acquire_for, DIFFUSION @@ -10085,8 +10085,7 @@ async def load_diffusion_model( @studio_router.post("/images/generate", response_model = DiffusionGenerateResponse) async def generate_diffusion_image( - request: DiffusionGenerateRequest, - current_subject: str = Depends(get_current_subject), + request: DiffusionGenerateRequest, current_subject: str = Depends(get_current_subject) ): from core.inference import image_gallery from core.inference.diffusion import get_diffusion_backend @@ -10119,20 +10118,25 @@ async def generate_diffusion_image( def _persist() -> list[dict]: records = [] for index, image in enumerate(result["images"]): - records.append(image_gallery.save(image, { - "prompt": request.prompt, - "negative_prompt": request.negative_prompt, - "width": request.width, - "height": request.height, - "steps": request.steps, - "guidance": request.guidance, - "seed": result["seed"], - # Position within the batch: images here share a seed + timestamp, - # so the export filename needs this to stay unique. - "batch_index": index, - "model": result.get("repo_id"), - "created_at": created_at, - })) + records.append( + image_gallery.save( + image, + { + "prompt": request.prompt, + "negative_prompt": request.negative_prompt, + "width": request.width, + "height": request.height, + "steps": request.steps, + "guidance": request.guidance, + "seed": result["seed"], + # Position within the batch: images here share a seed + timestamp, + # so the export filename needs this to stay unique. + "batch_index": index, + "model": result.get("repo_id"), + "created_at": created_at, + }, + ) + ) return records try: @@ -10165,8 +10169,7 @@ async def list_gallery_images( @studio_router.get("/images/gallery/{image_id}/file") async def get_gallery_image_file( - image_id: str, - current_subject: str = Depends(get_current_subject), + image_id: str, current_subject: str = Depends(get_current_subject) ): from core.inference import image_gallery @@ -10183,10 +10186,7 @@ async def get_gallery_image_file( @studio_router.delete("/images/gallery/{image_id}") -async def delete_gallery_image( - image_id: str, - current_subject: str = Depends(get_current_subject), -): +async def delete_gallery_image(image_id: str, current_subject: str = Depends(get_current_subject)): from core.inference import image_gallery deleted = await asyncio.to_thread(image_gallery.delete, image_id) @@ -10198,7 +10198,6 @@ async def delete_gallery_image( @studio_router.delete("/images/gallery") async def clear_gallery_images(current_subject: str = Depends(get_current_subject)): from core.inference import image_gallery - removed = await asyncio.to_thread(image_gallery.clear) return {"removed": removed} @@ -10216,19 +10215,16 @@ async def unload_diffusion_model(current_subject: str = Depends(get_current_subj @studio_router.get("/images/status", response_model = DiffusionStatusResponse) async def diffusion_status(current_subject: str = Depends(get_current_subject)): from core.inference.diffusion import get_diffusion_backend - return DiffusionStatusResponse(**get_diffusion_backend().status()) @studio_router.get("/images/load-progress", response_model = DiffusionLoadProgressResponse) async def diffusion_load_progress(current_subject: str = Depends(get_current_subject)): from core.inference.diffusion import get_diffusion_backend - return DiffusionLoadProgressResponse(**get_diffusion_backend().load_progress()) @studio_router.get("/images/generate-progress", response_model = DiffusionGenerateProgressResponse) async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)): from core.inference.diffusion import get_diffusion_backend - return DiffusionGenerateProgressResponse(**get_diffusion_backend().generate_progress()) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 2e87c947e3..98073231ff 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -844,8 +844,7 @@ async def list_local_models( models = [m for m in models if not _is_hidden_model(m.id, m.path)] # Tag each GGUF with its task so the Images picker can filter to diffusion. models = [ - m.model_copy(update = {"task": _local_model_task(m.path, m.model_format)}) - for m in models + m.model_copy(update = {"task": _local_model_task(m.path, m.model_format)}) for m in models ] return LocalModelListResponse( @@ -3063,11 +3062,24 @@ def _repo_gguf_last_modified(repo_info) -> float: # GGUF general.architecture values that denote a diffusion (image) model; # everything else is treated as a text model. Lets the Images picker show only # image GGUFs in its On Device list. -_DIFFUSION_GGUF_ARCHS = frozenset({ - "flux", "flux2", "sd1", "sd2", "sd3", "sdxl", "stable_diffusion", - "lumina2", "qwen_image", "qwenimage", "auraflow", "pixart", - "hunyuan_video", "wan", -}) +_DIFFUSION_GGUF_ARCHS = frozenset( + { + "flux", + "flux2", + "sd1", + "sd2", + "sd3", + "sdxl", + "stable_diffusion", + "lumina2", + "qwen_image", + "qwenimage", + "auraflow", + "pixart", + "hunyuan_video", + "wan", + } +) def _gguf_architecture(path: str) -> Optional[str]: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index df64528460..621ea6a899 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -280,12 +280,21 @@ def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch) fam = detect_family("unsloth/Qwen-Image-2512-GGUF") monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: "Qwen/Qwen-Image-2512") # Caller's explicit base wins and the HF tag is not consulted. - assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", "my/base", fam, None) == "my/base" + assert ( + diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", "my/base", fam, None) + == "my/base" + ) # No caller base: the repo's base_model tag (the variant base) is used. - assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", None, fam, None) == "Qwen/Qwen-Image-2512" + assert ( + diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", None, fam, None) + == "Qwen/Qwen-Image-2512" + ) # No caller base and no tag: the family fallback. monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: None) - assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", " ", fam, None) == fam.base_repo + assert ( + diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", " ", fam, None) + == fam.base_repo + ) def test_load_without_gguf_raises(): @@ -340,7 +349,9 @@ def test_base_file_downloaded_excludes_undownloaded(): assert _base_file_downloaded("vae/diffusion_pytorch_model.safetensors") # Excluded: the GGUF supplies the transformer; docs/assets and top-level files # are never downloaded, so counting them would peg the bar short of 100%. - assert not _base_file_downloaded("transformer/diffusion_pytorch_model-00001-of-00003.safetensors") + assert not _base_file_downloaded( + "transformer/diffusion_pytorch_model-00001-of-00003.safetensors" + ) assert not _base_file_downloaded("assets/Z-Image-Gallery.pdf") assert not _base_file_downloaded("README.md") assert not _base_file_downloaded(".gitattributes") @@ -374,7 +385,10 @@ def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path): (tmp_path / "model.gguf").write_bytes(b"weights") backend = DiffusionBackend() backend.load_pipeline( - str(tmp_path), gguf_filename = "model.gguf", base_repo = "Qwen/Qwen-Image", family_override = "qwen-image" + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "Qwen/Qwen-Image", + family_override = "qwen-image", ) backend.generate(prompt = "a sloth", guidance = 4.0) # Qwen-Image's distilled guidance is off; the real CFG must land on true_cfg_scale. @@ -386,9 +400,13 @@ def test_begin_load_rejects_concurrent(monkeypatch): backend = DiffusionBackend() # The worker resolves the base via a network lookup; stub it so the test is offline. monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None) - monkeypatch.setattr(DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0)) + monkeypatch.setattr( + DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0) + ) # Block the spawned worker so the load stays "in progress". - monkeypatch.setattr(DiffusionBackend, "load_pipeline", lambda self, **k: __import__("time").sleep(0.2)) + monkeypatch.setattr( + DiffusionBackend, "load_pipeline", lambda self, **k: __import__("time").sleep(0.2) + ) backend.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z-image-turbo-Q4_K_S.gguf") with pytest.raises(RuntimeError): backend.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z-image-turbo-Q4_K_S.gguf") diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 1309746dbb..e60fb3237a 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -51,7 +51,13 @@ class _FakeBackend: "error": None, } - def generate(self, *, seed = None, batch_size = 1, **kwargs): + def generate( + self, + *, + seed = None, + batch_size = 1, + **kwargs, + ): if not self.loaded: raise RuntimeError("No diffusion model is loaded.") # The real backend returns the PIL images; the route persists them. The @@ -110,13 +116,15 @@ def client(monkeypatch, tmp_path): monkeypatch.setattr(gallery_module, "save", _save) monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None) + def _list_images(limit = None, offset = 0): ordered = sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True) - return ordered[offset:] if limit is None else ordered[offset:offset + limit] + return ordered[offset:] if limit is None else ordered[offset : offset + limit] monkeypatch.setattr(gallery_module, "list_images", _list_images) monkeypatch.setattr( - gallery_module, "image_path", + gallery_module, + "image_path", lambda i: (tmp_path / f"{i}.png") if i in store else None, ) monkeypatch.setattr(gallery_module, "delete", lambda i: store.pop(i, None) is not None) @@ -129,11 +137,14 @@ def client(monkeypatch, tmp_path): def test_load_generate_status_unload_roundtrip(client): - loaded = 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": "base/repo", - }) + loaded = 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": "base/repo", + }, + ) assert loaded.status_code == 200 body = loaded.json() assert body["loaded"] is True and body["family"] == "z-image" @@ -161,7 +172,9 @@ def test_load_generate_status_unload_roundtrip(client): 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"}) + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) resp = client.post( "/api/inference/images/generate", json = {"prompt": "p", "batch_size": 3, "seed": 5}, @@ -175,7 +188,9 @@ def test_generate_batch_size_persists_each_image(client): def test_gallery_pagination(client): - client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}) + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) client.post("/api/inference/images/generate", json = {"prompt": "p", "batch_size": 5, "seed": 1}) page1 = client.get("/api/inference/images/gallery?limit=2&offset=0").json() assert len(page1["images"]) == 2 and page1["has_more"] is True @@ -184,7 +199,9 @@ def test_gallery_pagination(client): def test_generate_rejects_non_multiple_of_16(client): - client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}) + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) # Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since # Z-Image requires dimensions divisible by 16. for bad in (1001, 1000): @@ -213,7 +230,9 @@ def test_load_unknown_family_returns_400(client, monkeypatch): backend = _FakeBackend() backend.begin_load = _raise monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) - resp = client.post("/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"}) + resp = client.post( + "/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"} + ) assert resp.status_code == 400 assert "family" in resp.json()["detail"] @@ -223,7 +242,9 @@ def test_load_progress_route(client): idle = client.get("/api/inference/images/load-progress") assert idle.status_code == 200 and idle.json()["phase"] is None # After load: the fake reports ready. - client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}) + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) ready = client.get("/api/inference/images/load-progress") assert ready.json()["phase"] == "ready"