From 8c5a00e0ca40d2b1bfcdb3d6362deaca00d7564f Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 3 Jul 2026 07:07:43 -0300 Subject: [PATCH 01/17] Studio: add OpenAI-compatible /v1/images/generations endpoint (#6686) * Studio: add OpenAI-compatible /v1/images/generations endpoint * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fall back to the loaded base repo for image-generation defaults * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restrict the image-generation 503 to genuine unload races * Route /v1/images/generations through the active diffusion engine --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- .../core/inference/diffusion_families.py | 32 ++ studio/backend/models/inference.py | 65 +++ studio/backend/routes/inference.py | 167 ++++++++ .../test_openai_images_generations_route.py | 386 ++++++++++++++++++ 4 files changed, 650 insertions(+) create mode 100644 studio/backend/tests/test_openai_images_generations_route.py diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 5beb8e36b6..7f72743d35 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -209,6 +209,38 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: return base or fam.base_repo +# Default (steps, guidance) per model for callers that can't pass them — namely +# the OpenAI /v1/images/generations endpoint, whose spec has no step/guidance +# knobs. Distilled "turbo/schnell" models want few steps and no CFG; the full +# "dev" models want more steps and real CFG. Matched by substring, most specific +# first — the same scheme and values as the UI's MODEL_DEFAULTS table +# (studio/frontend/src/features/images/images-page.tsx); keep the two in sync. +_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( + ("z-image-turbo", 9, 0.0), + ("flux.1-schnell", 4, 0.0), + ("flux.1", 28, 3.5), + ("flux.2-klein", 4, 0.0), + ("qwen-image", 20, 4.0), + ("z-image", 20, 4.0), +) +# Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback. +_GENERATION_DEFAULT_FALLBACK = (9, 0.0) + + +def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]: + """Default ``(steps, guidance)`` for a loaded model. The first identifier that + names a known model wins (the repo id, then the resolved base repo), so a + local-path load — whose repo id is just a filesystem path that may not name + the model — still resolves via its base repo. Within an identifier, keys are + matched as substrings, most specific first (the same scheme as the UI).""" + for identifier in identifiers: + needle = (identifier or "").lower() + for key, steps, guidance in _GENERATION_DEFAULTS: + if key in needle: + return steps, guidance + return _GENERATION_DEFAULT_FALLBACK + + def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]: """The hosted pre-quantized transformer repo for ``scheme`` in this family, or None.""" for entry_scheme, repo_id in fam.prequant_repos: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 11d3e5025c..04b70e7e11 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1928,3 +1928,68 @@ class DiffusionStatusResponse(BaseModel): None, description = "Why diffusers was chosen over the native sd.cpp engine (null when none)", ) + + +# ── OpenAI-compatible images API (POST /v1/images/generations) ── +# +# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf +# OpenAI clients work unchanged. The loaded image GGUF stands in for the model; +# GPT-image-only knobs (quality, style, background, output_format, ...) are +# accepted and ignored, exactly as dall-e-2 ignores them. The size string is +# parsed and `stream` is rejected in the route, where the diffusion backend is +# in reach; everything Pydantic can check declaratively lives here. + + +class ImageGenerationRequest(BaseModel): + """OpenAI ``CreateImageRequest`` for ``POST /v1/images/generations``. + + ``prompt`` is the only required field, per the spec. Unlisted OpenAI fields + are ignored (Pydantic's default), matching dall-e-2's treatment of the + GPT-image-only parameters.""" + + prompt: str = Field(..., min_length = 1, description = "Text description of the image(s).") + model: Optional[str] = Field( + None, description = "Model id (informational; the loaded image model is used)." + ) + n: int = Field(1, ge = 1, le = 10, description = "Number of images to generate (1-10).") + size: str = Field( + "auto", description = "'auto' or 'x' (256-2048, each a multiple of 16)." + ) + response_format: Literal["url", "b64_json"] = Field( + "url", description = "Return each image as a URL or a base64-encoded PNG." + ) + user: Optional[str] = Field(None, description = "End-user identifier (accepted, unused).") + # gpt-image-only; declared so we can reject it with a clear error instead of + # silently returning JSON to a client that asked for an SSE stream. + stream: Optional[bool] = Field( + None, description = "Streaming image generation is not supported; omit or set false." + ) + + @field_validator("n", "size", "response_format", mode = "before") + @classmethod + def _null_means_default(cls, value, info): + # OpenAI marks these nullable WITH a default, so an explicit null means + # "use the default" — coalesce it instead of 400-ing a spec-valid body. + if value is None: + return cls.model_fields[info.field_name].default + return value + + +class ImageGenerationData(BaseModel): + """One image in an ``ImagesResponse`` (OpenAI ``Image``). Exactly one of + ``url`` / ``b64_json`` is set, per the request's ``response_format``; the + route serializes with ``exclude_none`` so the unused key is omitted.""" + + b64_json: Optional[str] = Field( + None, description = "Base64-encoded PNG (response_format=b64_json)." + ) + url: Optional[str] = Field(None, description = "URL to the PNG bytes (response_format=url).") + + +class ImageGenerationResponse(BaseModel): + """OpenAI ``ImagesResponse``. dall-e-shaped: the GPT-image-only top-level + fields (background/output_format/size/quality/usage) are omitted, since our + sizes wouldn't satisfy their fixed enums and we report no token usage.""" + + created: int = Field(..., description = "Unix timestamp (seconds) the images were created.") + data: list[ImageGenerationData] = Field(..., description = "The generated images.") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 67e541a8d8..7400a9f657 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1062,6 +1062,9 @@ from models.inference import ( DiffusionLoadProgressResponse, GalleryImage, GalleryListResponse, + ImageGenerationRequest, + ImageGenerationData, + ImageGenerationResponse, LoadResponse, LoadProgressResponse, UnloadResponse, @@ -11301,3 +11304,167 @@ async def diffusion_load_progress(current_subject: str = Depends(get_current_sub async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)): from core.inference.diffusion_engine_router import get_active_diffusion_engine return DiffusionGenerateProgressResponse(**get_active_diffusion_engine().generate_progress()) + + +# ────────────────────────────────────────────────────────────────────────── +# OpenAI-compatible images API (POST /v1/images/generations) +# +# The inference router is mounted at both /api/inference and /v1, so this also +# answers /v1/images/generations for off-the-shelf OpenAI clients. It maps +# OpenAI's CreateImageRequest onto the in-process diffusion backend and returns +# an ImagesResponse. Studio's own Image tab uses the richer /images/generate +# route above; this is the spec-shaped surface, and the single error boundary +# mapping backend exceptions to OpenAI error envelopes (the global /v1 handler +# wraps HTTPException detail into the envelope). +# ────────────────────────────────────────────────────────────────────────── + + +# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample +# x 2x patch); the named OpenAI sizes (1024x1024, 1536x1024, 256x256, ...) all +# satisfy this. Mirrors DiffusionGenerateRequest's width/height bounds so both +# generate paths accept the same geometry. +_IMAGE_SIZE_RE = _re.compile(r"^(\d{1,5})\s*x\s*(\d{1,5})$") +_IMAGE_DIM_MIN, _IMAGE_DIM_MAX = 256, 2048 +# Sanitized 503 detail shared by the pre-check and the unload-race branch, so both +# "no image model" responses stay identical. +_NO_IMAGE_MODEL_MSG = "No image model loaded. Load an image model first." + + +def _parse_openai_image_size(size: str) -> tuple[int, int]: + """OpenAI ``size`` -> (width, height). ``auto``/empty -> 1024x1024 (~1MP, what + these models target). Raises ValueError with a client-facing message.""" + text = (size or "").strip().lower() + if text in ("", "auto"): + return 1024, 1024 + match = _IMAGE_SIZE_RE.match(text) + if not match: + raise ValueError("size must be 'auto' or 'x', e.g. '1024x1024'.") + width, height = int(match.group(1)), int(match.group(2)) + for label, value in (("width", width), ("height", height)): + if not _IMAGE_DIM_MIN <= value <= _IMAGE_DIM_MAX: + raise ValueError(f"size {label} must be between {_IMAGE_DIM_MIN} and {_IMAGE_DIM_MAX}.") + if value % 16 != 0: + raise ValueError(f"size {label} must be a multiple of 16.") + return width, height + + +def _absolute_image_url(request: Request, relative: str) -> str: + """Join a relative gallery path onto the request's own scheme+host, for the + response_format=url links. Like every Studio route, the target needs the + bearer token; b64_json avoids that for clients that can't carry it.""" + return str(request.base_url).rstrip("/") + relative + + +@router.post( + "/images/generations", + response_model = ImageGenerationResponse, + response_model_exclude_none = True, +) +async def openai_image_generations( + body: ImageGenerationRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """OpenAI-compatible text-to-image (POST /v1/images/generations). + + Generates ``n`` images from ``prompt`` on the loaded diffusion model and + returns them as URLs (default) or base64 PNGs per ``response_format``. Steps + and guidance have no OpenAI knob, so they default per loaded model.""" + from core.inference import image_gallery + from core.inference.diffusion_engine_router import get_active_diffusion_engine + from core.inference.diffusion_families import default_generation_params + + if body.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Streaming image generation is not supported.", status = 400, param = "stream" + ), + ) + try: + width, height = _parse_openai_image_size(body.size) + except ValueError as exc: + raise HTTPException( + status_code = 400, detail = openai_error_body(str(exc), status = 400, param = "size") + ) + + # Use the active engine (diffusers OR native sd.cpp on a no-GPU host), the same + # accessor /images/generate uses, so a model loaded on the native engine isn't + # wrongly reported unloaded here. + backend = get_active_diffusion_engine() + status = backend.status() + if not status.get("loaded"): + # Mirror /v1/completions and /v1/embeddings, which 503 when their backend + # isn't loaded; the global handler turns this into the OpenAI envelope. + raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG) + + # Fall back to the resolved base repo so a local-path load (whose repo_id is a + # filesystem path) still gets the right per-model steps/guidance. + steps, guidance = default_generation_params(status.get("repo_id"), status.get("base_repo")) + try: + result = await asyncio.to_thread( + backend.generate, + prompt = body.prompt, + width = width, + height = height, + steps = steps, + guidance = guidance, + batch_size = body.n, + ) + except Exception as exc: # noqa: BLE001 (single boundary, sanitized envelope) + # A RuntimeError with the model now unloaded means it was evicted/unloaded + # between the readiness check above and the call (a transient race): 503. + # Every other failure (CUDA OOM, a diffusers shape/device error, both also + # RuntimeError) is a real 500, and its raw message must not reach the client. + if isinstance(exc, RuntimeError) and not backend.is_loaded: + raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG) + logger.error("openai_images.generate_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Image generation failed.") + + created = int(time.time()) + want_b64 = body.response_format == "b64_json" + # Persist each image with its full recipe embedded, like /images/generate, so + # the response_format=url links resolve and the images show up in the gallery. + recipe = { + "prompt": body.prompt, + "negative_prompt": None, + "width": width, + "height": height, + "steps": steps, + "guidance": guidance, + # The batch shares one base seed, so restoring a batch_index>0 sibling needs the + # original batch_size to replay it (same as /images/generate); persist it. + "batch_size": body.n, + "model": result.get("repo_id"), + "created_at": float(created), + } + # The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed + # per image (returned in ``seeds``), so record each image's own seed, like + # /images/generate, or a native batch_index>0 image shows the wrong seed. + per_image_seeds = result.get("seeds") + + def _persist() -> list[ImageGenerationData]: + items: list[ImageGenerationData] = [] + for index, image in enumerate(result["images"]): + seed = ( + per_image_seeds[index] + if per_image_seeds and index < len(per_image_seeds) + else result["seed"] + ) + record = image_gallery.save(image, {**recipe, "batch_index": index, "seed": seed}) + if want_b64: + encoded = image_gallery.image_b64(record["id"]) + if encoded is None: # vanished between write and read — fail the call + raise RuntimeError("generated image could not be read back for encoding") + items.append(ImageGenerationData(b64_json = encoded)) + else: + items.append(ImageGenerationData(url = _absolute_image_url(request, record["url"]))) + return items + + try: + data = await asyncio.to_thread(_persist) + except Exception as exc: # noqa: BLE001 + logger.error("openai_images.persist_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Failed to save the generated image.") + + return ImageGenerationResponse(created = created, data = data) diff --git a/studio/backend/tests/test_openai_images_generations_route.py b/studio/backend/tests/test_openai_images_generations_route.py new file mode 100644 index 0000000000..f82586a173 --- /dev/null +++ b/studio/backend/tests/test_openai_images_generations_route.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""FastAPI round-trip tests for the OpenAI-compatible POST /v1/images/generations. + +The diffusion backend and image gallery are replaced with light fakes, so these +exercise the route wiring, OpenAI param mapping, validation, error envelopes, and +response shape without torch, diffusers, weights, or a GPU. The pure helpers +(`_parse_openai_image_size`, `default_generation_params`) are unit-tested directly. +""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import core.inference.diffusion as diffusion_module +import core.inference.diffusion_engine_router as engine_router +import core.inference.image_gallery as gallery_module +from auth.authentication import get_current_subject +from core.inference.diffusion_families import default_generation_params +from routes.inference import router, _parse_openai_image_size +from utils.api_errors import install_api_error_handlers + + +# ── pure helpers ──────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "repo_id, expected", + [ + ("unsloth/Z-Image-Turbo-GGUF", (9, 0.0)), # turbo entry, before the z-image fallback + ("unsloth/Z-Image-GGUF", (20, 4.0)), + ("unsloth/FLUX.1-schnell-GGUF", (4, 0.0)), # schnell entry, before the flux.1 entry + ("black-forest-labs/FLUX.1-dev", (28, 3.5)), + ("unsloth/FLUX.2-klein-4B-GGUF", (4, 0.0)), + ("unsloth/Qwen-Image-2512-GGUF", (20, 4.0)), + ("some/unknown-model", (9, 0.0)), # fallback + ("", (9, 0.0)), + ], +) +def test_default_generation_params(repo_id, expected): + assert default_generation_params(repo_id) == expected + + +def test_default_generation_params_specificity_ordering(): + # The "-turbo" / "-schnell" entries must win over their broader siblings; a + # reorder that broke this would silently mis-default. + assert default_generation_params("x/Z-Image-Turbo") != default_generation_params("x/Z-Image") + assert default_generation_params("x/FLUX.1-schnell") != default_generation_params( + "x/FLUX.1-dev" + ) + + +def test_default_generation_params_falls_back_to_base_repo(): + # A local-path load: repo_id is a filesystem path that names no model, so the + # resolved base repo is what identifies it (and distinguishes dev from schnell). + assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-dev") == (28, 3.5) + assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-schnell") == ( + 4, + 0.0, + ) + assert default_generation_params("/models/my-ckpt", "Qwen/Qwen-Image") == (20, 4.0) + # repo_id wins when it already names the model; base repo is only a fallback. + assert default_generation_params("unsloth/Z-Image-Turbo-GGUF", "Tongyi-MAI/Z-Image") == (9, 0.0) + # Nothing identifiable -> fallback; None identifiers are skipped. + assert default_generation_params(None, None) == (9, 0.0) + assert default_generation_params("/models/x", None) == (9, 0.0) + + +@pytest.mark.parametrize( + "size, expected", + [ + ("auto", (1024, 1024)), + ("", (1024, 1024)), + ("AUTO", (1024, 1024)), + ("512x512", (512, 512)), + ("512x256", (512, 256)), + ("1792x1024", (1792, 1024)), # dall-e-3 named size: must pass the bounds + ("1024x1792", (1024, 1792)), + (" 256 x 256 ", (256, 256)), + ], +) +def test_parse_image_size_ok(size, expected): + assert _parse_openai_image_size(size) == expected + + +@pytest.mark.parametrize("size", ["abc", "100x100", "4096x4096", "300x300", "512", "x512", "0x0"]) +def test_parse_image_size_rejects(size): + with pytest.raises(ValueError): + _parse_openai_image_size(size) + + +# ── route round-trip ──────────────────────────────────────────────────── + + +class _FakeBackend: + def __init__( + self, + loaded = True, + repo_id = "unsloth/Z-Image-Turbo-GGUF", + base_repo = None, + generate_error = None, + unload_on_generate = False, + native_seeds = False, + ) -> None: + self._loaded = loaded + self._repo_id = repo_id + self._base_repo = base_repo + # Model the native sd.cpp engine, which returns a distinct seed per image. + self._native_seeds = native_seeds + # When set, generate() raises this; unload_on_generate flips is_loaded off + # first, to model the eviction/unload race vs an in-pipeline failure (OOM). + self._generate_error = generate_error + self._unload_on_generate = unload_on_generate + self.calls = [] + + @property + def is_loaded(self): + return self._loaded + + def status(self): + return { + "loaded": self._loaded, + "repo_id": self._repo_id if self._loaded else None, + "family": "z-image" if self._loaded else None, + "base_repo": self._base_repo if self._loaded else None, + "device": "cpu", + "dtype": "float32", + "cpu_offload": False, + } + + def generate( + self, + *, + prompt, + width, + height, + steps, + guidance, + batch_size = 1, + ): + if not self._loaded: + raise RuntimeError("No diffusion model is loaded.") + if self._generate_error is not None: + if self._unload_on_generate: + self._loaded = False + raise self._generate_error + self.calls.append( + dict( + prompt = prompt, + width = width, + height = height, + steps = steps, + guidance = guidance, + batch_size = batch_size, + ) + ) + out = { + "images": [object() for _ in range(batch_size)], + "seed": 4242, + "repo_id": self._repo_id, + } + if self._native_seeds: + out["seeds"] = [4242 + i for i in range(batch_size)] + return out + + +def _make_client(backend): + store = {} + + def _save(image, meta): + image_id = f"img{len(store)}" + record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"} + store[image_id] = record + return record + + app = FastAPI() + install_api_error_handlers(app) + app.include_router(router, prefix = "/v1") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + return TestClient(app), store, _save + + +@pytest.fixture +def client(monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None) + cli.backend = backend # type: ignore[attr-defined] + return cli + + +def _post(client, body): + return client.post("/v1/images/generations", json = body) + + +def test_url_response_shape(client): + resp = _post(client, {"prompt": "a sloth", "size": "256x256"}) + assert resp.status_code == 200 + body = resp.json() + assert set(body.keys()) == {"created", "data"} + assert isinstance(body["created"], int) and body["created"] > 0 + assert len(body["data"]) == 1 + item = body["data"][0] + assert "url" in item and "b64_json" not in item # exclude_none drops the unused key + assert item["url"].endswith("/file") + # Z-Image-Turbo defaults (9 steps, 0 guidance) flow into the backend call. + assert client.backend.calls[0] == dict( + prompt = "a sloth", width = 256, height = 256, steps = 9, guidance = 0.0, batch_size = 1 + ) + + +def test_b64_response_shape(client): + resp = _post(client, {"prompt": "a sloth", "size": "256x256", "response_format": "b64_json"}) + assert resp.status_code == 200 + item = resp.json()["data"][0] + assert "b64_json" in item and "url" not in item + assert item["b64_json"] == "QUJD" + + +def test_local_load_uses_base_repo_for_defaults(monkeypatch): + # repo_id is a local path that names no model; base_repo identifies FLUX.1-dev, + # so the route must pick 28 steps / 3.5 guidance, not the 9/0 fallback. + backend = _FakeBackend(repo_id = "/models/my-flux", base_repo = "black-forest-labs/FLUX.1-dev") + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 200 + assert backend.calls[0]["steps"] == 28 and backend.calls[0]["guidance"] == 3.5 + + +def test_pipeline_runtime_error_is_sanitized_500(monkeypatch): + # A RuntimeError raised inside the pipeline while the model stays loaded (e.g. + # CUDA OOM, a RuntimeError subclass) must be a sanitized 500, not a 503 that + # echoes the raw exception text. + oom = RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (GPU 0; 47.5 GiB total)") + backend = _FakeBackend(generate_error = oom) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 500 + assert resp.json()["error"]["message"] == "Image generation failed." + assert "CUDA" not in resp.text # raw exception text must not leak + + +def test_unload_race_returns_503(monkeypatch): + # The model is evicted/unloaded between the readiness check and the call: the + # RuntimeError with is_loaded now False is the one case that maps to 503. + backend = _FakeBackend( + generate_error = RuntimeError("No diffusion model is loaded."), + unload_on_generate = True, + ) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 503 + err = resp.json()["error"] + assert err["type"] == "api_error" + # The 503 carries the fixed sanitized message, not the raw exception text. + assert err["message"] == "No image model loaded. Load an image model first." + + +def test_non_runtime_pipeline_error_is_500(monkeypatch): + # A non-RuntimeError from the pipeline (the model stays loaded) must not route + # to the 503 branch (which is gated on isinstance RuntimeError) -> sanitized 500. + backend = _FakeBackend(generate_error = ValueError("bad tensor shape")) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 500 + assert "shape" not in resp.text + + +def test_n_maps_to_batch(client): + resp = _post(client, {"prompt": "p", "size": "256x256", "n": 3}) + assert resp.status_code == 200 + assert len(resp.json()["data"]) == 3 + assert client.backend.calls[0]["batch_size"] == 3 + + +def test_batch_persists_batch_size(monkeypatch): + # n>1 must persist batch_size in each gallery record so the Studio restore path + # can replay a batch_index>0 sibling (which shares the batch's single seed). + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3}) + assert resp.status_code == 200 + records = sorted(store.values(), key = lambda r: r["batch_index"]) + assert [r["batch_index"] for r in records] == [0, 1, 2] + assert all(r["batch_size"] == 3 for r in records) + + +def test_uses_active_engine_not_diffusers_singleton(monkeypatch): + # On a no-GPU host the loaded model lives behind the native sd_cpp engine, not the + # diffusers singleton. The route must query get_active_diffusion_engine (like + # /images/generate) or it 503s a model that is loaded and usable. + active = _FakeBackend(loaded = True) # the active (e.g. sd_cpp) engine, loaded + idle_diffusers = _FakeBackend(loaded = False) # diffusers singleton, empty + monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: active) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: idle_diffusers) + cli, store, _save = _make_client(active) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 200 + assert len(active.calls) == 1 # the active engine did the work, not the idle singleton + + +def test_native_batch_persists_per_image_seed(monkeypatch): + # The native sd.cpp engine returns a distinct seed per image (base+index) in + # "seeds"; each gallery record must store its own seed (like /images/generate), + # not the shared base, or a restored batch_index>0 image shows the wrong seed. + backend = _FakeBackend(native_seeds = True) + monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3}) + assert resp.status_code == 200 + records = sorted(store.values(), key = lambda r: r["batch_index"]) + assert [r["seed"] for r in records] == [4242, 4243, 4244] + + +def test_null_fields_coalesce_to_defaults(client): + # OpenAI marks n/size/response_format nullable-with-default: null -> default. + resp = _post(client, {"prompt": "p", "n": None, "size": None, "response_format": None}) + assert resp.status_code == 200 + assert len(resp.json()["data"]) == 1 + assert "url" in resp.json()["data"][0] + assert client.backend.calls[0]["width"] == 1024 # size null -> auto -> 1024 + + +@pytest.mark.parametrize( + "body, param", + [ + ({"size": "256x256"}, "prompt"), # missing prompt + ({"prompt": "", "size": "256x256"}, "prompt"), # empty prompt + ({"prompt": "p", "size": "300x300"}, "size"), # not multiple of 16 + ({"prompt": "p", "size": "abc"}, "size"), # unparseable + ({"prompt": "p", "stream": True}, "stream"), # streaming unsupported + ], +) +def test_validation_400_with_param(client, body, param): + resp = _post(client, body) + assert resp.status_code == 400 + err = resp.json()["error"] + assert err["type"] == "invalid_request_error" + assert err["param"] == param + for k in ("message", "code"): + assert k in err + + +@pytest.mark.parametrize("n", [0, 11, -1]) +def test_n_out_of_range_400(client, n): + resp = _post(client, {"prompt": "p", "size": "256x256", "n": n}) + assert resp.status_code == 400 + assert resp.json()["error"]["type"] == "invalid_request_error" + + +def test_no_model_loaded_503(monkeypatch): + backend = _FakeBackend(loaded = False) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p"}) + assert resp.status_code == 503 + # 503 still wears the OpenAI envelope (api_error) on the /v1 surface. + assert resp.json()["error"]["type"] == "api_error" + + +def test_auth_required(): + backend = _FakeBackend() + app = FastAPI() + install_api_error_handlers(app) + app.include_router(router, prefix = "/v1") + # No dependency override: the real auth dependency runs and rejects. + resp = TestClient(app).post("/v1/images/generations", json = {"prompt": "p"}) + assert resp.status_code in (401, 403) From 7d8b2db23639b47a9c103cce76efb30506dec642 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 08:37:32 -0700 Subject: [PATCH 02/17] Studio diffusion: persistent sd-server for the native engine (load once, serve many) (#6768) * Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict Phase 1 of porting the richer diffusion stack onto the image-generation backend. - Add a compartmentalized device/dtype policy module (diffusion_device.py) resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or fp32, never a silent fp16 that renders a black image. - Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved float16 to float32 for those families so they do not produce black images. - Split the backend locks: a generation holds only _generate_lock, so status, unload, and a new load are never blocked by a long denoise. Add per-generation cancellation via callback_on_step_end so an eviction or a superseding load preempts a running generation; a replacement load waits for it to stop before allocating, so two pipelines never sit in VRAM at once. - Validate a load request before the GPU handoff so an unloadable pick never evicts a working chat model, and reject missing local paths up front. - Add CPU-only tests for the device policy, dtype guard, lock split and cancellation, and validate-before-evict, plus a GPU benchmark/regression script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR against a saved reference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and VAE tiling/slicing from measured free device memory vs the model's estimated resident footprint, then applies it to the built pipeline. auto stays resident when the model fits (byte-identical to the prior resident path), and falls to whole-module offload when tight; fast/balanced/low_vram are explicit overrides. Sequential submodule offload is unreliable for GGUF transformers on diffusers 0.38, so it falls back to whole-module offload and status reports the policy actually engaged. Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM 47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost. 73 prior + 35 new CPU tests pass. * Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level, use_stream) that keeps the transformer flowing through the GPU a few blocks at a time while the text encoder / VAE stay resident, and fix VAE tiling to drive the VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so status never overstates either, and group falls back to whole-module offload when the transformer can't be streamed. Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 -> 2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names now match that tradeoff: balanced = stream the transformer, low_vram = offload every component. auto picks group when the companions fit resident, else model. 112 CPU tests pass. * Studio diffusion (Phase 5): image quality-vs-quant accuracy harness Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold prompt + seed fixed, render a grid with a reference quant (default BF16), then render each candidate quant and measure drift from the reference. Records mean PSNR + SSIM (pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity (transformers, --clip), plus file size, latency, and peak VRAM, then prints a quality-vs-cost table and recommends the smallest quant within a quality budget. --selftest validates the metrics on synthetic images with no GPU or model. Verified on Z-Image (B200): the table degrades monotonically with quant size (Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat (~0.34) -- quantization erodes fine detail far more than prompt adherence. * Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32) Add a speed_mode knob (off by default, so the render path stays bit-identical): default applies channels_last VAE + regional torch.compile of the denoiser's repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional compile is gated off for the GGUF transformer (dequantises per-op) and for families flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image), so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed optims run before placement/offload, per the diffusers composition order. status now reports speed_mode + the optims actually engaged. Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last', 'tf32'], compile correctly skipped for GGUF; generation works in every mode. 121 CPU tests pass. * Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3) storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16 compute dtype while normalisations and embeddings stay full precision. Applied before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder dense). status reports which encoders were cast. Verified on Z-Image (B200, balanced/group mode where the encoder stays resident): generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload) at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off by default and documented as such, with the Phase 5 harness to size the cost. 127 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob) Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant (fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8 stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run before placement; status reports the mode actually engaged. This is the lean realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the 3045-line port. Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default); size it per model with the Phase 5 quality harness. diffusion_bench gains --text-encoder-quant. 129 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm / XPU; this covers the hardware diffusers serves poorly, consuming the same split GGUF assets Studio already curates. - sd_cpp_args.py: pure sd-cli command builder. Maps the family to its text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1 CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential) to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu / --vae-tiling / --diffusion-fa), so one user knob drives both engines. - sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary() with the same precedence as the llama finder (env override, then the Studio install root, then in-tree, then PATH), an is_available/version probe, and a one-shot subprocess generate that streams progress and returns the PNG. runtime_env() prepends the binary's directory to the platform library path so a prebuilt's bundled libstable-diffusion.so resolves. select_diffusion_engine() is the pure routing decision (GPU backends to diffusers, CPU/MPS to native when present). - install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the Studio install root. resolve_release_asset() is a pure, unit-tested host-to-asset matrix. - scripts/sd_cpp_smoke.py: end-to-end native generation harness. Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine, routing, runtime env, and the installer resolver. Full diffusion suite 166 passing. Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears) The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with 'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2 *_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing min_features filter did not exclude them. Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed, guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers (M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and quantises these layers fine, so the exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note 'context_embedder' contains the substring 'text_embed', which is why the token is the specific 'time_text_embed', not 'text_embed'. Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s (5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already working) are unchanged. - diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes exclude_name_tokens; quantize_transformer passes it for int8 only. - hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps attention / FFN / sequence-embedder linears), while fp8 keeps them. - scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's Linear layers and derive the exclusion list. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears) The prequant-checkpoint builder applied the dense quant filter without the int8-only M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8 checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now produces a working checkpoint on every supported model, giving int8 (the consumer-preferred scheme) the same ~2x load-VRAM and download reduction fp8 already had. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the native stable-diffusion.cpp engine instead of diffusers, with diffusers as the guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM. - diffusion_engine_router: centralised engine selection (built on the existing select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason. - sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface backed by sd-cli, with lazy binary install, registry-driven asset fetch, step-progress parsing, and cancellation. - diffusion_families: per-family single-file VAE + text-encoder asset mapping. - sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled). - routes/inference + gpu_arbiter: drive the active engine via the router; the API now reports the active engine and any fallback reason. - tests for the backend, router, route selection, and cancellation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler Address review feedback on #6724: - engine router: unload the engine being deactivated on a switch, so the old model is not left resident-but-unreachable (the evictor only targets the active engine). - generate route: sd.cpp execution errors (nonzero exit / timeout / missing output) now map to 500, not 409 (which only means not-loaded / cancelled). - native batch: return per-image seeds and persist the actual seed for each image so every batch image is reproducible. - Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the stable-diffusion.cpp docs; other families keep sd-cli defaults. - honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload offload flags on the native load instead of hardcoding them off. - fail the load when the sd-cli binary is present but not runnable (version() now returns None on exec error / nonzero exit). - size estimate: only treat the transformer asset as a possible local path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF - apply_step_cache now engages only via the transformer's native enable_cache (the diffusers CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model now runs uncached per the best-effort contract. - GGUF transformers are quantized (the default Studio load path), so they now use the higher quantized FBCache threshold when the caller leaves it unset, instead of the dense default that could keep the cache from triggering. - fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so fullgraph=True failed warmup and silently measured an eager cached run); output dir is now relative to the script, not a hardcoded path. * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs without producing output (or without closing stdout) would never reach proc.wait and the wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the PROCESS, so the main thread always enforces the timeout and kills a hung process (which closes the pipe and ends the reader). Add a test that times out even when stdout blocks, and make the no-binary test hermetic so a host-installed sd-cli can't leak in. * Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn The filter callback can be invoked without a module name, so fqn.lower() would raise AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens, so the linear is kept) instead of crashing the quantise pass. * Studio diffusion (Phase 16) review fixes: native engine robustness - sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64); a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds aliased to the same image. Random seeds stay 53-bit (JS-safe). - sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub fall back to anonymous instead of failing auth on a blank token. - sd_cpp_backend: a superseding load now cancels the in-flight generation, so the old sd-cli can no longer return/persist an image from the previous model. - diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a slow 10+ GB free / CUDA sync does not block engine selection. - diffusion_engine_router: probe sd-cli runnability (version()) before committing to native, so a present-but-unrunnable binary falls back to diffusers at selection. - diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a CPU-only install can still reach the native sd.cpp engine instead of failing load. - tests updated for the runnability probe + a not-runnable fallback case. * Studio diffusion (Phase 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening - install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket); extract through a per-member containment check (Zip-Slip guard); expanduser the --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch the separately-published cudart runtime DLL archive so sd-cli.exe can start. - sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie. - tests: Zip-Slip rejection, normal extraction, studio-home discovery. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes the numbered files _ (base_0.png, base_1.png, ...) instead of the literal --output path. SdCppEngine.generate checked only the literal path, so a batch generation would exit 0 and then raise 'no image' (or return a stale file). generate now returns the literal path when present and otherwise falls back to the numbered siblings; single-image behavior is unchanged. Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected without error. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder Codex review: the M=1 modulation/embedder exclusion was wired only into the dense runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint quantised the AdaLN modulation and conditioning-embedder linears, and loading it via transformer_prequant_path (the load path only loads already-quantised tensors, it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes for the runtime path. Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 -> the M=1 exclusion, every other scheme -> none) and use it in both the runtime quantiser and the builder, so a prequant artifact's quantised-layer set always matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion). Test: int8_exclude_name_tokens returns the exclusion for int8 and () for fp8/nvfp4/mxfp8. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race Codex review on the native-engine routing: - The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat) unconditionally after engine selection. A native sd.cpp load on a pure-CPU host never touches the GPU, so that needlessly tore down the resident chat model. The handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it. - sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so the setting was unverifiable. status now derives them from state.offload_flags (still 'none' on CPU, where the flags are empty). - _run_load committed the new state without cancelling/waiting on a generation that started during the (slow) asset download, so a stale sd-cli run against the OLD model could finish afterward and persist an image from the previous model once the new load reported ready. The commit now signals the in-flight cancel and waits on _generate_lock before swapping _state (taken only at commit, so the download never serialises against generation), mirroring the diffusers load path. Tests: CPU native load skips the arbiter while a GPU native load takes it; status reports offload active when flags are set; _run_load cancels and waits for an in-flight generation before committing. * Studio diffusion (Phase 14) review round 2: align helper name with the stack Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the identical helper already present higher in the diffusion stack (Phase 16). The helper definition, the runtime quantiser call, and the offline builder are now byte-identical to that version, so the two branches no longer introduce a divergent name for the same single-source-of-truth and the stack merges without a conflict on this fix. No behavior change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion: persistent sd-server for the native engine (load once, serve many) The native sd.cpp tier ran sd-cli one-shot per image, so begin_load only resolved asset paths and every generation re-spawned sd-cli and reloaded the multi-GB GGUF from disk (a batch of N = N full reloads). This makes it a resident backend backed by stable-diffusion.cpp's persistent sd-server, mirroring the chat backend's llama-server lifecycle: - begin_load spawns sd-server once (the model loads there) and polls /v1/models until ready; unload kills it. - generate submits ONE async /sdcpp/v1/img_gen job for the whole batch (no reload), polls it to completion, and decodes the returned images. Step progress and ETA come from the server's stdout (the job JSON has no per-step field). - The one-shot sd-cli path is kept as an automatic fallback: it is used when sd-server is absent, and also when a present sd-server fails to start, so behavior is never worse than before. The public backend surface is unchanged, so routes/router need no change. New: sd_cpp_server.py (SdCppServer manager: spawn/readiness/job-submit-poll/cancel/stop, process spawned inside the drain thread so PR_SET_PDEATHSIG binds to the interpreter, not a transient thread; empty scratch dir for the server's per-request LoRA/upscaler/embd scans). Extended: sd_cpp_engine.py (find_sd_server_binary), sd_cpp_args.py (build_sd_cpp_server_command + build_img_gen_request), sd_cpp_backend.py (server/one-shot modes, ensure_sd_server_binary upgrades existing sd-cli-only installs), and the prebuilt installer (locate + chmod sd-server, which ships in the same archive as sd-cli). Verified on a B200 (Z-Image-Turbo-GGUF, CUDA sd-server): one model load across multiple generations (server pid stable, a single 'listening on:'), a batch served from one job with distinct per-image seeds, the second generation faster than the first, and unload/reload spawning a fresh process. 105 sd.cpp + 81 diffusion tests pass. Addresses the review of the Phase 16 native-engine PR. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio native diffusion: harden the persistent sd-server path Addresses review findings on the sd-server backend: - Router: treat a runnable sd-server as native availability, so an sd-server-only install (no sd-cli) still routes to the native engine instead of silently falling back to diffusers. - Backend: probe the sd-server binary before the multi-GB asset download, falling back to one-shot sd-cli up front when it cannot run. - Backend: a lazily cached one-shot fallback engine no longer pins the backend to one-shot; only an explicitly injected engine does, so a now-available server can be used on the next load. - Backend: mask explicit seeds to sd.cpp's signed int64 range before submitting a server job (large seeds were rejected/wrapped in server mode only), and split batches above the server's per-job limit into chunks, each with a timeout proportional to its image count. - Backend/server: make server startup cancellable. stop() signals an abort event before taking the lifecycle lock so a blocking readiness wait bails promptly; unload() stops a not-yet-committed pending server. - Backend: status() clears stale loaded state when the resident server has exited, so clients reload instead of hammering a dead process with 500s. - Server: abandon a poll whose best-effort cancel is not honored within a grace window (releasing the generate lock), report a pre-submit stop/cancel as cancellation (409, not 500), and verify JSON responses are the expected type before indexing. - Server: use a bounded deque for the stdout tail buffer. - Add native_mode to DiffusionStatusResponse so the field is not dropped by the response model. Adds regression tests for each behavioral change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * sd-server: harden native lifecycle and GPU install path - Treat a crashed sd-server probe (signal death / non-127 nonzero) as unavailable so a broken prebuilt falls back to diffusers instead of routing to a server that dies on startup. - Drop stale loaded state when a resident server has exited before a generate, returning the recoverable not-loaded path rather than a 500. - Reject incomplete server batches (fewer blobs than requested) like the one-shot path instead of silently dropping images. - Bound the server log tail in place (keep the deque(maxlen)) and bypass HTTP(S) proxies for the loopback client (trust_env=False). - Honor a stop() that arrives after the server is published but before start() takes the lock, so a cancelled load cannot leak a spawned model process. - Map a closed-client RuntimeError during poll to a cancellation when the generation is being cancelled, so unload races surface as 409 not 500. - Stop a timed-out server job (best-effort cancel then teardown) so an abandoned generation cannot keep denoising and block later loads. - Install the accelerator-matched sd-server build (ROCm/Vulkan/CUDA) and probe the resident server before auto-installing sd-cli, so a server-only or GPU host does not fetch the wrong or an unused binary. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../core/inference/diffusion_engine_router.py | 42 +- studio/backend/core/inference/sd_cpp_args.py | 130 +++++ .../backend/core/inference/sd_cpp_backend.py | 543 +++++++++++++++--- .../backend/core/inference/sd_cpp_engine.py | 94 ++- .../backend/core/inference/sd_cpp_server.py | 502 ++++++++++++++++ studio/backend/models/inference.py | 5 + .../tests/test_diffusion_engine_router.py | 14 + studio/backend/tests/test_sd_cpp_args.py | 103 ++++ studio/backend/tests/test_sd_cpp_backend.py | 313 ++++++++++ studio/backend/tests/test_sd_cpp_engine.py | 53 ++ studio/backend/tests/test_sd_cpp_server.py | 417 ++++++++++++++ studio/install_sd_cpp_prebuilt.py | 18 + 12 files changed, 2108 insertions(+), 126 deletions(-) create mode 100644 studio/backend/core/inference/sd_cpp_server.py create mode 100644 studio/backend/tests/test_sd_cpp_server.py diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 0c758501f8..ffe31d93a5 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -29,7 +29,12 @@ from typing import Any, Optional from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_families import DiffusionFamily, family_sd_cpp_supported -from core.inference.sd_cpp_backend import _install_allowed, ensure_sd_cpp_binary +from core.inference.sd_cpp_backend import ( + _install_allowed, + _server_binary_runnable, + ensure_sd_cpp_binary, + ensure_sd_server_binary, +) from core.inference.sd_cpp_engine import ( ENGINE_DIFFUSERS, ENGINE_SD_CPP, @@ -137,20 +142,37 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] fam_ok = family_sd_cpp_supported(fam) binary = None + server_binary = None if policy_eligible and fam_ok: - binary = ensure_sd_cpp_binary( + # Probe the resident sd-server FIRST: the backend PREFERS it, and an sd-server-only + # install (no sd-cli) must still route to native rather than silently falling back to + # diffusers. Checking it before the sd-cli install also means a server-only host does + # not pay an avoidable sd-cli download. Install the accelerator-matched build (ROCm / + # Vulkan / CUDA) so a forced-native GPU load gets the GPU server, not the CPU one. + server_binary = ensure_sd_server_binary( allow_install = _install_allowed(), accelerator = _install_accelerator_for(backend), ) - # Probe runnability here, before committing the route to native: a present but - # non-runnable binary (wrong arch, missing shared libs, no execute bit) would - # otherwise pass as available and only fail inside the background load, instead - # of falling back to diffusers now. + if server_binary and not _server_binary_runnable(server_binary): + logger.warning( + "sd-server at %s is present but not runnable; not using it", server_binary + ) + server_binary = None + # sd-cli is the one-shot fallback. Always LOCATE an existing binary, but only + # auto-INSTALL it when there is no usable server, so a server-only install is not + # forced to also download a CLI it will never use. Probe runnability before + # committing native: a present but non-runnable binary (wrong arch, missing shared + # libs, no execute bit) would otherwise pass as available and only fail inside the + # background load, instead of falling back to diffusers now. + binary = ensure_sd_cpp_binary( + allow_install = _install_allowed() and server_binary is None, + accelerator = _install_accelerator_for(backend), + ) if binary and SdCppEngine(binary = binary).version() is None: - logger.warning("sd-cli at %s is present but not runnable; using diffusers", binary) + logger.warning("sd-cli at %s is present but not runnable; not using it", binary) binary = None - native_available = bool(binary) and policy_eligible and fam_ok + native_available = bool(binary or server_binary) and policy_eligible and fam_ok choice = select_diffusion_engine( backend, native_available = native_available, prefer_native = prefer_native ) @@ -162,8 +184,8 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] reason = f"GPU backend '{backend}' uses diffusers" elif not fam_ok: reason = f"family '{fam.name}' has no native sd.cpp asset mapping" - elif not binary: - reason = "sd-cli binary unavailable" + elif not (binary or server_binary): + reason = "native sd.cpp binary unavailable" else: reason = "diffusers selected" return _activate(ENGINE_DIFFUSERS, reason) diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 5e0db4f0ce..40a5613273 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -325,6 +325,136 @@ def build_sd_cpp_upscale_command( return cmd +def build_sd_cpp_server_command( + binary: str, + files: SdCppModelFiles, + *, + host: str, + port: int, + vae_format: Optional[str] = None, + offload: Optional[list[str]] = None, + native_speed: Optional[str] = None, + threads: Optional[int] = None, + scratch_dir: Optional[str] = None, + verbose: bool = False, + extra_args: Optional[list[str]] = None, +) -> list[str]: + """Build the ``sd-server`` argv: model + hardware/server flags only. + + ``sd-server`` (stable-diffusion.cpp ``examples/server``) loads the model once at + spawn from the SAME flags ``sd-cli`` takes (``--diffusion-model`` / ``--vae`` / + the text encoders / ``--vae-format`` / offload + speed), and adds ``--listen-ip`` + / ``--listen-port``. Per-generation parameters (prompt, size, steps, seed, cfg, + sampler, batch) are NOT here -- they go in each ``/sdcpp/v1/img_gen`` request, so + one resident process serves many generations without reloading the weights. + + ``offload`` / ``native_speed`` map to the exact same sd.cpp flags as the one-shot + engine (``--offload-to-cpu`` / ``--diffusion-fa`` / ...), verified to be accepted + by ``sd-server --help``. ``scratch_dir`` (if given) is pointed at by the LoRA / + hires-upscaler / embeddings directory flags: sd-server's img_gen handler recursively + iterates those dirs, and an unset / missing dir makes it fail the request, so we give + it a real (empty) directory. ``extra_args`` is appended last so a power user can + override anything (sd.cpp's parser is last-wins). + """ + if not files.diffusion_model: + raise ValueError("diffusion_model path is required") + + cmd: list[str] = [binary, "--diffusion-model", files.diffusion_model] + for flag, value in ( + ("--vae", files.vae), + ("--clip_l", files.clip_l), + ("--clip_g", files.clip_g), + ("--t5xxl", files.t5xxl), + ("--llm", files.llm), + ("--qwen2vl", files.qwen2vl), + ): + if value: + cmd += [flag, value] + if vae_format: + cmd += ["--vae-format", vae_format] + cmd += ["--listen-ip", str(host), "--listen-port", str(int(port))] + if scratch_dir: + cmd += [ + "--lora-model-dir", + scratch_dir, + "--hires-upscalers-dir", + scratch_dir, + "--embd-dir", + scratch_dir, + ] + if threads is not None: + cmd += ["--threads", str(int(threads))] + + offload = list(offload or []) + if offload: + cmd += offload + # De-dup speed flags against offload (offload may already include --diffusion-fa). + cmd += [f for f in native_speed_flags(native_speed) if f not in offload] + if verbose: + cmd += ["-v"] + if extra_args: + cmd += list(extra_args) + return cmd + + +def build_img_gen_request( + *, + prompt: str, + negative_prompt: Optional[str] = None, + width: int = 1024, + height: int = 1024, + steps: Optional[int] = None, + seed: Optional[int] = None, + batch_count: int = 1, + sample_method: Optional[str] = None, + flow_shift: Optional[float] = None, + cfg_scale: Optional[float] = None, + distilled_guidance: Optional[float] = None, + output_format: str = "png", +) -> dict: + """Build the ``POST /sdcpp/v1/img_gen`` JSON body for one text-to-image request. + + The native ``sdcpp`` API takes the whole batch in one request (``batch_count``), + so a batch reuses the resident model with no reload. Sampling lives under + ``sample_params``; guidance is split exactly like the one-shot engine's + ``_map_guidance``: a FLUX distilled value goes to ``guidance.distilled_guidance``, + a real classifier-free scale goes to ``guidance.txt_cfg``. Only set keys are + emitted so the server applies its own defaults for the rest. + """ + if not str(prompt).strip(): + raise ValueError("prompt is required") + + guidance: dict = {} + if cfg_scale is not None: + guidance["txt_cfg"] = float(cfg_scale) + if distilled_guidance is not None: + guidance["distilled_guidance"] = float(distilled_guidance) + + sample_params: dict = {} + if steps is not None: + sample_params["sample_steps"] = int(steps) + if sample_method: + sample_params["sample_method"] = str(sample_method) + if flow_shift is not None: + sample_params["flow_shift"] = float(flow_shift) + if guidance: + sample_params["guidance"] = guidance + + req: dict = { + "prompt": prompt, + "negative_prompt": negative_prompt or "", + "width": int(width), + "height": int(height), + "batch_count": max(1, int(batch_count)), + "output_format": output_format, + } + if seed is not None: + req["seed"] = int(seed) + if sample_params: + req["sample_params"] = sample_params + return req + + def _fmt_float(value: float) -> str: """Compact float -> str: drop a trailing ``.0`` so ``1.0`` -> ``1`` (sd-cli accepts both, but the tidy form keeps logged commands readable).""" diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 3c79b14f2d..53f7131e16 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -49,13 +49,22 @@ from core.inference.diffusion_memory import ( OFFLOAD_NONE, OFFLOAD_SEQUENTIAL, ) -from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, offload_flags +from core.inference.sd_cpp_args import ( + SdCppGenParams, + SdCppModelFiles, + build_img_gen_request, + offload_flags, +) from core.inference.sd_cpp_engine import ( SdCppCancelled, SdCppEngine, find_sd_cpp_binary, + find_sd_server_binary, + runtime_env, ) +from core.inference.sd_cpp_server import SdCppServer from loggers import get_logger +from utils.subprocess_compat import windows_hidden_subprocess_kwargs logger = get_logger(__name__) @@ -68,6 +77,45 @@ _STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)") # download / extract / chmod. _install_lock = threading.Lock() +# sd-server accepts at most this many images per img_gen job; larger Studio batches +# (the request model allows up to 32) are split into chunks of this size, the way the +# one-shot path did them one image at a time. +_MAX_SERVER_BATCH = 8 + +# Per-image wall-clock budget for a server job, so a batch gets a timeout proportional to +# its image count (matching the one-shot path, where each image had its own budget) rather +# than one fixed deadline the whole batch has to finish within. +_SERVER_PER_IMAGE_TIMEOUT_S = 1800.0 + + +def _server_binary_runnable(binary: str) -> bool: + """Best-effort probe that ``binary`` can actually execute (not just exist). + + Runs `` --help`` with the same runtime env the server will use, so a present + but unrunnable build (wrong arch, missing shared libs, no execute bit) is caught before + a multi-GB asset download. Conservative: only a clear "cannot launch" signal (OSError, + or the dynamic-loader exit codes 126/127) returns False; anything else is treated as + runnable so a quirky ``--help`` exit code never blocks a working binary.""" + import subprocess + + try: + proc = subprocess.run( + [binary, "--help"], + capture_output = True, + timeout = 20, + env = runtime_env(binary), + **windows_hidden_subprocess_kwargs(), + ) + except OSError: + return False # cannot exec at all (wrong arch / no execute bit / missing loader) + except Exception: # noqa: BLE001 -- timeout or anything odd: don't block on a flaky probe + return True + # A negative return code is a signal death (e.g. -4 SIGILL from an incompatible + # prebuilt on an older CPU): the binary launches but immediately crashes, so treat it + # as unavailable and let the load fall back to diffusers instead of routing to a + # server that will die on startup. + return proc.returncode >= 0 and proc.returncode not in (126, 127) + def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu") -> Optional[str]: """Path to a usable ``sd-cli`` binary, installing the prebuilt once if needed. @@ -105,9 +153,51 @@ def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu" return None +def ensure_sd_server_binary( + *, allow_install: bool = True, accelerator: str = "cpu" +) -> Optional[str]: + """Path to a usable ``sd-server`` binary, installing the prebuilt once if needed. + + Unlike ``ensure_sd_cpp_binary``, this installs when *sd-server specifically* is + missing -- even if an ``sd-cli`` from an older install is already present -- so an + existing one-shot install is upgraded to the persistent server (the prebuilt archive + ships both). Returns None when it is absent and cannot be installed; the backend then + uses the one-shot fallback. Never raises. + """ + found = find_sd_server_binary() + if found: + return found + if not allow_install: + return None + with _install_lock: + found = find_sd_server_binary() + if found: + return found + try: + import sys + + studio_dir = Path(__file__).resolve().parents[3] # .../studio + if str(studio_dir) not in sys.path: + sys.path.insert(0, str(studio_dir)) + from install_sd_cpp_prebuilt import install as _install + except Exception as exc: # noqa: BLE001 -- import path / module issues are non-fatal + logger.warning("sd-server installer import failed: %s", exc) + return None + try: + _install(accelerator = accelerator) # extracts sd-cli AND sd-server + except Exception as exc: # noqa: BLE001 -- download/extract failure -> fall back + logger.warning("sd-server auto-install failed: %s", exc) + return None + return find_sd_server_binary() + + @dataclass(frozen = True) class _SdState: - """The loaded native checkpoint: resolved asset paths + run settings.""" + """The loaded native checkpoint: resolved asset paths + run settings. + + ``server`` is the resident ``sd-server`` process (the model is loaded once, inside + it) when ``mode == "server"``; in the ``"oneshot"`` fallback it is ``None`` and each + generation re-runs ``sd-cli``.""" repo_id: str base_repo: str @@ -120,6 +210,8 @@ class _SdState: threads: Optional[int] = None sampling_method: Optional[str] = None flow_shift: Optional[float] = None + server: Optional[SdCppServer] = None + mode: str = "server" def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str: @@ -191,11 +283,19 @@ class SdCppDiffusionBackend: self._lock = threading.Lock() self._generate_lock = threading.Lock() self._engine = engine # resolved lazily on first load so import stays cheap + # An engine passed in is an EXPLICIT injection (the test seam / escape hatch) and + # pins one-shot mode; an engine cached later by a runtime fallback must NOT, so a + # now-available server can still be used on the next load. + self._engine_injected = engine is not None self._state: Optional[_SdState] = None self._loading: Optional[_SdLoading] = None self._load_token = 0 self._cancel_event = threading.Event() self._active_generate_cancel: Optional[threading.Event] = None + # The sd-server being started for an in-flight load, before it is committed to + # _state. Tracked so an unload / superseding load can stop it mid-startup instead + # of leaving it loading (and holding the generate lock) for the whole timeout. + self._pending_server: Optional[SdCppServer] = None self._gen: Optional[_SdGen] = None @property @@ -212,6 +312,38 @@ class SdCppDiffusionBackend: self._engine = SdCppEngine(binary = binary) return self._engine + def _resolve_backend(self) -> tuple[str, Optional[str], Optional[SdCppEngine]]: + """Pick the native execution mode: ("server", binary, None) or ("oneshot", None, engine). + + The persistent ``sd-server`` is preferred (load once, serve many). The one-shot + ``sd-cli`` is the fallback for older / custom builds that lack the server target. + An explicitly injected engine forces one-shot (the unit-test seam and an escape + hatch), so a test never spawns a real server or triggers an install. A lazily + cached fallback engine does NOT force one-shot: once a resident server becomes + available (installed, or a per-model start that previously failed now works), the + next load can use it, instead of being pinned to one-shot for the whole session. + """ + if self._engine_injected and self._engine is not None: + return "oneshot", None, self._resolve_engine() + # Install the sd-server build matching the resolved device backend (ROCm / Vulkan / + # CUDA), not the default CPU build: a forced/enabled native load on a GPU host must + # not silently fetch the plain-CPU server. Lazy import avoids an import cycle with + # the router, which imports this backend during engine selection. + from core.inference.diffusion_engine_router import _install_accelerator_for + + accelerator = _install_accelerator_for( + getattr(resolve_diffusion_device_target(), "backend", "cpu") + ) + server_binary = ensure_sd_server_binary( + allow_install = _install_allowed(), accelerator = accelerator + ) + if server_binary is not None: + return "server", server_binary, None + logger.warning( + "sd-server not found; falling back to one-shot sd-cli (reloads the model per image)." + ) + return "oneshot", None, self._resolve_engine() + # ── Background load + progress ───────────────────────────────────────── def begin_load( @@ -298,10 +430,35 @@ class SdCppDiffusionBackend: _load_token: int, ) -> None: try: - # Ensure the binary up front so an install failure surfaces before the - # multi-GB asset pull (the router also pre-checks, but a forced reload here - # must not silently download then fail at generate). - engine = self._resolve_engine() + # Resolve the backend mode (persistent sd-server preferred, one-shot sd-cli + # fallback) and binary up front so an install / missing-binary failure + # surfaces before the multi-GB asset pull. + mode, server_binary, engine = self._resolve_backend() + if mode == "server": + # Probe the server binary before the multi-GB asset pull: a present but + # unrunnable build (wrong arch / missing libs) would otherwise download + # everything and only then fail to start. If it cannot run, fall back to + # the one-shot engine now (when it is usable), else surface the failure. + assert server_binary is not None + if not _server_binary_runnable(server_binary): + logger.warning( + "sd-server at %s is present but not runnable; trying one-shot sd-cli.", + server_binary, + ) + try: + usable = self._resolve_engine().version() is not None + except Exception: # noqa: BLE001 + usable = False + if not usable: + raise RuntimeError("sd-server binary is present but not runnable.") + mode, server_binary, engine = "oneshot", None, self._resolve_engine() + if mode == "oneshot": + # Probe the binary: version() returns None when the present binary cannot + # run (bad perms / missing libs), so fail now rather than commit a "ready" + # state that crashes on the first generation. + assert engine is not None + if engine.version() is None: + raise RuntimeError("sd-cli binary is present but not runnable.") assets = self._asset_specs(repo_id, gguf_filename, fam) self._set_expected_bytes(assets, hf_token) @@ -318,37 +475,21 @@ class SdCppDiffusionBackend: ) device = resolve_diffusion_device_target().device # Honor the requested speed everywhere; offload only off-CPU (forced - # sd_cpp / MPS), since on CPU sd-cli is resident in RAM and the offload - # flags are no-ops. + # sd_cpp / MPS), since on CPU the weights are resident in RAM and the + # offload flags are no-ops. offload: tuple[str, ...] = () if device != "cpu": offload = tuple(offload_flags(_memory_policy(memory_mode, cpu_offload))) - state = _SdState( - repo_id = repo_id, - base_repo = base, - family = fam, - device = device, - files = files, - vae_format = fam.sd_cpp_vae_format, - native_speed = _native_speed_for(speed_mode), - offload_flags = offload, - threads = None, - sampling_method = fam.sd_cpp_sampling_method, - flow_shift = fam.sd_cpp_flow_shift, - ) - # Probe the binary: version() returns None when the present binary cannot - # run (bad permissions / missing shared libs), so fail the load now rather - # than commit a "ready" state that crashes on the first generation. - if engine.version() is None: - raise RuntimeError("sd-cli binary is present but not runnable.") - # A generation that started during the (slow) asset download is still running - # against the OLD model. Abort it, then WAIT on _generate_lock for it to exit - # before publishing the new state -- otherwise that stale sd-cli run can finish - # afterward and persist an image from the previous model once this load reports - # ready (mirrors the diffusers load_pipeline commit). _generate_lock is taken - # only here, not during the download, so the long fetch never serialises against - # generation; the inner token re-check guards an unload/newer load arriving while - # we waited. + native_speed = _native_speed_for(speed_mode) + + # Tear down any previously-loaded model, then commit the new one. A generation + # that started during the (slow) asset download is still running against the OLD + # model: abort it and WAIT on _generate_lock for it to exit before swapping, or + # a stale run could finish afterward and persist an image from the previous + # model. For server mode we stop the old server and start (load) the new one + # HERE, under _generate_lock, so generation never races a half-loaded server and + # two resident models never coexist. _generate_lock is taken only now, not during + # the download, so the long fetch never serialises against generation. with self._lock: if self._load_token != _load_token: return # superseded / cancelled @@ -358,6 +499,78 @@ class SdCppDiffusionBackend: with self._lock: if self._load_token != _load_token: return # superseded / cancelled while waiting + old_state = self._state + self._state = None # the old model is being torn down + if old_state is not None and old_state.server is not None: + old_state.server.stop() + server: Optional[SdCppServer] = None + if mode == "server": + assert server_binary is not None + server = SdCppServer(server_binary) + # Publish the not-yet-committed server so unload() / a superseding load + # can stop it mid-startup (SdCppServer.stop aborts the readiness wait + # without waiting on the lifecycle lock), instead of it loading for the + # full startup timeout while holding the generate lock. + with self._lock: + self._pending_server = server + try: + # Blocks until the server has loaded the model and is answering + # (its readiness check); raises with the log tail on a failed load. + server.start( + files, + vae_format = fam.sd_cpp_vae_format, + offload = list(offload), + native_speed = native_speed, + threads = None, + ) + except SdCppCancelled: + # Startup was aborted by an unload / superseding load: stop the + # half-started server and bail (the outer handler returns cleanly). + server.stop() + raise + except Exception as start_exc: # noqa: BLE001 + # A present-but-unusable sd-server must be no worse than the + # one-shot engine: fall back to sd-cli when it is usable, else + # surface the server error. + logger.warning( + "sd-server failed to start (%s); falling back to one-shot sd-cli.", + start_exc, + ) + server.stop() + server = None + try: + usable = self._resolve_engine().version() is not None + except Exception: # noqa: BLE001 + usable = False + if not usable: + raise start_exc + mode = "oneshot" + finally: + with self._lock: + if self._pending_server is server: + self._pending_server = None + state = _SdState( + repo_id = repo_id, + base_repo = base, + family = fam, + device = device, + files = files, + vae_format = fam.sd_cpp_vae_format, + native_speed = native_speed, + offload_flags = offload, + threads = None, + sampling_method = fam.sd_cpp_sampling_method, + flow_shift = fam.sd_cpp_flow_shift, + server = server, + mode = mode, + ) + with self._lock: + if self._load_token != _load_token: + # Superseded / unloaded while we were loading: discard the server + # we just started so it doesn't leak (and keep _state unloaded). + if server is not None: + server.stop() + return self._state = state self._loading = None except SdCppCancelled: @@ -475,76 +688,62 @@ class SdCppDiffusionBackend: seed: Optional[int] = None, batch_size: int = 1, ) -> dict[str, Any]: - import tempfile - - from PIL import Image - cancel = threading.Event() with self._generate_lock: with self._lock: state = self._state if state is None: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) + # A resident server can exit while idle; if a client generates without first + # polling status, drop the stale loaded state and report not-loaded so it gets + # the recoverable reload path instead of a 500 from img_gen (not running). + if ( + state.mode == "server" + and state.server is not None + and not state.server.is_alive() + ): + self._state = None + raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) self._active_generate_cancel = cancel - engine = self._resolve_engine() try: if seed is None: seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1) else: seed = int(seed) cfg_scale, flux_guidance = _map_guidance(state.family, guidance) - extra_args: list[str] = [] - if state.vae_format: - extra_args += ["--vae-format", state.vae_format] - if state.flow_shift is not None: - extra_args += ["--flow-shift", repr(float(state.flow_shift))] - self._gen = _SdGen(total_steps = int(steps)) - images = [] - seeds: list[int] = [] - with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: - for index in range(max(1, int(batch_size))): - if cancel.is_set(): - raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # Distinct seed per batch image (sd-cli is one image/run here), - # so a batch is reproducible image-by-image from the base seed. - # Mask to sd-cli's int64 range, NOT 53 bits: the request model and - # the diffusers backend both accept large explicit seeds, so a tight - # 2**53 mask would silently truncate them (2**53 -> 0) and collide - # distinct requested seeds onto the same image. Randomly-drawn seeds - # above are already 53-bit (JS-safe); explicit seeds pass through. - seed_i = (seed + index) & ((1 << 63) - 1) - out_path = str(Path(tmpdir) / f"img_{index}.png") - params = SdCppGenParams( - prompt = prompt, - negative_prompt = negative_prompt or None, - width = int(width), - height = int(height), - steps = int(steps), - cfg_scale = cfg_scale, - guidance = flux_guidance, - seed = seed_i, - sampling_method = state.sampling_method, - batch_count = 1, - ) - engine.generate( - state.files, - params, - output_path = out_path, - offload = list(state.offload_flags) or None, - native_speed = state.native_speed, - threads = state.threads, - extra_args = extra_args or None, - on_log = self._on_log, - cancel_event = cancel, - ) - with Image.open(out_path) as im: - images.append(im.copy()) - seeds.append(seed_i) + if state.mode == "server" and state.server is not None: + images, seeds = self._generate_server( + state, + prompt = prompt, + negative_prompt = negative_prompt, + width = width, + height = height, + steps = steps, + seed = seed, + batch_size = batch_size, + cfg_scale = cfg_scale, + flux_guidance = flux_guidance, + cancel = cancel, + ) + else: + images, seeds = self._generate_oneshot( + state, + prompt = prompt, + negative_prompt = negative_prompt, + width = width, + height = height, + steps = steps, + seed = seed, + batch_size = batch_size, + cfg_scale = cfg_scale, + flux_guidance = flux_guidance, + cancel = cancel, + ) if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # ``seeds`` is the per-image seed (each sd-cli run used seed+index), so - # the route can persist the real seed for every image in the batch. + # ``seeds`` is the per-image seed (image i used seed+i), so the route can + # persist the real seed for every image in the batch. return { "images": images, "seed": int(seed), @@ -559,6 +758,144 @@ class SdCppDiffusionBackend: if self._active_generate_cancel is cancel: self._active_generate_cancel = None + def _generate_server( + self, + state: _SdState, + *, + prompt: str, + negative_prompt: Optional[str], + width: int, + height: int, + steps: int, + seed: int, + batch_size: int, + cfg_scale: Optional[float], + flux_guidance: Optional[float], + cancel: threading.Event, + ) -> tuple[list, list[int]]: + """Generate via the resident sd-server (no model reload). + + A batch larger than the server's per-job limit is split into chunks: the server + rejects a batch_count above _MAX_SERVER_BATCH, and the one-shot path served large + batches image-by-image, so preserve that. The base seed is masked to sd.cpp's + signed-int64 range (the request model / diffusers accept larger seeds), and each + chunk is submitted at base+offset so the per-image seeds stay reproducible. Each + chunk gets a timeout proportional to its image count so a slow CPU batch is not + cancelled partway through on one fixed deadline.""" + import io + + from PIL import Image + + assert state.server is not None + total = max(1, int(batch_size)) + # sd.cpp's image seed is signed int64; mask the base (and every derived seed) so a + # large explicit seed is not rejected / wrapped inconsistently by the server. + base_seed = int(seed) & ((1 << 63) - 1) + images: list = [] + seeds: list[int] = [] + for offset in range(0, total, _MAX_SERVER_BATCH): + if cancel.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + count = min(_MAX_SERVER_BATCH, total - offset) + chunk_seed = (base_seed + offset) & ((1 << 63) - 1) + payload = build_img_gen_request( + prompt = prompt, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + seed = chunk_seed, + batch_count = count, + sample_method = state.sampling_method, + flow_shift = state.flow_shift, + cfg_scale = cfg_scale, + distilled_guidance = flux_guidance, + ) + blobs = state.server.img_gen( + payload, + on_step = self._on_log, + cancel_event = cancel, + total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count, + ) + # All-or-nothing per chunk, like the one-shot path: if the server returns fewer + # blobs than requested (e.g. one image in the batch failed to encode), fail + # rather than silently dropping images from the user's requested batch. + if not cancel.is_set() and len(blobs) != count: + raise RuntimeError( + f"sd-server returned {len(blobs)} of {count} requested images in the batch." + ) + images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs) + # sd.cpp advances the seed per image within a job, so report chunk_seed+i. + seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs))) + return images, seeds + + def _generate_oneshot( + self, + state: _SdState, + *, + prompt: str, + negative_prompt: Optional[str], + width: int, + height: int, + steps: int, + seed: int, + batch_size: int, + cfg_scale: Optional[float], + flux_guidance: Optional[float], + cancel: threading.Event, + ) -> tuple[list, list[int]]: + """Fallback path: re-run one-shot sd-cli per image (reloads the model each time).""" + import tempfile + + from PIL import Image + + engine = self._resolve_engine() + extra_args: list[str] = [] + if state.vae_format: + extra_args += ["--vae-format", state.vae_format] + if state.flow_shift is not None: + extra_args += ["--flow-shift", repr(float(state.flow_shift))] + + images = [] + seeds: list[int] = [] + with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: + for index in range(max(1, int(batch_size))): + if cancel.is_set(): + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + # Distinct seed per batch image, reproducible image-by-image from the base + # seed. Mask to int64, NOT 53 bits: the request model and the diffusers + # backend both accept large explicit seeds, so a tight 2**53 mask would + # truncate them and collide distinct requested seeds onto the same image. + seed_i = (seed + index) & ((1 << 63) - 1) + out_path = str(Path(tmpdir) / f"img_{index}.png") + params = SdCppGenParams( + prompt = prompt, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + cfg_scale = cfg_scale, + guidance = flux_guidance, + seed = seed_i, + sampling_method = state.sampling_method, + batch_count = 1, + ) + engine.generate( + state.files, + params, + output_path = out_path, + offload = list(state.offload_flags) or None, + native_speed = state.native_speed, + threads = state.threads, + extra_args = extra_args or None, + on_log = self._on_log, + cancel_event = cancel, + ) + with Image.open(out_path) as im: + images.append(im.copy()) + seeds.append(seed_i) + return images, seeds + def _on_log(self, line: str) -> None: gen = self._gen if gen is None or gen.total_steps <= 0: @@ -596,13 +933,40 @@ class SdCppDiffusionBackend: with self._lock: if self._active_generate_cancel is not None: self._active_generate_cancel.set() + state = self._state self._state = None self._load_token += 1 self._loading = None + # A load may be mid server.start() with the server not yet committed to _state; + # grab it too so we can stop it (its startup is abortable) instead of leaving it + # loading for the full startup timeout. + pending = self._pending_server + self._pending_server = None + # Stop the resident server outside the lock (terminate can take a few seconds). A + # mid-flight generation had its cancel event set above, so its poll loop unwinds + # as the process goes away. + if state is not None and state.server is not None: + state.server.stop() + if pending is not None and pending is not (state.server if state else None): + pending.stop() return self.status() def status(self) -> dict[str, Any]: state = self._state + # A resident sd-server can exit after load (OOM-killed / crashed while idle). If so, + # drop the stale loaded state so status reports not-loaded and clients reload, + # instead of every generation failing with a 500 against a dead process. + if ( + state is not None + and state.mode == "server" + and state.server is not None + and not state.server.is_alive() + ): + logger.warning("sd-server exited after load; clearing loaded state") + with self._lock: + if self._state is state: + self._state = None + state = None if state is None: return { "loaded": False, @@ -622,6 +986,7 @@ class SdCppDiffusionBackend: "attention_backend": None, "transformer_cache": None, "engine": "sd_cpp", + "native_mode": None, } return { "loaded": True, @@ -645,6 +1010,8 @@ class SdCppDiffusionBackend: "attention_backend": None, "transformer_cache": None, "engine": "sd_cpp", + # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. + "native_mode": state.mode, } diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 182b25a1b7..1c8d8c7033 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -51,6 +51,9 @@ logger = logging.getLogger(__name__) # target is ``sd-cli``; older builds shipped ``sd`` -- both are probed on PATH. _BINARY_STEM = "sd-cli" _LEGACY_STEM = "sd" +# The persistent HTTP server target (stable-diffusion.cpp ``examples/server``). It +# ships next to ``sd-cli`` in both the prebuilt archives and the cmake build tree. +_SERVER_STEM = "sd-server" class SdCppCancelled(RuntimeError): @@ -119,11 +122,11 @@ def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[ return env -def _layout_candidates(root: Path) -> list[Path]: - """sd-cli locations under a stable-diffusion.cpp checkout/install ``root``, +def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]: + """``stem`` locations under a stable-diffusion.cpp checkout/install ``root``, highest priority first: the cmake ``build/bin`` tree, then a Windows Release subdir, then the root itself.""" - name = _binary_name(_BINARY_STEM) + name = _binary_name(stem) cands = [ root / "build" / "bin" / name, root / "build" / "bin" / "Release" / name, @@ -133,66 +136,101 @@ def _layout_candidates(root: Path) -> list[Path]: return cands -def find_sd_cpp_binary() -> Optional[str]: - """Locate the ``sd-cli`` binary, or None. +def _first_file(paths: list[Path]) -> Optional[str]: + for p in paths: + try: + if p.is_file(): + return str(p) + except OSError: + continue + return None - Search order (mirrors the llama.cpp finder so a Studio install lands where - both engines look): - 1. ``SD_CLI_PATH`` env -- a direct path to the binary. + +def _find_binary( + *, direct_env: str, path_stems: tuple[str, ...], layout_stem: str +) -> Optional[str]: + """Shared finder for the stable-diffusion.cpp binaries. + + Search order (mirrors the llama.cpp finder so a Studio install lands where every + binary is looked for): + 1. ``direct_env`` -- a direct path to the binary. 2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir. 3. the installer target: ``/../stable-diffusion.cpp`` when that env (or ``STUDIO_HOME``) is set, else ``~/.unsloth/stable-diffusion.cpp``. 4. ``./stable-diffusion.cpp`` in-tree build (developer checkout). - 5. ``sd-cli`` (then legacy ``sd``) on PATH. + 5. ``path_stems`` on PATH (in order). """ - - def _first_file(paths: list[Path]) -> Optional[str]: - for p in paths: - try: - if p.is_file(): - return str(p) - except OSError: - continue - return None - # 1. Direct binary path. - env_bin = os.environ.get("SD_CLI_PATH") + env_bin = os.environ.get(direct_env) if env_bin and Path(env_bin).is_file(): return env_bin # 2. Custom install dir. custom = os.environ.get("UNSLOTH_SD_CPP_PATH") if custom: - hit = _first_file(_layout_candidates(Path(custom))) + hit = _first_file(_layout_candidates(Path(custom), layout_stem)) if hit: return hit - # 3. Default install root: the installer's default_install_dir() -- a sibling of - # the llama.cpp install under UNSLOTH_STUDIO_HOME / STUDIO_HOME when set, else - # ~/.unsloth. Mirror that env resolution or a custom Studio home never resolves. + # 3. Default install root. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME the same way + # the installer's default_install_dir does (base = the Studio home's parent), so + # a binary installed under a custom Studio root is discovered and side-by-side + # Studios stay isolated; falls back to the sibling of ~/.unsloth/llama.cpp. studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") - default_base = Path(studio_home).parent if studio_home else Path.home() / ".unsloth" - hit = _first_file(_layout_candidates(default_base / "stable-diffusion.cpp")) + default_root = ( + Path(studio_home).parent / "stable-diffusion.cpp" + if studio_home + else Path.home() / ".unsloth" / "stable-diffusion.cpp" + ) + hit = _first_file(_layout_candidates(default_root, layout_stem)) if hit: return hit # 4. In-tree developer build: /stable-diffusion.cpp. try: project_root = Path(__file__).resolve().parents[4] - hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp")) + hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp", layout_stem)) if hit: return hit except (OSError, IndexError): pass # 5. PATH. - for stem in (_BINARY_STEM, _LEGACY_STEM): + for stem in path_stems: on_path = shutil.which(stem) if on_path: return on_path return None +def find_sd_cpp_binary() -> Optional[str]: + """Locate the one-shot ``sd-cli`` binary (env ``SD_CLI_PATH``), or None. + + Probes ``sd-cli`` then legacy ``sd`` on PATH. This is the fallback engine once the + persistent ``sd-server`` exists; it also still backs the ESRGAN upscale mode. + """ + return _find_binary( + direct_env = "SD_CLI_PATH", + path_stems = (_BINARY_STEM, _LEGACY_STEM), + layout_stem = _BINARY_STEM, + ) + + +def find_sd_server_binary() -> Optional[str]: + """Locate the persistent ``sd-server`` binary (env ``SD_SERVER_PATH``), or None. + + Same precedence as ``find_sd_cpp_binary`` but keyed to the ``sd-server`` stem, so + a Studio install (prebuilt archive or cmake build, both of which ship ``sd-server`` + next to ``sd-cli``) is found in the same places. Preferred over the one-shot CLI: + it loads the model once and serves many generations without reloading from disk. + """ + return _find_binary( + direct_env = "SD_SERVER_PATH", + path_stems = (_SERVER_STEM,), + layout_stem = _SERVER_STEM, + ) + + class SdCppEngine: """A thin handle over a located ``sd-cli`` binary. diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py new file mode 100644 index 0000000000..17b269e71d --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -0,0 +1,502 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persistent ``sd-server`` (stable-diffusion.cpp) process manager. + +The native diffusion tier used to shell out to one-shot ``sd-cli`` per image, which +reloaded the multi-GB GGUF from disk every generation. ``sd-server`` (the upstream +``examples/server`` target) loads the model once at spawn and serves many generations +over HTTP, exactly like the chat backend's persistent ``llama-server``. This manager +owns ONLY the process + HTTP lifecycle; the backend (``sd_cpp_backend.py``) still owns +asset resolution, request validation, and the public Studio surface. + +Shape mirrors ``core/rag/embed_llama_server.py``: + * ``start`` -- pick a free loopback port, spawn the server (model loads here), + drain stdout on a daemon thread, poll until ready. + * ``img_gen`` -- POST ``/sdcpp/v1/img_gen`` (the whole batch in one request), + poll the async job to a terminal state, return image bytes. + * ``stop`` -- SIGTERM -> wait -> SIGKILL, join the drain thread (idempotent). + +Readiness is real: upstream ``main.cpp`` loads the model BEFORE it binds the port and +prints ``listening on:``, so a 200 from ``GET /v1/models`` (a trivial handler) means the +model is loaded; a load failure exits the process before listening and is surfaced with +the captured log tail. (The richer ``/sdcpp/v1/capabilities`` handler can block in some +builds, so it is not used for readiness.) The job JSON has no per-step field, so step progress is recovered by +parsing the server's stdout (the same ``N/M`` lines ``sd-cli`` emits), routed to the +active generation's callback. + +Import-light on purpose (no torch / diffusers / PIL), so selecting the native tier on +a CPU box never drags the GPU stack into the process. +""" + +from __future__ import annotations + +import atexit +import base64 +import logging +import shutil +import socket +import subprocess +import tempfile +import threading +import time +from collections import deque +from typing import Any, Callable, Optional + +import httpx + +from core.inference.sd_cpp_args import SdCppModelFiles, build_sd_cpp_server_command +from core.inference.sd_cpp_engine import SdCppCancelled, runtime_env +from utils.native_path_leases import child_env_without_native_path_secret +from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid +from utils.subprocess_compat import windows_hidden_subprocess_kwargs + +logger = logging.getLogger(__name__) + +# httpx transport errors meaning "the server is gone / connection refused" -- treated +# as "not ready yet" while polling readiness, and as a fatal "server died" mid-request. +_TRANSPORT_ERRORS = ( + httpx.ConnectError, + httpx.ReadError, + httpx.RemoteProtocolError, + httpx.WriteError, +) + +# Readiness probe. Upstream binds the port only AFTER the model is loaded, so any 200 +# means ready. We use /v1/models (a trivial, always-fast handler) rather than +# /sdcpp/v1/capabilities: the capabilities handler can block in some builds (it enumerates +# model metadata), which would stall readiness even though the server is up. +_READY_PATH = "/v1/models" +# Native async sdcpp API. +_IMG_GEN_PATH = "/sdcpp/v1/img_gen" +_JOBS_PATH = "/sdcpp/v1/jobs" + +_TERMINAL_OK = "completed" +_TERMINAL_FAIL = "failed" +_TERMINAL_CANCELLED = "cancelled" + +# After a cancel is requested, how long to let the server reflect it in job status before +# abandoning the poll. The native cancel is best-effort, so without this cap a server that +# ignores/loses the cancel would keep this call (and the backend's generate lock) alive +# until the job finishes naturally, blocking a superseding load from swapping the model. +_CANCEL_GRACE_S = 5.0 + + +class SdCppServer: + """A resident ``sd-server`` subprocess plus the HTTP client that drives it.""" + + def __init__( + self, + binary: str, + *, + host: str = "127.0.0.1", + ) -> None: + self.binary = binary + self.host = host + self.port: Optional[int] = None + self._process: Optional[subprocess.Popen] = None + # Fixed-size, thread-safe tail buffer: the drain thread appends while lifecycle / + # request threads read it for diagnostics, so a deque(maxlen) is safer and cheaper + # than a list with manual slicing. + self._tail: deque[str] = deque(maxlen = 200) + self._stdout_thread: Optional[threading.Thread] = None + self._lifecycle_lock = threading.Lock() + # Set (lock-free) by stop() so a blocking start()/readiness wait can be aborted + # promptly without waiting on the lifecycle lock start() holds. + self._abort = threading.Event() + # Set for the duration of a generation so the continuous stdout drain can feed + # the active request's step-progress callback; cleared in img_gen's finally. + self._step_listener: Optional[Callable[[str], None]] = None + # trust_env=False: this client only ever talks to the loopback sd-server, so it must + # not route through HTTP_PROXY/HTTPS_PROXY (a proxy without 127.0.0.1 in NO_PROXY + # would break readiness/generation). Matches the local llama-server clients. + self._client = httpx.Client(timeout = 30.0, trust_env = False) + self._scratch_dir: Optional[str] = None + self._stopped = False + atexit.register(self.stop) + + # ── lifecycle ──────────────────────────────────────────────────────────── + + @property + def base_url(self) -> str: + return f"http://{self.host}:{self.port}" + + def is_alive(self) -> bool: + return self._process is not None and self._process.poll() is None + + @staticmethod + def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + def start( + self, + files: SdCppModelFiles, + *, + vae_format: Optional[str] = None, + offload: Optional[list[str]] = None, + native_speed: Optional[str] = None, + threads: Optional[int] = None, + env: Optional[dict[str, str]] = None, + startup_timeout: float = 600.0, + ) -> None: + """Spawn the server (which loads the model) and block until it is ready. + + Raises ``RuntimeError`` (with the captured log tail) if the process exits during + startup or never answers within ``startup_timeout``. Holds the lifecycle lock so + a concurrent start/stop can't interleave. + """ + with self._lifecycle_lock: + # A stop()/unload that raced in AFTER the backend published this server as + # _pending_server but BEFORE start() took the lock has already set _abort and + # closed the httpx client. Honor that delivered stop instead of clearing the + # abort and spawning a model process the cancelled load would then leak. + if self._stopped or self._abort.is_set(): + raise SdCppCancelled("sd-server start was cancelled before launch.") + self._abort.clear() + port = self._find_free_port() + # An empty scratch dir for sd-server's LoRA / upscaler / embeddings scans + # (it recursively iterates them per request and errors on a missing dir). + self._scratch_dir = tempfile.mkdtemp(prefix = "sdcpp_dirs_") + cmd = build_sd_cpp_server_command( + self.binary, + files, + host = self.host, + port = port, + vae_format = vae_format, + offload = list(offload or []), + native_speed = native_speed, + threads = threads, + scratch_dir = self._scratch_dir, + verbose = True, # sd-server prints the per-step sampling lines we parse + ) + run_env = runtime_env(self.binary, child_env_without_native_path_secret()) + if env: + run_env.update(env) + logger.info("starting sd-server: %s", " ".join(cmd)) + # Clear in place: reassigning to [] drops the deque(maxlen=200) bound, so the + # continuous stdout drain would then grow the tail without limit for the whole + # resident-server lifetime. + self._tail.clear() + self._spawn_error: Optional[Exception] = None + spawned = threading.Event() + + # Spawn INSIDE the drain thread, which then reads stdout for the process's whole + # lifetime. child_popen_kwargs() sets PR_SET_PDEATHSIG, which on Linux is bound to + # the CREATING THREAD -- so the child must be created by a thread that outlives it, + # or a transient spawner thread ending would kill the server. The drain thread is + # exactly that long-lived owner; it dies only when the process exits or the + # interpreter goes away (the case we DO want to reap the GPU-resident server). + def _own_process() -> None: + try: + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + errors = "replace", + env = run_env, + **windows_hidden_subprocess_kwargs(), + **child_popen_kwargs(), + ) + except Exception as exc: # noqa: BLE001 -- surface the spawn failure to start() + self._spawn_error = exc + spawned.set() + return + self._process = proc + self.port = port + adopt_pid(proc.pid) # so a global shutdown sweep also reaps it + spawned.set() + self._drain_stdout(proc) + # stdout closed == the process exited; reap it so it is not left a zombie + # until the next stop()/reload. + try: + proc.wait(timeout = 5) + except Exception: # noqa: BLE001 + pass + + self._stdout_thread = threading.Thread( + target = _own_process, daemon = True, name = "sd-server-owner" + ) + self._stdout_thread.start() + spawned.wait() + if self._spawn_error is not None: + self._dispose() + raise RuntimeError(f"failed to spawn sd-server: {self._spawn_error}") + if not self._wait_ready(startup_timeout): + tail = "\n".join(list(self._tail)[-30:]) + aborted = self._abort.is_set() + self._kill_locked() + self._dispose() + if aborted: + raise SdCppCancelled("sd-server startup was cancelled.") + raise RuntimeError("sd-server failed to become ready. Last output:\n" + tail[:2000]) + + def _wait_ready( + self, + timeout: float, + interval: float = 0.5, + ) -> bool: + """Poll ``/v1/models`` until 200; bail early if the process exits. + + Upstream binds the port only AFTER the model is loaded, so a 200 here is a true + ready signal (no half-loaded race).""" + deadline = time.monotonic() + timeout + url = f"{self.base_url}{_READY_PATH}" + while time.monotonic() < deadline: + # A concurrent stop() (unload / superseding load) sets _abort so this wait can + # bail without holding the model-load hostage for the full startup_timeout. + if self._abort.is_set(): + logger.info("sd-server startup aborted before ready") + return False + if not self.is_alive(): + code = None if self._process is None else self._process.returncode + logger.error("sd-server exited early during load (code %s)", code) + return False + try: + if self._client.get(url, timeout = 2.0).status_code == 200: + return True + except (*_TRANSPORT_ERRORS, httpx.TimeoutException): + pass + time.sleep(interval) + logger.error("sd-server readiness timed out after %ss", timeout) + return False + + def _drain_stdout(self, proc: subprocess.Popen) -> None: + """Drain stdout so the pipe never deadlocks; keep a tail for diagnostics and + feed each line to the active generation's step callback.""" + try: + assert proc.stdout is not None + for raw in proc.stdout: + line = raw.rstrip() + if not line: + continue + self._tail.append(line) # deque(maxlen) discards the oldest automatically + logger.debug("[sd-server] %s", line) + cb = self._step_listener + if cb is not None: + try: + cb(line) + except Exception: # noqa: BLE001 -- a progress callback must never break drain + pass + except Exception: # noqa: BLE001 -- drain thread must never raise (pipe closed at teardown) + pass + + def stop(self) -> None: + """Terminate the server (SIGTERM -> SIGKILL), join the drain, and release the HTTP + client + atexit handler. Idempotent.""" + # Signal abort BEFORE contending for the lifecycle lock: a concurrent start() holds + # that lock for the whole (up to startup_timeout) readiness wait, so setting the + # event lets that wait bail immediately instead of stop() blocking behind it. + self._abort.set() + self._stopped = True + with self._lifecycle_lock: + self._kill_locked() + self._dispose() + + def _dispose(self) -> None: + """Release per-instance resources (on stop / failed start). The backend never + reuses a disposed server, so this closes the pooled httpx client and drops the + atexit handler that would otherwise pin every reloaded instance for the session.""" + try: + atexit.unregister(self.stop) + except Exception: # noqa: BLE001 + pass + try: + self._client.close() + except Exception: # noqa: BLE001 + pass + if self._scratch_dir: + shutil.rmtree(self._scratch_dir, ignore_errors = True) + self._scratch_dir = None + + def _kill_locked(self) -> None: + proc = self._process + if proc is None: + return + pid = proc.pid + try: + proc.terminate() + proc.wait(timeout = 5) + except subprocess.TimeoutExpired: + logger.warning("sd-server did not exit on SIGTERM; killing") + try: + proc.kill() + proc.wait(timeout = 5) + except Exception: # noqa: BLE001 -- best-effort teardown + pass + except Exception as exc: # noqa: BLE001 + logger.warning("error terminating sd-server: %s", exc) + finally: + forget_pid(pid) + self._process = None + self.port = None + if self._stdout_thread is not None: + self._stdout_thread.join(timeout = 2) + self._stdout_thread = None + + # ── generation ─────────────────────────────────────────────────────────── + + def img_gen( + self, + payload: dict[str, Any], + *, + on_step: Optional[Callable[[str], None]] = None, + cancel_event: Optional[threading.Event] = None, + poll_interval: float = 0.4, + submit_timeout: float = 60.0, + total_timeout: float = 1800.0, + ) -> list[bytes]: + """Submit one async ``img_gen`` job, poll it to completion, return image bytes. + + ``on_step`` receives each server stdout line (for the step bar). ``cancel_event``, + when set, cancels the job via the native endpoint and raises ``SdCppCancelled``. + Raises ``RuntimeError`` on submit/poll failures (including the server dying), with + the log tail attached. + """ + # If the server was already stopped for a cancel/unload/superseding load that set + # the cancel event before this submit began, report it as a cancellation (which the + # route maps to a client-state 409) rather than a generic "server died" 500. + if self._stopped or not self.is_alive(): + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + raise RuntimeError("sd-server is not running.") + + self._step_listener = on_step + job_id: Optional[str] = None + try: + # Submit -> 202 Accepted + job id. + try: + resp = self._client.post( + f"{self.base_url}{_IMG_GEN_PATH}", json = payload, timeout = submit_timeout + ) + except (*_TRANSPORT_ERRORS, httpx.TimeoutException) as exc: + raise RuntimeError(self._died_message("img_gen submit", exc)) from exc + if resp.status_code == 429: + raise RuntimeError("sd-server job queue is full (HTTP 429).") + if resp.status_code not in (200, 202): + raise RuntimeError( + f"sd-server img_gen submit -> {resp.status_code}: {resp.text[:500]}" + ) + try: + job = resp.json() + except ValueError as exc: + raise RuntimeError( + f"sd-server img_gen returned a non-JSON submit response: {exc}" + ) from exc + if not isinstance(job, dict): + raise RuntimeError( + f"sd-server img_gen returned an unexpected submit response type: {type(job)}" + ) + job_id = job.get("id") + if not job_id: + raise RuntimeError(f"sd-server img_gen returned no job id: {job}") + + # Poll the job to a terminal state. + deadline = time.monotonic() + total_timeout + cancel_sent_at: Optional[float] = None + while True: + if cancel_event is not None and cancel_event.is_set(): + if cancel_sent_at is None: + self.cancel(job_id) + cancel_sent_at = time.monotonic() + elif time.monotonic() - cancel_sent_at > _CANCEL_GRACE_S: + # The best-effort cancel was not reflected in job status within the + # grace window; abandon the poll so the caller can stop the server + # instead of holding the generate lock until the job finishes. + raise SdCppCancelled("sd-server generation was cancelled.") + if not self.is_alive(): + # If we're unwinding a cancel (e.g. unload killed the server), surface a + # clean cancellation rather than a generic "server died" error. + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + raise RuntimeError(self._died_message("img_gen poll", None)) + if time.monotonic() > deadline: + # Best-effort cancel, then tear the server down: current sd-server does + # not interrupt an already-generating job (cancel_generating=false / 409), + # so leaving it up would keep denoising the abandoned job and block later + # generations/reloads behind it. Stopping frees the slot; the backend sees + # the dead server on the next generate and takes the recoverable reload path. + self.cancel(job_id) + self.stop() + raise RuntimeError(f"sd-server generation timed out after {total_timeout}s") + try: + jr = self._client.get(f"{self.base_url}{_JOBS_PATH}/{job_id}", timeout = 10.0) + except (*_TRANSPORT_ERRORS, httpx.TimeoutException): + time.sleep(poll_interval) + continue + except RuntimeError as exc: + # A concurrent stop()/unload closes the shared httpx client; httpx then + # raises a plain RuntimeError ("client has been closed") that is NOT a + # transport error. When we are being cancelled, report it as a clean + # cancellation (route -> 409) instead of a generic 500 generation failure. + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") from exc + raise + if jr.status_code in (404, 410): + raise RuntimeError(f"sd-server job {job_id} is gone (HTTP {jr.status_code}).") + if jr.status_code != 200: + time.sleep(poll_interval) + continue + try: + jd = jr.json() + except ValueError as exc: + raise RuntimeError(f"sd-server job status was not JSON: {exc}") from exc + if not isinstance(jd, dict): + raise RuntimeError( + f"sd-server job status returned an unexpected response type: {type(jd)}" + ) + status = jd.get("status") + if status == _TERMINAL_OK: + return self._decode_images(jd) + if status == _TERMINAL_FAIL: + err = jd.get("error") or {} + raise RuntimeError( + "sd-server generation failed: " + f"{err.get('code', 'error')}: {err.get('message', '')}".strip() + ) + if status == _TERMINAL_CANCELLED: + raise SdCppCancelled("sd-server generation was cancelled.") + time.sleep(poll_interval) + finally: + self._step_listener = None + + def cancel(self, job_id: str) -> None: + """Best-effort native cancel of an in-flight job.""" + try: + self._client.post(f"{self.base_url}{_JOBS_PATH}/{job_id}/cancel", timeout = 5.0) + except Exception: # noqa: BLE001 -- cancel is best-effort + pass + + @staticmethod + def _decode_images(job: dict[str, Any]) -> list[bytes]: + # Defensive against an unexpected response shape (a misbehaving/older server): + # verify each level is the type we index before calling dict/list methods. + result = job.get("result") if isinstance(job, dict) else None + images = result.get("images") if isinstance(result, dict) else None + items = [it for it in images if isinstance(it, dict)] if isinstance(images, list) else [] + out: list[bytes] = [] + for item in sorted(items, key = lambda d: d.get("index", 0)): + b64 = item.get("b64_json") + if not b64: + continue + try: + out.append(base64.b64decode(b64)) + except Exception as exc: # noqa: BLE001 + raise RuntimeError(f"sd-server returned an undecodable image: {exc}") from exc + if not out: + raise RuntimeError("sd-server completed the job but returned no images.") + return out + + def _died_message(self, where: str, exc: Optional[Exception]) -> str: + tail = "\n".join(list(self._tail)[-20:]) + base = f"sd-server connection lost during {where}" + if not self.is_alive(): + code = None if self._process is None else self._process.returncode + base += f" (process exited, code {code})" + if exc is not None: + base += f": {exc}" + if tail: + base += "\nLast output:\n" + tail[:1500] + return base diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 04b70e7e11..57ce68b1cb 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1924,6 +1924,11 @@ class DiffusionStatusResponse(BaseModel): ) transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp") + native_mode: Optional[str] = Field( + None, + description = "Native sd.cpp execution mode: server (resident sd-server) | oneshot " + "(per-image sd-cli) | null (diffusers engine)", + ) fallback_reason: Optional[str] = Field( None, description = "Why diffusers was chosen over the native sd.cpp engine (null when none)", diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 898711e8bd..f4cebc893d 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -32,6 +32,10 @@ def _clean_env_and_state(monkeypatch): "get_active_diffusion_engine", lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}), ) + # Default: no resident sd-server (so existing tests exercise the sd-cli path only) and + # a stubbed runnability probe, so neither reaches the real install/exec path. + monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: None) + monkeypatch.setattr(r, "_server_binary_runnable", lambda *_a, **_k: True) yield @@ -70,6 +74,16 @@ def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch): assert r.active_engine_name() == ENGINE_SD_CPP +def test_cpu_with_only_sd_server_picks_sd_cpp(monkeypatch): + # An sd-server-only install (no runnable sd-cli) must still route to native: the + # backend prefers the resident server, so a runnable sd-server is native availability. + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, None) # no sd-cli + monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None)) + monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: "/usr/bin/sd-server") + assert _select() == ENGINE_SD_CPP + + def test_present_but_not_runnable_binary_falls_back(monkeypatch): # A binary that exists but cannot run (version() -> None) must fall back to # diffusers at selection, not commit native and fail inside the load. diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index d3133b08d2..126eed4855 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -21,7 +21,9 @@ from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, SdCppUpscaleParams, + build_img_gen_request, build_sd_cpp_command, + build_sd_cpp_server_command, build_sd_cpp_upscale_command, native_speed_flags, offload_flags, @@ -364,3 +366,104 @@ def test_build_upscale_requires_input_and_model(): SdCppUpscaleParams(input_image = "/i.png", upscale_model = ""), output_path = "/o.png", ) + + +# ── sd-server spawn command ────────────────────────────────────────────────── + + +def test_server_command_has_model_and_listen_but_no_request_params(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf") + cmd = build_sd_cpp_server_command( + "/bin/sd-server", files, host = "127.0.0.1", port = 5678, vae_format = "flux2" + ) + assert _pair(cmd, "--diffusion-model") == "/m/z.gguf" + assert _pair(cmd, "--vae") == "/m/ae.sft" + assert _pair(cmd, "--llm") == "/m/q.gguf" + assert _pair(cmd, "--vae-format") == "flux2" + assert _pair(cmd, "--listen-ip") == "127.0.0.1" + assert _pair(cmd, "--listen-port") == "5678" + # Per-request parameters must NOT be baked into the spawn command. + for flag in ( + "--prompt", + "--seed", + "--steps", + "--cfg-scale", + "--guidance", + "--width", + "--height", + "--batch-count", + ): + assert flag not in cmd + + +def test_server_command_maps_offload_and_speed_and_dedupes(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_server_command( + "/bin/sd-server", + files, + host = "127.0.0.1", + port = 1, + offload = ["--offload-to-cpu", "--diffusion-fa"], + native_speed = "default", # would add --diffusion-fa again + threads = 8, + ) + assert _pair(cmd, "--threads") == "8" + assert cmd.count("--diffusion-fa") == 1 # de-duped against offload + assert "--offload-to-cpu" in cmd + + +def test_server_command_scratch_dir_expands_to_lora_upscaler_embd(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_server_command( + "/bin/sd-server", files, host = "127.0.0.1", port = 1, scratch_dir = "/tmp/scratch" + ) + assert _pair(cmd, "--lora-model-dir") == "/tmp/scratch" + assert _pair(cmd, "--hires-upscalers-dir") == "/tmp/scratch" + assert _pair(cmd, "--embd-dir") == "/tmp/scratch" + # Absent when not requested. + bare = build_sd_cpp_server_command("/bin/sd-server", files, host = "127.0.0.1", port = 1) + assert "--lora-model-dir" not in bare and "--hires-upscalers-dir" not in bare + + +def test_server_command_requires_diffusion_model(): + with pytest.raises(ValueError): + build_sd_cpp_server_command( + "/bin/sd-server", SdCppModelFiles(diffusion_model = ""), host = "127.0.0.1", port = 1 + ) + + +# ── img_gen request body ───────────────────────────────────────────────────── + + +def test_img_gen_request_maps_core_fields(): + req = build_img_gen_request( + prompt = "a fox", + negative_prompt = "blurry", + width = 512, + height = 768, + steps = 8, + seed = 42, + batch_count = 3, + sample_method = "euler", + cfg_scale = 4.0, + ) + assert req["prompt"] == "a fox" and req["negative_prompt"] == "blurry" + assert req["width"] == 512 and req["height"] == 768 + assert req["seed"] == 42 and req["batch_count"] == 3 + assert req["sample_params"]["sample_steps"] == 8 + assert req["sample_params"]["sample_method"] == "euler" + assert req["sample_params"]["guidance"]["txt_cfg"] == 4.0 + assert req["output_format"] == "png" + + +def test_img_gen_request_flux_uses_distilled_guidance(): + req = build_img_gen_request(prompt = "x", steps = 4, distilled_guidance = 3.5, flow_shift = 3.0) + g = req["sample_params"]["guidance"] + assert g["distilled_guidance"] == 3.5 + assert "txt_cfg" not in g + assert req["sample_params"]["flow_shift"] == 3.0 + + +def test_img_gen_request_requires_prompt(): + with pytest.raises(ValueError): + build_img_gen_request(prompt = " ", steps = 4) diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 46a0514332..fb64f1324b 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -77,10 +77,69 @@ def _loaded_backend(fam_name = "z-image", engine = None): vae_format = fam.sd_cpp_vae_format, sampling_method = fam.sd_cpp_sampling_method, flow_shift = fam.sd_cpp_flow_shift, + mode = "oneshot", # this fixture injects an engine, so it exercises the one-shot path ) return b +class _FakeServer: + """Stands in for SdCppServer: records the spawn + one img_gen per whole batch.""" + + def __init__(self, binary): + self.binary = binary + self.started = None + self.stopped = False + self.payloads = [] + self.timeouts = [] + self.alive = True + + def is_alive(self): + return self.alive and not self.stopped + + def start( + self, + files, + *, + vae_format = None, + offload = None, + native_speed = None, + threads = None, + ): + self.started = dict( + files = files, + vae_format = vae_format, + offload = offload, + native_speed = native_speed, + threads = threads, + ) + + def img_gen( + self, + payload, + *, + on_step = None, + cancel_event = None, + total_timeout = None, + ): + import io as _io + + self.payloads.append(payload) + self.timeouts.append(total_timeout) + if on_step is not None: + steps = payload.get("sample_params", {}).get("sample_steps", 0) + on_step(f" {steps}/{steps}") + n = int(payload.get("batch_count", 1)) + blobs = [] + for i in range(n): + buf = _io.BytesIO() + Image.new("RGB", (1, 1), (i, i, i)).save(buf, format = "PNG") + blobs.append(buf.getvalue()) + return blobs + + def stop(self): + self.stopped = True + + # ── asset resolution ────────────────────────────────────────────────────────── @@ -321,6 +380,260 @@ def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF" +# ── persistent sd-server mode ────────────────────────────────────────────────── + + +def test_resolve_backend_prefers_server(monkeypatch): + b = SdCppDiffusionBackend() # no injected engine + monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server") + mode, binary, engine = b._resolve_backend() + assert mode == "server" and binary == "/x/sd-server" and engine is None + + +def test_resolve_backend_injected_engine_forces_oneshot(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + mode, binary, engine = b._resolve_backend() + assert mode == "oneshot" and binary is None and engine is not None + + +def test_resolve_backend_falls_back_to_oneshot_without_server(monkeypatch): + b = SdCppDiffusionBackend() + monkeypatch.setattr(bk, "find_sd_server_binary", lambda: None) + monkeypatch.setattr(bk, "_install_allowed", lambda: False) # don't attempt a real install + monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli") + mode, binary, engine = b._resolve_backend() + assert mode == "oneshot" and engine is not None + + +def test_resolve_backend_cached_fallback_engine_does_not_pin_oneshot(monkeypatch): + # A lazily cached fallback engine (NOT an explicit injection) must not force one-shot: + # once a server is available again, the next load can use it. + b = SdCppDiffusionBackend() # no injected engine + b._engine = _FakeEngine() # simulate a prior lazy one-shot fallback caching the engine + monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server") + mode, binary, engine = b._resolve_backend() + assert mode == "server" and binary == "/x/sd-server" and engine is None + + +def _run_server_load( + monkeypatch, + b, + servers, + fam_name = "z-image", +): + fam = detect_family(fam_name) + monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server") + # The fake binary path is not a real executable; skip the up-front runnability probe. + monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True) + + def _factory(binary): + s = _FakeServer(binary) + servers.append(s) + return s + + monkeypatch.setattr(bk, "SdCppServer", _factory) + monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: []) + monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None) + monkeypatch.setattr( + b, + "_fetch_assets", + lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"}, + ) + monkeypatch.setattr( + bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu") + ) + b._load_token = 1 + b._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z.gguf", + base = fam.base_repo, + fam = fam, + hf_token = None, + _load_token = 1, + ) + + +def test_server_load_spawns_once_and_status_reports_mode(monkeypatch): + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + assert len(servers) == 1 + assert servers[0].started is not None # the model is loaded once, at spawn + assert b._state is not None and b._state.mode == "server" and b._state.server is servers[0] + assert b.status()["native_mode"] == "server" + + +def test_server_generate_uses_one_request_for_whole_batch(monkeypatch): + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 7, batch_size = 3) + assert len(out["images"]) == 3 + assert all(isinstance(im, Image.Image) for im in out["images"]) + # ONE job for the whole batch (no per-image model reload), unlike the one-shot path. + assert len(servers[0].payloads) == 1 + assert servers[0].payloads[0]["batch_count"] == 3 + assert out["seed"] == 7 and out["seeds"] == [7, 8, 9] + # step progress was driven from the server's stdout line. + assert b._gen is None # cleared after generate + + +def test_server_generate_splits_batches_above_server_limit(monkeypatch): + # A batch above the server's per-job limit is chunked (the one-shot path did these + # image-by-image); each chunk gets a timeout proportional to its image count. + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 100, batch_size = 10) + assert len(out["images"]) == 10 + counts = [p["batch_count"] for p in servers[0].payloads] + assert counts == [bk._MAX_SERVER_BATCH, 10 - bk._MAX_SERVER_BATCH] # [8, 2] + # Each chunk's timeout scales with its image count, not one fixed batch deadline. + assert servers[0].timeouts == [ + bk._SERVER_PER_IMAGE_TIMEOUT_S * 8, + bk._SERVER_PER_IMAGE_TIMEOUT_S * 2, + ] + # Seeds run contiguously across chunks (chunk 2 submitted at base + 8). + assert out["seeds"] == list(range(100, 110)) + assert servers[0].payloads[1]["seed"] == 108 + + +def test_server_generate_masks_large_seed(monkeypatch): + # sd.cpp's image seed is signed int64; a larger explicit seed must be masked before it + # reaches the server (the request model / diffusers accept up to 2**64 - 1). + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 2**64 - 1, batch_size = 1) + assert servers[0].payloads[0]["seed"] <= (1 << 63) - 1 + assert all(s <= (1 << 63) - 1 for s in out["seeds"]) + + +def test_status_clears_when_server_died(monkeypatch): + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + assert b.status()["loaded"] is True + servers[0].alive = False # the resident server crashed / was OOM-killed + st = b.status() + assert st["loaded"] is False + assert b._state is None # stale state was dropped so clients reload + + +def test_server_generate_progress_from_stdout(monkeypatch): + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + + seen = {} + + class _WatchServer(_FakeServer): + def img_gen( + self, + payload, + *, + on_step = None, + cancel_event = None, + total_timeout = None, + ): + on_step(" 4/8") + seen["mid"] = b.generate_progress() + return super().img_gen( + payload, on_step = on_step, cancel_event = cancel_event, total_timeout = total_timeout + ) + + b._state = bk._SdState( + repo_id = b._state.repo_id, + base_repo = b._state.base_repo, + family = b._state.family, + device = b._state.device, + files = b._state.files, + vae_format = b._state.vae_format, + sampling_method = b._state.sampling_method, + flow_shift = b._state.flow_shift, + server = _WatchServer("/x/sd-server"), + mode = "server", + ) + b.generate(prompt = "x", steps = 8, seed = 1) + assert seen["mid"]["step"] == 4 and seen["mid"]["total_steps"] == 8 + + +def test_server_unload_stops_server(monkeypatch): + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + st = b.unload() + assert st["loaded"] is False + assert servers[0].stopped is True + assert b._state is None + + +def test_server_reload_stops_old_server_before_new(monkeypatch): + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + # A second load must tear down the first server and start a fresh one. + b._load_token = 2 + fam = detect_family("z-image") + b._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z.gguf", + base = fam.base_repo, + fam = fam, + hf_token = None, + _load_token = 2, + ) + assert len(servers) == 2 + assert servers[0].stopped is True # old server stopped + assert b._state.server is servers[1] and servers[1].stopped is False + + +def test_server_start_failure_falls_back_to_oneshot(monkeypatch): + # A present-but-broken sd-server must not fail the load when sd-cli works. + b = SdCppDiffusionBackend() + monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server") + # Probe passes; the failure we exercise here is in start(), not the up-front probe. + monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True) + + class _BadServer: + def __init__(self, binary): + self.stopped = False + + def start(self, *a, **k): + raise RuntimeError("sd-server broken") + + def stop(self): + self.stopped = True + + monkeypatch.setattr(bk, "SdCppServer", _BadServer) + fake = _FakeEngine() + monkeypatch.setattr(b, "_resolve_engine", lambda: fake) + monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: []) + monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None) + monkeypatch.setattr( + b, + "_fetch_assets", + lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"}, + ) + monkeypatch.setattr( + bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu") + ) + fam = detect_family("z-image") + b._load_token = 1 + b._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z.gguf", + base = fam.base_repo, + fam = fam, + hf_token = None, + _load_token = 1, + ) + assert b._state is not None and b._state.mode == "oneshot" and b._state.server is None + # and it can still generate via the one-shot engine + out = b.generate(prompt = "x", steps = 4, seed = 1) + assert len(out["images"]) == 1 and len(fake.calls) == 1 + + def test_run_load_redacts_paths_in_progress_error(monkeypatch): # A load failure surfaced via load_progress() must run through redact_native_paths, the # same scrub the diffusers load path applies, so a registered native path can't leak. diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 5c9d3256af..daaf5223a8 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -24,6 +24,7 @@ from core.inference.sd_cpp_engine import ( ENGINE_SD_CPP, SdCppEngine, find_sd_cpp_binary, + find_sd_server_binary, runtime_env, select_diffusion_engine, ) @@ -75,6 +76,58 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch): assert find_sd_cpp_binary() is None +# ── sd-server discovery ────────────────────────────────────────────────────── + + +def _clear_server_env(monkeypatch): + monkeypatch.delenv("SD_SERVER_PATH", raising = False) + monkeypatch.delenv("SD_CLI_PATH", raising = False) + monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False) + + +def test_find_server_prefers_sd_server_path_env(tmp_path, monkeypatch): + _clear_server_env(monkeypatch) + binary = tmp_path / "sd-server" + binary.write_text("x") + monkeypatch.setenv("SD_SERVER_PATH", str(binary)) + monkeypatch.setattr(eng.shutil, "which", lambda *_a: None) + assert find_sd_server_binary() == str(binary) + + +def test_find_server_build_layout(tmp_path, monkeypatch): + _clear_server_env(monkeypatch) + root = tmp_path / "sdcpp" + built = root / "build" / "bin" / "sd-server" + built.parent.mkdir(parents = True) + built.write_text("x") + monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root)) + monkeypatch.setattr(eng.shutil, "which", lambda *_a: None) + assert find_sd_server_binary() == str(built) + + +def test_find_server_path_fallback(tmp_path, monkeypatch): + _clear_server_env(monkeypatch) + monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome")) + monkeypatch.setattr( + eng.shutil, "which", lambda stem: "/usr/bin/sd-server" if stem == "sd-server" else None + ) + assert find_sd_server_binary() == "/usr/bin/sd-server" + + +def test_find_server_not_confused_with_sd_cli(tmp_path, monkeypatch): + # A tree that has only sd-cli must NOT be reported as an sd-server (and vice versa), + # so the backend correctly falls back to one-shot when only the CLI is present. + _clear_server_env(monkeypatch) + root = tmp_path / "sdcpp" + (root / "build" / "bin").mkdir(parents = True) + (root / "build" / "bin" / "sd-cli").write_text("x") + monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root)) + monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome")) + monkeypatch.setattr(eng.shutil, "which", lambda *_a: None) + assert find_sd_server_binary() is None + assert find_sd_cpp_binary() == str(root / "build" / "bin" / "sd-cli") + + # ── availability / version ────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_sd_cpp_server.py b/studio/backend/tests/test_sd_cpp_server.py new file mode 100644 index 0000000000..9e040814e9 --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_server.py @@ -0,0 +1,417 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the persistent sd-server process manager (SdCppServer). + +Hermetic: subprocess.Popen and the httpx client are faked, so nothing spawns a real +binary or opens a socket beyond the free-port probe.""" + +from __future__ import annotations + +import base64 +import io +import threading + +import pytest +from PIL import Image + +from core.inference import sd_cpp_server as srv +from core.inference.sd_cpp_args import SdCppModelFiles +from core.inference.sd_cpp_engine import SdCppCancelled +from core.inference.sd_cpp_server import SdCppServer + +_FILES = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/vae.sft", llm = "/m/llm.sft") + + +def _png_b64(shade: int) -> str: + buf = io.BytesIO() + Image.new("RGB", (1, 1), (shade, shade, shade)).save(buf, format = "PNG") + return base64.b64encode(buf.getvalue()).decode() + + +class _FakePopen: + """Minimal Popen stand-in. stdout yields the scripted lines then BLOCKS until the + process is terminated/killed/exited -- mirroring a real child that holds its pipe + open for its lifetime (so the owner/drain thread stays alive, as in production).""" + + def __init__( + self, + lines = (), + exit_code = None, + ): + self.pid = 4242 + self._lines = list(lines) + self._exit = exit_code # None == alive + self.returncode = exit_code + self.terminated = False + self.killed = False + self._done = threading.Event() + if exit_code is not None: + self._done.set() + + @property + def stdout(self): + def _gen(): + for ln in self._lines: + yield ln + self._done.wait() # hold the pipe open until the process ends + + return _gen() + + def poll(self): + return self._exit + + def terminate(self): + self.terminated = True + self._exit = 0 + self.returncode = 0 + self._done.set() + + def wait(self, timeout = None): + self._done.wait(timeout) + if self._exit is None: + self._exit = 0 + self.returncode = 0 + return self.returncode + + def kill(self): + self.killed = True + self._exit = -9 + self.returncode = -9 + self._done.set() + + +class _Resp: + def __init__( + self, + status_code, + payload = None, + text = "", + bad_json = False, + ): + self.status_code = status_code + self._payload = payload if payload is not None else {} + self.text = text + self._bad_json = bad_json + + def json(self): + if self._bad_json: + raise ValueError("not json") + return self._payload + + +class _FakeClient: + def __init__( + self, + *, + get = None, + post = None, + ): + self._get = get or (lambda url: _Resp(200, {})) + self._post = post or (lambda url, json: _Resp(202, {"id": "job1"})) + self.get_urls = [] + self.post_calls = [] + self.closed = False + + def get( + self, + url, + timeout = None, + ): + self.get_urls.append(url) + return self._get(url) + + def post( + self, + url, + json = None, + timeout = None, + ): + self.post_calls.append((url, json)) + return self._post(url, json) + + def close(self): + self.closed = True + + +@pytest.fixture +def patched(monkeypatch): + """Neutralise process-lifetime side effects for the manager under test.""" + monkeypatch.setattr(srv, "adopt_pid", lambda pid: None) + monkeypatch.setattr(srv, "forget_pid", lambda pid: None) + monkeypatch.setattr(srv, "child_popen_kwargs", lambda: {}) + monkeypatch.setattr(srv, "windows_hidden_subprocess_kwargs", lambda: {}) + return monkeypatch + + +def _server_with(popen, client): + s = SdCppServer("/x/sd-server") + s._client = client + # Attach the fake process + port so generation tests can run without start(). + s._process = popen + s.port = 1234 + return s + + +# ── start / readiness ────────────────────────────────────────────────────────── + + +def test_start_becomes_ready_when_capabilities_200(patched): + popen = _FakePopen(lines = ["loading model", "listening on: http://127.0.0.1:1"]) + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + s = _server_with( + popen, _FakeClient(get = lambda url: _Resp(200, {"model": {"path": "/m/z.gguf"}})) + ) + s.start(_FILES, startup_timeout = 5.0) + assert s.is_alive() is True + assert s.port is not None + + +def test_start_fails_fast_when_process_exits(patched): + # Model load failed -> process exits before listening; start must raise with the tail. + popen = _FakePopen(lines = ["error: bad model"], exit_code = 1) + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + # Capabilities never answers (connection refused) -> readiness relies on exit detection. + s = _server_with( + popen, _FakeClient(get = lambda url: (_ for _ in ()).throw(srv.httpx.ConnectError("refused"))) + ) + with pytest.raises(RuntimeError, match = "failed to become ready"): + s.start(_FILES, startup_timeout = 2.0) + + +# ── generation ─────────────────────────────────────────────────────────────── + + +def _completed_job(images_b64): + return _Resp( + 200, + { + "status": "completed", + "result": {"images": [{"index": i, "b64_json": b} for i, b in enumerate(images_b64)]}, + }, + ) + + +def test_img_gen_returns_image_bytes_in_index_order(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobA"}), + # result images deliberately out of order -> manager must sort by index. + get = lambda url: _Resp( + 200, + { + "status": "completed", + "result": { + "images": [ + {"index": 1, "b64_json": _png_b64(200)}, + {"index": 0, "b64_json": _png_b64(50)}, + ] + }, + }, + ), + ), + ) + blobs = s.img_gen({"prompt": "x", "batch_count": 2, "sample_params": {"sample_steps": 4}}) + assert len(blobs) == 2 + first = Image.open(io.BytesIO(blobs[0])).convert("RGB").getpixel((0, 0)) + assert first == (50, 50, 50) # index 0 first + + +def test_img_gen_failed_job_raises(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobF"}), + get = lambda url: _Resp( + 200, {"status": "failed", "error": {"code": "x", "message": "boom"}} + ), + ), + ) + with pytest.raises(RuntimeError, match = "generation failed.*boom"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_queue_full_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(429, text = "busy"))) + with pytest.raises(RuntimeError, match = "queue is full"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_cancel_posts_cancel_and_raises(patched): + popen = _FakePopen() + cancel = threading.Event() + cancel.set() # already cancelled before the first poll + client = _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobC"}), + get = lambda url: _Resp( + 200, {"status": "cancelled", "error": {"code": "cancelled", "message": "c"}} + ), + ) + s = _server_with(popen, client) + with pytest.raises(SdCppCancelled): + s.img_gen({"prompt": "x"}, cancel_event = cancel) + assert any(url.endswith("/cancel") for url, _ in client.post_calls) + + +def test_img_gen_detects_server_death(patched): + popen = _FakePopen() + + def _die_get(url): + popen._exit = 137 # the process died between submit and poll + return _Resp(200, {"status": "generating"}) + + s = _server_with( + popen, _FakeClient(post = lambda url, json: _Resp(202, {"id": "jobD"}), get = _die_get) + ) + with pytest.raises(RuntimeError, match = "connection lost|process exited"): + s.img_gen({"prompt": "x"}) + + +# ── stdout routing + stop ────────────────────────────────────────────────────── + + +def test_drain_routes_lines_to_step_listener_and_tail(patched): + s = SdCppServer("/x/sd-server") + seen = [] + s._step_listener = seen.append + # exit_code set so stdout ends after the scripted lines (a live fake would block). + s._drain_stdout(_FakePopen(lines = ["sampling 1/8", "", "sampling 8/8", "done"], exit_code = 0)) + assert "sampling 1/8" in seen and "sampling 8/8" in seen + assert "" not in seen # blank lines skipped + assert s._tail[-1] == "done" + + +def test_stop_is_idempotent_and_terminates(patched): + popen = _FakePopen() + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + client = _FakeClient(get = lambda url: _Resp(200, {})) + s = _server_with(popen, client) + s.start(_FILES, startup_timeout = 5.0) + s.stop() + assert popen.terminated is True + assert s.is_alive() is False + assert client.closed is True # stop() releases the pooled HTTP client + s.stop() # second call must not raise + + +def test_img_gen_submit_error_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(400, text = "bad params"))) + with pytest.raises(RuntimeError, match = "submit -> 400"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_malformed_submit_json_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, bad_json = True))) + with pytest.raises(RuntimeError, match = "non-JSON submit"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_empty_result_raises(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobE"}), + get = lambda url: _Resp(200, {"status": "completed", "result": {"images": []}}), + ), + ) + with pytest.raises(RuntimeError, match = "no images"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_rejected_after_stop(patched): + popen = _FakePopen() + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {}))) + s.start(_FILES, startup_timeout = 5.0) + s.stop() + with pytest.raises(RuntimeError, match = "not running"): + s.img_gen({"prompt": "x"}) + + +# ── cancellation + defensive parsing (review follow-ups) ─────────────────────── + + +def test_img_gen_cancelled_before_submit_reports_cancellation(patched): + # The server was stopped for a cancel/unload before submit; with the cancel event set + # this must surface as a cancellation (route -> 409), not a generic "not running" 500. + popen = _FakePopen() + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {}))) + s.start(_FILES, startup_timeout = 5.0) + s.stop() + cancel = threading.Event() + cancel.set() + with pytest.raises(SdCppCancelled): + s.img_gen({"prompt": "x"}, cancel_event = cancel) + + +def test_img_gen_abandons_when_cancel_not_honored(patched): + # A best-effort cancel the server ignores must not pin this call (and the generate + # lock) until natural completion: after the grace window it raises cancellation. + patched.setattr(srv, "_CANCEL_GRACE_S", 0.0) + popen = _FakePopen() + cancel = threading.Event() + cancel.set() + client = _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobG"}), + get = lambda url: _Resp(200, {"status": "generating"}), # never terminal + ) + s = _server_with(popen, client) + with pytest.raises(SdCppCancelled): + s.img_gen({"prompt": "x"}, cancel_event = cancel, poll_interval = 0.01) + + +def test_img_gen_non_dict_submit_json_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, ["not", "a", "dict"]))) + with pytest.raises(RuntimeError, match = "unexpected submit response"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_non_dict_status_json_raises(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobH"}), + get = lambda url: _Resp(200, ["unexpected"]), + ), + ) + with pytest.raises(RuntimeError, match = "unexpected response type"): + s.img_gen({"prompt": "x"}, poll_interval = 0.01) + + +def test_decode_images_tolerates_unexpected_shapes(): + # A misbehaving/older server can return non-dict result/images/items; _decode_images + # must raise a clean "no images" rather than an AttributeError on .get(). + for job in ({"result": ["x"]}, {"result": {"images": "nope"}}, {"result": {"images": [1, 2]}}): + with pytest.raises(RuntimeError, match = "no images"): + SdCppServer._decode_images(job) + + +def test_start_aborted_by_concurrent_stop(patched): + # A stop() during the readiness wait must abort start() promptly (without waiting out + # the startup timeout) and surface as a cancellation. + popen = _FakePopen(lines = ["loading model"]) + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + + def _never_ready(url): + raise srv.httpx.ConnectError("refused") + + s = _server_with(popen, _FakeClient(get = _never_ready)) + + def _stop_soon(): + import time as _t + _t.sleep(0.2) + s.stop() + + threading.Thread(target = _stop_soon, daemon = True).start() + with pytest.raises(SdCppCancelled): + s.start(_FILES, startup_timeout = 30.0) diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 4b2b36f95a..c76aee11f8 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -140,6 +140,17 @@ def _locate_sd_cli(root: Path) -> Optional[Path]: return None +def _locate_sd_server(root: Path) -> Optional[Path]: + """The persistent ``sd-server`` binary in the extracted tree, if the archive ships + one (modern stable-diffusion.cpp releases do). Best-effort: the native backend + falls back to one-shot ``sd-cli`` when it is absent.""" + name = "sd-server.exe" if sys.platform == "win32" else "sd-server" + for p in root.rglob(name): + if p.is_file(): + return p + return None + + def _download( url: str, dest: Path, @@ -237,6 +248,13 @@ def install( if sys.platform != "win32": _make_executable(sd_cli) print(f"installed sd-cli -> {sd_cli}", flush = True) + # The same archive ships the persistent sd-server; make it runnable too so the + # native backend can prefer it (load once, serve many) over one-shot sd-cli. + sd_server = _locate_sd_server(target) + if sd_server is not None and sys.platform != "win32": + _make_executable(sd_server) + if sd_server is not None: + print(f"installed sd-server -> {sd_server}", flush = True) return sd_cli From 467e74baeedb5a3e9b4a70ae223911887e2b957e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:45:14 -0700 Subject: [PATCH 03/17] Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI redesign (#6769) * Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict Phase 1 of porting the richer diffusion stack onto the image-generation backend. - Add a compartmentalized device/dtype policy module (diffusion_device.py) resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or fp32, never a silent fp16 that renders a black image. - Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved float16 to float32 for those families so they do not produce black images. - Split the backend locks: a generation holds only _generate_lock, so status, unload, and a new load are never blocked by a long denoise. Add per-generation cancellation via callback_on_step_end so an eviction or a superseding load preempts a running generation; a replacement load waits for it to stop before allocating, so two pipelines never sit in VRAM at once. - Validate a load request before the GPU handoff so an unloadable pick never evicts a working chat model, and reject missing local paths up front. - Add CPU-only tests for the device policy, dtype guard, lock split and cancellation, and validate-before-evict, plus a GPU benchmark/regression script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR against a saved reference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and VAE tiling/slicing from measured free device memory vs the model's estimated resident footprint, then applies it to the built pipeline. auto stays resident when the model fits (byte-identical to the prior resident path), and falls to whole-module offload when tight; fast/balanced/low_vram are explicit overrides. Sequential submodule offload is unreliable for GGUF transformers on diffusers 0.38, so it falls back to whole-module offload and status reports the policy actually engaged. Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM 47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost. 73 prior + 35 new CPU tests pass. * Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level, use_stream) that keeps the transformer flowing through the GPU a few blocks at a time while the text encoder / VAE stay resident, and fix VAE tiling to drive the VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so status never overstates either, and group falls back to whole-module offload when the transformer can't be streamed. Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 -> 2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names now match that tradeoff: balanced = stream the transformer, low_vram = offload every component. auto picks group when the companions fit resident, else model. 112 CPU tests pass. * Studio diffusion (Phase 5): image quality-vs-quant accuracy harness Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold prompt + seed fixed, render a grid with a reference quant (default BF16), then render each candidate quant and measure drift from the reference. Records mean PSNR + SSIM (pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity (transformers, --clip), plus file size, latency, and peak VRAM, then prints a quality-vs-cost table and recommends the smallest quant within a quality budget. --selftest validates the metrics on synthetic images with no GPU or model. Verified on Z-Image (B200): the table degrades monotonically with quant size (Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat (~0.34) -- quantization erodes fine detail far more than prompt adherence. * Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32) Add a speed_mode knob (off by default, so the render path stays bit-identical): default applies channels_last VAE + regional torch.compile of the denoiser's repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional compile is gated off for the GGUF transformer (dequantises per-op) and for families flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image), so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed optims run before placement/offload, per the diffusers composition order. status now reports speed_mode + the optims actually engaged. Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last', 'tf32'], compile correctly skipped for GGUF; generation works in every mode. 121 CPU tests pass. * Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3) storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16 compute dtype while normalisations and embeddings stay full precision. Applied before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder dense). status reports which encoders were cast. Verified on Z-Image (B200, balanced/group mode where the encoder stays resident): generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload) at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off by default and documented as such, with the Phase 5 harness to size the cost. 127 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob) Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant (fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8 stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run before placement; status reports the mode actually engaged. This is the lean realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the 3045-line port. Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default); size it per model with the Phase 5 quality harness. diffusion_bench gains --text-encoder-quant. 129 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm / XPU; this covers the hardware diffusers serves poorly, consuming the same split GGUF assets Studio already curates. - sd_cpp_args.py: pure sd-cli command builder. Maps the family to its text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1 CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential) to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu / --vae-tiling / --diffusion-fa), so one user knob drives both engines. - sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary() with the same precedence as the llama finder (env override, then the Studio install root, then in-tree, then PATH), an is_available/version probe, and a one-shot subprocess generate that streams progress and returns the PNG. runtime_env() prepends the binary's directory to the platform library path so a prebuilt's bundled libstable-diffusion.so resolves. select_diffusion_engine() is the pure routing decision (GPU backends to diffusers, CPU/MPS to native when present). - install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the Studio install root. resolve_release_asset() is a pure, unit-tested host-to-asset matrix. - scripts/sd_cpp_smoke.py: end-to-end native generation harness. Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine, routing, runtime env, and the installer resolver. Full diffusion suite 166 passing. Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears) The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with 'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2 *_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing min_features filter did not exclude them. Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed, guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers (M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and quantises these layers fine, so the exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note 'context_embedder' contains the substring 'text_embed', which is why the token is the specific 'time_text_embed', not 'text_embed'. Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s (5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already working) are unchanged. - diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes exclude_name_tokens; quantize_transformer passes it for int8 only. - hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps attention / FFN / sequence-embedder linears), while fp8 keeps them. - scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's Linear layers and derive the exclusion list. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears) The prequant-checkpoint builder applied the dense quant filter without the int8-only M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8 checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now produces a working checkpoint on every supported model, giving int8 (the consumer-preferred scheme) the same ~2x load-VRAM and download reduction fp8 already had. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the native stable-diffusion.cpp engine instead of diffusers, with diffusers as the guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM. - diffusion_engine_router: centralised engine selection (built on the existing select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason. - sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface backed by sd-cli, with lazy binary install, registry-driven asset fetch, step-progress parsing, and cancellation. - diffusion_families: per-family single-file VAE + text-encoder asset mapping. - sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled). - routes/inference + gpu_arbiter: drive the active engine via the router; the API now reports the active engine and any fallback reason. - tests for the backend, router, route selection, and cancellation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler Address review feedback on #6724: - engine router: unload the engine being deactivated on a switch, so the old model is not left resident-but-unreachable (the evictor only targets the active engine). - generate route: sd.cpp execution errors (nonzero exit / timeout / missing output) now map to 500, not 409 (which only means not-loaded / cancelled). - native batch: return per-image seeds and persist the actual seed for each image so every batch image is reproducible. - Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the stable-diffusion.cpp docs; other families keep sd-cli defaults. - honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload offload flags on the native load instead of hardcoding them off. - fail the load when the sd-cli binary is present but not runnable (version() now returns None on exec error / nonzero exit). - size estimate: only treat the transformer asset as a possible local path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF - apply_step_cache now engages only via the transformer's native enable_cache (the diffusers CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model now runs uncached per the best-effort contract. - GGUF transformers are quantized (the default Studio load path), so they now use the higher quantized FBCache threshold when the caller leaves it unset, instead of the dense default that could keep the cache from triggering. - fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so fullgraph=True failed warmup and silently measured an eager cached run); output dir is now relative to the script, not a hardcoded path. * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs without producing output (or without closing stdout) would never reach proc.wait and the wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the PROCESS, so the main thread always enforces the timeout and kills a hung process (which closes the pipe and ends the reader). Add a test that times out even when stdout blocks, and make the no-binary test hermetic so a host-installed sd-cli can't leak in. * Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn The filter callback can be invoked without a module name, so fqn.lower() would raise AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens, so the linear is kept) instead of crashing the quantise pass. * Studio diffusion (Phase 16) review fixes: native engine robustness - sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64); a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds aliased to the same image. Random seeds stay 53-bit (JS-safe). - sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub fall back to anonymous instead of failing auth on a blank token. - sd_cpp_backend: a superseding load now cancels the in-flight generation, so the old sd-cli can no longer return/persist an image from the previous model. - diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a slow 10+ GB free / CUDA sync does not block engine selection. - diffusion_engine_router: probe sd-cli runnability (version()) before committing to native, so a present-but-unrunnable binary falls back to diffusers at selection. - diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a CPU-only install can still reach the native sd.cpp engine instead of failing load. - tests updated for the runnability probe + a not-runnable fallback case. * Studio diffusion (Phase 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening - install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket); extract through a per-member containment check (Zip-Slip guard); expanduser the --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch the separately-published cudart runtime DLL archive so sd-cli.exe can start. - sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie. - tests: Zip-Slip rejection, normal extraction, studio-home discovery. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes the numbered files _ (base_0.png, base_1.png, ...) instead of the literal --output path. SdCppEngine.generate checked only the literal path, so a batch generation would exit 0 and then raise 'no image' (or return a stale file). generate now returns the literal path when present and otherwise falls back to the numbered siblings; single-image behavior is unchanged. Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected without error. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder Codex review: the M=1 modulation/embedder exclusion was wired only into the dense runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint quantised the AdaLN modulation and conditioning-embedder linears, and loading it via transformer_prequant_path (the load path only loads already-quantised tensors, it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes for the runtime path. Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 -> the M=1 exclusion, every other scheme -> none) and use it in both the runtime quantiser and the builder, so a prequant artifact's quantised-layer set always matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion). Test: int8_exclude_name_tokens returns the exclusion for int8 and () for fp8/nvfp4/mxfp8. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race Codex review on the native-engine routing: - The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat) unconditionally after engine selection. A native sd.cpp load on a pure-CPU host never touches the GPU, so that needlessly tore down the resident chat model. The handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it. - sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so the setting was unverifiable. status now derives them from state.offload_flags (still 'none' on CPU, where the flags are empty). - _run_load committed the new state without cancelling/waiting on a generation that started during the (slow) asset download, so a stale sd-cli run against the OLD model could finish afterward and persist an image from the previous model once the new load reported ready. The commit now signals the in-flight cancel and waits on _generate_lock before swapping _state (taken only at commit, so the download never serialises against generation), mirroring the diffusers load path. Tests: CPU native load skips the arbiter while a GPU native load takes it; status reports offload active when flags are set; _run_load cancels and waits for an in-flight generation before committing. * Studio diffusion (Phase 14) review round 2: align helper name with the stack Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the identical helper already present higher in the diffusion stack (Phase 16). The helper definition, the runtime quantiser call, and the offline builder are now byte-identical to that version, so the two branches no longer introduce a divergent name for the same single-source-of-truth and the stack merges without a conflict on this fix. No behavior change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion: eager patches + torch.compile cache speed phase Adds the opt-in speed path for the GGUF diffusion transformer behind a selectable speed mode (default off, so output is unchanged until a profile is chosen): - diffusion_eager_patches.py: shared eager fast-paths (channels_last, attention/backend selection, fused norms and QKV) installed at load and rolled back on unload or failed load. - diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent torch.compile cache and the GGUF-transformer compile wiring. - diffusion_arch_patches.py: architecture-specific patches. - diffusion_patch_backend.py: shared install/restore plumbing. - diffusion_speed.py: speed-profile planning. Tests for each module plus the benchmarking and probe scripts used to measure speed, memory, and accuracy of the path. * Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI Backend: - Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file fp8 transformers, gated to the unsloth org plus a curated allowlist. - Image-conditioned workflows built with Pipeline.from_pipe so they reuse the loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint, outpaint, and a hires-fix upscale pass. - Instruction editing as its own family kind (Qwen-Image-Edit-2511, FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and multi-reference) plus klein inpaint. - Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the decoded image size and cap upscale output to avoid OOM on large inputs. - Fixes: from_pipe defaulting to a float32 recast that crashed torchao quantized transformers; image-conditioned calls forcing the slider size onto the input image. Native sd.cpp engine rejects image-conditioned and reference requests it cannot serve. Frontend: - Redesigned Images page with capability-gated workflow tabs (Create, Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor, client-side outpaint, and a multi-reference picker. - Advanced options moved to a right-docked panel mirroring Chat: closed by default, toggled by a single fixed top-bar button that stays in place. sd.cpp installer: pin the release, verify each download's sha256, add a download timeout, and make the source repo configurable for a future mirror. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio Images: correct the Advanced panel comment (closed by default, fixed toggle) * Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a stray "gguf" tag on the Hub but ship no .gguf files. The model search classified them as GGUF from the bare tag, so the picker rendered the GGUF variant expander, which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata (populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected. * Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline load when the id is unsloth-hosted or on-device (the backend infers the family + base repo and gates loads to unsloth/* or local paths), and show a clear message otherwise instead of silently ignoring the click. Curated and GGUF paths are unchanged. * Studio Images: keep curated safetensors models in Recommended after download The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet), so they vanished from the picker after the first load and could only be found by typing an exact search. The row already renders a downloaded badge, matching how GGUF Recommended rows stay visible when cached. Drop the exclusion so the curated safetensors always list. * Studio Images: clarify the GGUF transformer-quant Advanced control Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode" with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/ FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF, falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only. * Studio Images: list on-device unsloth diffusion models in the picker The Images picker's On Device tab hid every non-GGUF cached repo whenever a task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit and FP8 safetensors) never showed up there. List cached repos that pass the task gate, limited under a filter to unsloth-hosted ones so base repos (which fail the diffusion load trust gate) don't appear only to dead-end on click. Chat behavior is unchanged: the task gate still drops image repos there. * Studio: hide single-file image checkpoints from the chat model picker The chat picker treats a cached repo as an image model, and hides it, only when it ships a diffusers model_index.json. Single-file, ComfyUI, and ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models. Fall back to resolving the repo id against the known diffusion families, the same resolver the Images backend loads from, so these checkpoints are tagged text-to-image and stay in the Images picker only. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio Images: add the FLUX.2-dev model family Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline) resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline + Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated, reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror. text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from klein's turbo defaults. Verified live: GGUF load resolves the family + gated base repo and generates a real 1024x1024 image on GPU. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio Images: clearer error for an unsupported diffusion model When a repo id resolves to no diffusion family the load raised 'Could not infer a diffusion family... Pass family_override (z-image)', which points at an unrelated family and doesn't say what is supported. Replace it with a message that lists the supported families (from a new supported_family_names helper) and notes that video models and image models whose diffusers transformer has no single-file loader are not supported. Applies to both the diffusers and native sd.cpp load paths. Also refreshes two stale family-registry comments that still called FLUX.2-dev omitted. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove stray async task scratch outputs committed by mistake * Diffusion: guard trust check against OSError and validate conditioning inputs - _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid characters (or a bare owner/name id) can't raise OSError; treat any failure as not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request still raises the clear FileNotFoundError for a genuinely missing local pick. - generate(): reject mask_image / upscale / reference_images supplied without an input image, and reject reference_images on a family that does not support reference conditioning, instead of silently degrading to txt2img / img2img. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address Codex review findings on the image-workflows PR Keep diffusion.py importable without torch: the compile/arch patch modules import torch at module level, so import them lazily at their load/unload call sites instead of at module load. This restores the torchless contract so get_diffusion_backend() works on a CPU/native sd.cpp install. Match family reject keywords and aliases as whole path/name segments, not raw substrings, so an unrelated word like edited, edition, or kontextual no longer misroutes or hides a valid base image model, while supported edit families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same segment matching in the picker task filter. Route FLUX.2-dev native guidance through --guidance like the other FLUX families rather than --cfg-scale. Reject native upscale requests that have no input image. Read image header dimensions and reject over-limit inputs before decoding pixels, so a crafted small-payload image cannot spike memory. Reject an upscale that would shrink the source below its input size. Validate the model_kind against the filename extension before the GPU handoff. Estimate a local diffusers pipeline's size from its on-disk weights so auto memory planning does not skip offload and OOM. Report workflows: [txt2img] from the native backend status so the Create tab stays enabled for a loaded native model. Clamp the outpaint canvas to the backend's 4096px decode limit. Adds regression tests for segment matching and kind/extension validation. * Address further Codex findings on the image-workflows PR - Persist the actual output image size in the gallery recipe instead of the request sliders: Transform/Inpaint/Edit derive the size from the uploaded image, Extend grows the canvas, and Upscale resizes it, so the sliders recorded (and later restored) the wrong dimensions for those workflows. - Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file name) in validate_load_request, so the unloadable pick fails before chat is evicted rather than deep in from_pretrained. - Only publish an image-conditioned from_pipe wrapper to the shared aux cache when the load is still current: from_pipe runs under the generate lock but not the state lock, so an unload racing its construction could otherwise cache a wrapper over torn-down modules that a later load would reuse. - Verify the Windows CUDA runtime archive checksum before extracting it, like the main sd-cli archive, so a corrupt or tampered runtime is rejected rather than extracted next to the binary. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/core/inference/diffusion.py | 890 +++++++++++-- .../core/inference/diffusion_arch_patches.py | 530 ++++++++ .../core/inference/diffusion_compile_cache.py | 351 +++++ .../core/inference/diffusion_eager_patches.py | 214 ++++ .../core/inference/diffusion_engine_router.py | 13 +- .../core/inference/diffusion_families.py | 181 ++- .../core/inference/diffusion_gguf_compile.py | 131 ++ .../core/inference/diffusion_memory.py | 11 + .../core/inference/diffusion_patch_backend.py | 68 + .../backend/core/inference/diffusion_speed.py | 91 +- studio/backend/core/inference/sd_cpp_args.py | 1 + .../backend/core/inference/sd_cpp_backend.py | 47 +- studio/backend/models/inference.py | 79 +- studio/backend/routes/inference.py | 41 +- studio/backend/routes/models.py | 18 +- studio/backend/tests/conftest.py | 6 + .../tests/test_diffusion_arch_patches.py | 274 ++++ .../backend/tests/test_diffusion_backend.py | 735 ++++++++++- .../tests/test_diffusion_compile_cache.py | 238 ++++ .../tests/test_diffusion_eager_patches.py | 190 +++ .../tests/test_diffusion_gguf_compile.py | 73 ++ studio/backend/tests/test_diffusion_routes.py | 32 +- studio/backend/tests/test_diffusion_speed.py | 67 +- studio/backend/tests/test_sd_cpp_install.py | 142 ++ .../assistant-ui/model-selector/pickers.tsx | 90 +- .../hub/hooks/use-hub-model-search.ts | 14 +- studio/frontend/src/features/images/api.ts | 41 +- .../src/features/images/images-page.tsx | 1139 ++++++++++++++++- studio/install_sd_cpp_prebuilt.py | 128 +- 29 files changed, 5556 insertions(+), 279 deletions(-) create mode 100644 studio/backend/core/inference/diffusion_arch_patches.py create mode 100644 studio/backend/core/inference/diffusion_compile_cache.py create mode 100644 studio/backend/core/inference/diffusion_eager_patches.py create mode 100644 studio/backend/core/inference/diffusion_gguf_compile.py create mode 100644 studio/backend/core/inference/diffusion_patch_backend.py create mode 100644 studio/backend/tests/test_diffusion_arch_patches.py create mode 100644 studio/backend/tests/test_diffusion_compile_cache.py create mode 100644 studio/backend/tests/test_diffusion_eager_patches.py create mode 100644 studio/backend/tests/test_diffusion_gguf_compile.py diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 1680ce0288..eb4bb33034 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -3,12 +3,17 @@ """Local diffusion (text-to-image) backend. -A torch-only singleton: it dequantises a single-file GGUF on-device via -``GGUFQuantizationConfig`` and pulls the rest of the pipeline (VAE, text -encoders, scheduler) from the matching base repo. torch/diffusers are imported -lazily so this stays importable in a no-torch runtime. ``begin_load`` runs on a -background thread; poll ``load_progress`` for the download bar. GPU-handoff -policy lives in the arbiter the routes call, not here. +A torch-only singleton that loads one of three "kinds" (see ``resolve_model_kind``): +a single-file GGUF transformer dequantised on-device via ``GGUFQuantizationConfig``, +a single-file safetensors transformer (e.g. fp8), or a full diffusers pipeline via +``from_pretrained`` (which re-applies an embedded quant config such as bnb-4bit). The +single-file kinds pull the rest of the pipeline (VAE, text encoders, scheduler) from +the matching base repo; the pipeline kind pulls everything from the repo itself. +Non-GGUF kinds are gated to the ``unsloth/*`` org (or a local path) for safety. + +torch/diffusers are imported lazily so this stays importable in a no-torch runtime. +``begin_load`` runs on a background thread; poll ``load_progress`` for the download +bar. GPU-handoff policy lives in the arbiter the routes call, not here. """ from __future__ import annotations @@ -30,6 +35,7 @@ from .diffusion_families import ( detect_family_for_pick, resolve_base_repo, resolve_local_gguf_child, + supported_family_names, ) from .diffusion_device import ( DiffusionDeviceTarget, @@ -41,14 +47,17 @@ from .diffusion_memory import ( apply_memory_plan, estimate_gguf_resident_mib, estimate_image_runtime_mib, + estimate_safetensors_dense_mib, file_size_mib, plan_diffusion_memory, snapshot_device_memory, ) from .diffusion_speed import ( SPEED_DEFAULT, + SPEED_MAX, SPEED_OFF, apply_speed_optims, + compile_eligible, resolve_speed_mode, restore_backend_flags, snapshot_backend_flags, @@ -57,6 +66,8 @@ from .diffusion_attention import ( apply_attention_backend, select_attention_backend, ) +from . import diffusion_compile_cache as compile_cache +from . import diffusion_gguf_compile as gguf_compile from .diffusion_cache import apply_step_cache from .diffusion_precision import quantize_text_encoders from .diffusion_prequant import ( @@ -74,6 +85,124 @@ from .diffusion_transformer_quant import ( logger = get_logger(__name__) +# A load resolves to exactly one of these "kinds", which decide how the transformer +# (and the rest of the pipeline) is built: +# "gguf" -- a single-file GGUF transformer dequantised on-device via +# GGUFQuantizationConfig; the VAE / text encoders / scheduler come +# from the companion base diffusers repo. The original behaviour. +# "single_file" -- a single-file *.safetensors transformer loaded with from_single_file +# WITHOUT the GGUF dequant config (e.g. an fp8 checkpoint); companions +# still come from the base repo. +# "pipeline" -- a full diffusers repo loaded with pipeline_cls.from_pretrained(repo_id), +# which pulls every component (transformer included) and re-applies any +# embedded quantization_config (e.g. a bnb-4bit pipeline) automatically. +_MODEL_KINDS = frozenset({"gguf", "single_file", "pipeline"}) + + +def resolve_model_kind(gguf_filename: Optional[str], model_kind: Optional[str] = None) -> str: + """Classify a load request into one of ``_MODEL_KINDS``. + + An explicit ``model_kind`` wins (validated). Otherwise the kind is inferred from + the single-file name: a ``.gguf`` name is ``"gguf"``, any other single-file name is + ``"single_file"``, and the absence of a name is a full ``"pipeline"`` load. Pure and + network-free, so the route, validation, and load paths all agree on the kind.""" + if model_kind: + kind = model_kind.strip().lower() + if kind not in _MODEL_KINDS: + raise ValueError( + f"Unknown model_kind '{model_kind}'. Expected one of {sorted(_MODEL_KINDS)}." + ) + return kind + name = (gguf_filename or "").strip() + if not name: + return "pipeline" + if name.lower().endswith(".gguf"): + return "gguf" + return "single_file" + + +def _decode_b64_image(data: str, *, mode: str = "RGB") -> Any: + """Decode a base64 (optionally ``data:`` URL) image string to a PIL image. + + The image-conditioned workflows (img2img / inpaint / edit) transport the input + image and mask as base64 in the JSON request, so this is the single decode path. + A mask is decoded as single-channel ``L``; the source image as ``RGB``.""" + import base64 + import binascii + import io + + from PIL import Image + + raw = data.strip() + if raw.startswith("data:"): + # data:[][;base64], + _, _, raw = raw.partition(",") + try: + blob = base64.b64decode(raw, validate = False) + except (binascii.Error, ValueError) as exc: + raise ValueError(f"Invalid base64 image data: {exc}") from exc + # Bound the decoded size. Every image-conditioned workflow (img2img / inpaint / upscale / + # reference / edit) decodes through here, so this single guard protects init, mask, and + # each reference image uniformly. PIL only WARNS in its 89-178MP "decompression bomb" soft + # zone and still loads (~0.5 GB RGB each, times up to 4 with multi-reference); cap the side + # well below that. 4096px covers txt2img's 2048 max, upscales, and normal outpaint canvases; + # anything larger is rejected with a clear 400 instead of risking an OOM. + max_side = 4096 + try: + img = Image.open(io.BytesIO(blob)) + # Read the declared dimensions from the header (Image.open is lazy) and reject an + # over-limit image BEFORE img.load() decompresses its pixels, so a crafted + # small-payload/huge-dimension file can't spike memory before the guard runs. + w, h = img.size + if w > max_side or h > max_side: + raise ValueError(f"Image is too large ({w}x{h}); maximum is {max_side}px per side.") + img.load() + except ValueError: + raise # the size guard's own message; don't wrap it as a decode error + except Exception as exc: # noqa: BLE001 — surfaced as a 400 to the client + raise ValueError(f"Could not decode image: {exc}") from exc + return img.convert(mode) + + +def _snap_to_multiple(img: Any, multiple: int = 16) -> Any: + """Resize a PIL image so both sides are multiples of ``multiple`` (rounded to nearest, + minimum one multiple), preserving content with a high-quality resample. + + Image-conditioned pipelines (Z-Image / Qwen / FLUX: 8x VAE downsample + 2x patch) reject + sizes that are not divisible by 16. Rather than error on an odd-sized upload, snap it so + the workflow just works; rounding to nearest keeps the rescale minimal/accurate.""" + from PIL import Image + + w, h = img.size + nw = max(multiple, int(round(w / multiple)) * multiple) + nh = max(multiple, int(round(h / multiple)) * multiple) + if (nw, nh) != (w, h): + img = img.resize((nw, nh), Image.LANCZOS) + return img + + +def _is_trusted_diffusion_repo(repo_id: str) -> bool: + """Whether a NON-GGUF load is allowed for ``repo_id``. + + Making ``gguf_filename`` optional opens a ``from_pretrained`` / ``from_single_file`` + on an arbitrary repo, which fetches and deserialises third-party weights. So the + non-GGUF paths are gated to the ``unsloth/*`` org (the curated safetensors models) and + to local paths the user explicitly pointed at (already on their disk). The GGUF path + is unchanged and stays open to any repo, as before. + + A bare ``owner/name`` HF id is never a real filesystem path, and an id with invalid + characters makes ``Path.exists()`` raise OSError; treat any such failure as "not a + local path" so the trust decision falls through to the unsloth/ check (the loader's + validate_load_request raises the clear FileNotFoundError for a genuinely missing + local pick).""" + try: + if Path(repo_id).expanduser().exists(): + return True + except OSError: + pass + return repo_id.strip().lower().startswith("unsloth/") + + @dataclass(frozen = True) class _LoadState: """Everything about the currently-loaded pipeline, swapped as one unit.""" @@ -90,6 +219,10 @@ class _LoadState: offload_policy: str = OFFLOAD_NONE vae_tiling: bool = False memory_mode: str = "auto" + # The resolved load kind: "gguf" | "single_file" | "pipeline". Surfaced in status so the + # UI can gate GGUF-only controls (the dense transformer_quant fast path only engages on + # the gguf kind; on single_file/pipeline it is a silent no-op). + kind: str = "gguf" # The opt-in speed profile (Phase 3). speed_mode: str = SPEED_OFF speed_optims: tuple = () @@ -107,6 +240,12 @@ class _LoadState: attention_backend: Optional[str] = None # Step cache engaged ("fbcache") or None. Opt-in, for many-step models. transformer_cache: Optional[str] = None + # Shared eager monkey-patches (diffusion_eager_patches) installed for this load (any + # non-off speed tier). Uninstalled on unload so a later `off` load is bit-identical. + eager_patched: bool = False + # Pre-warmed torch.compile cache context (diffusion_compile_cache.CacheContext) when a + # compiled tier ran, else None. Carries the per-key inductor dir + bundle for save/restore. + compile_cache_ctx: Any = None @dataclass @@ -183,6 +322,11 @@ class DiffusionBackend: # The callback mutates _gen and generate_progress() reads it, both lock-free, # so per-step progress polling stays live during a generation. self._gen: Optional[_GenState] = None + # Cache of image-conditioned workflow pipelines (img2img / inpaint) built via + # Pipeline.from_pipe around the loaded text-to-image pipe. They share its already + # resident modules (no extra VRAM, no reload), so we build each once per load and + # reuse it. Keyed by pipeline class name; cleared on unload with the base pipe. + self._aux_pipes: dict[str, Any] = {} @property def is_loaded(self) -> bool: @@ -280,28 +424,30 @@ class DiffusionBackend: *, gguf_filename: Optional[str] = None, family_override: Optional[str] = None, + model_kind: Optional[str] = None, ) -> DiffusionFamily: """Cheap, network-free validation shared by the route (before it evicts the - chat model) and both load paths, so an unloadable pick fails BEFORE the GPU - handoff. Raises ValueError for a missing gguf_filename or undetectable - family, and ValueError/FileNotFoundError for a bad local GGUF path. Touches - no GPU, network, or state.""" - if not gguf_filename: - raise ValueError( - "gguf_filename is required: this backend loads single-file GGUF checkpoints only." - ) - # Reject a non-GGUF single-file name (e.g. README.md, config.json) here, before - # the route hands the GPU over: without this a family-looking repo_id paired with - # a non-GGUF filename passes preflight, evicts the chat model, and only fails in - # the background from_single_file -- exactly the eviction this validation prevents. - if not gguf_filename.lower().endswith(".gguf"): - raise ValueError( - f"gguf_filename must name a .gguf single-file checkpoint; got '{gguf_filename}'." - ) + chat model) and the load paths, so an unloadable pick fails BEFORE the GPU + handoff. Resolves the load kind (gguf / single_file / pipeline), then raises + ValueError for a missing single-file name, a non-unsloth non-GGUF repo, or an + undetectable family, and ValueError/FileNotFoundError for a bad local path. + Touches no GPU, network, or state.""" + kind = resolve_model_kind(gguf_filename, model_kind) fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: raise ValueError( - f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)." + f"'{repo_id}' is not a supported diffusion image model. Supported families: " + f"{', '.join(supported_family_names())}. If this is a variant of one of them, " + f"pass family_override with that family name. (Video models and image models " + f"whose diffusers transformer has no single-file loader are not supported.)" + ) + # Non-GGUF loads (a single-file safetensors transformer, or a full pipeline) + # are gated to the unsloth org or a local path -- they fetch + deserialise + # weights, so an arbitrary remote repo is rejected here, before any work. + if kind != "gguf" and not _is_trusted_diffusion_repo(repo_id): + raise ValueError( + 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." ) # 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. @@ -309,15 +455,58 @@ class DiffusionBackend: # missing one is an error here; a bare "org/name" id is a remote HF repo and # is left for the background load to resolve. local_root = Path(repo_id).expanduser() - if local_root.exists(): - resolve_local_gguf_child(local_root, gguf_filename) - elif ( - # POSIX path-shaped, a "."/".." prefix (covers ./ ../ and their Windows .\ ..\ - # forms), a Windows separator anywhere (never present in a bare "org/name" HF - # id), or an absolute path on this OS. + # POSIX path-shaped, a "."/".." prefix (covers ./ ../ and their Windows .\ ..\ + # forms), a Windows separator anywhere (never present in a bare "org/name" HF + # id), or an absolute path on this OS. + path_shaped = ( repo_id.startswith(("/", "\\", "~", ".")) or "\\" in repo_id or local_root.is_absolute() - ): - raise FileNotFoundError(f"Local model path does not exist: {repo_id}") + ) + if kind in ("gguf", "single_file"): + if not gguf_filename: + raise ValueError(f"a single-file checkpoint name is required for a '{kind}' load.") + # Fail a kind/extension mismatch here (before the route evicts chat and grabs the + # GPU), instead of deep in the background from_single_file: a "gguf" load needs a + # .gguf file, and a "single_file" load must not be handed a .gguf. + is_gguf_name = gguf_filename.lower().endswith(".gguf") + if kind == "gguf" and not is_gguf_name: + raise ValueError("a 'gguf' load requires a .gguf checkpoint name.") + if kind == "single_file" and is_gguf_name: + raise ValueError("a .gguf checkpoint needs model_kind 'gguf', not 'single_file'.") + # A single-file load must name an actual checkpoint: an arbitrary repo file + # (README.md, config.json) would pass preflight, evict the chat model, and + # only fail in the background from_single_file -- the eviction this + # validation exists to prevent. + if kind == "single_file" and not gguf_filename.lower().endswith(".safetensors"): + raise ValueError( + f"'{gguf_filename}' is not a loadable single-file checkpoint " + f"(expected a .safetensors name; use a .gguf name for a GGUF load)." + ) + if local_root.exists(): + resolve_local_gguf_child(local_root, gguf_filename) + elif path_shaped: + raise FileNotFoundError(f"Local model path does not exist: {repo_id}") + else: # pipeline + if gguf_filename: + raise ValueError( + "a 'pipeline' load takes a full diffusers repo, not a single-file name." + ) + if local_root.exists(): + if not (local_root / "model_index.json").exists(): + raise FileNotFoundError( + f"Local pipeline directory has no model_index.json: {repo_id}" + ) + elif path_shaped: + raise FileNotFoundError(f"Local model path does not exist: {repo_id}") + elif repo_id.upper().endswith("-GGUF"): + # A remote "*-GGUF" id is a single-file GGUF repo, not a full diffusers + # pipeline: loading it as a pipeline passes the trusted-repo check, evicts + # chat, then fails in the background when from_pretrained finds no + # model_index.json. Reject the certain case here (no network round-trip) + # so the bad pick fails before the GPU handoff, as the route expects. + raise ValueError( + f"'{repo_id}' is a single-file GGUF repo; load it with model_kind 'gguf' " + f"and a .gguf filename, not as a full pipeline." + ) return fam # ── Background load + progress ───────────────────────────────────────── @@ -340,13 +529,17 @@ class DiffusionBackend: attention_backend: Optional[str] = None, transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, + model_kind: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" # 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 fam = self.validate_load_request( - repo_id, gguf_filename = gguf_filename, family_override = family_override + repo_id, + gguf_filename = gguf_filename, + family_override = family_override, + model_kind = model_kind, ) with self._lock: @@ -380,6 +573,7 @@ class DiffusionBackend: attention_backend = attention_backend, transformer_cache = transformer_cache, transformer_cache_threshold = transformer_cache_threshold, + model_kind = model_kind, _load_token = token, ), daemon = True, @@ -395,21 +589,30 @@ class DiffusionBackend: fam = detect_family_for_pick( kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override") ) - base = _resolve_base_repo( - kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token") - ) + kind = resolve_model_kind(kwargs.get("gguf_filename"), kwargs.get("model_kind")) + if kind == "pipeline": + # The full pipeline IS the repo: from_pretrained pulls every component + # (transformer included) from it, so the base repo is the repo itself. + base = kwargs["repo_id"] + else: + base = _resolve_base_repo( + kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token") + ) kwargs["base_repo"] = base expected, base_files = self._estimate_download_bytes( kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token"), + kind = kind, # The dense transformer-quant path downloads the base repo's # transformer/ shards via from_pretrained(subfolder="transformer") # INSIDE the locked finalize phase, where unload/cancellation cannot # preempt the multi-GB pull. When that path can actually run, pull the - # shards here in the preemptible prefetch instead. - include_transformer = self._dense_quant_prefetch_needed(fam, kwargs), + # shards here in the preemptible prefetch instead. (Pipeline loads + # already include transformer/ via their own filter.) + include_transformer = kind == "gguf" + and self._dense_quant_prefetch_needed(fam, kwargs), ) with self._lock: # Stamp progress only if this load is still current; a superseding @@ -454,7 +657,11 @@ class DiffusionBackend: if loading is None: return _progress("ready" if self._state is not None else None) - downloaded = self._cache_bytes(loading.repo_id) + self._cache_bytes(loading.base_repo) + # Sum the checkpoint repo + companion base cache. For a full-pipeline load the + # base IS the repo, so count it once (else the bar double-counts to "finalizing"). + downloaded = self._cache_bytes(loading.repo_id) + if loading.base_repo and loading.base_repo != loading.repo_id: + downloaded += self._cache_bytes(loading.base_repo) expected = loading.expected_bytes # Downloads done but pipeline still dequantising / moving to GPU. The cache # scan can slightly exceed the estimate (extra cached quants, blob padding), @@ -483,16 +690,29 @@ class DiffusionBackend: base_repo: str, hf_token: Optional[str], *, + kind: str = "gguf", include_transformer: bool = False, ) -> tuple[int, list[str]]: """Total download size for the progress bar, plus the base-repo files to - fetch (the prefetch reuses this list, so the base is listed only once).""" + fetch (the prefetch reuses this list, so the base is listed only once). + + For a ``pipeline`` load the whole repo IS the pipeline (``base_repo`` is the + repo itself), so the transformer/ subfolder is INCLUDED -- unlike the GGUF / + single-file paths, where the transformer is the single file and the base repo + supplies only the companions.""" from huggingface_hub import HfApi api = HfApi() total = 0 base_files: list[str] = [] try: + if kind == "pipeline": + info = api.model_info(repo_id, files_metadata = True, token = hf_token) + for s in info.siblings: + if _pipeline_file_downloaded(s.rfilename): + base_files.append(s.rfilename) + total += s.size or 0 + return total, base_files # Skip the Hub size lookup for a LOCAL gguf path: model_info(repo_id) would # raise on a filesystem path and (caught below) skip the base-repo lookup too, # so the companion VAE/text-encoder files would never be prefetched and would @@ -525,6 +745,30 @@ class DiffusionBackend: return 0 # repo not in cache yet return total + @staticmethod + def _local_dir_weight_bytes(path: Path, *, exclude_transformer: bool) -> int: + """Sum the on-disk weight files under a local diffusers directory. The HF blob + cache is empty for a local path, so this is the only size signal for auto memory + planning; without it a large local model folds to zero and the planner skips + offload and OOMs. ``exclude_transformer`` drops the ``transformer/`` subfolder + for GGUF/single-file loads (their transformer is the single file, not resident + here); a full pipeline load keeps it (the whole repo is resident).""" + total = 0 + for f in path.rglob("*"): + if f.suffix.lower() not in (".safetensors", ".bin", ".pt", ".ckpt"): + continue + try: + rel = f.relative_to(path) + except ValueError: + continue + if exclude_transformer and rel.parts and rel.parts[0] == "transformer": + continue + try: + total += f.stat().st_size + except OSError: + continue + return total + @staticmethod def _companion_cache_bytes(base: str) -> int: """Resident companion (VAE + text-encoder) size for the memory plan. @@ -536,21 +780,7 @@ class DiffusionBackend: weights to zero and auto planning can pick a resident placement that OOMs.""" local = Path(base).expanduser() if local.is_dir(): - total = 0 - for f in local.rglob("*"): - if f.suffix.lower() not in (".safetensors", ".bin", ".pt", ".ckpt"): - continue - try: - rel = f.relative_to(local) - except ValueError: - continue - if rel.parts and rel.parts[0] == "transformer": - continue # supplied by the GGUF single-file; not resident here - try: - total += f.stat().st_size - except OSError: - continue - return total + return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True) return DiffusionBackend._cache_bytes(base) # ── Synchronous load / generate / unload ─────────────────────────────── @@ -573,6 +803,7 @@ class DiffusionBackend: attention_backend: Optional[str] = None, transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, + model_kind: Optional[str] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -581,9 +812,17 @@ class DiffusionBackend: # load anonymously, not 401 as an explicit empty credential. hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None fam = self.validate_load_request( - repo_id, gguf_filename = gguf_filename, family_override = family_override + repo_id, + gguf_filename = gguf_filename, + family_override = family_override, + model_kind = model_kind, + ) + kind = resolve_model_kind(gguf_filename, model_kind) + # For a full pipeline the repo itself supplies every component, so it is its + # own base; the single-file kinds resolve the companion base diffusers repo. + base = ( + repo_id if kind == "pipeline" else _resolve_base_repo(repo_id, base_repo, fam, hf_token) ) - base = _resolve_base_repo(repo_id, base_repo, fam, hf_token) target = self._resolve_device_target(fam) device, dtype = target.device, target.dtype @@ -614,7 +853,13 @@ class DiffusionBackend: # checkpoints never sit in VRAM at once. self._unload_locked() - gguf_path = self._resolve_gguf_path(repo_id, gguf_filename, hf_token) + # The single-file kinds resolve a checkpoint path (GGUF or safetensors); + # the pipeline kind has none (from_pretrained pulls the repo directly). + single_file_path = ( + self._resolve_gguf_path(repo_id, gguf_filename, hf_token) + if kind in ("gguf", "single_file") + else None + ) transformer_cls = getattr(diffusers, fam.transformer_class) pipeline_cls = getattr(diffusers, fam.pipeline_class) @@ -622,17 +867,30 @@ class DiffusionBackend: # the real budget) -- this also doubles as the dense-quant preflight: the # dense bf16 transformer must fit resident, so the fast path is offered only # when the plan is `none`. - plan = self._plan_memory(target, gguf_path, base, fam, memory_mode, cpu_offload) + plan = self._plan_memory( + target, + single_file_path, + base, + fam, + memory_mode, + cpu_offload, + kind = kind, + repo_id = repo_id, + ) # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it # (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul # dequant on both speed and quality, at the cost of a higher-memory dense # load. Gated on CUDA + bf16 + a resident fit; ANY failure (unsupported arch - # / scheme, OOM, partial quant) falls back to the GGUF build below. + # / scheme, OOM, partial quant) falls back to the GGUF build below. Only the + # GGUF kind offers it: it materialises the dense bf16 transformer from the + # base repo, which the safetensors kinds (a single-file or already-quantized + # pipeline) do not have. pipe = None transformer_quant_engaged = None if ( - normalize_transformer_quant(transformer_quant) is not None + kind == "gguf" + and normalize_transformer_quant(transformer_quant) is not None and dense_transformer_supported(target) and plan.offload_policy == OFFLOAD_NONE ): @@ -665,30 +923,47 @@ class DiffusionBackend: clear_gpu_cache() if pipe is None: - # Default: dequantise the single-file GGUF transformer on-device; the - # VAE / text-encoder / scheduler come from the base diffusers repo - # (GGUF is transformer-only). - transformer = transformer_cls.from_single_file( - gguf_path, - quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), - torch_dtype = dtype, - config = base, - subfolder = "transformer", - # Forward the token: the config is fetched from the (possibly gated) - # base repo before from_pretrained gets a chance to authenticate. - token = hf_token, - ) + if kind == "pipeline": + # Full diffusers repo: from_pretrained pulls every component + # (transformer + VAE + text encoders + scheduler) from the repo + # and re-applies any embedded quantization_config (e.g. bnb-4bit), + # so a pre-quantized pipeline reloads quantized with no extra config. + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} + if hf_token: + pipe_kwargs["token"] = hf_token + pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs) + else: + # Single-file transformer; the VAE / text-encoder / scheduler come + # from the base diffusers repo (the single file is transformer-only). + sf_kwargs: dict[str, Any] = { + "torch_dtype": dtype, + "config": base, + "subfolder": "transformer", + # Forward the token: the config is fetched from the (possibly + # gated) base repo before from_pretrained can authenticate. + "token": hf_token, + } + if kind == "gguf": + # Dequantise the GGUF transformer on-device at the compute dtype. + sf_kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig( + compute_dtype = dtype + ) + # A safetensors single-file (e.g. fp8) carries its own dtype, so no + # GGUF dequant config is passed. + transformer = transformer_cls.from_single_file( + single_file_path, **sf_kwargs + ) - pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} - if hf_token: - pipe_kwargs["token"] = hf_token - pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer} + if hf_token: + pipe_kwargs["token"] = hf_token + pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) # Resolve the effective speed mode: GGUF models default to the # near-lossless `default` profile (compile is ~2.2x and sits below # the quant noise floor), dense models stay bit-identical `off`. An # explicit speed_mode (incl. "off") is honored verbatim. - effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename)) + effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf") # A torchao-quantized dense transformer runs its matmuls through the # regional torch.compile; UNcompiled (eager) it is ~30x slower and would # lose to the GGUF fallback. A dense model otherwise resolves to `off`, so @@ -732,17 +1007,86 @@ class DiffusionBackend: quant_active = transformer_quant_engaged is not None or bool(gguf_filename), logger = logger, ) - # apply_speed_optims flips the process-global TF32 / cudnn.benchmark - # flags. If a later step here (text-encoder quant, memory plan) then - # raises -- e.g. OOM -- those flags would leak flipped and a subsequent - # `off` load would no longer be bit-identical. Restore the snapshot - # unless we reach the commit (unload restores on the happy path). - committed = False + # Install the shared compile-safe eager patches (fused RMSNorm / + # AdaLayerNorm) for any active speed tier. They are class-level, idempotent + # and math-equivalent (FMA / fused -> neutral under compile, equal-or-more + # accurate), so they help eager AND compiled runs. The bit-identical `off` + # reference path must run with them UNINSTALLED, so uninstall there. + # + # Everything from here to the _LoadState commit mutates PROCESS-WIDE state + # (class patches, TORCHINDUCTOR_CACHE_DIR, backend flags). _unload_locked only + # reverses it via _state, so a failure BEFORE the commit would leak it (and + # break the next `off` load's bit-identity). Guard the whole block: on any + # pre-commit failure, restore everything; on success the commit transfers + # ownership to _state and _unload_locked takes over. + # The GGUF-specific speed lever (compiled dequant) applies only when the + # GGUF transformer was ACTUALLY loaded. On the dense torchao-quant + # fast path (fp8 / int8 / fp4) `gguf_filename` is still set as the fallback, + # but `pipe.transformer` is dense (no GGUFLinear), and those schemes need the + # REGIONAL block compile (dynamic quant is ~30x slower eager), not the GGUF + # dequant compile -- so treat the transformer as non-GGUF here. The + # safetensors kinds (single_file / pipeline) likewise have no GGUFLinear. + gguf_transformer = kind == "gguf" and transformer_quant_engaged is None + + eager_patched = False + compile_ctx = None + state_committed = False + # Lazy import: these patch modules import torch at module level, so + # importing them here (not at module load) keeps diffusion.py torch-free + # to import, letting get_diffusion_backend() run on a torchless native install. + from .diffusion_eager_patches import ( + install_compile_safe_patches, + uninstall_patches, + ) + from .diffusion_arch_patches import ( + install_arch_patches, + uninstall_arch_patches, + ) + try: + if effective_speed != SPEED_OFF: + install_compile_safe_patches() + # Per-arch compile-safe fusions (qwen _modulate / z-image residual + # addcmul, etc.). Also neutral under compile, so on for every active + # tier; tracked by the same eager_patched flag for uninstall. + install_arch_patches() + eager_patched = True + else: + uninstall_patches() + uninstall_arch_patches() + + # Pre-warmed torch.compile cache (Mega-cache): when a compiled tier will + # run, point inductor at a per-fingerprint dir and load a matching bundle + # BEFORE the first compiled forward, so the one-time 25-58s compile can be + # paid once (by us / a first run) and reused. A miss is silent -> local + # compile, exactly as today. + if effective_speed in (SPEED_DEFAULT, SPEED_MAX) and compile_eligible( + target, is_gguf = gguf_transformer, family = fam + ): + compile_ctx = compile_cache.begin( + family = fam.name, + transformer = getattr(pipe, "transformer", None), + dtype = getattr(target, "dtype", None), + quant = transformer_quant_engaged, + attention_backend = attention_engaged, + compile_kwargs = { + # Mirrors apply_speed_optims' fullgraph decision: an active + # step cache OR a planned offload graph-breaks, so the cached + # bundle must be keyed on the same fullgraph setting. + "fullgraph": cache_engaged is None + and plan.offload_policy == OFFLOAD_NONE, + "dynamic": effective_speed != SPEED_MAX, + "mode": "max-autotune-no-cudagraphs" + if effective_speed == SPEED_MAX + else "default", + }, + logger = logger, + ) + speed_applied = apply_speed_optims( pipe, target, - is_gguf = bool(gguf_filename), + is_gguf = gguf_transformer, family = fam, speed_mode = effective_speed, cache_active = cache_engaged is not None, @@ -774,8 +1118,8 @@ class DiffusionBackend: # the model's estimated resident size). apply_memory_plan returns the # (policy, tiling) ACTUALLY engaged (it may fall back to whole-module # offload, and tiling is a no-op on a pipeline with no tiling control), so - # status stays honest. The dense fast path already placed the pipe resident; - # for the `none` policy this is an idempotent re-placement. + # status stays honest. The dense fast path already placed the pipe + # resident; for the `none` policy this is an idempotent re-placement. effective_policy, effective_tiling = apply_memory_plan( pipe, plan, device = device, logger = logger ) @@ -787,6 +1131,7 @@ class DiffusionBackend: base_repo = base, device = device, dtype = str(dtype).replace("torch.", ""), + kind = kind, cpu_offload = effective_policy != OFFLOAD_NONE, offload_policy = effective_policy, vae_tiling = effective_tiling, @@ -798,14 +1143,24 @@ class DiffusionBackend: transformer_quant = transformer_quant_engaged, attention_backend = attention_engaged, transformer_cache = cache_engaged, + eager_patched = eager_patched, + compile_cache_ctx = compile_ctx, ) - committed = True + state_committed = True finally: - if not committed: - # Restore the flags AND free the half-built pipe's VRAM: the - # failed load never commits _state, so nothing else reclaims it - # until the next unload. + # Pre-commit failure: nothing owns the process-wide mutations yet, so + # roll them back here (symmetric with _unload_locked). + if not state_committed: restore_backend_flags(backend_flags_before) + compile_cache.restore(compile_ctx) + # apply_speed_optims may have installed the compiled GGUF dequant + # before a later step failed; uninstall is idempotent. + gguf_compile.uninstall_all() + if eager_patched: + uninstall_patches() + uninstall_arch_patches() + # Also free the half-built pipe's VRAM: the failed load never + # commits _state, so nothing else reclaims it until the next unload. clear_gpu_cache() logger.info( @@ -915,34 +1270,68 @@ class DiffusionBackend: def _plan_memory( self, target: DiffusionDeviceTarget, - gguf_path: str, + single_file_path: Optional[str], base: str, fam: DiffusionFamily, memory_mode: Optional[str], cpu_offload: bool, + *, + kind: str = "gguf", + repo_id: Optional[str] = None, ): """Build the memory plan for this load: snapshot free device memory and estimate the model's resident footprint, then let the planner pick an offload policy + VAE memory savers. Kept on the backend so the cached base - repo (companion text-encoder / VAE) feeds the size estimate.""" + repo (companion text-encoder / VAE) feeds the size estimate. + + The size estimate is per-kind: diffusers keeps GGUF weights packed (per-matmul + transient dequant), so a GGUF loads near its on-disk size; a safetensors + single-file loads near its on-disk size (it carries its dtype); and a full + pipeline is one cached download (transformer + companions), already compressed.""" device_memory = snapshot_device_memory(target) - transformer_resident = estimate_gguf_resident_mib(file_size_mib(gguf_path)) - # The companion components (VAE + text encoders) load near their on-disk - # size; sum whatever the prefetch placed in the base-repo cache, or -- for a - # LOCAL diffusers base -- the on-disk component weights (the blob cache is - # empty for a local path, which would otherwise fold multi-GB companions to 0 - # and let auto planning pick a resident placement that OOMs). - companion = self._companion_cache_bytes(base) - companion_mib = int(companion // (1024 * 1024)) if companion else None - model_dense_mib = None - if transformer_resident is not None: - model_dense_mib = transformer_resident + (companion_mib or 0) - # Feed the variant hint (gguf filename + base repo) next to the family name so - # estimate_image_runtime_mib sees distilled markers ("turbo"/"schnell") that + if kind == "pipeline": + # The whole repo (transformer + companions) is one cached download; the + # cached bytes are the resident estimate (bnb-4bit / fp8 stay compressed). + # A LOCAL pipeline path isn't in the HF blob cache, so sum its on-disk weights + # (transformer included) instead of folding to zero and skipping offload. + local_repo = Path(repo_id).expanduser() if repo_id else None + if local_repo is not None and local_repo.is_dir(): + cached = self._local_dir_weight_bytes(local_repo, exclude_transformer = False) + else: + cached = self._cache_bytes(repo_id) if repo_id else 0 + cached_mib = int(cached // (1024 * 1024)) if cached else None + model_dense_mib = estimate_safetensors_dense_mib(cached_mib) + companion_mib = None + else: + if kind == "single_file": + # Safetensors single-file: no dequant expansion (it carries its dtype). + transformer_resident = estimate_safetensors_dense_mib( + file_size_mib(single_file_path) + ) + else: + transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path)) + # The companion components (VAE + text encoders) load near their on-disk + # size; sum whatever the prefetch placed in the base-repo cache, or -- for a + # LOCAL diffusers base -- the on-disk component weights (the blob cache is + # empty for a local path, which would otherwise fold multi-GB companions to 0 + # and let auto planning pick a resident placement that OOMs). + companion = self._companion_cache_bytes(base) + companion_mib = int(companion // (1024 * 1024)) if companion else None + model_dense_mib = None + if transformer_resident is not None: + model_dense_mib = transformer_resident + (companion_mib or 0) + # Feed the variant hint (single-file basename + base/repo) next to the family name + # so estimate_image_runtime_mib sees distilled markers ("turbo"/"schnell") that # detect_family normalizes out of fam.name -- distilled models need ~15% less # activation headroom, and over-reserving can force needless offload / tiling. variant_hint = " ".join( - p for p in (fam.name, Path(gguf_path).name if gguf_path else "", base or "") if p + p + for p in ( + fam.name, + Path(single_file_path).name if single_file_path else "", + repo_id or base or "", + ) + if p ) runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = variant_hint) return plan_diffusion_memory( @@ -955,6 +1344,60 @@ class DiffusionBackend: explicit_offload = cpu_offload, ) + def _workflow_pipe(self, state: _LoadState, class_name: Optional[str], workflow: str) -> Any: + """The diffusers pipeline for an image-conditioned ``workflow``, built once and + cached. ``Pipeline.from_pipe`` re-wires the loaded text-to-image pipe's resident + modules (transformer/VAE/text-encoder, incl. any compiled/quantised state) into + the workflow pipeline class, so there is no extra VRAM and no reload. Raises a + clear ValueError when the family does not support the workflow.""" + if not class_name: + raise ValueError( + f"{workflow} is not supported for the '{state.family.name}' model family." + ) + cached = self._aux_pipes.get(class_name) + if cached is not None: + return cached + import diffusers + + # torch_dtype=None is load-bearing: diffusers' from_pipe defaults torch_dtype to + # torch.float32 and then runs new_pipeline.to(dtype=float32) over EVERY component. + # That recast (a) needlessly upcasts the reused bf16 modules and (b) hard-crashes + # on the dense-quant fast path -- a torchao-quantized + torch.compiled transformer + # has tensor-subclass Linear weights that torch.nn.Module._apply cannot swap_tensors + # ("Couldn't swap Linear.weight"). Passing None makes from_pipe skip the cast and + # reuse the resident modules AT THEIR LOADED dtype, which is the whole point of + # from_pipe (component reuse, no reload, no extra VRAM). + pipe = getattr(diffusers, class_name).from_pipe(state.pipe, torch_dtype = None) + # Only publish to the shared aux cache if THIS load is still current. from_pipe runs + # under _generate_lock but NOT _lock, so an unload()/superseding load can clear + # _aux_pipes and null _state while it builds; caching unconditionally would re-insert + # a wrapper over now-stale modules that a later same-workflow load would reuse (or + # keep the old VRAM pinned). This generation still uses the returned pipe. + with self._lock: + if self._state is state: + self._aux_pipes[class_name] = pipe + return pipe + + @staticmethod + def _align_vae_dtype(pipe: Any) -> None: + """Cast the VAE to the transformer's compute dtype before an image-conditioned + call. The img2img/inpaint pipelines VAE-encode the input image at the text- + encoder dtype (bf16), but a prior txt2img DECODE may have left the shared VAE + upcast to fp32 (its ``force_upcast`` path), so the encode would mismatch + (bf16 image vs fp32 VAE). Re-aligning here is safe: our families run bf16 or + fp32 only (the fp16 guard promotes fp16), and a later txt2img decode re-upcasts + as needed. Best-effort; a no-op when already aligned.""" + transformer = getattr(pipe, "transformer", None) + vae = getattr(pipe, "vae", None) + if transformer is None or vae is None: + return + try: + target_dtype = transformer.dtype + if next(vae.parameters()).dtype != target_dtype: + vae.to(dtype = target_dtype) + except (StopIteration, AttributeError, RuntimeError): + pass + @staticmethod def _reset_step_cache(pipe: Any) -> None: """Clear the transformer's stateful step cache (FBCache) before a generation. @@ -989,8 +1432,22 @@ class DiffusionBackend: guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, + # Image-conditioned workflows (base64 / data-URL): an init image alone selects + # img2img; an init image + mask selects inpaint. ``strength`` is the img2img/ + # inpaint denoise strength (0 = keep source, 1 = full redraw). None = txt2img. + init_image: Optional[str] = None, + mask_image: Optional[str] = None, + strength: Optional[float] = None, + # Upscale (hires fix): a factor > 1 with an init image enlarges the input and + # re-denoises it at low strength to paint detail at the higher resolution. + upscale: Optional[float] = None, + # Reference workflow (FLUX.2): ADDITIONAL reference images beyond ``init_image``. The + # pipeline accepts a list, so multiple references can be combined (subject + style, + # character + scene). Ignored by non-reference workflows. + reference_images: Optional[list[str]] = None, ) -> dict[str, Any]: import torch + from PIL import Image # A per-generation cancel Event: unload()/a superseding load set THIS event # (registered under _lock below) to abort just this denoise. _generate_lock @@ -1020,10 +1477,121 @@ class DiffusionBackend: seed = int(seed) generator.manual_seed(seed) + # Select the pipeline for this workflow. txt2img uses the loaded pipe; + # img2img/inpaint reuse its resident modules via from_pipe (no reload); + # an edit model's OWN loaded pipe is already the edit pipeline. + pipe = state.pipe + init_pil = mask_pil = None + ref_extra: list = [] + # Validate parameter dependencies up front: mask / upscale / reference all + # need an input image, and reference conditioning needs a family that + # supports it. Without these guards an unsupported combination would be + # silently ignored and quietly fall back to txt2img / img2img. + if init_image is None: + if mask_image is not None: + raise ValueError("mask_image requires an input image (init_image).") + if upscale is not None and upscale > 1.0: + raise ValueError("upscale requires an input image (init_image).") + if reference_images: + raise ValueError("reference_images require an input image (init_image).") + if reference_images and not getattr(state.family, "reference", False): + raise ValueError( + f"Reference images are not supported for the '{state.family.name}' " + "model family." + ) + if getattr(state.family, "edit", False): + # Instruction editing: the loaded pipe is the edit pipeline. It always + # needs an input image; the prompt is the edit instruction. No mask, no + # from_pipe (the model has no plain text-to-image mode). + if init_image is None: + raise ValueError( + f"{state.family.name} is an image-editing model: provide an input image." + ) + workflow = "edit" + init_pil = _decode_b64_image(init_image, mode = "RGB") + elif mask_image is not None and init_image is not None: + workflow = "inpaint" + pipe = self._workflow_pipe(state, state.family.inpaint_pipeline_class, workflow) + init_pil = _decode_b64_image(init_image, mode = "RGB") + mask_pil = _decode_b64_image(mask_image, mode = "L") + elif init_image is not None and upscale is not None and upscale > 1.0: + # Upscale (hires fix): enlarge the input with Lanczos, then re-run the + # img2img pipeline on it at a low denoise strength so the transformer + # adds high-frequency detail without redrawing the content. Shares the + # img2img pipeline/modules via from_pipe (no extra VRAM, no reload). + workflow = "upscale" + pipe = self._workflow_pipe(state, state.family.img2img_pipeline_class, workflow) + init_pil = _decode_b64_image(init_image, mode = "RGB") + iw, ih = init_pil.size + # Cap the factor, THEN cap the absolute output: a large input times the + # factor (e.g. 1024 at 4x = 4096, or a big upload) would otherwise OOM the + # VAE/transformer. Bound the longest side to 2048 (txt2img's own max), + # scaling both dims to keep the aspect ratio; round to a multiple of 16 + # (VAE downsample + patch size require it for our families). + factor = max(1.0, min(float(upscale), 4.0)) + tw_f, th_f = iw * factor, ih * factor + max_side = 2048 + fit = min(1.0, max_side / max(tw_f, th_f)) + tw = max(16, int(round(tw_f * fit / 16.0)) * 16) + th = max(16, int(round(th_f * fit / 16.0)) * 16) + # After the absolute cap, the target must still exceed the input, or + # "upscale" would shrink it (e.g. a 3000px source at 2x clamps to 2048). + # Reject rather than silently return a smaller image than uploaded. + if max(tw, th) <= max(iw, ih): + raise ValueError( + f"Upscale would not enlarge this image: its longest side " + f"({max(iw, ih)}px) already meets the {max_side}px output limit. " + f"Use a smaller source image." + ) + init_pil = init_pil.resize((tw, th), Image.LANCZOS) + if strength is None: + # Hires-fix default: low enough to preserve content, high enough to + # synthesise new detail at the higher resolution. + strength = 0.35 + elif getattr(state.family, "reference", False) and init_image is not None: + # FLUX.2-style reference conditioning: the loaded pipe (Flux2KleinPipeline) + # takes the reference image directly via its `image` arg and generates a + # fresh image at the REQUESTED size, guided by both the prompt and the + # reference. No from_pipe (the loaded pipe already supports it), no strength + # (reference-conditioning, not a denoise blend), and the output size comes + # from the sliders (the pipeline resizes the reference to ~1MP itself). + # Checked AFTER inpaint/upscale so a mask/upscale request on a reference + # family (FLUX.2-klein also has an inpaint pipeline) still routes correctly. + workflow = "reference" + init_pil = _decode_b64_image(init_image, mode = "RGB") + # Additional references (FLUX.2 accepts a list): decode them so the + # conditioning combines all of them. Capped to keep VRAM bounded. + ref_extra = [ + _decode_b64_image(x, mode = "RGB") for x in (reference_images or [])[:3] + ] + elif init_image is not None: + workflow = "img2img" + pipe = self._workflow_pipe(state, state.family.img2img_pipeline_class, workflow) + init_pil = _decode_b64_image(init_image, mode = "RGB") + else: + workflow = "txt2img" + # Auto-resize odd-sized inputs to a multiple of 16 for the workflows whose + # OUTPUT size is taken from the input image (img2img / inpaint / extend / edit), + # so an upload like 186px tall no longer fails the pipeline's divisibility check. + # txt2img/reference use the validated slider size; upscale already produced a /16 + # target. The mask is matched to the snapped image so inpaint stays aligned. + if init_pil is not None and workflow in ("img2img", "inpaint", "edit"): + init_pil = _snap_to_multiple(init_pil, 16) + if mask_pil is not None and mask_pil.size != init_pil.size: + from PIL import Image as _PILImage + mask_pil = mask_pil.resize(init_pil.size, _PILImage.NEAREST) + if init_pil is not None: + # Keep the VAE encode dtype consistent with the input image. + self._align_vae_dtype(pipe) + + # Pipelines vary in which kwargs they accept (img2img derives size from the + # input image and may reject width/height; a distilled pipe may take no + # negative prompt or step callback), so gate every optional kwarg on the + # actual signature. + call_params = inspect.signature(pipe.__call__).parameters + kwargs: dict[str, Any] = { "prompt": prompt, - "width": width, - "height": height, "num_inference_steps": steps, # Most pipelines take guidance via "guidance_scale"; Qwen-Image # uses "true_cfg_scale" (its distilled guidance is off). @@ -1033,10 +1601,33 @@ class DiffusionBackend: # share this call's seed, drawn sequentially from one generator. "num_images_per_prompt": batch_size, } - # Pipelines vary in which kwargs they accept (a distilled pipeline may - # take neither a negative prompt nor a step callback), so only pass - # those where the signature has them. - call_params = inspect.signature(state.pipe.__call__).parameters + if init_pil is not None: + # Reference with extra images passes the whole list (FLUX.2 combines them); + # every other workflow takes the single image. + kwargs["image"] = [init_pil, *ref_extra] if ref_extra else init_pil + if mask_pil is not None and "mask_image" in call_params: + kwargs["mask_image"] = mask_pil + if strength is not None and "strength" in call_params: + kwargs["strength"] = strength + # width/height. txt2img uses the requested slider size. Image-conditioned + # pipes must use the INPUT IMAGE's own size, NOT the slider: the output is + # the redrawn/extended input, and the denoise builds latents from the image, + # so a slider size that differs from the image mismatches (e.g. a 1536px + # outpaint vs a 1024 slider -> "tensor a (128) must match tensor b (192)"). + # Many img2img/inpaint pipelines drop width/height entirely; pass them only + # when accepted, derived from the image so they are always consistent. + if workflow in ("txt2img", "reference"): + # txt2img and FLUX.2 reference both generate at the REQUESTED size; the + # reference pipe resizes the conditioning image itself, so it must not be + # pinned to the input image's size like img2img/inpaint/upscale are. + kwargs["width"] = width + kwargs["height"] = height + elif init_pil is not None: + iw, ih = init_pil.size + if "width" in call_params: + kwargs["width"] = iw + if "height" in call_params: + kwargs["height"] = ih if negative_prompt and "negative_prompt" in call_params: kwargs["negative_prompt"] = negative_prompt @@ -1071,13 +1662,20 @@ class DiffusionBackend: # inference_mode is strictly faster than the no_grad diffusers # uses internally and numerically identical for inference. with torch.inference_mode(): - images = state.pipe(**kwargs).images + images = pipe(**kwargs).images finally: self._gen = None # A cancelled denoise returns early with a partial/garbage image; # don't hand it back to be persisted. if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) + # The first compiled generation just paid the compile cost; persist the + # warm torch.compile cache bundle when saving is enabled (distributor / + # first-run warm). Idempotent + best-effort -- never fails a generation. + try: + compile_cache.save(state.compile_cache_ctx, logger = logger) + except Exception: # noqa: BLE001 — cache persistence is best-effort + pass # Return the PIL images (not yet encoded): the route embeds each # image's recipe and persists it via the gallery. return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id} @@ -1133,6 +1731,24 @@ class DiffusionBackend: # Restore the process-wide backend flags (TF32 / cudnn.benchmark) this load # may have flipped, so the next `off` load is bit-identical again. restore_backend_flags(state.backend_flags_before) + # Restore TORCHINDUCTOR_CACHE_DIR and uninstall the shared eager patches, so a + # later `off` load runs the bit-identical reference path. Both are idempotent. + compile_cache.restore(state.compile_cache_ctx) + # Uninstall the GGUF dequant accelerators (compiled dequant / global weight + # buffer) this load may have installed, so a later `off` load runs the stock, + # bit-identical dequant. Idempotent. + gguf_compile.uninstall_all() + if state.eager_patched: + # Lazy import (torch at module level) to keep diffusion.py torch-free to import. + from .diffusion_eager_patches import uninstall_patches + from .diffusion_arch_patches import uninstall_arch_patches + + uninstall_patches() + uninstall_arch_patches() + # Drop the workflow pipes built around this load's modules so they don't pin the + # freed pipeline (they only re-wire its components, but holding the wrappers + # would keep the modules alive past unload). + self._aux_pipes.clear() self._state = None del state clear_gpu_cache() @@ -1147,6 +1763,7 @@ class DiffusionBackend: "base_repo": None, "device": None, "dtype": None, + "model_kind": None, "cpu_offload": False, "offload_policy": None, "vae_tiling": False, @@ -1157,6 +1774,7 @@ class DiffusionBackend: "transformer_quant": None, "attention_backend": None, "transformer_cache": None, + "workflows": [], } return { "loaded": True, @@ -1165,6 +1783,7 @@ class DiffusionBackend: "base_repo": state.base_repo, "device": state.device, "dtype": state.dtype, + "model_kind": state.kind, "cpu_offload": state.cpu_offload, "offload_policy": state.offload_policy, "vae_tiling": state.vae_tiling, @@ -1175,9 +1794,37 @@ class DiffusionBackend: "transformer_quant": state.transformer_quant, "attention_backend": state.attention_backend, "transformer_cache": state.transformer_cache, + # Image-conditioned workflows the loaded family supports, so the UI can gate + # its tabs. txt2img is always available on the diffusers engine. + "workflows": _family_workflows(state.family), } +def _family_workflows(fam: DiffusionFamily) -> list[str]: + """The workflow ids the diffusers engine can run for ``fam`` (drives UI gating).""" + # Instruction-editing families have no plain text-to-image mode: their pipeline always + # takes an input image + instruction, so they expose only the "edit" workflow. + if getattr(fam, "edit", False): + return ["edit"] + workflows = ["txt2img"] + # Reference families (FLUX.2) keep txt2img and add reference conditioning via their own + # pipeline's optional image arg (no img2img/inpaint classes needed). + if getattr(fam, "reference", False): + workflows.append("reference") + if getattr(fam, "img2img_pipeline_class", None): + # Upscale (hires fix) runs on the img2img pipeline, so it is available exactly + # when img2img is. + workflows.append("img2img") + workflows.append("upscale") + if getattr(fam, "inpaint_pipeline_class", None): + workflows.append("inpaint") + # Outpaint (extend) reuses the inpaint pipeline with a padded canvas + border mask, + # so it needs an inpaint pipeline that preserves the (larger) canvas size. + if getattr(fam, "inpaint_preserves_size", True): + workflows.append("outpaint") + return workflows + + def _resolve_base_repo( repo_id: str, base_repo: Optional[str], fam: DiffusionFamily, hf_token: Optional[str] ) -> str: @@ -1223,6 +1870,19 @@ def _base_file_downloaded(rfilename: str, *, include_transformer: bool = False) return not rfilename.startswith("assets/") +def _pipeline_file_downloaded(rfilename: str) -> bool: + """True for files a full-pipeline ``from_pretrained`` fetches. + + Like ``_base_file_downloaded`` but for the ``pipeline`` kind, where the repo + supplies its OWN transformer weights, so the ``transformer/`` subfolder is kept. + Top-level docs (README/PDF/images) and ``assets/`` are still skipped so the + progress estimate matches what actually lands on disk. + """ + if "/" not in rfilename: # top-level: only the pipeline manifest is fetched + return rfilename == "model_index.json" + return not rfilename.startswith("assets/") + + def _progress( phase: Optional[str], bytes_downloaded: int = 0, diff --git a/studio/backend/core/inference/diffusion_arch_patches.py b/studio/backend/core/inference/diffusion_arch_patches.py new file mode 100644 index 0000000000..e079e09603 --- /dev/null +++ b/studio/backend/core/inference/diffusion_arch_patches.py @@ -0,0 +1,530 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Per-architecture eager fusions for the diffusion DiT blocks. + +The shared patches in ``diffusion_eager_patches.py`` cover what diffusers factors into +SHARED classes (``RMSNorm``, the ``AdaLayerNorm*`` modulation classes -- the latter only +used by flux.1). The remaining fusible ops -- the gated residual ``x = x + gate * out`` and +the inline modulation ``norm * (1 + scale) + shift`` -- are written longhand in each +family's PER-MODEL block ``forward`` (flux.2 / qwen / z-image), so they need per-arch +patches. This module supplies them, in the ``unsloth_zoo.temporary_patches`` style: one +small patch function per target, applied through the shared, fingerprint-checked, +reversible backend (``diffusion_patch_backend`` -> ``patch_function`` / ``restore_original``). + +The single actionable fused op is ``torch.addcmul(a, b, c) == a + b * c`` (one FMA kernel, +1-ULP vs mul+add, MORE accurate). Two forms exist: + * out-of-place ``torch.addcmul(...)`` -- COMPILE-SAFE: lowers to plain ops, no aliasing, + neutral under ``torch.compile``. Used here for ALL non-``off`` tiers. It captures the + real eager win (fewer kernel launches); it does NOT save the output allocation. + * in-place ``x.addcmul_(...)`` -- would also save the allocation, but the allocation part + measured NEUTRAL on this stack (CUDA caching allocator already recycles; cf. the GGUF + weight-buffer result), and in-place residual mutation is compile-unsafe + aliasing-risky. + So it is deliberately NOT used; the out-of-place form is the whole, safe win. + +Each patch is guarded TWICE: ``can_safely_patch`` checks the forward SIGNATURE, and a local +source-body check (``_body_has``) confirms the exact lines we rewrite are still present, so a +future diffusers that changed the block body simply leaves the block UNPATCHED (correctness +first) rather than running a stale copy. Kill-switch: ``UNSLOTH_DIFFUSION_ARCH_PATCHES=0``. + +Implemented for all four families (extend by adding entries to ``_SPECS``): + * qwen-image ``QwenImageTransformerBlock._modulate`` -- modulation addcmul (all 4 sites). + * z-image ``ZImageTransformerBlock.forward`` -- the 2 gated-residual addcmuls. + * flux.1 ``FluxTransformerBlock.forward`` (inline norm2 modulation + 4 gated residuals) + + ``FluxSingleTransformerBlock.forward`` (residual + gate*proj_out). + * flux.2-klein ``Flux2TransformerBlock.forward`` (4 inline modulations + 4 gated residuals) + + ``Flux2SingleTransformerBlock.forward`` (inline modulation + gated residual). +""" + +from __future__ import annotations + +import inspect +import logging +import os +from typing import Any, Callable, Optional + +import torch + +from .diffusion_patch_backend import apply_patch, revert_patch + +logger = logging.getLogger(__name__) + +_ENV_ENABLE = "UNSLOTH_DIFFUSION_ARCH_PATCHES" + + +def _patches_enabled() -> bool: + return (os.environ.get(_ENV_ENABLE) or "").strip().lower() not in ("0", "off", "false", "no") + + +def _body_has(fn: Callable, *needles: str) -> bool: + """True iff every ``needle`` appears in ``fn``'s source -- a body-drift guard so a patch + self-disables if diffusers changed the lines it rewrites.""" + try: + src = inspect.getsource(fn) + except (OSError, TypeError): + return False + return all(n in src for n in needles) + + +# ===================================================================================== +# qwen-image: QwenImageTransformerBlock._modulate (modulation addcmul, all 4 call sites) +# ===================================================================================== +def _qwen_modulate( + self, + x, + mod_params, + index = None, +): + """diffusers 0.38 ``QwenImageTransformerBlock._modulate`` with the final + ``x*(1+scale)+shift`` fused to ``torch.addcmul`` (covers both the global and the + per-token ``index`` branches, since both end in that same expression).""" + shift, scale, gate = mod_params.chunk(3, dim = -1) + + if index is not None: + actual_batch = shift.size(0) // 2 + shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] + scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:] + gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:] + index_expanded = index.unsqueeze(-1) + shift_0_exp = shift_0.unsqueeze(1) + shift_1_exp = shift_1.unsqueeze(1) + scale_0_exp = scale_0.unsqueeze(1) + scale_1_exp = scale_1.unsqueeze(1) + gate_0_exp = gate_0.unsqueeze(1) + gate_1_exp = gate_1.unsqueeze(1) + shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp) + scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp) + gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp) + else: + shift_result = shift.unsqueeze(1) + scale_result = scale.unsqueeze(1) + gate_result = gate.unsqueeze(1) + + # fused: x * (1 + scale_result) + shift_result + return torch.addcmul(shift_result, x, 1 + scale_result), gate_result + + +def _spec_qwen_modulate(): + try: + from diffusers.models.transformers.transformer_qwenimage import ( + QwenImageTransformerBlock as cls, + ) + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "_modulate", None) + if orig is None or not _body_has(orig, "x * (1 + scale_result) + shift_result"): + return None + return (cls, "_modulate", _qwen_modulate) + + +# ===================================================================================== +# z-image: ZImageTransformerBlock.forward (the 2 gated-residual addcmuls) +# ===================================================================================== +def _zimage_forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + freqs_cis: torch.Tensor, + adaln_input: torch.Tensor | None = None, + noise_mask: torch.Tensor | None = None, + adaln_noisy: torch.Tensor | None = None, + adaln_clean: torch.Tensor | None = None, +): + """diffusers 0.38 ``ZImageTransformerBlock.forward`` with the two gated residuals + ``x = x + gate * sublayer`` fused to ``torch.addcmul`` (out-of-place). The shift-free + ``*scale`` modulation and the non-gated (``else``) residuals are left as-is.""" + from diffusers.models.transformers.transformer_z_image import select_per_token + + if self.modulation: + seq_len = x.shape[1] + + if noise_mask is not None: + mod_noisy = self.adaLN_modulation(adaln_noisy) + mod_clean = self.adaLN_modulation(adaln_clean) + + scale_msa_noisy, gate_msa_noisy, scale_mlp_noisy, gate_mlp_noisy = mod_noisy.chunk( + 4, dim = 1 + ) + scale_msa_clean, gate_msa_clean, scale_mlp_clean, gate_mlp_clean = mod_clean.chunk( + 4, dim = 1 + ) + + gate_msa_noisy, gate_mlp_noisy = gate_msa_noisy.tanh(), gate_mlp_noisy.tanh() + gate_msa_clean, gate_mlp_clean = gate_msa_clean.tanh(), gate_mlp_clean.tanh() + + scale_msa_noisy, scale_mlp_noisy = 1.0 + scale_msa_noisy, 1.0 + scale_mlp_noisy + scale_msa_clean, scale_mlp_clean = 1.0 + scale_msa_clean, 1.0 + scale_mlp_clean + + scale_msa = select_per_token(scale_msa_noisy, scale_msa_clean, noise_mask, seq_len) + scale_mlp = select_per_token(scale_mlp_noisy, scale_mlp_clean, noise_mask, seq_len) + gate_msa = select_per_token(gate_msa_noisy, gate_msa_clean, noise_mask, seq_len) + gate_mlp = select_per_token(gate_mlp_noisy, gate_mlp_clean, noise_mask, seq_len) + else: + mod = self.adaLN_modulation(adaln_input) + scale_msa, gate_msa, scale_mlp, gate_mlp = mod.unsqueeze(1).chunk(4, dim = 2) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + # Attention block -- fused gated residual: x + gate_msa * attention_norm2(attn_out) + attn_out = self.attention( + self.attention_norm1(x) * scale_msa, attention_mask = attn_mask, freqs_cis = freqs_cis + ) + x = torch.addcmul(x, gate_msa, self.attention_norm2(attn_out)) + + # FFN block -- fused gated residual: x + gate_mlp * ffn_norm2(feed_forward(...)) + x = torch.addcmul( + x, gate_mlp, self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp)) + ) + else: + attn_out = self.attention( + self.attention_norm1(x), attention_mask = attn_mask, freqs_cis = freqs_cis + ) + x = x + self.attention_norm2(attn_out) + x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x))) + + return x + + +def _spec_zimage_forward(): + try: + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "x = x + gate_msa * self.attention_norm2(attn_out)", + "x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))", + ): + return None + return (cls, "forward", _zimage_forward) + + +# ===================================================================================== +# flux.1: FluxTransformerBlock / FluxSingleTransformerBlock +# (block modulation goes through AdaLayerNormZero -- already handled by the shared patch; +# here we fuse the inline norm2 modulation + the gated residual adds.) +# ===================================================================================== +def _flux_double_forward( + self, + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb = None, + joint_attention_kwargs = None, +): + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, emb = temb + ) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = ( + self.norm1_context(encoder_hidden_states, emb = temb) + ) + joint_attention_kwargs = joint_attention_kwargs or {} + attention_outputs = self.attn( + hidden_states = norm_hidden_states, + encoder_hidden_states = norm_encoder_hidden_states, + image_rotary_emb = image_rotary_emb, + **joint_attention_kwargs, + ) + if len(attention_outputs) == 2: + attn_output, context_attn_output = attention_outputs + elif len(attention_outputs) == 3: + attn_output, context_attn_output, ip_attn_output = attention_outputs + + # fused: hidden_states + gate_msa * attn_output + hidden_states = torch.addcmul(hidden_states, gate_msa.unsqueeze(1), attn_output) + + norm_hidden_states = self.norm2(hidden_states) + # fused: norm * (1 + scale_mlp) + shift_mlp + norm_hidden_states = torch.addcmul( + shift_mlp[:, None], norm_hidden_states, 1 + scale_mlp[:, None] + ) + + ff_output = self.ff(norm_hidden_states) + hidden_states = torch.addcmul(hidden_states, gate_mlp.unsqueeze(1), ff_output) + if len(attention_outputs) == 3: + hidden_states = hidden_states + ip_attn_output + + encoder_hidden_states = torch.addcmul( + encoder_hidden_states, c_gate_msa.unsqueeze(1), context_attn_output + ) + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = torch.addcmul( + c_shift_mlp[:, None], norm_encoder_hidden_states, 1 + c_scale_mlp[:, None] + ) + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = torch.addcmul( + encoder_hidden_states, c_gate_mlp.unsqueeze(1), context_ff_output + ) + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +def _spec_flux_double(): + try: + from diffusers.models.transformers.transformer_flux import FluxTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "hidden_states = hidden_states + attn_output", + "norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]", + "encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output", + ): + return None + return (cls, "forward", _flux_double_forward) + + +def _flux_single_forward( + self, + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb = None, + joint_attention_kwargs = None, +): + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim = 1) + + residual = hidden_states + norm_hidden_states, gate = self.norm(hidden_states, emb = temb) + mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states = norm_hidden_states, + image_rotary_emb = image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim = 2) + gate = gate.unsqueeze(1) + # fused: residual + gate * proj_out(hidden_states) + hidden_states = torch.addcmul(residual, gate, self.proj_out(hidden_states)) + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + encoder_hidden_states, hidden_states = ( + hidden_states[:, :text_seq_len], + hidden_states[:, text_seq_len:], + ) + return encoder_hidden_states, hidden_states + + +def _spec_flux_single(): + try: + from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "hidden_states = gate * self.proj_out(hidden_states)", + "hidden_states = residual + hidden_states", + ): + return None + return (cls, "forward", _flux_single_forward) + + +# ===================================================================================== +# flux.2-klein: Flux2TransformerBlock / Flux2SingleTransformerBlock +# (modulation is INLINE here -- not via AdaLayerNorm -- so we fuse both the modulation and +# the gated residuals; scale/shift/gate are [B,1,dim] so no [:, None] is needed.) +# ===================================================================================== +def _flux2_double_forward( + self, + hidden_states, + encoder_hidden_states, + temb_mod_img, + temb_mod_txt, + image_rotary_emb = None, + joint_attention_kwargs = None, +): + from diffusers.models.transformers.transformer_flux2 import Flux2Modulation + + joint_attention_kwargs = joint_attention_kwargs or {} + (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = Flux2Modulation.split( + temb_mod_img, 2 + ) + (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = ( + Flux2Modulation.split(temb_mod_txt, 2) + ) + + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = torch.addcmul(shift_msa, norm_hidden_states, 1 + scale_msa) + + norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) + norm_encoder_hidden_states = torch.addcmul( + c_shift_msa, norm_encoder_hidden_states, 1 + c_scale_msa + ) + + attention_outputs = self.attn( + hidden_states = norm_hidden_states, + encoder_hidden_states = norm_encoder_hidden_states, + image_rotary_emb = image_rotary_emb, + **joint_attention_kwargs, + ) + attn_output, context_attn_output = attention_outputs + + hidden_states = torch.addcmul(hidden_states, gate_msa, attn_output) + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = torch.addcmul(shift_mlp, norm_hidden_states, 1 + scale_mlp) + + ff_output = self.ff(norm_hidden_states) + hidden_states = torch.addcmul(hidden_states, gate_mlp, ff_output) + + encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_msa, context_attn_output) + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = torch.addcmul( + c_shift_mlp, norm_encoder_hidden_states, 1 + c_scale_mlp + ) + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_mlp, context_ff_output) + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +def _spec_flux2_double(): + try: + from diffusers.models.transformers.transformer_flux2 import Flux2TransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa", + "hidden_states = hidden_states + gate_mlp * ff_output", + "encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output", + ): + return None + return (cls, "forward", _flux2_double_forward) + + +def _flux2_single_forward( + self, + hidden_states, + encoder_hidden_states, + temb_mod, + image_rotary_emb = None, + joint_attention_kwargs = None, + split_hidden_states = False, + text_seq_len = None, +): + from diffusers.models.transformers.transformer_flux2 import Flux2Modulation + + if encoder_hidden_states is not None: + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim = 1) + + mod_shift, mod_scale, mod_gate = Flux2Modulation.split(temb_mod, 1)[0] + + norm_hidden_states = self.norm(hidden_states) + norm_hidden_states = torch.addcmul(mod_shift, norm_hidden_states, 1 + mod_scale) + + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states = norm_hidden_states, + image_rotary_emb = image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = torch.addcmul(hidden_states, mod_gate, attn_output) + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + if split_hidden_states: + encoder_hidden_states, hidden_states = ( + hidden_states[:, :text_seq_len], + hidden_states[:, text_seq_len:], + ) + return encoder_hidden_states, hidden_states + else: + return hidden_states + + +def _spec_flux2_single(): + try: + from diffusers.models.transformers.transformer_flux2 import ( + Flux2SingleTransformerBlock as cls, + ) + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift", + "hidden_states = hidden_states + mod_gate * attn_output", + ): + return None + return (cls, "forward", _flux2_single_forward) + + +# ===================================================================================== +# registry + lifecycle +# ===================================================================================== +# Each entry is a zero-arg resolver returning (cls, attr, new_fn) or None (target absent / +# body drifted). All entries here are COMPILE-SAFE (out-of-place addcmul). +_SPECS: tuple[Callable[[], Optional[tuple]], ...] = ( + _spec_qwen_modulate, + _spec_zimage_forward, + _spec_flux_double, + _spec_flux_single, + _spec_flux2_double, + _spec_flux2_single, +) + +# (cls, attr) pairs we successfully patched, for an exact reverse. +_patched: list[tuple] = [] + + +def install_arch_patches() -> int: + """Install the per-arch compile-safe fusions (idempotent). Returns the count applied. + + Safe for every non-``off`` tier: the fusions lower to plain ops and are neutral under + ``torch.compile`` (so they also help the ``max`` regionally-compiled blocks).""" + if not _patches_enabled(): + uninstall_arch_patches() + return 0 + if _patched: + return len(_patched) + for resolve in _SPECS: + try: + spec = resolve() + except Exception as exc: # noqa: BLE001 + logger.warning( + "arch-patch: resolver %s failed: %s", getattr(resolve, "__name__", resolve), exc + ) + spec = None + if spec is None: + continue + cls, attr, new_fn = spec + if apply_patch(cls, attr, new_fn, match_level = "relaxed"): + _patched.append((cls, attr)) + else: + logger.warning( + "arch-patch: skipping %s.%s (signature mismatch / unavailable)", + getattr(cls, "__name__", cls), + attr, + ) + logger.info("arch-patch: installed %d/%d per-arch fusions", len(_patched), len(_SPECS)) + return len(_patched) + + +def uninstall_arch_patches() -> None: + """Restore every per-arch patched method/forward (idempotent).""" + for cls, attr in list(_patched): + revert_patch(cls, attr) + _patched.clear() + + +def is_installed() -> bool: + return bool(_patched) diff --git a/studio/backend/core/inference/diffusion_compile_cache.py b/studio/backend/core/inference/diffusion_compile_cache.py new file mode 100644 index 0000000000..97812ba9f9 --- /dev/null +++ b/studio/backend/core/inference/diffusion_compile_cache.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pre-warmed ``torch.compile`` cache for the diffusion denoiser (Mega-cache). + +The regional ``torch.compile`` of the repeated denoiser block (``diffusion_speed.py``) +pays a one-time 25-58s compile on the FIRST image after a load. This module lets that +cost be paid ONCE -- by us (the distributor) ahead of time, or by the user on a first +run -- and reused on every later load via torch's portable Mega-cache +(``torch.compiler.save_cache_artifacts`` / ``load_cache_artifacts``, torch >= 2.7). + +PORTABILITY IS NOT UNIVERSAL. A compiled artifact is only valid for the SAME torch +version, Triton version, CUDA build, and GPU architecture it was produced on (and the +same model graph: family, dtype, quant scheme, attention backend, compile kwargs, shape +bucket). torch validates these on load and a mismatch simply yields no cache hit -- it +does NOT error. So this layer is built around an EXACT-MATCH fingerprint with a SILENT +FALLBACK to local compile: a miss is normal and never fatal. We therefore ship per-arch +bundles keyed by the full fingerprint, never one universal cache. See +``outputs/compile_cache/DISTRIBUTION.md``. + +Lifecycle (driven by the caller, around ``_compile_repeated_blocks``): + 1. ``begin(...)`` -> build the fingerprint, point ``TORCHINDUCTOR_CACHE_DIR`` at a + per-key dir, and ``load_cache_artifacts`` if a matching bundle + exists. Must run BEFORE the first compiled forward. + 2. (compile + one warmup forward happen as usual; on a hit they reuse the cache.) + 3. ``save(...)`` -> ``save_cache_artifacts`` to the bundle + manifest, AFTER the + warmup forward, when in distributor/save mode. + 4. ``restore(...)`` -> put ``TORCHINDUCTOR_CACHE_DIR`` back on unload. + +Everything is env-gated and best-effort; torch is imported lazily. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import os +import time +from pathlib import Path +from typing import Any, Optional + +# ----------------------------------------------------------------------------- env knobs +# UNSLOTH_DIFFUSION_COMPILE_CACHE: auto (default) | 0 | 1 +# auto -> load a matching bundle if present (no automatic save). +# 1 -> load AND save (distributor / first-run warm). +# 0 -> disabled (plain local compile, no cache dir override). +# UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR: root dir for bundles (default under the workspace). +# UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE: 1 -> force-enable save even in "auto". +_ENV_MODE = "UNSLOTH_DIFFUSION_COMPILE_CACHE" +_ENV_DIR = "UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR" +_ENV_SAVE = "UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE" + +_DEFAULT_ROOT = Path.home() / ".cache" / "unsloth" / "diffusion_compile_cache" + +_MANIFEST_NAME = "manifest.json" +_BUNDLE_NAME = "cache.bin" +_FORMAT_VERSION = 1 + + +def cache_mode() -> str: + """``off`` | ``auto`` | ``on`` from the environment. ``auto`` is the default.""" + raw = (os.environ.get(_ENV_MODE) or "auto").strip().lower() + if raw in ("0", "off", "false", "no"): + return "off" + if raw in ("1", "on", "true", "yes"): + return "on" + return "auto" + + +def _save_enabled(mode: str) -> bool: + if mode == "off": + return False + if mode == "on": + return True + # auto: save only if explicitly opted in. + return (os.environ.get(_ENV_SAVE) or "").strip().lower() in ("1", "on", "true", "yes") + + +def cache_root() -> Path: + root = os.environ.get(_ENV_DIR) + return Path(root) if root else _DEFAULT_ROOT + + +# --------------------------------------------------------------------------- fingerprint +def _triton_version() -> Optional[str]: + try: + import triton # noqa: PLC0415 + return str(getattr(triton, "__version__", None)) + except Exception: # noqa: BLE001 — triton optional + return None + + +def _diffusers_version() -> Optional[str]: + try: + import diffusers # noqa: PLC0415 + return str(getattr(diffusers, "__version__", None)) + except Exception: # noqa: BLE001 + return None + + +def environment_fingerprint() -> dict[str, Any]: + """The HARD-portability dimensions: any difference here invalidates a bundle. + + These mirror what torch's inductor cache itself keys on (torch + triton + CUDA + + GPU type), plus diffusers (the graph source). We surface them explicitly so the + manifest is self-describing and a mismatch is obvious to a human, not just to torch. + """ + fp: dict[str, Any] = { + "format": _FORMAT_VERSION, + "torch": None, + "torch_cuda": None, + "triton": _triton_version(), + "diffusers": _diffusers_version(), + "gpu_name": None, + "gpu_capability": None, + } + try: + import torch # noqa: PLC0415 + + fp["torch"] = str(torch.__version__) + fp["torch_cuda"] = str(torch.version.cuda) + if torch.cuda.is_available(): + fp["gpu_name"] = torch.cuda.get_device_name(0) + cap = torch.cuda.get_device_capability(0) + fp["gpu_capability"] = f"sm_{cap[0]}{cap[1]}" + except Exception: # noqa: BLE001 — best-effort + pass + return fp + + +def model_fingerprint( + *, + family: Any, + transformer: Any, + dtype: Any, + quant: Any, + attention_backend: Any, + compile_kwargs: dict[str, Any], + shape_bucket: Any = None, +) -> dict[str, Any]: + """The MODEL-graph dimensions that change the compiled artifact. + + ``family`` is the Unsloth family name; ``transformer`` is the live module (we read + its class name + ``_repeated_blocks`` so the key tracks exactly what gets compiled). + """ + blocks = list(getattr(transformer, "_repeated_blocks", []) or []) + return { + "family": str(family), + "transformer_cls": type(transformer).__name__ if transformer is not None else None, + "repeated_blocks": sorted(str(b) for b in blocks), + "dtype": str(dtype), + "quant": str(quant) if quant is not None else "none", + "attention_backend": str(attention_backend) if attention_backend is not None else "default", + "compile_kwargs": {k: compile_kwargs[k] for k in sorted(compile_kwargs)}, + "shape_bucket": shape_bucket, + } + + +def cache_key(env_fp: dict[str, Any], model_fp: dict[str, Any]) -> str: + payload = json.dumps({"env": env_fp, "model": model_fp}, sort_keys = True, default = str) + return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:32] + + +# ----------------------------------------------------------------------------- lifecycle +@dataclasses.dataclass +class CacheContext: + """Carries the per-load cache state between ``begin`` and ``save``/``restore``.""" + + key: str + dir: Path + bundle: Path + manifest_path: Path + env_fp: dict[str, Any] + model_fp: dict[str, Any] + mode: str + hit: bool = False + saved: bool = False + prev_inductor_dir: Optional[str] = None + prev_inductor_dir_set: bool = False + + +def begin( + *, + family: Any, + transformer: Any, + dtype: Any, + quant: Any, + attention_backend: Any, + compile_kwargs: dict[str, Any], + shape_bucket: Any = None, + logger: Any = None, +) -> Optional[CacheContext]: + """Point inductor at a per-key dir and load a matching bundle, BEFORE compile. + + Returns a ``CacheContext`` to pass to ``save``/``restore``, or ``None`` when the + cache is disabled or torch lacks the Mega-cache API. Never raises. + """ + mode = cache_mode() + if mode == "off": + return None + try: + import torch # noqa: PLC0415 + if not ( + hasattr(torch.compiler, "save_cache_artifacts") + and hasattr(torch.compiler, "load_cache_artifacts") + ): + _warn(logger, "Mega-cache API unavailable (need torch >= 2.7); skipping") + return None + except Exception as exc: # noqa: BLE001 + _warn(logger, f"torch import failed: {exc}") + return None + + env_fp = environment_fingerprint() + model_fp = model_fingerprint( + family = family, + transformer = transformer, + dtype = dtype, + quant = quant, + attention_backend = attention_backend, + compile_kwargs = compile_kwargs, + shape_bucket = shape_bucket, + ) + key = cache_key(env_fp, model_fp) + cdir = cache_root() / key + ctx = CacheContext( + key = key, + dir = cdir, + bundle = cdir / _BUNDLE_NAME, + manifest_path = cdir / _MANIFEST_NAME, + env_fp = env_fp, + model_fp = model_fp, + mode = mode, + ) + + # Isolate inductor's on-disk cache per key so bundles never cross-contaminate. + try: + cdir.mkdir(parents = True, exist_ok = True) + ctx.prev_inductor_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR") + ctx.prev_inductor_dir_set = True + os.environ["TORCHINDUCTOR_CACHE_DIR"] = str(cdir / "inductor") + except Exception as exc: # noqa: BLE001 + _warn(logger, f"could not set TORCHINDUCTOR_CACHE_DIR: {exc}") + + # Try an exact-match load. A miss/mismatch is normal and non-fatal. + if ctx.bundle.exists() and ctx.manifest_path.exists(): + ctx.hit = _try_load(ctx, logger) + else: + _info(logger, f"compile-cache: no bundle for key {key} (will compile locally)") + return ctx + + +def _try_load(ctx: CacheContext, logger: Any) -> bool: + try: + manifest = json.loads(ctx.manifest_path.read_text()) + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: unreadable manifest: {exc}") + return False + + # Exact-match guard (defence in depth: torch also validates internally on load). + if manifest.get("env") != ctx.env_fp or manifest.get("model") != ctx.model_fp: + _warn(logger, "compile-cache: fingerprint mismatch; falling back to local compile") + return False + + try: + data = ctx.bundle.read_bytes() + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: cannot read bundle: {exc}") + return False + + # Integrity check (corruption / truncation; not a security signature). + digest = hashlib.sha256(data).hexdigest() + if manifest.get("sha256") and manifest["sha256"] != digest: + _warn(logger, "compile-cache: bundle checksum mismatch; ignoring") + return False + + try: + import torch # noqa: PLC0415 + + info = torch.compiler.load_cache_artifacts(data) + if info is None: + _warn(logger, "compile-cache: load_cache_artifacts returned None (no hit)") + return False + _info(logger, f"compile-cache: loaded bundle for key {ctx.key}") + return True + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: load failed: {exc}") + return False + + +def save(ctx: Optional[CacheContext], *, logger: Any = None) -> bool: + """Persist the compiled artifacts to the bundle + manifest, AFTER a warmup forward. + + No-op unless save is enabled (mode ``on`` or ``UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE``). + Returns True if a bundle was written. + """ + if ctx is None or not _save_enabled(ctx.mode) or ctx.saved: + return False + try: + import torch # noqa: PLC0415 + result = torch.compiler.save_cache_artifacts() + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: save_cache_artifacts failed: {exc}") + return False + if not result or result[0] is None: + _warn(logger, "compile-cache: nothing to save (empty artifacts)") + return False + + data = result[0] + try: + ctx.dir.mkdir(parents = True, exist_ok = True) + ctx.bundle.write_bytes(data) + manifest = { + "format": _FORMAT_VERSION, + "key": ctx.key, + "created": time.time(), + "bytes": len(data), + "sha256": hashlib.sha256(data).hexdigest(), + "env": ctx.env_fp, + "model": ctx.model_fp, + } + ctx.manifest_path.write_text(json.dumps(manifest, indent = 2, sort_keys = True, default = str)) + ctx.saved = True + _info(logger, f"compile-cache: saved bundle ({len(data)} bytes) for key {ctx.key}") + return True + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: could not write bundle: {exc}") + return False + + +def restore(ctx: Optional[CacheContext]) -> None: + """Restore ``TORCHINDUCTOR_CACHE_DIR`` to its pre-load value. Call on unload.""" + if ctx is None or not ctx.prev_inductor_dir_set: + return + try: + if ctx.prev_inductor_dir is None: + os.environ.pop("TORCHINDUCTOR_CACHE_DIR", None) + else: + os.environ["TORCHINDUCTOR_CACHE_DIR"] = ctx.prev_inductor_dir + except Exception: # noqa: BLE001 + pass + + +def _warn(logger: Any, msg: str) -> None: + if logger is not None: + logger.warning("diffusion.compile_cache: %s", msg) + + +def _info(logger: Any, msg: str) -> None: + if logger is not None: + logger.info("diffusion.compile_cache: %s", msg) diff --git a/studio/backend/core/inference/diffusion_eager_patches.py b/studio/backend/core/inference/diffusion_eager_patches.py new file mode 100644 index 0000000000..593f545984 --- /dev/null +++ b/studio/backend/core/inference/diffusion_eager_patches.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Reversible, failure-safe monkey-patches that speed up the diffusion denoiser. + +These patch a couple of SHARED diffusers building-block classes (so a tiny patch surface +covers all four Studio families) to make inference faster WITHOUT reducing accuracy and +WITHOUT regressing ``torch.compile``. Both patches are *compile-safe*: they lower to plain +torch ops, are bit-identical or FMA-1-ULP (i.e. equal-or-MORE accurate) vs stock diffusers, +and ``torch.compile`` fuses either form to the same kernel -- so they help the EAGER path +and are neutral under compile. + +What it patches (measured on a B200, bf16, DiT shapes): + +* ``normalization.RMSNorm.forward`` -> fused ``F.rms_norm`` on the common path (non-NPU, + no bias, weight None/fp16/bf16; else the exact original, incl. its fp32 quirk). This is + the standout win (~6-12x per call) because QK-norm runs every attention block on + Qwen-Image + Z-Image. Bit-identical in bf16. + +* ``AdaLayerNormContinuous`` / ``AdaLayerNormZero`` / ``AdaLayerNormZeroSingle`` -> + the ``norm(x)*(1+scale)+shift`` modulation fused via ``torch.addcmul`` (~1.2x). fp32 + exact; bf16 within 1 ULP (FMA, a single rounding -> more accurate than mul+add). Covers + flux.1 / flux.2-klein / qwen-image. + +Deliberately NOT patched (evidence-based): +* ``FeedForward`` (skipping its eval no-op Dropout) -- measured a small REGRESSION because + the per-module ``isinstance`` check costs more than the skipped identity dispatch. +* GEGLU / SwiGLU / GELU -- already mul-bound; a real win needs a custom Triton/CUDA kernel + (out of scope, correctness risk). +* Per-family custom MLP/norm classes and attention -- not shared (no leverage), and + attention already routes through ``F.scaled_dot_product_attention`` via the existing + ``set_attention_backend`` dispatcher. + +Lifecycle: ``install_compile_safe_patches()`` is idempotent and patches at the class level. +Install it for any active speed tier; the bit-identical ``off`` reference path must run with +the patches UNINSTALLED, so the caller uninstalls on an ``off`` load and on unload. The +Studio CHAT<->DIFFUSION arbiter guarantees a single active diffusion pipe, so class-level +state is safe. +""" + +from __future__ import annotations + +import logging +import os +from typing import Callable, Optional + +import torch +import torch.nn.functional as F + +from .diffusion_patch_backend import apply_patch, revert_patch + +logger = logging.getLogger(__name__) + +# Kill-switch: set UNSLOTH_DIFFUSION_EAGER_PATCHES=0 to disable the patches entirely (for +# A/B benchmarking or to rule them out while debugging). Enabled by default. +_ENV_ENABLE = "UNSLOTH_DIFFUSION_EAGER_PATCHES" + + +def _patches_enabled() -> bool: + return (os.environ.get(_ENV_ENABLE) or "").strip().lower() not in ("0", "off", "false", "no") + + +# --------------------------------------------------------------------------- # +# Resolve the diffusers classes we patch. Any import failure -> that patch is +# simply unavailable (None) and is skipped at install time. +# --------------------------------------------------------------------------- # +try: + from diffusers.models.normalization import ( + AdaLayerNormContinuous as _AdaLayerNormContinuous, + AdaLayerNormZero as _AdaLayerNormZero, + AdaLayerNormZeroSingle as _AdaLayerNormZeroSingle, + RMSNorm as _RMSNorm, + ) +except Exception: # noqa: BLE001 + _AdaLayerNormContinuous = _AdaLayerNormZero = _AdaLayerNormZeroSingle = _RMSNorm = None + +try: + from diffusers.utils.import_utils import is_torch_npu_available as _is_npu + _NPU = bool(_is_npu()) +except Exception: # noqa: BLE001 + _NPU = False + + +# --------------------------------------------------------------------------- # +# Patched forwards. Each mirrors diffusers 0.38 semantics, with the documented +# fused fast path. ``addcmul(input, t1, t2) == input + t1 * t2`` in one fused kernel. +# --------------------------------------------------------------------------- # +def _adaln_continuous_forward(self, x, conditioning_embedding): + emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim = 1) + # original: self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return torch.addcmul(shift[:, None, :], self.norm(x), 1 + scale[:, None, :]) + + +def _adaln_zero_forward( + self, + x, + timestep = None, + class_labels = None, + hidden_dtype = None, + emb = None, +): + if self.emb is not None: + emb = self.emb(timestep, class_labels, hidden_dtype = hidden_dtype) + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim = 1) + # original: self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + x = torch.addcmul(shift_msa[:, None], self.norm(x), 1 + scale_msa[:, None]) + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +def _adaln_zero_single_forward( + self, + x, + emb = None, +): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa = emb.chunk(3, dim = 1) + x = torch.addcmul(shift_msa[:, None], self.norm(x), 1 + scale_msa[:, None]) + return x, gate_msa + + +# Filled in at install time with the ORIGINAL RMSNorm.forward so the guarded fast path +# can fall back for the uncommon (NPU / bias / fp32-weight) cases. +_orig_rmsnorm_forward: Optional[Callable] = None + + +def _rmsnorm_forward(self, hidden_states): + # Fall back to the exact original for cases where F.rms_norm is NOT equivalent to + # diffusers' implementation: + # * NPU / bias / fp32-weight -> the original has special handling / an fp32 quirk; + # * tuple `dim` -> diffusers always reduces the LAST dim (`mean(-1)`) while + # F.rms_norm reduces every dim in `self.dim` (differs for a multi-dim shape); + # * dtype mismatch (e.g. fp32 activations into an fp16/bf16-weight norm) -> diffusers + # computes the variance in fp32 from the ORIGINAL tensor and only casts before the + # weight multiply, so casting first would change the variance. + if _NPU or self.bias is not None or _orig_rmsnorm_forward is None or len(tuple(self.dim)) != 1: + return _orig_rmsnorm_forward(self, hidden_states) # type: ignore[misc] + weight = self.weight + if weight is None: + return F.rms_norm(hidden_states, self.dim, None, self.eps) + if weight.dtype in (torch.float16, torch.bfloat16) and hidden_states.dtype == weight.dtype: + # Common DiT path (bf16 activations + bf16 weight): F.rms_norm matches diffusers + # bit-for-bit (both reduce the variance in fp32 internally), just fused. + return F.rms_norm(hidden_states, self.dim, weight, self.eps) + # Mixed dtype / fp32 weight -> keep the exact original behaviour. + return _orig_rmsnorm_forward(self, hidden_states) # type: ignore[misc] + + +# --------------------------------------------------------------------------- # +# Install / uninstall. All swaps go through the shared patch backend +# (unsloth_zoo patch_function / restore_original): the live original is fingerprinted +# (can_safely_patch, relaxed) so a diffusers that renamed/reordered a forward's params is +# left UNPATCHED instead of miscompiled, and the original is stashed for an exact restore. +# --------------------------------------------------------------------------- # +def _specs(): + # (class, patched_fn) + return [ + (_AdaLayerNormContinuous, _adaln_continuous_forward), + (_AdaLayerNormZero, _adaln_zero_forward), + (_AdaLayerNormZeroSingle, _adaln_zero_single_forward), + (_RMSNorm, _rmsnorm_forward), + ] + + +# Classes whose `forward` we successfully patched, so uninstall reverts exactly those. +_patched: list[type] = [] + + +def install_compile_safe_patches() -> int: + """Install the shared compile-safe speedup patches (idempotent). + + Returns the number of patches applied. A second call while installed is a no-op. + """ + global _orig_rmsnorm_forward + if not _patches_enabled(): + uninstall_patches() # ensure OFF even if a prior call installed them + return 0 + if _patched: + return len(_patched) + for cls, new_fn in _specs(): + if cls is None: + continue + # Capture the live original BEFORE patching so the RMSNorm fast path can fall back + # to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases. + if cls is _RMSNorm: + _orig_rmsnorm_forward = cls.forward + if apply_patch(cls, "forward", new_fn, match_level = "relaxed"): + _patched.append(cls) + else: + logger.warning( + "eager-patch: skipping %s (signature mismatch / unavailable)", + getattr(cls, "__name__", cls), + ) + if cls is _RMSNorm: + _orig_rmsnorm_forward = None + logger.info( + "eager-patch: installed %d/%d shared diffusion patches", len(_patched), len(_specs()) + ) + return len(_patched) + + +def uninstall_patches() -> None: + """Restore every patched class to its exact original ``forward`` (idempotent).""" + global _orig_rmsnorm_forward + for cls in list(_patched): + revert_patch(cls, "forward") + _patched.clear() + _orig_rmsnorm_forward = None + + +def is_installed() -> bool: + """True if any compile-safe patch is currently installed.""" + return bool(_patched) diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index ffe31d93a5..957df10ef7 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -117,7 +117,12 @@ def _activate(name: str, reason: Optional[str]) -> Any: return get_active_diffusion_engine() -def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] = None) -> Any: +def select_and_activate_engine( + fam: DiffusionFamily, + *, + hf_token: Optional[str] = None, + model_kind: Optional[str] = None, +) -> Any: """Pick + activate the engine for loading ``fam`` on this host; return the engine. Falls back to diffusers (recording a reason) whenever the native route is @@ -125,6 +130,12 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] native asset mapping, or the sd-cli binary is unavailable -- always BEFORE the slow load begins, so a fallback never strands a half-native load. """ + # Non-GGUF loads (a single-file safetensors transformer, or a full diffusers + # pipeline) only run on diffusers: the native sd.cpp engine consumes single-file + # GGUF checkpoints only, so force diffusers before the device/native checks below. + if model_kind and model_kind != "gguf": + return _activate(ENGINE_DIFFUSERS, f"non-GGUF load ({model_kind}) requires diffusers") + forced, sd_cpp_pref, mps_enabled = _engine_config() if forced == ENGINE_DIFFUSERS: diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 7f72743d35..2f9ba12c62 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -37,6 +37,32 @@ class DiffusionFamily: # Pipeline kwarg carrying the guidance value. Most use "guidance_scale"; # Qwen-Image's distilled guidance is off, so its real CFG is "true_cfg_scale". cfg_kwarg: str = "guidance_scale" + # Optional diffusers pipeline classes for image-conditioned workflows. The backend + # builds these around the ALREADY-loaded transformer/VAE/text-encoder via + # ``Pipeline.from_pipe`` (no extra weights, no reload), so a family only needs the + # class name here to gain the workflow. None = the family does not support it (the + # UI gates the workflow off). The base text-to-image pipeline is ``pipeline_class``. + img2img_pipeline_class: Optional[str] = None + inpaint_pipeline_class: Optional[str] = None + # True when the inpaint pipeline keeps the input canvas size, so it can also drive + # outpaint (extend), where the padded canvas is LARGER than the original. False for + # FLUX.2 (its pipelines scale any >1MP input down to ~1MP, which shrinks an outpaint + # canvas back and defeats the extend). Such families get Inpaint but not Extend. + inpaint_preserves_size: bool = True + # True for instruction-editing families (Qwen-Image-Edit / FLUX Kontext): the model's + # OWN pipeline (``pipeline_class``) is the edit pipeline -- it takes an input image plus + # a text instruction and has no plain text-to-image mode. So these expose only the + # "edit" workflow, require an input image at generate time, and the loaded pipe is used + # directly (no from_pipe). ``base_repo`` here is the matching diffusers repo that + # supplies the VAE / text-encoder / processor / scheduler for the GGUF transformer. + edit: bool = False + # True for families whose OWN text-to-image pipeline ALSO accepts reference image(s) + # (FLUX.2: Flux2KleinPipeline takes an optional ``image`` arg). Unlike ``edit`` these + # families still do plain text-to-image (no image), and unlike img2img the conditioning + # is reference-based, not a denoise blend: there is no ``strength`` and the output size + # comes from the requested width/height, not the reference's size. The loaded pipe is + # used directly (no from_pipe). Exposes a "reference" workflow alongside "txt2img". + reference: bool = False # Extra lowercased substrings (besides ``name``) that map a repo id here. aliases: tuple[str, ...] = field(default_factory = tuple) # True for families whose activations overflow float16's finite range @@ -78,9 +104,9 @@ class DiffusionFamily: # Keyed by architecture, not per model variant: a checkpoint's specific base repo # is read from its HF base_model tag at load time, so one entry covers Turbo/full, # schnell/dev, etc. base_repo here is only a fallback. Only archs whose diffusers -# transformer supports from_single_file load here (ERNIE-Image does not, yet). -# FLUX.2-dev and FLUX.2-klein-9B are left out only because their base diffusers -# repos are gated; the open klein-4B base stands in for the klein family below. +# transformer supports from_single_file load here (ERNIE-Image does not, yet; LTX +# video models are out of scope). FLUX.2-klein-9B shares the klein family (its base +# repo is resolved per-variant), and FLUX.2-dev has its own family below. _FAMILIES: tuple[DiffusionFamily, ...] = ( DiffusionFamily( name = "flux.1", @@ -88,6 +114,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "FluxTransformer2DModel", base_repo = "black-forest-labs/FLUX.1-schnell", aliases = ("flux1", "flux-1"), + img2img_pipeline_class = "FluxImg2ImgPipeline", + inpaint_pipeline_class = "FluxInpaintPipeline", sd_cpp_vae = ("black-forest-labs/FLUX.1-schnell", "ae.safetensors"), sd_cpp_text_encoders = ( ("comfyanonymous/flux_text_encoders", "clip_l.safetensors", "clip_l"), @@ -95,14 +123,21 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( ), ), # FLUX.2-klein is a distinct pipeline (Flux2KleinPipeline) with a Qwen3 text - # encoder, not the Mistral-based Flux2Pipeline; it must precede a generic - # flux match. The base Flux2Pipeline (FLUX.2-dev) is gated, so it's omitted. + # encoder, not the Mistral-based Flux2Pipeline; it must precede a generic flux + # match. The Mistral-based Flux2Pipeline is the separate flux.2-dev family below. DiffusionFamily( name = "flux.2-klein", pipeline_class = "Flux2KleinPipeline", transformer_class = "Flux2Transformer2DModel", base_repo = "black-forest-labs/FLUX.2-klein-4B", aliases = ("flux2-klein",), + # Flux2KleinPipeline natively accepts reference image(s) via its `image` arg, so it + # exposes a "reference" workflow on top of plain text-to-image. It has a dedicated + # inpaint pipeline too (no img2img one), so it also gets inpaint + extend (outpaint). + reference = True, + inpaint_pipeline_class = "Flux2KleinInpaintPipeline", + # FLUX.2 scales >1MP inputs down to ~1MP, so outpaint (a larger canvas) can't grow. + inpaint_preserves_size = False, # FLUX.2 uses a distinct 32-channel autoencoder; sd-cli needs the latent # format override. The single-file VAE ships in Comfy-Org/flux2-dev (the # klein-4B repo only has a sharded diffusers VAE). Shares Qwen3-4B with z-image. @@ -112,6 +147,61 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"), ), ), + # FLUX.2-dev is the full (non-distilled) FLUX.2. It uses the Mistral-based + # Flux2Pipeline, distinct from klein's Qwen3-based Flux2KleinPipeline, so it needs + # its own entry. Its base diffusers repo is gated (gated=auto) but reachable with an + # HF token. text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint + # pipeline for dev. VAE + Mistral text encoder come from the open Comfy-Org/flux2-dev + # mirror for the sd-cli path (shares the FLUX.2 32-channel AE with klein). + DiffusionFamily( + name = "flux.2-dev", + pipeline_class = "Flux2Pipeline", + transformer_class = "Flux2Transformer2DModel", + base_repo = "black-forest-labs/FLUX.2-dev", + aliases = ("flux2-dev", "flux2dev"), + sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"), + sd_cpp_vae_format = "flux2", + sd_cpp_text_encoders = ( + ( + "Comfy-Org/flux2-dev", + "split_files/text_encoders/mistral_3_small_flux2_bf16.safetensors", + "llm", + ), + ), + ), + DiffusionFamily( + # Instruction editing with FLUX. FluxKontextPipeline takes an input image + an edit + # instruction; the GGUF transformer is the standard FluxTransformer2DModel, with the + # T5/CLIP text encoders + VAE from the base diffusers repo. cfg defaults to + # guidance_scale (FLUX). Most-specific aliases first so detect_family prefers this + # over the plain "flux.1" family and un-rejects the "kontext" keyword for it. + name = "flux.1-kontext", + pipeline_class = "FluxKontextPipeline", + transformer_class = "FluxTransformer2DModel", + base_repo = "black-forest-labs/FLUX.1-Kontext-dev", + aliases = ("flux.1-kontext-dev", "flux1-kontext", "flux-kontext", "kontext"), + edit = True, + ), + DiffusionFamily( + # Instruction editing (image-in + text-instruction-out). The 2511 checkpoint ships + # as QwenImageEditPlusPipeline (multi-image-capable); the GGUF transformer is the + # standard QwenImageTransformer2DModel, with the VAE / Qwen2.5-VL text-encoder / + # image processor / scheduler coming from the base diffusers repo. Most-specific + # aliases first so detect_family prefers this over the plain "qwen-image" family. + name = "qwen-image-edit", + pipeline_class = "QwenImageEditPlusPipeline", + transformer_class = "QwenImageTransformer2DModel", + base_repo = "Qwen/Qwen-Image-Edit-2511", + cfg_kwarg = "true_cfg_scale", + aliases = ( + "qwen-image-edit-2511", + "qwen-image-edit-2509", + "qwen-image-edit", + "qwen_image_edit", + "qwenimageedit", + ), + edit = True, + ), DiffusionFamily( name = "qwen-image", pipeline_class = "QwenImagePipeline", @@ -119,6 +209,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( base_repo = "Qwen/Qwen-Image", cfg_kwarg = "true_cfg_scale", aliases = ("qwen_image", "qwenimage"), + img2img_pipeline_class = "QwenImageImg2ImgPipeline", + inpaint_pipeline_class = "QwenImageInpaintPipeline", sd_cpp_vae = ("Comfy-Org/Qwen-Image_ComfyUI", "split_files/vae/qwen_image_vae.safetensors"), # The Qwen2.5-VL text encoder as a Q4_K_M GGUF keeps the CPU RAM win (the # bf16 safetensors encoder is ~15 GB). sd-cli's --qwen2vl is an alias of --llm. @@ -139,6 +231,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "ZImageTransformer2DModel", base_repo = "Tongyi-MAI/Z-Image-Turbo", aliases = ("zimage", "z_image"), + img2img_pipeline_class = "ZImageImg2ImgPipeline", + inpaint_pipeline_class = "ZImageInpaintPipeline", # Z-Image's MLP down-projections peak near 9e5, which overflows float16. fp16_incompatible = True, sd_cpp_vae = ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"), @@ -150,15 +244,45 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( # Editing / inpaint checkpoints share an arch keyword but need a different # pipeline and an input image, which this text-to-image backend doesn't drive. -_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "inpainting") +# "layered" rejects Qwen-Image-Layered: its transformer sets additional_t_cond=True +# and expects an extra addition_t_cond input that the standard QwenImagePipeline +# never supplies, so it loads but crashes at the first denoise step. Rejecting it +# here fails the load fast with a clear message and hides it from the picker. +_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "layered") + + +def _token_in_needle(token: str, needle: str) -> bool: + """True when ``token`` appears in ``needle`` as a whole path/name segment, i.e. + delimited by a separator (``- _ . / \\``) or a string boundary, not merely as a + raw substring. This keeps multi-part tokens matching where they should + ('qwen-image-edit' in 'qwen-image-edit-2511') while preventing a short token from + matching inside an unrelated word ('kontext' must not match 'kontextual', 'edit' + must not match 'edition').""" + return re.search(r"(?:^|[-_./\\])" + re.escape(token) + r"(?:$|[-_./\\])", needle) is not None + + +def _best_family_match(needle: str) -> Optional[DiffusionFamily]: + """The family whose name/alias is the LONGEST whole-segment token of ``needle``. + Longest = most specific, so an edit checkpoint ('...qwen-image-edit-2511...') + matches the 'qwen-image-edit' family rather than the generic 'qwen-image' one. + Segment matching (not raw substring) stops a short alias like 'kontext' from + hijacking an unrelated path such as '.../kontextual/z-image-...gguf'.""" + best: Optional[tuple[DiffusionFamily, int]] = None + for fam in _FAMILIES: + for token in (fam.name, *fam.aliases): + if _token_in_needle(token, needle) and (best is None or len(token) > best[1]): + best = (fam, len(token)) + return best[0] if best else None def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[DiffusionFamily]: """Resolve a ``DiffusionFamily`` from a repo id, or an explicit override. - ``override`` matches a family ``name`` or alias exactly; otherwise the repo - id is scanned for the first family whose name/alias appears in it. Image - editing checkpoints are rejected (None) since this backend is text-to-image. + ``override`` matches a family ``name`` or alias exactly. Otherwise the most-specific + family whose name/alias is a substring of the repo id wins. Supported editing families + (Qwen-Image-Edit) match here; unsupported editing/inpaint/layered checkpoints that only + share a base family's arch keyword are still rejected (None), because they need a + different pipeline + input this backend's base text-to-image path doesn't drive. """ if override: key = override.strip().lower() @@ -167,24 +291,33 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff return fam return None needle = repo_id.lower() - # Match edit keywords as whole id segments, not raw substrings, so a normal - # text-to-image repo like ".../some-image-edition" isn't misread as an editing - # checkpoint. Qwen-Image-Edit / FLUX.1-Kontext still match (edit/kontext are - # whole tokens there). Scope the check to the LAST path component (the model id - # or filename), not arbitrary parent directories: a valid file selected as - # repo_id `/models/edit` + filename `Z-Image-Turbo-Q4.gguf` must not be rejected - # just because a parent folder happens to be named `edit`. The combined - # `repo_id/gguf_filename` fallback passes the filename as that last segment. - basename = re.split(r"[/\\]+", needle)[-1] - segments = set(re.split(r"[-_.]+", basename)) - if any(kw in segments for kw in _EDIT_KEYWORDS): - return None - for fam in _FAMILIES: - if fam.name in needle or any(alias in needle for alias in fam.aliases): - return fam + match = _best_family_match(needle) + if match is not None: + # Don't let a generic base family (e.g. qwen-image) swallow a variant it can't run + # (qwen-image-LAYERED, ...-Inpaint): if the id still carries a reject keyword the + # matched family does not itself declare, reject so the load fails fast + clearly. + # Scope the keyword check to the LAST path component (the model id or + # filename), not arbitrary parent directories: a valid file selected as + # repo_id `/models/edit` + filename `Z-Image-Turbo-Q4.gguf` must not be + # rejected because a parent folder happens to be named `edit`. The + # combined `repo_id/gguf_filename` fallback passes the filename last. + basename = re.split(r"[/\\]+", needle)[-1] + matched_tokens = (match.name, *match.aliases) + if any( + _token_in_needle(kw, basename) and not any(kw in tok for tok in matched_tokens) + for kw in _EDIT_KEYWORDS + ): + return None + return match return None +def supported_family_names() -> tuple[str, ...]: + """Family names accepted as ``family_override`` and shown in the unknown-model + error. Kept in registry order so the message lists what the backend can load.""" + return tuple(fam.name for fam in _FAMILIES) + + def detect_family_for_pick( repo_id: str, gguf_filename: Optional[str] = None, diff --git a/studio/backend/core/inference/diffusion_gguf_compile.py b/studio/backend/core/inference/diffusion_gguf_compile.py new file mode 100644 index 0000000000..aef086e734 --- /dev/null +++ b/studio/backend/core/inference/diffusion_gguf_compile.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGUF dequant accelerator for the light-compile (``default``) diffusion path. + +Profiling (outputs/profile_eager/) showed that ~70-80% of EAGER GGUF denoise CUDA +time is the per-forward weight dequant: every ``GGUFLinear.forward`` calls +``diffusers.quantizers.gguf.utils.dequantize_gguf_tensor`` -> ``dequantize_blocks_Q4_K``, +a ~20-op pure-PyTorch chain (nibble shifts + masks + block-scale mul + zero-point sub + +assembly), once per linear per step. + +COMPILED DEQUANT (``install_compiled_dequant``): swap ``dequantize_gguf_tensor`` for +``torch.compile(orig, dynamic=True)``. Inductor fuses the op chain into a few kernels. +Measured 1.24-1.64x warm with a small one-time compile (~7.5-10.4s) and ZERO extra VRAM -- +the weights stay quantized. ``dynamic=True`` is key: the dequant inputs are the WEIGHT +tensors (fixed shapes, independent of image resolution / batch), so it compiles once and +never recompiles on a resolution change. ``GGUFLinear.forward_native`` resolves the +function as a module global, so replacing the module attribute reroutes every linear. + +It is the ``default`` tier's lever (the transformer block stays eager). It is deliberately +NOT used under ``max`` (full regional block compile): there the block is compiled as one +graph which fuses the dequant inline, and a separately-compiled dequant would be traced +into that graph and break it -- so ``max`` runs the stock dequant and lets the block +compile fuse it. + +The swap goes through the shared, fingerprint-checked, reversible patch backend +(``diffusion_patch_backend``); ``uninstall_*`` restores the exact original so a later +bit-identical ``off`` load runs the stock dequant. Kill-switch: +``UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT=0``. torch / diffusers are imported lazily. + +(A global weight-buffer accelerator lived here too but was removed: it measured neutral +end-to-end -- the CUDA caching allocator already serves the per-forward cast allocation +from its pool with zero ``cudaMalloc`` churn, so reusing one buffer saved nothing on a +DiT's compute-bound forward. See outputs/arch_patch/SUMMARY.md.) +""" + +from __future__ import annotations + +import os +from typing import Any + +from .diffusion_patch_backend import apply_patch, revert_patch + +# --- kill-switch ------------------------------------------------------------------- + +_ENV_COMPILE_DEQUANT = "UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT" +_DISABLED = {"0", "off", "false", "no"} + + +def _enabled(env_name: str) -> bool: + """Enabled unless explicitly disabled (default ON).""" + return str(os.environ.get(env_name, "1")).strip().lower() not in _DISABLED + + +def _gguf_utils(): + """The diffusers GGUF utils module, or None if this diffusers build lacks it.""" + try: + from diffusers.quantizers.gguf import utils as gguf_utils # noqa: PLC0415 + return gguf_utils + except Exception: # noqa: BLE001 — old/!GGUF diffusers -> accelerator is a no-op + return None + + +# --- compiled dequant -------------------------------------------------------------- + +# True while our compiled wrapper is installed (the shared patch backend stashes the +# original on the module for an exact restore). +_compiled_dequant_installed = False +_DEQUANT_ATTR = "dequantize_gguf_tensor" + + +def is_compiled_dequant_installed() -> bool: + return _compiled_dequant_installed + + +def install_compiled_dequant(logger: Any = None) -> bool: + """Replace ``dequantize_gguf_tensor`` with ``torch.compile(orig, dynamic=True)`` via the + shared patch backend (original stashed for restore). + + Idempotent (a second call is a no-op while installed). Returns True if the compiled + dequant is in place afterwards, False if disabled / unavailable / it failed.""" + global _compiled_dequant_installed + if not _enabled(_ENV_COMPILE_DEQUANT): + return False + if _compiled_dequant_installed: + return True + gguf_utils = _gguf_utils() + if gguf_utils is None or not hasattr(gguf_utils, _DEQUANT_ATTR): + return False + try: + import torch # noqa: PLC0415 + + compiled = torch.compile(gguf_utils.dequantize_gguf_tensor, dynamic = True) + # force=True: the new callable is the SAME function compiled, so its fingerprint + # differs from the original and can_safely_patch would (correctly) reject it. + if apply_patch(gguf_utils, _DEQUANT_ATTR, compiled, force = True): + _compiled_dequant_installed = True + return True + return False + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "install_compiled_dequant", exc) + _compiled_dequant_installed = False + return False + + +def uninstall_compiled_dequant() -> None: + """Restore the original ``dequantize_gguf_tensor``. Idempotent.""" + global _compiled_dequant_installed + if not _compiled_dequant_installed: + return + gguf_utils = _gguf_utils() + if gguf_utils is not None: + revert_patch(gguf_utils, _DEQUANT_ATTR) + _compiled_dequant_installed = False + + +# --- convenience ------------------------------------------------------------------- + + +def uninstall_all() -> None: + """Uninstall the GGUF accelerator. Idempotent; safe to call on every unload.""" + uninstall_compiled_dequant() + + +def is_installed() -> bool: + return is_compiled_dequant_installed() + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.gguf_compile: %s failed: %s", what, exc) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 656e6abfcf..c56d1592ce 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -236,6 +236,17 @@ def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]: return int(storage_mib * 1.05) # small margin for allocator + bf16 norms/biases +def estimate_safetensors_dense_mib(storage_mib: Optional[int]) -> Optional[int]: + """Resident size of a safetensors checkpoint, in MiB. + + Unlike a GGUF (which is dequantised to bf16/fp16 on load, so a 4-bit file + expands ~4x), a safetensors checkpoint loads near its on-disk size: a dense + bf16 file is already bf16, and a bnb-4bit / fp8 file stays compressed in VRAM. + So the on-disk size is the estimate, returned unchanged (None passes through). + """ + return storage_mib + + def estimate_image_runtime_mib( *, width: Optional[int], diff --git a/studio/backend/core/inference/diffusion_patch_backend.py b/studio/backend/core/inference/diffusion_patch_backend.py new file mode 100644 index 0000000000..c133bfaf47 --- /dev/null +++ b/studio/backend/core/inference/diffusion_patch_backend.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""One vetted path for every diffusion monkey-patch. + +Thin wrappers over ``unsloth_zoo.temporary_patches.utils`` ``patch_function`` / +``restore_original`` so all of the backend's runtime patching (eager fusions, GGUF +accelerators, per-arch block rewrites) goes through the SAME fingerprint-checked, +reversible mechanism: + +* ``patch_function`` stores the live original under a unique attribute and, unless + ``force=True``, runs ``can_safely_patch`` (a parameter-name/kind/required fingerprint; + ``match_level="relaxed"`` ignores type-annotation drift but still rejects a real + signature change) -- so a future diffusers/transformers that renamed or reordered a + forward's parameters is simply left unpatched instead of silently miscompiled. +* ``restore_original`` puts the stored original back -- exact, idempotent uninstall. + +``unsloth_zoo`` is imported LAZILY inside each call, never at module import: it runs GPU +detection at import time and raises without an accelerator (set ``UNSLOTH_ALLOW_CPU=1`` to +bypass, as the test conftest does), and the diffusion backend must stay importable on a +CPU-only host. This mirrors the backend's existing deferred unsloth_zoo imports (see +``core/training/trainer.py``, ``core/export/export.py``). If the import fails (unsloth_zoo +absent, or a no-GPU host without the bypass), patching is a best-effort no-op: the stock +forward runs, correctness is preserved, only the optimisation is skipped. +""" + +from __future__ import annotations + +from typing import Any + + +def apply_patch( + target: Any, + attr: str, + new_fn: Any, + *, + match_level: str = "relaxed", + force: bool = False, +) -> bool: + """Patch ``target.attr -> new_fn`` via ``unsloth_zoo`` ``patch_function`` (the original is + stashed for ``revert_patch``). Returns True iff the swap was applied. Returns False + (never raises) if unsloth_zoo is unavailable or ``can_safely_patch`` rejects the swap. + + ``force=True`` skips the safety check -- use it only when the new callable is the SAME + function transformed (e.g. its ``torch.compile`` wrapper), where a fingerprint mismatch + is expected and benign.""" + try: + from unsloth_zoo.temporary_patches.utils import patch_function + except Exception: # noqa: BLE001 — no unsloth_zoo / no-GPU host -> optimisation skipped + return False + try: + return bool(patch_function(target, attr, new_fn, match_level = match_level, force = force)) + except Exception: # noqa: BLE001 — best-effort; leave the original in place + return False + + +def revert_patch(target: Any, attr: str) -> bool: + """Restore ``target.attr`` from the original stashed by ``apply_patch``. Idempotent and + best-effort: returns False (never raises) if there is nothing stored or unsloth_zoo is + unavailable.""" + try: + from unsloth_zoo.temporary_patches.utils import restore_original + except Exception: # noqa: BLE001 + return False + try: + return bool(restore_original(target, attr)) + except Exception: # noqa: BLE001 + return False diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 99f506450a..9c99de9406 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -6,22 +6,34 @@ Off by default, so the default render path stays bit-identical to a plain run (the property the regression harness checks). When the operator opts in, this applies the near-lossless speedups in the order the diffusers guides recommend -(channels_last + cudnn.benchmark -> regional compile, with TF32 / fused-QKV under -"max"): +(channels_last + cudnn.benchmark -> compile, with TF32 / fused-QKV under "max"): off - nothing (default; bit-identical reference). - default - near-lossless: channels_last VAE memory format + cudnn.benchmark conv - autotune + regional torch.compile of the denoiser's repeated block WHERE - eligible (bf16, CUDA, a compile-friendly family). Compile is the big win - (~2.3x denoise on the GGUF Z-Image transformer, PSNR ~36 dB vs eager, - well above the Q4 quantisation noise floor, so it does not meaningfully - move output quality). - max - default plus near-lossless TF32 matmul and fused QKV projections. + eager - everything lossless EXCEPT torch.compile: channels_last VAE + + cudnn.benchmark + the attention backend + the shared eager monkey-patches + (fused RMSNorm / AdaLayerNorm + per-arch addcmul fusions, see + diffusion_eager_patches.py / diffusion_arch_patches.py). The fast first-image + / casual-use path -- no compile tax to amortise. + default - LIGHT compile. For a GGUF model: channels_last + cudnn.benchmark + + torch.compile of ONLY the dequant op chain + (``torch.compile(dequantize_gguf_tensor, dynamic=True)``) -- the dequant is + ~70-80% of eager GGUF time, so fusing it gives ~1.24-1.64x for a small + one-time compile (~7.5-10.4s) and ZERO extra VRAM, resolution-invariant + (the dequant inputs are fixed-shape weights). For a dense (non-GGUF) model + there is no dequant, so ``default`` falls back to regional torch.compile of + the denoiser's repeated block (the only compile lever a dense model has). + max - the FULL torch.compile: regional max-autotune compile of the denoiser's + repeated block (which fuses the GGUF dequant AND the matmul/norm/elementwise + in one graph -- ~3.2x on the GGUF Z-Image transformer, PSNR ~36 dB vs eager, + well above the Q4 noise floor) plus TF32 matmul and fused QKV projections. -Regional compile used to be gated off for the GGUF transformer, but it compiles and -runs faster on the current diffusers/torch (measured; the GGUF dequant ops stay -eager and the rest of the repeated block compiles), so the GGUF gate is removed; the -per-family ``supports_torch_compile`` flag and the bf16/CUDA checks still apply. +Tier rationale: ``default`` is the cheap, always-amortising compile (compile just the +hot GGUF dequant; the block stays eager) so the first image is fast and VRAM is +untouched; ``max`` pays the larger regional-compile tax for the bigger warm speedup. +The compiled dequant is deliberately skipped under ``max`` -- the regional block compile +subsumes the dequant fusion (a separately-compiled dequant would be traced into that +graph and break it), so ``max`` runs the stock dequant and lets the block compile it. The +per-family ``supports_torch_compile`` flag and the bf16/CUDA checks gate regional compile. The backend flags this layer flips (TF32, cudnn.benchmark) are PROCESS-WIDE, so ``snapshot_backend_flags`` / ``restore_backend_flags`` let the caller capture the @@ -34,10 +46,13 @@ from __future__ import annotations from typing import Any, Optional +from . import diffusion_gguf_compile as gguf_compile + SPEED_OFF = "off" +SPEED_EAGER = "eager" SPEED_DEFAULT = "default" SPEED_MAX = "max" -SPEED_MODES = (SPEED_OFF, SPEED_DEFAULT, SPEED_MAX) +SPEED_MODES = (SPEED_OFF, SPEED_EAGER, SPEED_DEFAULT, SPEED_MAX) def snapshot_backend_flags() -> Optional[dict]: @@ -105,12 +120,13 @@ def normalize_speed_mode(value: Optional[str]) -> str: def resolve_speed_mode(value: Optional[str], *, is_gguf: bool) -> str: """The effective speed mode when the caller leaves it UNSET (``None``). - A GGUF model defaults to ``default``: regional compile is ~2.2x faster and its - numeric perturbation sits well below the quantisation noise floor (measured - PSNR ~37 dB compile-vs-eager versus ~21 dB Q4-vs-bf16), so it does not reduce - output quality relative to the dense reference. A dense (non-GGUF) model stays - ``off`` / bit-identical, since there compile would be the only source of drift. - An explicit value -- including ``"off"`` -- is always honored verbatim.""" + A GGUF model defaults to ``default``: it compiles only the hot dequant op chain + (~70-80% of eager GGUF time) for ~1.24-1.64x at a small one-time compile and zero + extra VRAM -- a cheap, always-amortising win whose numeric perturbation sits well + below the quantisation noise floor (the dequant graph is unchanged, just + Inductor-fused). A dense (non-GGUF) model stays ``off`` / bit-identical, since there + compile would be the only source of drift. An explicit value -- including ``"off"`` + -- is always honored verbatim.""" if value is None: return SPEED_DEFAULT if is_gguf else SPEED_OFF return normalize_speed_mode(value) @@ -165,6 +181,7 @@ def apply_speed_optims( "tf32": False, "fused_qkv": False, "compiled": False, + "compiled_dequant": False, } mode = normalize_speed_mode(speed_mode) # TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max, @@ -175,29 +192,49 @@ def apply_speed_optims( if mode == SPEED_OFF: return applied + on_cuda = getattr(target, "device", None) == "cuda" + family_allows_compile = bool(getattr(family, "supports_torch_compile", True)) + # Lossless: a channels-last VAE speeds up its convolutions with no numeric change. applied["channels_last"] = _vae_channels_last(pipe, logger) # Near-lossless: let cuDNN autotune the fixed-shape VAE convs (CUDA only). It may # pick a different conv algorithm, so it is a "default"-tier (not bit-identical) win. - if getattr(target, "device", None) == "cuda": + if on_cuda: applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger) - # Near-lossless and the largest win: regional compile of the repeated denoiser - # block, where eligible (now incl. the GGUF transformer). `max` opts into - # max-autotune (longer compile, autotuned kernels). - if compile_eligible(target, is_gguf = is_gguf, family = family): + # --- the compile lever, remapped per tier ---------------------------------------- + # default = LIGHT compile: for a GGUF model, compile ONLY the dequant op chain + # (~70-80% of eager GGUF time) -- cheap, VRAM-free, resolution-invariant; the + # transformer block stays eager. A dense model has no dequant, so default falls + # back to the regional block compile (its only compile lever). + # max = FULL compile: regional max-autotune compile of the repeated denoiser block + # (fuses dequant + matmul + norm + elementwise in one graph). It subsumes the + # dequant fusion, so we do NOT also install the standalone compiled dequant here. + # eager = no compile at all. + if mode == SPEED_DEFAULT: + if is_gguf and on_cuda and family_allows_compile: + applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger) + elif compile_eligible(target, is_gguf = is_gguf, family = family): + applied["compiled"] = _compile_repeated_blocks( + pipe, + logger, + max_autotune = False, + cache_active = cache_active, + offload_active = offload_active, + ) + elif mode == SPEED_MAX and compile_eligible(target, is_gguf = is_gguf, family = family): applied["compiled"] = _compile_repeated_blocks( pipe, logger, - max_autotune = mode == SPEED_MAX, + max_autotune = True, cache_active = cache_active, offload_active = offload_active, ) if mode == SPEED_MAX: # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed. - if getattr(target, "device", None) == "cuda": + if on_cuda: applied["tf32"] = _enable_tf32(logger) applied["fused_qkv"] = _fuse_qkv(pipe, logger) diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 40a5613273..ae729a881a 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -40,6 +40,7 @@ from core.inference.diffusion_memory import ( _TE_FLAGS_BY_FAMILY: dict[str, tuple[str, ...]] = { "z-image": ("--llm",), "flux.2-klein": ("--llm",), + "flux.2-dev": ("--llm",), "qwen-image": ("--qwen2vl",), "flux.1": ("--clip_l", "--t5xxl"), } diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 53f7131e16..757a28a3a0 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -42,6 +42,7 @@ from core.inference.diffusion_families import ( family_sd_cpp_supported, resolve_base_repo, resolve_local_gguf_child, + supported_family_names, ) from core.inference.diffusion_memory import ( OFFLOAD_GROUP, @@ -270,7 +271,7 @@ def _map_guidance( classifier-free ``--cfg-scale``. A distilled 0/1 means CFG off (sd-cli's 1.0); a value > 1 is real CFG. Mirrors the engine mapping validated in the CPU benchmark. """ - if fam.name in ("flux.1", "flux.2-klein"): + if fam.name in ("flux.1", "flux.2-klein", "flux.2-dev"): return None, (float(guidance) if guidance is not None else None) cfg = float(guidance) if (guidance is not None and guidance > 1.0) else 1.0 return cfg, None @@ -366,6 +367,9 @@ class SdCppDiffusionBackend: attention_backend: Optional[str] = None, transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, + # Accepted for a uniform engine interface; the native engine is GGUF-only, so a + # non-GGUF kind never routes here (the router forces diffusers for those). + model_kind: Optional[str] = None, ) -> dict[str, Any]: """Validate, then fetch assets on a daemon thread. Returns at once.""" # An empty / whitespace token is "no token": passing "" verbatim to HfApi / @@ -381,7 +385,11 @@ class SdCppDiffusionBackend: # validation and then dead-end here on a no-GPU (native-routed) host. fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: - raise ValueError(f"Could not infer a diffusion family for '{repo_id}'.") + raise ValueError( + f"'{repo_id}' is not a supported diffusion image model. Supported families: " + f"{', '.join(supported_family_names())}. If this is a variant of one of them, " + f"pass family_override with that family name." + ) if not family_sd_cpp_supported(fam): raise ValueError(f"Family '{fam.name}' has no native sd.cpp asset mapping.") @@ -687,7 +695,36 @@ class SdCppDiffusionBackend: guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, + # Accepted for a uniform engine interface. The native engine is text-to-image + # only for now (sd-cli's init-img/mask plumbing is not wired), so an image- + # conditioned request is rejected clearly rather than silently dropping the input. + init_image: Optional[str] = None, + mask_image: Optional[str] = None, + strength: Optional[float] = None, + # Accepted for the uniform engine interface; upscale needs an init image, so the + # init_image guard below rejects it on the native engine like img2img/inpaint. + upscale: Optional[float] = None, + # Reference workflow is GPU/diffusers-only (FLUX.2); accepted for interface parity. + reference_images: Optional[list[str]] = None, ) -> dict[str, Any]: + import tempfile + + from PIL import Image + + if ( + init_image is not None + or mask_image is not None + or reference_images + or (upscale is not None and upscale > 1) + ): + # upscale needs an input image, so a direct API call with upscale > 1 but no + # init_image must be rejected too rather than silently returning a plain, + # un-upscaled text-to-image result (the diffusers backend rejects the same). + raise ValueError( + "img2img / inpaint / reference / upscale are not yet supported on the native " + "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows." + ) + cancel = threading.Event() with self._generate_lock: with self._lock: @@ -987,6 +1024,7 @@ class SdCppDiffusionBackend: "transformer_cache": None, "engine": "sd_cpp", "native_mode": None, + "workflows": [], } return { "loaded": True, @@ -1012,6 +1050,11 @@ class SdCppDiffusionBackend: "engine": "sd_cpp", # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. "native_mode": state.mode, + # The native engine supports plain text-to-image only (generate() rejects + # img2img / inpaint / reference / upscale), so advertise just txt2img. Without + # this the status omits workflows, the UI reads [], and it disables the Create + # tab for a loaded native model, stranding the user on an image-only tab. + "workflows": ["txt2img"], } diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 57ce68b1cb..41faf29303 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1688,9 +1688,19 @@ class AnthropicMessagesResponse(BaseModel): 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" + model_path: str = Field(..., description = "Diffusion repo id or local path") + gguf_filename: Optional[str] = Field( + None, + description = "The chosen single-file checkpoint (GGUF or safetensors) inside " + "model_path. Required for the gguf / single_file kinds; omit for a full pipeline.", + ) + model_kind: Optional[Literal["gguf", "single_file", "pipeline"]] = Field( + None, + description = "How to load the model (null = auto-detect from gguf_filename): gguf " + "(single-file GGUF transformer, dequantised on-device), single_file (single-file " + "safetensors transformer, e.g. fp8), or pipeline (a full diffusers repo via " + "from_pretrained, embedded quant auto-applied). Non-GGUF kinds are restricted to " + "unsloth/* repos (or a local path).", ) base_repo: Optional[str] = Field( None, description = "Companion diffusers repo for VAE/text-encoders (default: family base)" @@ -1707,10 +1717,12 @@ class DiffusionLoadRequest(BaseModel): "cut), low_vram (offload every component, lowest VRAM, slower). " "Overrides cpu_offload when set.", ) - speed_mode: Optional[Literal["off", "default", "max"]] = Field( + speed_mode: Optional[Literal["off", "eager", "default", "max"]] = Field( None, description = "Opt-in speed optims (default off -> bit-identical output): " - "default (channels_last + regional torch.compile where eligible), " + "eager (channels_last + cudnn + attention + fused RMSNorm/AdaLayerNorm patches, " + "NO torch.compile -> fast first image, no compile tax), " + "default (also regional torch.compile where eligible), " "max (also TF32 + fused QKV).", ) text_encoder_quant: Optional[Literal["fp8", "nvfp4"]] = Field( @@ -1823,6 +1835,55 @@ class DiffusionGenerateRequest(BaseModel): batch_size: int = Field( 1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)" ) + # Image-conditioned workflows (base64 or data-URL). An init_image alone runs img2img; + # init_image + mask_image runs inpaint. Both require a model family with the matching + # pipeline (img2img/inpaint) or the load is rejected with a clear message. + # Cap each base64 image string so a single request can't buffer a multi-GB payload (the + # decoded dimensions are bounded separately in the backend). ~32 MiB comfortably fits a + # full 4096px image yet rejects abuse. + init_image: Optional[str] = Field( + None, + max_length = 32 * 1024 * 1024, + description = "Base64/data-URL source image for img2img or inpaint (omit for txt2img)", + ) + mask_image: Optional[str] = Field( + None, + max_length = 32 * 1024 * 1024, + description = "Base64/data-URL mask for inpaint (white = repaint, black = keep). " + "Requires init_image.", + ) + strength: Optional[float] = Field( + None, + ge = 0.0, + le = 1.0, + description = "img2img/inpaint denoise strength: 0 keeps the source, 1 fully " + "redraws it. Ignored for txt2img.", + ) + upscale: Optional[float] = Field( + None, + ge = 1.0, + le = 4.0, + description = "Upscale (hires fix) factor for an init_image: enlarges the source " + "by this multiple and re-denoises at low strength. Requires init_image; " + "ignored for txt2img/inpaint/edit.", + ) + reference_images: Optional[list[str]] = Field( + None, + max_length = 3, + description = "Additional reference images (base64/data-URL) for the FLUX.2 reference " + "workflow, combined with init_image. Up to 3; ignored by other workflows.", + ) + + @field_validator("reference_images") + @classmethod + def _bounded_reference_items(cls, value: Optional[list[str]]) -> Optional[list[str]]: + # Each reference is a base64 image; bound its length like init_image/mask_image so a + # request carrying several references can't buffer a multi-GB payload. + if value is not None: + for item in value: + if len(item) > 32 * 1024 * 1024: + raise ValueError("each reference image must be at most 32 MiB (base64)") + return value @field_validator("width", "height") @classmethod @@ -1899,6 +1960,9 @@ class DiffusionStatusResponse(BaseModel): base_repo: Optional[str] = Field(None, description = "Companion diffusers base repo") device: Optional[str] = Field(None, description = "Device the pipeline is on") dtype: Optional[str] = Field(None, description = "Compute dtype") + model_kind: Optional[str] = Field( + None, description = "Resolved load kind: gguf | single_file | pipeline (gates GGUF-only UI)" + ) cpu_offload: bool = Field(False, description = "Whether CPU offload is engaged") offload_policy: Optional[str] = Field( None, description = "Resolved offload policy: none | group | model | sequential" @@ -1923,6 +1987,11 @@ class DiffusionStatusResponse(BaseModel): "_native_cudnn), or null for the default SDPA", ) transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") + workflows: list[str] = Field( + default_factory = list, + description = "Image workflows the loaded family supports (drives UI tab gating): " + "txt2img, img2img, inpaint. Empty when nothing is loaded or on the native engine.", + ) engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp") native_mode: Optional[str] = Field( None, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7400a9f657..2c0c2bf6b0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11051,7 +11051,7 @@ def _guard_diffusion_load_against_training() -> None: async def load_diffusion_model( request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject) ): - from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion import get_diffusion_backend, resolve_model_kind from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_engine_router import ( annotate_status, @@ -11062,14 +11062,18 @@ async def load_diffusion_model( backend = get_diffusion_backend() try: + # Resolve the load kind once (gguf / single_file / pipeline) so validation, + # engine selection, and the load all agree. A bad explicit kind raises here -> 400. + kind = resolve_model_kind(request.gguf_filename, request.model_kind) # Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family, - # missing local GGUF) must not evict a working chat model and then 400. The - # validated family also drives engine selection below. + # missing local GGUF, a non-unsloth non-GGUF repo) must not evict a working chat + # model and then 400. The validated family also drives engine selection below. fam = await asyncio.to_thread( backend.validate_load_request, request.model_path, gguf_filename = request.gguf_filename, family_override = request.family_override, + model_kind = kind, ) # Refuse while training is running: a multi-GB diffusion pipeline would # compete with the training subprocess for VRAM. The chat path does the @@ -11077,15 +11081,17 @@ async def load_diffusion_model( _guard_diffusion_load_against_training() # Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU), # installing the sd-cli binary if needed -- all BEFORE evicting chat, so a - # native fallback never strands a half-loaded state. - engine = await asyncio.to_thread(select_and_activate_engine, fam, hf_token = request.hf_token) + # native fallback never strands a half-loaded state. Non-GGUF kinds force diffusers. + engine = await asyncio.to_thread( + select_and_activate_engine, fam, hf_token = request.hf_token, model_kind = kind + ) # Take the GPU from the chat backend only when this load will actually use it, # which is exactly the resolved device being non-CPU. diffusers on an accelerator # and a force-native sd.cpp load on CUDA/XPU/MPS both resolve to that device; a # native sd.cpp load on a pure-CPU host does not. Crucially, a CPU-only host with # no usable sd-cli falls back to diffusers ON CPU -- that also never touches GPU # memory, so keying off the engine name (not the device) would wrongly evict a - # resident chat model for a load that can't use the GPU. Gate on the device. + # resident chat model for a load that cannot use the GPU. Gate on the device. device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) needs_gpu = device != "cpu" if needs_gpu: @@ -11117,6 +11123,7 @@ async def load_diffusion_model( attention_backend = request.attention_backend, transformer_cache = request.transformer_cache, transformer_cache_threshold = request.transformer_cache_threshold, + model_kind = kind, ) return DiffusionStatusResponse(**annotate_status(status_dict)) except (ValueError, FileNotFoundError) as exc: @@ -11149,7 +11156,16 @@ async def generate_diffusion_image( guidance = request.guidance, seed = request.seed, batch_size = request.batch_size, + init_image = request.init_image, + mask_image = request.mask_image, + strength = request.strength, + upscale = request.upscale, + reference_images = request.reference_images, ) + except ValueError as exc: + # Bad client input (undecodable image/mask, or a workflow the loaded family + # doesn't support) — a 400 with the reason, not a generic 500. + raise HTTPException(status_code = 400, detail = str(exc)) except RuntimeError as exc: # Only "no model loaded" / user-cancelled are client-state (409); both engines # raise these two EXACT messages. The native sd.cpp engine also raises @@ -11161,10 +11177,10 @@ async def generate_diffusion_image( msg = str(exc) if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG): raise HTTPException(status_code = 409, detail = msg) - logger.error("diffusion.generate_failed: %s", exc) + logger.error("diffusion.generate_failed: %s", exc, exc_info = True) raise HTTPException(status_code = 500, detail = "Image generation failed.") except Exception as exc: - logger.error("diffusion.generate_failed: %s", exc) + logger.error("diffusion.generate_failed: %s", exc, exc_info = True) raise HTTPException(status_code = 500, detail = "Image generation failed.") # Persist each image with its full recipe embedded. The diffusers batch shares @@ -11187,8 +11203,13 @@ async def generate_diffusion_image( { "prompt": request.prompt, "negative_prompt": request.negative_prompt, - "width": request.width, - "height": request.height, + # Persist the ACTUAL output size, not the request sliders: Transform/ + # Inpaint/Edit derive it from the uploaded image, Extend grows the + # canvas, and Upscale resizes it, so request.width/height would record + # (and later restore) the wrong dimensions for those workflows. For + # plain txt2img the image size equals the sliders anyway. + "width": getattr(image, "width", None) or request.width, + "height": getattr(image, "height", None) or request.height, "steps": request.steps, "guidance": request.guidance, "seed": seed, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index da23205f07..cfba8a1181 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3188,9 +3188,15 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): def _repo_is_diffusers(repo_info) -> bool: - """A diffusers pipeline repo (e.g. a Z-Image / FLUX base) carries a top-level - model_index.json. These render images rather than chat, so the chat picker - hides them — mirroring how cached diffusion GGUFs are classified by arch.""" + """True for an image-diffusion repo, so the chat picker hides it (it renders + images, not chat) and the Images picker claims it — mirroring how cached + diffusion GGUFs are classified by arch. + + Two signals: a full diffusers pipeline carries a top-level model_index.json, + while single-file / ComfyUI / ControlNet image checkpoints (e.g. an FP8 + Qwen-Image or a z-image .safetensors) ship none. For those, fall back to the + repo id resolving to a known diffusion family — the same resolver the Images + backend loads from — so they don't surface as loadable chat models.""" try: for rev in repo_info.revisions: for f in rev.files: @@ -3198,6 +3204,12 @@ def _repo_is_diffusers(repo_info) -> bool: return True except Exception: pass + try: + from core.inference.diffusion_families import detect_family + if detect_family(getattr(repo_info, "repo_id", "") or "") is not None: + return True + except Exception: + pass return False diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index b0b9ee309c..435bafe484 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -22,6 +22,12 @@ _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) +# Let the diffusion patch backend lazily import unsloth_zoo on a CPU-only / no-GPU test +# host: unsloth_zoo runs accelerator detection at import and raises without a GPU unless +# this is set (device_type.get_device_type checks torch.cuda first, so it is a no-op on a +# real GPU run). setdefault so an explicit override wins. +os.environ.setdefault("UNSLOTH_ALLOW_CPU", "1") + # Pytest CLI options diff --git a/studio/backend/tests/test_diffusion_arch_patches.py b/studio/backend/tests/test_diffusion_arch_patches.py new file mode 100644 index 0000000000..77276a0fb1 --- /dev/null +++ b/studio/backend/tests/test_diffusion_arch_patches.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Numerical + lifecycle tests for the per-arch eager fusions (``diffusion_arch_patches``). + +Each per-arch patch only fuses ``a + b*c`` -> ``torch.addcmul`` (1-ULP, more accurate), so +the patched method/forward must match the stock diffusers one within fp tolerance. We also +check install/uninstall reversibility + idempotency, the kill-switch, and the body-drift +guard (a diffusers whose block body changed is left unpatched). Runs on CPU. +""" + +from __future__ import annotations + +import types + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("diffusers") + +from core.inference import diffusion_arch_patches as ap # noqa: E402 + + +@pytest.fixture(autouse = True) +def _clean(): + ap.uninstall_arch_patches() + yield + ap.uninstall_arch_patches() + + +# ── qwen-image _modulate (modulation addcmul) ─────────────────────────────────── + + +def test_qwen_modulate_matches_stock_global_and_indexed(): + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q + + B, L, D = 2, 16, 64 + x = torch.randn(B, L, D) + mod = torch.randn(B, 3 * D) + # _modulate uses no real `self` state, so call it unbound with self=None. + ref_x, ref_g = Q._modulate(None, x, mod) + got_x, got_g = ap._qwen_modulate(None, x, mod) + torch.testing.assert_close(got_x, ref_x, atol = 1e-5, rtol = 1e-4) + assert torch.equal(got_g, ref_g) + + # per-token `index` branch (mod batch is 2*B). + idx = torch.randint(0, 2, (B, L)) + mod2 = torch.randn(2 * B, 3 * D) + ref2_x, ref2_g = Q._modulate(None, x, mod2, idx) + got2_x, got2_g = ap._qwen_modulate(None, x, mod2, idx) + torch.testing.assert_close(got2_x, ref2_x, atol = 1e-5, rtol = 1e-4) + assert torch.equal(got2_g, ref2_g) + + +# ── z-image block forward (gated-residual addcmul) ────────────────────────────── + + +class _AttnStub(torch.nn.Module): + """Deterministic stand-in for ZImageAttention so the block forward runs without RoPE / + freqs_cis (the patch only changes the residual adds, which is what we validate).""" + + def __init__(self, dim): + super().__init__() + self.proj = torch.nn.Linear(dim, dim, bias = False) + + def forward(self, h, **kwargs): + return self.proj(h) + + +def _zimage_block(dim = 64, heads = 4): + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock + + blk = ZImageTransformerBlock( + layer_id = 0, + dim = dim, + n_heads = heads, + n_kv_heads = heads, + norm_eps = 1e-5, + qk_norm = True, + modulation = True, + ).eval() + blk.attention = _AttnStub(dim).eval() + return blk + + +def _adaln_dim(dim): + from diffusers.models.transformers.transformer_z_image import ADALN_EMBED_DIM + return min(dim, ADALN_EMBED_DIM) + + +def test_zimage_forward_matches_stock_global_modulation(): + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock + + torch.manual_seed(0) + blk = _zimage_block() + B, L, D = 2, 16, 64 + x = torch.randn(B, L, D) + adaln = torch.randn(B, _adaln_dim(D)) + + with torch.inference_mode(): + ref = ZImageTransformerBlock.forward(blk, x, None, None, adaln_input = adaln).clone() + got = ap._zimage_forward(blk, x, None, None, adaln_input = adaln) + torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4) + + +def test_zimage_forward_matches_stock_per_token_modulation(): + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock + + torch.manual_seed(1) + blk = _zimage_block() + B, L, D = 2, 16, 64 + x = torch.randn(B, L, D) + ad = _adaln_dim(D) + adaln_noisy = torch.randn(B, ad) + adaln_clean = torch.randn(B, ad) + noise_mask = torch.randint(0, 2, (B, L)) + + with torch.inference_mode(): + ref = ZImageTransformerBlock.forward( + blk, + x, + None, + None, + noise_mask = noise_mask, + adaln_noisy = adaln_noisy, + adaln_clean = adaln_clean, + ).clone() + got = ap._zimage_forward( + blk, + x, + None, + None, + noise_mask = noise_mask, + adaln_noisy = adaln_noisy, + adaln_clean = adaln_clean, + ) + torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4) + + +# ── flux.1 / flux.2 block forwards (modulation + gated-residual addcmul) ───────── + + +class _Tuple2AttnStub(torch.nn.Module): + """Double-stream attention stub -> (img_out, ctx_out).""" + + def __init__(self, dim): + super().__init__() + self.pi = torch.nn.Linear(dim, dim, bias = False) + self.pc = torch.nn.Linear(dim, dim, bias = False) + + def forward( + self, + hidden_states, + encoder_hidden_states = None, + **kwargs, + ): + return self.pi(hidden_states), self.pc(encoder_hidden_states) + + +class _SingleAttnStub(torch.nn.Module): + """Single-stream attention stub -> tensor.""" + + def __init__(self, dim): + super().__init__() + self.p = torch.nn.Linear(dim, dim, bias = False) + + def forward(self, hidden_states, **kwargs): + return self.p(hidden_states) + + +def _close_any(got, ref): + if isinstance(ref, tuple): + assert len(got) == len(ref) + for g, r in zip(got, ref): + torch.testing.assert_close(g, r, atol = 1e-5, rtol = 1e-4) + else: + torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4) + + +D, H = 64, 4 +B, L, LC = 2, 16, 8 + + +def test_flux_double_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux import FluxTransformerBlock + + torch.manual_seed(0) + blk = FluxTransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval() + blk.attn = _Tuple2AttnStub(D).eval() + hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D) + with torch.inference_mode(): + ref = FluxTransformerBlock.forward(blk, hs, ehs, temb) + got = ap._flux_double_forward(blk, hs, ehs, temb) + _close_any(got, ref) + + +def test_flux_single_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock + + torch.manual_seed(1) + blk = FluxSingleTransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval() + blk.attn = _SingleAttnStub(D).eval() + hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D) + with torch.inference_mode(): + ref = FluxSingleTransformerBlock.forward(blk, hs, ehs, temb) + got = ap._flux_single_forward(blk, hs, ehs, temb) + _close_any(got, ref) + + +def test_flux2_double_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux2 import Flux2TransformerBlock + + torch.manual_seed(2) + blk = Flux2TransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval() + blk.attn = _Tuple2AttnStub(D).eval() + hs, ehs = torch.randn(B, L, D), torch.randn(B, LC, D) + tmi, tmt = torch.randn(B, 6 * D), torch.randn(B, 6 * D) + with torch.inference_mode(): + ref = Flux2TransformerBlock.forward(blk, hs, ehs, tmi, tmt) + got = ap._flux2_double_forward(blk, hs, ehs, tmi, tmt) + _close_any(got, ref) + + +def test_flux2_single_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux2 import Flux2SingleTransformerBlock + + torch.manual_seed(3) + blk = Flux2SingleTransformerBlock( + dim = D, num_attention_heads = H, attention_head_dim = D // H + ).eval() + blk.attn = _SingleAttnStub(D).eval() + hs, ehs, tm = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, 3 * D) + with torch.inference_mode(): + ref = Flux2SingleTransformerBlock.forward(blk, hs, ehs, tm) + got = ap._flux2_single_forward(blk, hs, ehs, tm) + _close_any(got, ref) + + +# ── lifecycle ─────────────────────────────────────────────────────────────────── + + +def test_install_idempotent_and_reversible(): + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock as Z + + q_orig, z_orig = Q._modulate, Z.forward + n1 = ap.install_arch_patches() + n2 = ap.install_arch_patches() # idempotent + assert n1 == 6 and n2 == n1 # qwen + z-image + flux.1 x2 + flux.2 x2 + assert Q._modulate is not q_orig and Z.forward is not z_orig + assert ap.is_installed() + + ap.uninstall_arch_patches() + assert not ap.is_installed() + assert Q._modulate is q_orig and Z.forward is z_orig # exact restore + ap.uninstall_arch_patches() # idempotent + + +def test_kill_switch(monkeypatch): + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q + + monkeypatch.setenv("UNSLOTH_DIFFUSION_ARCH_PATCHES", "0") + orig = Q._modulate + assert ap.install_arch_patches() == 0 + assert not ap.is_installed() + assert Q._modulate is orig + + +def test_body_drift_guard_skips_changed_block(monkeypatch): + # If a resolver's body-check fails (diffusers changed the lines we rewrite), that patch + # is skipped. Force the qwen resolver to see a drifted body. + monkeypatch.setattr(ap, "_body_has", lambda fn, *needles: False) + assert ap.install_arch_patches() == 0 + assert not ap.is_installed() diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 2801e96c3c..2b258df003 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -22,10 +22,19 @@ from core.inference.diffusion import ( _base_file_downloaded, _resolve_diffusion_compute_dtype, ) + +# diffusion.py imports the compile/arch patch modules LAZILY (they pull torch at module +# level, and diffusion.py must stay importable on a torchless native install). Import them +# here at collection time -- under the real torch -- so they are cached in sys.modules +# before the fake-torch fixtures swap it out; otherwise the lazy import inside load_pipeline +# would try to build them against the incomplete stub torch. +import core.inference.diffusion_eager_patches # noqa: E402,F401 +import core.inference.diffusion_arch_patches # noqa: E402,F401 from core.inference.diffusion_families import ( detect_family, resolve_base_repo, resolve_local_gguf_child, + supported_family_names, ) @@ -45,29 +54,63 @@ def test_detect_family_from_repo_id(): assert klein.cfg_kwarg == "guidance_scale" # Both klein sizes share the one family (base repo resolved per-variant). assert detect_family("unsloth/FLUX.2-klein-9B-GGUF").name == "flux.2-klein" - # Only klein is wired up; the Mistral-based FLUX.2-dev base repo is gated. - assert detect_family("unsloth/FLUX.2-dev-GGUF") is None + # FLUX.2-dev is the Mistral-based Flux2Pipeline, a distinct family from klein; its + # gated base repo is reachable with an HF token. It must not collide with klein. + dev = detect_family("unsloth/FLUX.2-dev-GGUF") + assert dev.name == "flux.2-dev" + assert dev.pipeline_class == "Flux2Pipeline" + assert dev.base_repo == "black-forest-labs/FLUX.2-dev" + assert detect_family("black-forest-labs/FLUX.2-dev").name == "flux.2-dev" # Qwen-Image guides via true_cfg_scale, not guidance_scale. assert detect_family("unsloth/Qwen-Image-2512-GGUF").cfg_kwarg == "true_cfg_scale" assert detect_family("unsloth/Z-Image-GGUF").cfg_kwarg == "guidance_scale" - # Image-editing checkpoints are rejected (text-to-image backend only). - assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None - assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") is None + # Qwen-Image-Edit is a SUPPORTED instruction-editing family (its own edit pipeline); + # the most-specific match wins so it doesn't fall back to the generic qwen-image. + edit = detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") + assert edit.name == "qwen-image-edit" + assert edit.pipeline_class == "QwenImageEditPlusPipeline" + assert edit.edit is True + assert detect_family("unsloth/Qwen-Image-Edit-2509-GGUF").name == "qwen-image-edit" + # FLUX Kontext is a SUPPORTED editing family (FluxKontextPipeline); the "kontext" + # keyword is un-rejected for it, and it must win over the generic "flux.1" match. + kontext = detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") + assert kontext.name == "flux.1-kontext" + assert kontext.pipeline_class == "FluxKontextPipeline" + assert kontext.edit is True + assert kontext.cfg_kwarg == "guidance_scale" + # A plain FLUX.1 checkpoint must still resolve to the base flux.1 family, not kontext. + assert detect_family("unsloth/FLUX.1-dev-GGUF").name == "flux.1" + # A plain Qwen-Image checkpoint must still resolve to the base family, not edit. + assert detect_family("unsloth/Qwen-Image-2512-GGUF").name == "qwen-image" assert detect_family("meta-llama/Llama-3-8B") is None +def test_detect_family_matches_reject_and_alias_by_segment(): + # Reject keywords and short aliases must match whole path/name segments, not raw + # substrings, so an unrelated word that merely CONTAINS one does not misroute a + # valid base model (regression: substring matching broke these). + assert detect_family("/models/edited/z-image-turbo-Q4_K_M.gguf").name == "z-image" + assert detect_family("unsloth/Z-Image-Edition-GGUF").name == "z-image" + assert detect_family("/models/kontextual/z-image-turbo-Q4_K_M.gguf").name == "z-image" + # Supported edit families still resolve (edit / kontext are whole tokens there). + assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF").name == "qwen-image-edit" + assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF").name == "flux.1-kontext" + # Unsupported variants sharing only a base arch keyword are still rejected. + assert detect_family("unsloth/Qwen-Image-Layered-GGUF") is None + assert detect_family("unsloth/Qwen-Image-2512-Inpaint") is None + + def test_detect_family_edit_keyword_scoped_to_basename(): from core.inference.diffusion_families import detect_family_for_pick - # A parent directory named `edit`/`kontext`/`inpaint` must NOT reject a valid - # text-to-image file: only the model id / filename basename is scanned for the - # edit keyword. A direct local pick arrives as (parent_dir, filename). - assert detect_family("/models/edit") is None # dir alone is ambiguous + # A parent directory named `edit`/`inpaint` must NOT poison a valid pick: only + # the model id / filename basename is scanned for reject keywords. A direct + # local pick arrives as (parent_dir, filename). + assert detect_family("/models/edit") is None # the dir alone is ambiguous assert detect_family_for_pick("/models/edit", "Z-Image-Turbo-Q4.gguf").name == "z-image" - assert detect_family_for_pick("/models/kontext", "qwen-image-2512-Q4.gguf").name == "qwen-image" - # But a genuine editing checkpoint (keyword in the id/filename) is still rejected. - assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None - assert detect_family_for_pick("/models/misc", "Qwen-Image-Edit-2511-Q4.gguf") is None + assert detect_family_for_pick("/models/inpaint", "qwen-image-2512-Q4.gguf").name == "qwen-image" + # A genuinely unsupported variant keyword in the FILENAME still rejects. + assert detect_family_for_pick("/models/misc", "Qwen-Image-Layered-Q4.gguf") is None def test_detect_family_override(): @@ -76,6 +119,16 @@ def test_detect_family_override(): assert detect_family("local/path", override = "not-a-family") is None +def test_supported_family_names(): + names = supported_family_names() + # The unknown-model error lists these, so the key families must be present. + for expected in ("flux.1", "flux.2-klein", "flux.2-dev", "qwen-image", "z-image"): + assert expected in names + # Every listed name is a valid family_override (round-trips through detect_family). + for name in names: + assert detect_family("some/unknown-repo", override = name) is not None + + def test_resolve_base_repo(): fam = detect_family("x", override = "z-image") assert resolve_base_repo(fam, None) == fam.base_repo @@ -210,6 +263,86 @@ class _FakeTransformer: return object() +class _FakeImg2ImgPipe: + """An img2img pipeline call: records the image-conditioned kwargs. Its signature + declares image/strength but NOT width/height, mirroring real img2img pipelines + (which derive the output size from the input image).""" + + last_kwargs: dict = {} + + def __call__( + self, + *, + prompt = None, + image = None, + strength = None, + negative_prompt = None, + callback_on_step_end = None, + guidance_scale = None, + true_cfg_scale = None, + **kwargs, + ): + _FakeImg2ImgPipe.last_kwargs = { + "prompt": prompt, + "image": image, + "strength": strength, + **kwargs, + } + n = kwargs.get("num_images_per_prompt", 1) + return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)]) + + +class _FakeImg2ImgPipeline: + built_from: object = None + from_pipe_kwargs: dict = {} + + @classmethod + def from_pipe(cls, base_pipe, **kwargs): + _FakeImg2ImgPipeline.built_from = base_pipe + _FakeImg2ImgPipeline.from_pipe_kwargs = kwargs + return _FakeImg2ImgPipe() + + +class _FakeInpaintPipe: + """An inpaint pipeline call: records image + mask_image + strength. Real inpaint + pipelines take both an init image and a grayscale mask and derive output size from + the input, so width/height are not in its signature.""" + + last_kwargs: dict = {} + + def __call__( + self, + *, + prompt = None, + image = None, + mask_image = None, + strength = None, + negative_prompt = None, + callback_on_step_end = None, + guidance_scale = None, + true_cfg_scale = None, + **kwargs, + ): + _FakeInpaintPipe.last_kwargs = { + "prompt": prompt, + "image": image, + "mask_image": mask_image, + "strength": strength, + **kwargs, + } + n = kwargs.get("num_images_per_prompt", 1) + return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)]) + + +class _FakeInpaintPipeline: + built_from: object = None + + @classmethod + def from_pipe(cls, base_pipe, **kwargs): + _FakeInpaintPipeline.built_from = base_pipe + return _FakeInpaintPipe() + + @pytest.fixture def fake_runtime(monkeypatch): torch = types.ModuleType("torch") @@ -226,9 +359,15 @@ def fake_runtime(monkeypatch): diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype) diffusers.ZImagePipeline = _FakePipeline diffusers.ZImageTransformer2DModel = _FakeTransformer + diffusers.ZImageImg2ImgPipeline = _FakeImg2ImgPipeline + diffusers.ZImageInpaintPipeline = _FakeInpaintPipeline # Qwen-Image too, so the true_cfg_scale cfg-kwarg path is exercisable. diffusers.QwenImagePipeline = _FakePipeline diffusers.QwenImageTransformer2DModel = _FakeTransformer + diffusers.QwenImageImg2ImgPipeline = _FakeImg2ImgPipeline + diffusers.QwenImageInpaintPipeline = _FakeInpaintPipeline + # Instruction-editing pipeline (Qwen-Image-Edit): its own pipeline IS the loaded one. + diffusers.QwenImageEditPlusPipeline = _FakePipeline monkeypatch.setitem(sys.modules, "torch", torch) monkeypatch.setitem(sys.modules, "diffusers", diffusers) @@ -237,6 +376,10 @@ def fake_runtime(monkeypatch): monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: None) _FakePipeline.last = {} _FakeTransformer.last = {} + _FakeImg2ImgPipeline.built_from = None + _FakeImg2ImgPipe.last_kwargs = {} + _FakeInpaintPipeline.built_from = None + _FakeInpaintPipe.last_kwargs = {} yield @@ -289,6 +432,533 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): assert backend.is_loaded is False +def _tiny_png_b64() -> str: + import base64 + import io + + from PIL import Image + + buf = io.BytesIO() + Image.new("RGB", (64, 64), (120, 30, 30)).save(buf, format = "PNG") + return base64.b64encode(buf.getvalue()).decode() + + +def test_generate_img2img_uses_from_pipe(fake_runtime, tmp_path): + """An init_image routes generate() through the family's img2img pipeline, built via + Pipeline.from_pipe around the loaded pipe (no reload), with image + strength passed + and width/height dropped (the img2img pipe derives size from the input image).""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + # The loaded family advertises the image-conditioned workflows for UI gating + # (upscale rides the img2img pipeline, so it appears whenever img2img does). + assert backend.status()["workflows"] == ["txt2img", "img2img", "upscale", "inpaint", "outpaint"] + + loaded_pipe = backend._state.pipe + out = backend.generate( + prompt = "a car at sunset", + steps = 4, + guidance = 0.0, + seed = 3, + init_image = _tiny_png_b64(), + strength = 0.5, + ) + assert len(out["images"]) == 1 + # from_pipe was handed the loaded text-to-image pipe (component reuse, no reload). + assert _FakeImg2ImgPipeline.built_from is loaded_pipe + # ...and with torch_dtype=None so from_pipe SKIPS its default float32 recast, which + # both upcasts the reused bf16 modules and crashes on torchao-quantized weights. + assert _FakeImg2ImgPipeline.from_pipe_kwargs.get("torch_dtype", "MISSING") is None + call = _FakeImg2ImgPipe.last_kwargs + assert call["image"] is not None # decoded source image passed through + assert call["strength"] == 0.5 + assert "width" not in call and "height" not in call # img2img derives size from image + + # A txt2img call after it still uses the base pipe (no image kwarg). + backend.generate(prompt = "plain", steps = 4, seed = 1) + assert backend._state.pipe.last_kwargs.get("image") is None + + +def test_generate_img2img_unsupported_family_raises(fake_runtime, tmp_path, monkeypatch): + """A family with no image-conditioning at all (no img2img/inpaint/edit/reference) rejects + an init_image with a clear error rather than failing deep in the pipeline.""" + from core.inference.diffusion_families import DiffusionFamily + + # A synthetic txt2img-only family: no img2img/inpaint pipeline, not edit, not reference. + # (Every shipped family now supports some image workflow, so build one for this case.) + plain = DiffusionFamily( + name = "plain-test", + pipeline_class = "ZImagePipeline", + transformer_class = "ZImageTransformer2DModel", + base_repo = "base/repo", + ) + monkeypatch.setattr( + "core.inference.diffusion.detect_family_for_pick", + lambda repo_id, gguf_filename = None, override = None: plain, + ) + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline(str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo") + assert backend.status()["workflows"] == ["txt2img"] + with pytest.raises(ValueError, match = "img2img"): + backend.generate(prompt = "x", steps = 4, init_image = _tiny_png_b64()) + + +def test_generate_rejects_conditioning_without_init_image(fake_runtime, tmp_path): + """mask / upscale / reference all need an input image; without one they must raise a + clear ValueError rather than silently degrading to txt2img.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + with pytest.raises(ValueError, match = "mask_image requires"): + backend.generate(prompt = "x", steps = 4, mask_image = _mask_b64(64)) + with pytest.raises(ValueError, match = "upscale requires"): + backend.generate(prompt = "x", steps = 4, upscale = 2.0) + with pytest.raises(ValueError, match = "reference_images require"): + backend.generate(prompt = "x", steps = 4, reference_images = [_tiny_png_b64()]) + + +def test_generate_rejects_reference_on_unsupported_family(fake_runtime, tmp_path): + """A non-reference family rejects reference_images instead of silently dropping them.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + with pytest.raises(ValueError, match = "Reference images are not supported"): + backend.generate( + prompt = "x", + steps = 4, + init_image = _tiny_png_b64(), + reference_images = [_tiny_png_b64()], + ) + + +def test_generate_upscale_enlarges_and_low_strength(fake_runtime, tmp_path): + """An init_image + upscale factor routes generate() through the family's img2img + pipeline (hires fix): the source is enlarged to size*factor (rounded to /16) before the + denoise, the strength defaults low, and the factor is capped so a huge value can't OOM.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + # Upscale rides the img2img pipeline, so it is advertised alongside img2img. + assert "upscale" in backend.status()["workflows"] + + loaded_pipe = backend._state.pipe + out = backend.generate( + prompt = "a crisp photo", + steps = 4, + guidance = 0.0, + seed = 3, + init_image = _tiny_png_b64(), + upscale = 2.0, # 64 -> 128, no explicit strength + ) + assert len(out["images"]) == 1 + # Reuses the resident modules via from_pipe (no reload, no extra VRAM). + assert _FakeImg2ImgPipeline.built_from is loaded_pipe + call = _FakeImg2ImgPipe.last_kwargs + # The image handed to the pipe is the ENLARGED source (64 * 2 = 128, already /16). + assert call["image"].size == (128, 128) + # Strength defaults to the hires-fix value when the caller sends none. + assert call["strength"] == 0.35 + + # The factor is capped at 4x so a large request can't blow up the VAE/transformer. + backend.generate( + prompt = "x", + steps = 4, + seed = 1, + init_image = _tiny_png_b64(), + upscale = 99.0, + ) + assert _FakeImg2ImgPipe.last_kwargs["image"].size == (256, 256) # 64 * 4 (capped) + + # An explicit strength overrides the hires-fix default. + backend.generate( + prompt = "x", + steps = 4, + seed = 1, + init_image = _tiny_png_b64(), + upscale = 1.5, + strength = 0.2, + ) + assert _FakeImg2ImgPipe.last_kwargs["strength"] == 0.2 + # 64 * 1.5 = 96, already a multiple of 16. + assert _FakeImg2ImgPipe.last_kwargs["image"].size == (96, 96) + + +def _png_b64(side: int) -> str: + import base64 + import io + + from PIL import Image + + buf = io.BytesIO() + Image.new("RGB", (side, side), (10, 20, 30)).save(buf, format = "PNG") + return base64.b64encode(buf.getvalue()).decode() + + +def test_decode_image_rejects_oversized(fake_runtime, tmp_path): + """An input image larger than the per-side cap is rejected with a clear error (protects + img2img / inpaint / reference from decompression-bomb / OOM inputs), not a 500.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + with pytest.raises(ValueError, match = "too large"): + backend.generate(prompt = "x", steps = 4, init_image = _png_b64(4112)) # > 4096/side + + +def test_upscale_output_is_capped(fake_runtime, tmp_path): + """Upscale bounds the absolute output side to 2048 even when input*factor exceeds it, so a + large upload at 4x can't OOM the VAE/transformer.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + backend.generate(prompt = "x", steps = 4, seed = 1, init_image = _png_b64(1024), upscale = 4.0) + # 1024 * 4 = 4096 -> clamped to 2048 (longest side), still a multiple of 16. + assert _FakeImg2ImgPipe.last_kwargs["image"].size == (2048, 2048) + + +def _mask_b64(side: int) -> str: + import base64 + import io + + from PIL import Image + + buf = io.BytesIO() + img = Image.new("L", (side, side), 0) + for y in range(side // 4, 3 * side // 4): + for x in range(side // 4, 3 * side // 4): + img.putpixel((x, y), 255) + img.save(buf, format = "PNG") + return base64.b64encode(buf.getvalue()).decode() + + +def test_img2img_snaps_non_multiple_of_16(fake_runtime, tmp_path): + """An odd-sized img2img upload (not divisible by 16) is auto-resized to the nearest + multiple of 16 so the pipeline's divisibility check passes instead of erroring.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + backend.generate(prompt = "x", steps = 4, seed = 1, init_image = _png_b64(186), strength = 0.5) + # 186 / 16 = 11.625 -> round to 12 -> 192. + assert _FakeImg2ImgPipe.last_kwargs["image"].size == (192, 192) + + +def test_inpaint_snaps_image_and_mask_together(fake_runtime, tmp_path): + """Inpaint snaps the odd-sized input to /16 AND resizes the mask to match, so the image + and mask stay aligned (a mismatch would crash the inpaint pipeline).""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + backend.generate( + prompt = "x", + steps = 4, + seed = 1, + init_image = _png_b64(186), + mask_image = _mask_b64(186), + strength = 0.5, + ) + assert _FakeInpaintPipe.last_kwargs["image"].size == (192, 192) + assert _FakeInpaintPipe.last_kwargs["mask_image"].size == (192, 192) + + +def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_path): + """A reference family (FLUX.2-klein) advertises txt2img + reference, and a generate with + an init_image passes it as the loaded pipe's `image` arg (no from_pipe, no strength) while + the output size stays the REQUESTED slider size (the pipe resizes the reference itself).""" + import diffusers + + diffusers.Flux2KleinPipeline = _FakePipeline + diffusers.Flux2KleinInpaintPipeline = _FakeInpaintPipeline + diffusers.Flux2Transformer2DModel = _FakeTransformer + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "flux.2-klein", + ) + # FLUX.2-klein: txt2img + reference (own pipe) + inpaint (dedicated pipe). No img2img class, + # so no img2img/upscale. + assert backend.status()["workflows"] == ["txt2img", "reference", "inpaint"] + + loaded_pipe = backend._state.pipe + out = backend.generate( + prompt = "a portrait in this style", + steps = 6, + guidance = 4.0, + seed = 5, + width = 768, + height = 512, + init_image = _tiny_png_b64(), + strength = 0.5, + ) + assert len(out["images"]) == 1 + call = loaded_pipe.last_kwargs + assert call["image"] is not None # reference handed to the loaded pipe + assert call["width"] == 768 and call["height"] == 512 # OUTPUT size = sliders, not input + assert "strength" not in call # reference conditioning has no strength + assert "mask_image" not in call + # Guidance flows via guidance_scale (FLUX.2 default behaviour). + assert call["guidance_scale"] == 4.0 + + # Multi-reference: extra reference_images are combined with init_image into a LIST so the + # model can blend several references (subject + style). + backend.generate( + prompt = "combine these", + steps = 6, + seed = 9, + width = 1024, + height = 1024, + init_image = _tiny_png_b64(), + reference_images = [_tiny_png_b64(), _tiny_png_b64()], + ) + img_arg = loaded_pipe.last_kwargs["image"] + assert isinstance(img_arg, list) and len(img_arg) == 3 # primary + 2 extras + + # Branch ordering: an init image + MASK on a reference family must route to inpaint (the + # dedicated pipeline), NOT be swallowed by the reference branch (which ignores the mask). + backend.generate( + prompt = "repaint here", + steps = 6, + seed = 2, + init_image = _tiny_png_b64(), + mask_image = _tiny_mask_b64(), + strength = 0.8, + ) + assert _FakeInpaintPipeline.built_from is loaded_pipe # built via from_pipe off the load + assert _FakeInpaintPipe.last_kwargs["mask_image"] is not None + assert _FakeInpaintPipe.last_kwargs["strength"] == 0.8 + + # Without an init image the same family does plain txt2img (no image arg). + backend.generate(prompt = "just text", steps = 6, seed = 1) + assert backend._state.pipe.last_kwargs.get("image") is None + + +def _tiny_mask_b64() -> str: + import base64 + import io + + from PIL import Image + + buf = io.BytesIO() + # A grayscale mask: white square (repaint) on black (keep). + img = Image.new("L", (64, 64), 0) + for y in range(16, 48): + for x in range(16, 48): + img.putpixel((x, y), 255) + img.save(buf, format = "PNG") + return base64.b64encode(buf.getvalue()).decode() + + +def test_generate_inpaint_uses_from_pipe(fake_runtime, tmp_path): + """An init_image + mask_image routes generate() through the family's inpaint pipeline, + built via Pipeline.from_pipe around the loaded pipe (no reload), with the decoded image + + mask + strength passed through and width/height dropped (size derives from the input).""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + loaded_pipe = backend._state.pipe + out = backend.generate( + prompt = "a red door", + steps = 4, + guidance = 0.0, + seed = 5, + init_image = _tiny_png_b64(), + mask_image = _tiny_mask_b64(), + strength = 0.7, + ) + assert len(out["images"]) == 1 + # The inpaint pipe (not img2img) was selected and built from the loaded pipe. + assert _FakeInpaintPipeline.built_from is loaded_pipe + assert _FakeImg2ImgPipeline.built_from is None + call = _FakeInpaintPipe.last_kwargs + assert call["image"] is not None and call["mask_image"] is not None + assert call["strength"] == 0.7 + assert "width" not in call and "height" not in call # inpaint derives size from image + + +def test_image_conditioned_passes_image_size_not_slider(fake_runtime, tmp_path): + """When the workflow pipe DOES accept width/height, an image-conditioned call must pass + the INPUT IMAGE's size, never the txt2img slider size -- otherwise a non-slider-sized + input (e.g. a 1536px outpaint canvas with a 1024 slider) mismatches the latents + ("tensor a (128) must match tensor b (192)"). Covers Transform + Extend with any size.""" + import base64 + import io + + from PIL import Image + + class _SizePipe: + last: dict = {} + + def __call__( + self, + *, + prompt = None, + image = None, + strength = None, + width = None, + height = None, + negative_prompt = None, + callback_on_step_end = None, + guidance_scale = None, + true_cfg_scale = None, + **kwargs, + ): + _SizePipe.last = {"width": width, "height": height} + n = kwargs.get("num_images_per_prompt", 1) + return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)]) + + class _SizePipeline: + @classmethod + def from_pipe(cls, base_pipe, **kwargs): + return _SizePipe() + + import diffusers + + diffusers.ZImageImg2ImgPipeline = _SizePipeline + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" + ) + buf = io.BytesIO() + Image.new("RGB", (96, 64), (10, 20, 30)).save(buf, format = "PNG") # non-square, non-slider + b64 = base64.b64encode(buf.getvalue()).decode() + backend.generate(prompt = "x", steps = 4, width = 1024, height = 1024, init_image = b64, strength = 0.5) + # The pipe got the IMAGE's 96x64, not the 1024x1024 slider. + assert _SizePipe.last == {"width": 96, "height": 64} + + +def test_edit_family_uses_own_pipeline_and_requires_image(fake_runtime, tmp_path): + """An instruction-editing family (Qwen-Image-Edit) exposes only the 'edit' workflow, + runs the image through its OWN loaded pipeline (no from_pipe), and rejects a call with + no input image.""" + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "Qwen/Qwen-Image-Edit-2511", + family_override = "qwen-image-edit", + ) + # Edit families advertise only the edit workflow (no txt2img / img2img / inpaint). + assert backend.status()["workflows"] == ["edit"] + loaded_pipe = backend._state.pipe + + out = backend.generate( + prompt = "make it night", + steps = 8, + guidance = 4.0, + seed = 1, + init_image = _tiny_png_b64(), + ) + assert len(out["images"]) == 1 + # The loaded pipe handled it directly -- no from_pipe img2img/inpaint was built. + assert backend._state.pipe is loaded_pipe + assert _FakeImg2ImgPipeline.built_from is None and _FakeInpaintPipeline.built_from is None + assert loaded_pipe.last_kwargs.get("image") is not None + + # An edit model with no input image fails fast with a clear message. + with pytest.raises(ValueError, match = "image"): + backend.generate(prompt = "make it night", steps = 8) + + +def test_load_pipeline_kind_uses_from_pretrained(fake_runtime): + """A full-pipeline (no single-file) load on an unsloth/* repo builds the pipe with + pipeline_cls.from_pretrained(repo_id) -- NO single-file transformer build, NO GGUF + quant config -- so an embedded bnb-4bit config is reloaded by diffusers itself.""" + backend = DiffusionBackend() + status = backend.load_pipeline( + "unsloth/Z-Image-Turbo-unsloth-bnb-4bit", family_override = "z-image" + ) + assert status["loaded"] is True + assert status["family"] == "z-image" + # from_pretrained pointed at the repo itself (it IS its own base), with no transformer. + assert _FakePipeline.last["base"] == "unsloth/Z-Image-Turbo-unsloth-bnb-4bit" + assert "transformer" not in _FakePipeline.last + # The GGUF single-file build path was never taken. + assert _FakeTransformer.last == {} + + +def test_load_single_file_safetensors_no_gguf_config(fake_runtime, tmp_path): + """A single-file *.safetensors transformer is built with from_single_file WITHOUT the + GGUF dequant config (it carries its own dtype), then assembled from the base repo.""" + (tmp_path / "model.safetensors").write_bytes(b"weights") + backend = DiffusionBackend() + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.safetensors", + base_repo = "base/repo", + family_override = "qwen-image", + ) + assert status["loaded"] is True + assert _FakeTransformer.last["path"] == str((tmp_path / "model.safetensors").resolve()) + assert _FakeTransformer.last["subfolder"] == "transformer" + # No GGUF quant config on the safetensors path (the GGUF path sets one). + assert "quantization_config" not in _FakeTransformer.last + assert _FakePipeline.last["base"] == "base/repo" + assert "transformer" in _FakePipeline.last + + +def test_load_pipeline_rejects_non_unsloth_repo(fake_runtime): + backend = DiffusionBackend() + with pytest.raises(ValueError, match = "unsloth"): + backend.load_pipeline("randomorg/Z-Image-bnb-4bit", family_override = "z-image") + + +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. + assert detect_family("unsloth/Qwen-Image-Layered-GGUF") is None + assert detect_family("unsloth/qwen_image_layered") is None + + +def test_failed_load_rolls_back_eager_patches(fake_runtime, tmp_path, monkeypatch): + """A load failure AFTER the eager patches install but BEFORE the _LoadState commit must + roll the process-wide patches back, so the next bit-identical `off` load is not + contaminated (the asymmetric-cleanup bug the reviewers flagged).""" + from core.inference import diffusion as diff_mod + from core.inference import diffusion_eager_patches as ep + + (tmp_path / "model.gguf").write_bytes(b"x") + ep.uninstall_patches() # clean slate + + def _boom(*_a, **_k): + raise RuntimeError("placement boom") + + # apply_memory_plan runs AFTER the patches are installed, before _LoadState commits. + monkeypatch.setattr(diff_mod, "apply_memory_plan", _boom) + backend = DiffusionBackend() + with pytest.raises(RuntimeError): + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + family_override = "z-image", + base_repo = "base/repo", + speed_mode = "eager", # != off -> installs the shared patches + ) + assert ep.is_installed() is False # rolled back by the load-failure finally + assert backend.is_loaded is False + + def test_cpu_offload_ignored_off_cuda(fake_runtime, tmp_path): (tmp_path / "model.gguf").write_bytes(b"x") backend = DiffusionBackend() @@ -381,8 +1051,10 @@ def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch) def test_load_without_gguf_raises(): backend = DiffusionBackend() - with pytest.raises(ValueError): - backend.load_pipeline("unsloth/Z-Image-Turbo-GGUF") # no gguf_filename + # No gguf_filename -> a full-pipeline load, gated to unsloth/*; a non-unsloth repo + # is rejected before any GPU/network work. + with pytest.raises(ValueError, match = "unsloth"): + backend.load_pipeline("some-org/Z-Image-bnb-4bit") def test_load_unknown_family_raises(): @@ -770,8 +1442,22 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime): def test_validate_load_request(tmp_path): backend = DiffusionBackend() - with pytest.raises(ValueError, match = "gguf_filename"): - backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF") + # No filename + unsloth repo -> a full-pipeline load (allowed for unsloth/*). + assert backend.validate_load_request("unsloth/Z-Image-Turbo-unsloth-bnb-4bit").name == "z-image" + # No filename + non-unsloth repo -> a pipeline load, gated to unsloth/* -> rejected. + with pytest.raises(ValueError, match = "unsloth"): + backend.validate_load_request("some-org/Z-Image-bnb-4bit") + # An explicit gguf/single_file kind still requires a single-file name. + with pytest.raises(ValueError, match = "single-file"): + backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", model_kind = "gguf") + # A pipeline kind must NOT carry a single-file name. + with pytest.raises(ValueError, match = "pipeline"): + backend.validate_load_request( + "unsloth/Z-Image-Turbo-bnb-4bit", gguf_filename = "q.gguf", model_kind = "pipeline" + ) + # A single-file safetensors load is also gated to unsloth/* repos. + with pytest.raises(ValueError, match = "unsloth"): + backend.validate_load_request("some-org/Z-Image", gguf_filename = "model.safetensors") with pytest.raises(ValueError, match = "family"): backend.validate_load_request("meta/Llama-3", gguf_filename = "q.gguf") # A family-looking repo paired with a non-GGUF single-file name is rejected here, @@ -783,6 +1469,21 @@ def test_validate_load_request(tmp_path): backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf").name == "z-image" ) + # A kind/extension mismatch fails fast here, before the route evicts chat + grabs the + # GPU only to fail in the background from_single_file path. + with pytest.raises(ValueError, match = ".gguf"): + backend.validate_load_request( + "unsloth/Z-Image-Turbo-GGUF", gguf_filename = "model.safetensors", model_kind = "gguf" + ) + with pytest.raises(ValueError, match = "gguf"): + backend.validate_load_request( + "unsloth/Qwen-Image-2512-FP8", gguf_filename = "q.gguf", model_kind = "single_file" + ) + # A remote "*-GGUF" repo loaded as a full pipeline (no single-file name) is a single-file + # GGUF repo, so from_pretrained would find no pipeline manifest and fail after chat is + # already evicted; reject it here before the GPU handoff. + with pytest.raises(ValueError, match = "GGUF"): + backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", model_kind = "pipeline") # A local path with a missing child fails here (before any GPU/network work). with pytest.raises(FileNotFoundError): backend.validate_load_request( diff --git a/studio/backend/tests/test_diffusion_compile_cache.py b/studio/backend/tests/test_diffusion_compile_cache.py new file mode 100644 index 0000000000..7193e9c0f4 --- /dev/null +++ b/studio/backend/tests/test_diffusion_compile_cache.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the pre-warmed torch.compile cache (``diffusion_compile_cache.py``). + +The Mega-cache API (``torch.compiler.save_cache_artifacts`` / ``load_cache_artifacts``) +is monkeypatched with deterministic in-memory fakes so the fingerprint / exact-match / +integrity / fallback / lifecycle logic is exercised without a real compile. The +fingerprint helpers run against the real torch on this box. +""" + +from __future__ import annotations + +import json +import types + +import pytest + +from core.inference import diffusion_compile_cache as cc + + +def _transformer(blocks = ("FluxTransformerBlock", "FluxSingleTransformerBlock")): + return types.SimpleNamespace(_repeated_blocks = list(blocks)) + + +_BEGIN_KW = dict( + family = "flux.1", + dtype = "torch.bfloat16", + quant = None, + attention_backend = "_native_cudnn", + compile_kwargs = {"fullgraph": True, "dynamic": True}, + shape_bucket = "1024x1024", +) + + +# --------------------------------------------------------------------------- fingerprint +def test_environment_fingerprint_has_hard_dimensions(): + fp = cc.environment_fingerprint() + for k in ("torch", "torch_cuda", "triton", "diffusers", "gpu_name", "gpu_capability"): + assert k in fp + + +def test_cache_key_stable_across_kwarg_order(): + efp = cc.environment_fingerprint() + t = _transformer() + a = cc.model_fingerprint( + family = "flux.1", + transformer = t, + dtype = "bf16", + quant = None, + attention_backend = "x", + compile_kwargs = {"fullgraph": True, "dynamic": True}, + ) + b = cc.model_fingerprint( + family = "flux.1", + transformer = t, + dtype = "bf16", + quant = None, + attention_backend = "x", + compile_kwargs = {"dynamic": True, "fullgraph": True}, + ) + assert cc.cache_key(efp, a) == cc.cache_key(efp, b) + + +@pytest.mark.parametrize( + "field,value", + [ + ("family", "qwen-image"), + ("dtype", "torch.float16"), + ("quant", "int8"), + ("attention_backend", "native"), + ("shape_bucket", "512x512"), + ], +) +def test_cache_key_sensitive_to_model_dims(field, value): + efp = cc.environment_fingerprint() + t = _transformer() + base = dict( + family = "flux.1", + transformer = t, + dtype = "bf16", + quant = None, + attention_backend = "x", + compile_kwargs = {"fullgraph": True}, + shape_bucket = "1024x1024", + ) + k0 = cc.cache_key(efp, cc.model_fingerprint(**base)) + base[field] = value + assert cc.cache_key(efp, cc.model_fingerprint(**base)) != k0 + + +def test_repeated_blocks_change_key(): + efp = cc.environment_fingerprint() + k1 = cc.cache_key( + efp, + cc.model_fingerprint( + family = "f", + transformer = _transformer(("A",)), + dtype = "bf16", + quant = None, + attention_backend = "x", + compile_kwargs = {}, + ), + ) + k2 = cc.cache_key( + efp, + cc.model_fingerprint( + family = "f", + transformer = _transformer(("B",)), + dtype = "bf16", + quant = None, + attention_backend = "x", + compile_kwargs = {}, + ), + ) + assert k1 != k2 + + +# ----------------------------------------------------------------------------- env knobs +@pytest.mark.parametrize( + "raw,expected", + [ + ("0", "off"), + ("off", "off"), + ("1", "on"), + ("on", "on"), + ("auto", "auto"), + ("", "auto"), + ("garbage", "auto"), + ], +) +def test_cache_mode(monkeypatch, raw, expected): + monkeypatch.setenv(cc._ENV_MODE, raw) + assert cc.cache_mode() == expected + + +def test_cache_mode_default_auto(monkeypatch): + monkeypatch.delenv(cc._ENV_MODE, raising = False) + assert cc.cache_mode() == "auto" + + +# ------------------------------------------------------------------------------ disabled +def test_begin_returns_none_when_disabled(monkeypatch): + monkeypatch.setenv(cc._ENV_MODE, "0") + assert cc.begin(transformer = _transformer(), **_BEGIN_KW) is None + + +def test_begin_returns_none_without_megacache_api(monkeypatch): + monkeypatch.setenv(cc._ENV_MODE, "auto") + fake_torch = types.ModuleType("torch") + fake_torch.compiler = types.SimpleNamespace() # no save/load attrs + monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch) + assert cc.begin(transformer = _transformer(), **_BEGIN_KW) is None + + +# ----------------------------------------------------------------- megacache fake + flow +@pytest.fixture +def fake_megacache(monkeypatch): + """Patch torch.compiler save/load with deterministic in-memory behaviour.""" + import torch + + state = {"saved": None, "loaded_with": None} + + def fake_save(): + return (b"ARTIFACT-BYTES", None) + + def fake_load(data: bytes): + state["loaded_with"] = data + return object() if data == b"ARTIFACT-BYTES" else None + + monkeypatch.setattr(torch.compiler, "save_cache_artifacts", fake_save, raising = False) + monkeypatch.setattr(torch.compiler, "load_cache_artifacts", fake_load, raising = False) + return state + + +def test_save_then_load_roundtrip(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "on") # load + save + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + + # First load: cold (no bundle yet). + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx is not None and ctx.hit is False + assert cc.save(ctx) is True + assert ctx.bundle.exists() and ctx.manifest_path.exists() + + # Second load with the SAME fingerprint: warm hit. + ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx2 is not None and ctx2.hit is True + assert fake_megacache["loaded_with"] == b"ARTIFACT-BYTES" + assert ctx2.key == ctx.key + + +def test_no_save_in_auto_mode(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "auto") + monkeypatch.delenv(cc._ENV_SAVE, raising = False) + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert cc.save(ctx) is False # auto without SAVE opt-in does not write + assert not ctx.bundle.exists() + + +def test_fingerprint_mismatch_falls_back(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "on") + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + cc.save(ctx) + + # Tamper the manifest's env fingerprint -> exact-match guard must reject the bundle. + manifest = json.loads(ctx.manifest_path.read_text()) + manifest["env"]["torch"] = "0.0.0-other" + ctx.manifest_path.write_text(json.dumps(manifest)) + + ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx2.hit is False # mismatch -> local compile, non-fatal + + +def test_corrupt_bundle_rejected(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "on") + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + cc.save(ctx) + ctx.bundle.write_bytes(b"CORRUPTED") # manifest sha256 no longer matches + + ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx2.hit is False + + +# ------------------------------------------------------------------------------- restore +def test_restore_inductor_dir(monkeypatch, tmp_path, fake_megacache): + import os + + monkeypatch.setenv(cc._ENV_MODE, "auto") + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", "/tmp/prior-inductor") + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert os.environ["TORCHINDUCTOR_CACHE_DIR"] != "/tmp/prior-inductor" # redirected + cc.restore(ctx) + assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == "/tmp/prior-inductor" # restored diff --git a/studio/backend/tests/test_diffusion_eager_patches.py b/studio/backend/tests/test_diffusion_eager_patches.py new file mode 100644 index 0000000000..a32dbec2ee --- /dev/null +++ b/studio/backend/tests/test_diffusion_eager_patches.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Numerical + lifecycle tests for the shared eager speedup patches. + +Builds the REAL diffusers 0.38 modules, captures the stock output, installs the patches, +and asserts the patched output matches within tolerance (fp32 on CPU always; bf16 on CUDA +when available). Also checks install/uninstall reversibility + idempotency, the +signature-guard no-op, and that a patched block compiles ``fullgraph=True`` (no graph break). +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("diffusers") + +import torch.nn as nn # noqa: E402 + +from core.inference import diffusion_eager_patches as ep # noqa: E402 +from diffusers.models.normalization import ( # noqa: E402 + AdaLayerNormContinuous, + AdaLayerNormZero, + AdaLayerNormZeroSingle, + RMSNorm, +) + +B, S, D, COND = 2, 16, 64, 32 + + +@pytest.fixture(autouse = True) +def _clean_patches(): + ep.uninstall_patches() + yield + ep.uninstall_patches() + + +def _devices_dtypes(): + cases = [("cpu", torch.float32)] + if torch.cuda.is_available(): + cases.append(("cuda", torch.bfloat16)) + return cases + + +def _build(cls, device, dtype): + torch.manual_seed(0) + if cls is RMSNorm: + m = RMSNorm(D, eps = 1e-6, elementwise_affine = True) + elif cls is AdaLayerNormContinuous: + m = AdaLayerNormContinuous( + D, COND, elementwise_affine = False, eps = 1e-6, norm_type = "layer_norm" + ) + elif cls is AdaLayerNormZero: + m = AdaLayerNormZero(D, num_embeddings = None, norm_type = "layer_norm") + elif cls is AdaLayerNormZeroSingle: + m = AdaLayerNormZeroSingle(D, norm_type = "layer_norm") + return m.to(device = device, dtype = dtype).eval() + + +def _inputs(cls, device, dtype): + torch.manual_seed(1) + x = torch.randn(B, S, D, device = device, dtype = dtype) + if cls is RMSNorm: + return (x,) + if cls is AdaLayerNormContinuous: + return (x, torch.randn(B, COND, device = device, dtype = dtype)) + # AdaLayerNormZero / Single take the conditioning emb of width D + return (x, torch.randn(B, D, device = device, dtype = dtype)) + + +def _call(cls, m, args): + if cls is AdaLayerNormZero: + return m(args[0], emb = args[1]) + return m(*args) + + +def _first(out): + return out[0] if isinstance(out, tuple) else out + + +@pytest.mark.parametrize( + "cls", [RMSNorm, AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle] +) +@pytest.mark.parametrize("device,dtype", _devices_dtypes()) +def test_patched_matches_original(cls, device, dtype): + m = _build(cls, device, dtype) + args = _inputs(cls, device, dtype) + + with torch.inference_mode(): + ref = _first(_call(cls, m, args)).clone() + + assert ep.install_compile_safe_patches() >= 1 + with torch.inference_mode(): + got = _first(_call(cls, m, args)) + + # The fused ops are FMA-based (addcmul) / fused (F.rms_norm): within ~1 ULP of the + # stock mul+add (and more accurate, single rounding), NOT bit-identical in fp32. + atol, rtol = (1e-5, 1e-4) if dtype == torch.float32 else (8e-3, 8e-3) + torch.testing.assert_close(got, ref, atol = atol, rtol = rtol) + + +def test_rmsnorm_mixed_dtype_falls_back(): + """fp32 activations into a bf16-weight RMSNorm: diffusers reduces variance in fp32 from + the original tensor, so the fused path must FALL BACK (identical output, not divergent).""" + m = RMSNorm(D, eps = 1e-6, elementwise_affine = True).to(torch.bfloat16).eval() + x = torch.randn(B, S, D, dtype = torch.float32) + with torch.inference_mode(): + ref = m(x).clone() + ep.install_compile_safe_patches() + with torch.inference_mode(): + got = m(x) + torch.testing.assert_close(got, ref, atol = 0.0, rtol = 0.0) # exact fallback + + +def test_rmsnorm_tuple_dim_falls_back(): + """diffusers RMSNorm always reduces the LAST dim even for a tuple `dim`; F.rms_norm + would reduce all of them, so a multi-dim `dim` must FALL BACK to the original.""" + m = RMSNorm((2, D), eps = 1e-6, elementwise_affine = True).eval() + x = torch.randn(B, 2, D) + with torch.inference_mode(): + ref = m(x).clone() + ep.install_compile_safe_patches() + with torch.inference_mode(): + got = m(x) + torch.testing.assert_close(got, ref, atol = 0.0, rtol = 0.0) # exact fallback + + +def test_install_idempotent_and_reversible(): + rms = RMSNorm(D, eps = 1e-6) + orig = RMSNorm.forward + n1 = ep.install_compile_safe_patches() + n2 = ep.install_compile_safe_patches() # second call is a no-op + assert n1 >= 1 and n2 == n1 + assert RMSNorm.forward is not orig + assert ep.is_installed() + ep.uninstall_patches() + assert RMSNorm.forward is orig # exact restore + assert not ep.is_installed() + ep.uninstall_patches() # idempotent uninstall + del rms + + +def test_kill_switch_disables_patches(monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_EAGER_PATCHES", "0") + orig = RMSNorm.forward + assert ep.install_compile_safe_patches() == 0 # no-op + assert not ep.is_installed() + assert RMSNorm.forward is orig # untouched + + +def test_signature_guard_skips_changed_class(monkeypatch): + """A diffusers class whose forward signature differs must be left untouched.""" + + class WeirdRMS(nn.Module): + def forward(self, x, extra): # not (self, hidden_states) + return x + + orig = WeirdRMS.forward + monkeypatch.setattr(ep, "_RMSNorm", WeirdRMS) + monkeypatch.setattr(ep, "_AdaLayerNormContinuous", None) + monkeypatch.setattr(ep, "_AdaLayerNormZero", None) + monkeypatch.setattr(ep, "_AdaLayerNormZeroSingle", None) + applied = ep.install_compile_safe_patches() + assert applied == 0 # nothing matched -> nothing patched + assert WeirdRMS.forward is orig # left untouched + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason = "compile graph-break check needs CUDA") +def test_no_graph_break_under_fullgraph(): + ep.install_compile_safe_patches() + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.rms = RMSNorm(D, eps = 1e-6) + self.ada = AdaLayerNormContinuous( + D, COND, elementwise_affine = False, norm_type = "layer_norm" + ) + + def forward(self, x, cond): + return self.ada(self.rms(x), cond) + + m = Block().to("cuda", torch.bfloat16).eval() + x = torch.randn(B, S, D, device = "cuda", dtype = torch.bfloat16) + cond = torch.randn(B, COND, device = "cuda", dtype = torch.bfloat16) + compiled = torch.compile(m, fullgraph = True) # raises if a graph break occurs + with torch.inference_mode(): + out = compiled(x, cond) + assert out.shape == (B, S, D) diff --git a/studio/backend/tests/test_diffusion_gguf_compile.py b/studio/backend/tests/test_diffusion_gguf_compile.py new file mode 100644 index 0000000000..c453713d29 --- /dev/null +++ b/studio/backend/tests/test_diffusion_gguf_compile.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the compiled GGUF dequant accelerator (``diffusion_gguf_compile.py``). + +Covers install/uninstall idempotency + exact reversibility, the kill-switch, and the +on-by-default behaviour. Runs on CPU -- patching the module attribute is lazy +(torch.compile only traces on the first real call). +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") +gguf_utils = pytest.importorskip("diffusers.quantizers.gguf.utils") + +from core.inference import diffusion_gguf_compile as gc # noqa: E402 + + +@pytest.fixture(autouse = True) +def _clean(): + # Always start and end from a clean, unpatched state so tests do not leak the + # process-wide patch into each other. + gc.uninstall_all() + yield + gc.uninstall_all() + + +def test_compiled_dequant_install_uninstall_reversible(): + orig = gguf_utils.dequantize_gguf_tensor + assert gc.is_compiled_dequant_installed() is False + + assert gc.install_compiled_dequant() is True + assert gc.is_compiled_dequant_installed() is True + # The module attribute is now a different (compiled) callable... + assert gguf_utils.dequantize_gguf_tensor is not orig + # ...idempotent: a second install is a no-op, attribute unchanged. + patched = gguf_utils.dequantize_gguf_tensor + assert gc.install_compiled_dequant() is True + assert gguf_utils.dequantize_gguf_tensor is patched + + gc.uninstall_compiled_dequant() + assert gc.is_compiled_dequant_installed() is False + # Exact original restored. + assert gguf_utils.dequantize_gguf_tensor is orig + # Uninstall is idempotent. + gc.uninstall_compiled_dequant() + assert gguf_utils.dequantize_gguf_tensor is orig + + +def test_compiled_dequant_kill_switch(monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT", "0") + orig = gguf_utils.dequantize_gguf_tensor + assert gc.install_compiled_dequant() is False + assert gc.is_compiled_dequant_installed() is False + assert gguf_utils.dequantize_gguf_tensor is orig + + +def test_compiled_dequant_on_by_default(monkeypatch): + # The compiled dequant is the real win, so it is ON without any env opt-in. + monkeypatch.delenv("UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT", raising = False) + assert gc.install_compiled_dequant() is True + assert gc.is_compiled_dequant_installed() is True + + +def test_uninstall_all(monkeypatch): + orig = gguf_utils.dequantize_gguf_tensor + gc.install_compiled_dequant() + assert gc.is_installed() is True + gc.uninstall_all() + assert gc.is_installed() is False + assert gguf_utils.dequantize_gguf_tensor is orig diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index a3c78a21aa..f6a9ea2766 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -35,13 +35,21 @@ class _FakeBackend: *, gguf_filename = None, family_override = None, + model_kind = None, ): # Mirror the real backend's cheap validation so the route's # validate-before-evict ordering is exercised. + from core.inference.diffusion import resolve_model_kind from core.inference.diffusion_families import detect_family - if not gguf_filename: - raise ValueError("gguf_filename is required.") + kind = resolve_model_kind(gguf_filename, model_kind) + if kind in ("gguf", "single_file") and not gguf_filename: + raise ValueError("a single-file checkpoint name is required.") + # Non-GGUF loads are gated to unsloth/* (or a local path), like the real backend. + if kind != "gguf" and not model_path.lower().startswith("unsloth/"): + raise ValueError( + f"Non-GGUF diffusion loads are restricted to unsloth/* repos; got '{model_path}'." + ) fam = detect_family(model_path, family_override) if fam is None: raise ValueError(f"Could not infer a diffusion family for '{model_path}'.") @@ -254,10 +262,24 @@ def test_generate_rejects_non_multiple_of_16(client): assert ok.status_code == 200 -def test_load_requires_gguf_filename(client): - # gguf_filename is now mandatory — a load without it is a 422. +def test_non_gguf_load_restricted_to_unsloth(client): + # gguf_filename is optional now; with none, the load is a full-pipeline kind, which + # is gated to unsloth/* repos. A non-unsloth repo (no filename) is rejected -> 400. resp = client.post("/api/inference/images/load", json = {"model_path": "x/z-image"}) - assert resp.status_code == 422 + assert resp.status_code == 400 + assert "unsloth" in resp.json()["detail"].lower() + + +def test_pipeline_load_allowed_for_unsloth_repo(client): + # An unsloth/* repo with no filename loads as a full diffusers pipeline (kind auto + # = pipeline); the route forwards model_kind="pipeline" to begin_load. + resp = client.post( + "/api/inference/images/load", json = {"model_path": "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"} + ) + assert resp.status_code == 200 + backend = diffusion_module.get_diffusion_backend() + assert backend.last_load_kwargs["model_kind"] == "pipeline" + assert backend.last_load_kwargs.get("gguf_filename") is None def test_generate_without_load_returns_409(client): diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 1f64b31b7a..de73121e00 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -14,8 +14,10 @@ import types import pytest +from core.inference import diffusion_speed as ds_mod from core.inference.diffusion_speed import ( SPEED_DEFAULT, + SPEED_EAGER, SPEED_MAX, SPEED_OFF, apply_speed_optims, @@ -27,6 +29,20 @@ from core.inference.diffusion_speed import ( ) +def _stub_gguf_accel(monkeypatch): + """Replace the real compiled-dequant installer (which touches torch.compile / + diffusers) with a recorder, so the tier-gating logic in apply_speed_optims is tested + in isolation. Returns a dict of how many times it was called.""" + called = {"compiled_dequant": 0} + + def _install(logger = None): + called["compiled_dequant"] += 1 + return True + + monkeypatch.setattr(ds_mod.gguf_compile, "install_compiled_dequant", _install) + return called + + def _target( *, device = "cuda", @@ -194,14 +210,18 @@ def test_speed_off_applies_nothing(monkeypatch): "tf32": False, "fused_qkv": False, "compiled": False, + "compiled_dequant": False, } assert pipe.vae.mem_format is None and pipe.compiled is False # off must not touch any process-wide flag (bit-identical reference path). assert torch.backends.cudnn.benchmark is False -def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): +def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch): + # A DENSE model has no GGUF dequant to compile, so `default` falls back to the + # regional block compile (its only compile lever) -- and no GGUF accelerators. torch = _stub_torch(monkeypatch) + called = _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT @@ -214,18 +234,22 @@ def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): # default also autotunes the VAE convs but does NOT flip TF32 or fuse QKV. assert applied["cudnn_benchmark"] is True and torch.backends.cudnn.benchmark is True assert applied["tf32"] is False and applied["fused_qkv"] is False + # No GGUF dequant on a dense model. + assert applied["compiled_dequant"] is False + assert called == {"compiled_dequant": 0} def test_offload_active_drops_fullgraph(monkeypatch): # Group/model/sequential offload installs a torch.compiler.disable'd onload hook; # compiling with fullgraph=True then crashes at the first denoise step. Same reason # as an active step cache -> fullgraph must drop to False when offload is planned. + # (Dense model: on this branch GGUF `default` takes the compiled-dequant path.) _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), - is_gguf = True, + is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT, offload_active = True, @@ -234,15 +258,50 @@ def test_offload_active_drops_fullgraph(monkeypatch): assert pipe.compile_kwargs["fullgraph"] is False -def test_speed_default_compiles_gguf(monkeypatch): +def test_speed_default_gguf_compiles_only_dequant(monkeypatch): + # GGUF `default` is the LIGHT path: compile ONLY the dequant op chain, NOT the + # regional block compile. _stub_torch(monkeypatch) + called = _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["channels_last"] is True - # GGUF now compiles (the big near-lossless win). + assert applied["compiled_dequant"] is True + # The transformer block is NOT regionally compiled under GGUF default. + assert applied["compiled"] is False and pipe.compiled is False + assert called == {"compiled_dequant": 1} + + +def test_speed_eager_gguf_installs_no_accelerator(monkeypatch): + # eager = lossless-but-no-compile: neither the compiled dequant nor the regional + # block compile run; only the process-wide lossless levers (channels_last, cudnn) + # and the shared/per-arch eager monkey-patches (installed elsewhere) engage. + _stub_torch(monkeypatch) + called = _stub_gguf_accel(monkeypatch) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_EAGER + ) + assert applied["compiled_dequant"] is False and applied["compiled"] is False + assert pipe.compiled is False + assert called == {"compiled_dequant": 0} + + +def test_speed_max_gguf_regional_compile_not_dequant(monkeypatch): + # GGUF `max` = the FULL regional block compile (which fuses the dequant inline), so + # the standalone compiled dequant is deliberately OFF. + _stub_torch(monkeypatch) + called = _stub_gguf_accel(monkeypatch) + pipe = _Pipe(with_compile = True, with_fuse = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_MAX + ) assert applied["compiled"] is True and pipe.compiled is True + assert pipe.compile_kwargs["mode"] == "max-autotune-no-cudagraphs" + assert applied["compiled_dequant"] is False + assert called == {"compiled_dequant": 0} def test_speed_default_cudnn_benchmark_only_on_cuda(monkeypatch): diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index b009cd37a6..4e06a3cdba 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -17,13 +17,25 @@ _STUDIO = Path(__file__).resolve().parents[2] if str(_STUDIO) not in sys.path: sys.path.insert(0, str(_STUDIO)) +import hashlib # noqa: E402 +import io # noqa: E402 +import json # noqa: E402 +import urllib.error # noqa: E402 import zipfile # noqa: E402 import pytest # noqa: E402 +import install_sd_cpp_prebuilt as sdmod # noqa: E402 from install_sd_cpp_prebuilt import ( # noqa: E402 + DEFAULT_REPO, + DEFAULT_TAG, + _fetch_release, + _pinned_tag, + _repo, _safe_extractall, + _verify_sha256, default_install_dir, + install, resolve_release_asset, ) @@ -124,6 +136,136 @@ def test_default_install_dir_is_sibling_of_llama(monkeypatch): assert d.parent.name == ".unsloth" +# ── version pin + source repo (reproducibility) ───────────────────────────── + + +def test_pinned_tag_default_and_override(monkeypatch): + monkeypatch.delenv("UNSLOTH_SD_CPP_TAG", raising = False) + assert _pinned_tag() == DEFAULT_TAG # pinned, not "latest" + monkeypatch.setenv("UNSLOTH_SD_CPP_TAG", "master-999-deadbee") + assert _pinned_tag() == "master-999-deadbee" + monkeypatch.setenv("UNSLOTH_SD_CPP_TAG", "") # explicit empty -> track latest + assert _pinned_tag() is None + + +def test_repo_default_and_mirror_override(monkeypatch): + monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False) + assert _repo() == DEFAULT_REPO == "leejet/stable-diffusion.cpp" + monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", "unslothai/stable-diffusion.cpp") + assert _repo() == "unslothai/stable-diffusion.cpp" + + +# ── sha256 integrity check ────────────────────────────────────────────────── + + +def test_verify_sha256_accepts_matching_digest(tmp_path): + f = tmp_path / "asset.zip" + f.write_bytes(b"hello sd-cli") + digest = "sha256:" + hashlib.sha256(b"hello sd-cli").hexdigest() + _verify_sha256(f, digest) # no raise + + +def test_verify_sha256_rejects_mismatch(tmp_path): + f = tmp_path / "asset.zip" + f.write_bytes(b"tampered") + bad = "sha256:" + hashlib.sha256(b"original").hexdigest() + with pytest.raises(RuntimeError, match = "sha256 mismatch"): + _verify_sha256(f, bad) + + +def test_verify_sha256_skips_when_absent_or_unknown(tmp_path): + f = tmp_path / "asset.zip" + f.write_bytes(b"x") + _verify_sha256(f, None) # no digest published -> warn + proceed (no raise) + _verify_sha256(f, "md5:abc") # unrecognised algo -> skip (no raise) + + +# ── _fetch_release: pinned-tag 404 -> latest fallback ─────────────────────── + + +def test_fetch_release_falls_back_to_latest_on_404(monkeypatch): + calls: list[str] = [] + + class _Resp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return json.dumps({"tag_name": "latest-xyz", "assets": []}).encode() + + def fake_urlopen(req, timeout = 30.0): + url = getattr(req, "full_url", req) + calls.append(url) + if "/tags/" in url: + raise urllib.error.HTTPError(url, 404, "not found", None, None) + return _Resp() + + monkeypatch.setattr(sdmod.urllib.request, "urlopen", fake_urlopen) + rel = _fetch_release("gone-tag", repo = "leejet/stable-diffusion.cpp") + assert rel["tag_name"] == "latest-xyz" + assert any("/tags/gone-tag" in c for c in calls) and any(c.endswith("/latest") for c in calls) + + +def test_fetch_release_propagates_non_404(monkeypatch): + def fake_urlopen(req, timeout = 30.0): + url = getattr(req, "full_url", req) + raise urllib.error.HTTPError(url, 403, "rate limited", None, None) + + monkeypatch.setattr(sdmod.urllib.request, "urlopen", fake_urlopen) + with pytest.raises(urllib.error.HTTPError): + _fetch_release("any-tag") + + +# ── install(): download -> verify -> extract -> locate (offline) ──────────── + + +def _zip_with_sd_cli() -> bytes: + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("build/bin/sd-cli", b"#!/bin/sh\necho sd-cli\n") + return buf.getvalue() + + +def _stub_release(monkeypatch, *, zip_bytes: bytes, digest: str): + name = "sd-master-deadbee-bin-Linux-Ubuntu-24.04-x86_64.zip" + release = { + "tag_name": "master-1-deadbee", + "assets": [ + { + "name": name, + "browser_download_url": f"https://example.invalid/{name}", + "digest": digest, + } + ], + } + monkeypatch.setattr(sdmod, "_fetch_release", lambda *a, **k: release) + monkeypatch.setattr(sdmod, "_download", lambda url, dest, **k: dest.write_bytes(zip_bytes)) + monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux") + monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64") + return name + + +def test_install_downloads_verifies_extracts(tmp_path, monkeypatch): + zb = _zip_with_sd_cli() + name = _stub_release( + monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest() + ) + sd_cli = install(install_dir = tmp_path) + assert sd_cli.name == "sd-cli" and sd_cli.is_file() + assert not (tmp_path / name).exists() # archive cleaned up after extract + + +def test_install_sha256_mismatch_raises_and_cleans_up(tmp_path, monkeypatch): + zb = _zip_with_sd_cli() + name = _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + "0" * 64) + with pytest.raises(RuntimeError, match = "sha256 mismatch"): + install(install_dir = tmp_path) + assert not (tmp_path / name).exists() # the finally: drops the bad archive + + # ── safe extraction (Zip-Slip guard) ───────────────────────────────────────── 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 622dd39d69..b8d75b2eb6 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1065,15 +1065,26 @@ const NON_CHAT_TASKS: readonly string[] = [...IMAGE_GEN_TASKS, UNSUPPORTED_DIFFU // which the text-to-image backend rejects (mirrors its _EDIT_KEYWORDS). Hidden by // id so they don't show in the Images picker only to 400 on load. Keeping the // image-to-image task itself is required: some supported models (FLUX.2-klein) -// carry that tag too. -const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint", "inpainting"] as const; +// carry that tag too. "layered" hides Qwen-Image-Layered, which needs a dedicated +// pipeline (additional_t_cond) the standard text-to-image path can't drive. +const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint", "layered"] as const; +// Editing families the backend now SUPPORTS (their own Edit workflow) -- must not be +// hidden even though their id contains an edit keyword. Mirrors the backend's +// qwen-image-edit family in diffusion_families.py. +const SUPPORTED_EDIT_KEYWORDS = ["qwen-image-edit", "kontext"] as const; +// Match a keyword as a whole path/name segment (bounded by a separator or a string +// edge), not a raw substring, so "edit" does not hide ".../edited/..." or an +// "*-edition" repo and "kontext" does not hide ".../kontextual/...". These keywords +// are literals of [a-z-], so no regex escaping is needed. Mirrors _token_in_needle in +// diffusion_families.py. +function idHasSegment(id: string, keyword: string): boolean { + return new RegExp(`(?:^|[-_./\\\\])${keyword}(?:$|[-_./\\\\])`).test(id); +} function isImageEditModel(repoId: string | null | undefined): boolean { if (!repoId) return false; - // Whole-segment match (not substring) so a normal model like "...-edition" - // isn't hidden; mirrors the backend detect_family segment check. Split on - // both path separators so a Windows local path is segmented too. - const segments = new Set(repoId.toLowerCase().split(/[-_./\\]+/)); - return IMAGE_EDIT_KEYWORDS.some((kw) => segments.has(kw)); + const id = repoId.toLowerCase(); + if (SUPPORTED_EDIT_KEYWORDS.some((kw) => idHasSegment(id, kw))) return false; + return IMAGE_EDIT_KEYWORDS.some((kw) => idHasSegment(id, kw)); } // Gate an on-device model by the picker's task scope. With a filter (the Images @@ -1829,6 +1840,21 @@ export function HubModelPicker({ task, ]); + // Curated non-GGUF (safetensors) models for the Images picker. The HF listing + + // Recommended gate only surface GGUF on a GPU host (isRecommendableFormat), so a + // bnb-4bit / fp8 safetensors model would never appear there. These curated entries + // (the non-GGUF ModelOptions passed in) are shown explicitly above the GGUF rows so + // the user can pick a full diffusers pipeline. Only the Images picker (task set) + // curates them; already-downloaded ones show under Downloaded instead. + const curatedSafetensorsRows = useMemo(() => { + if (!task) return []; + // Always list the curated safetensors (bnb-4bit / fp8) diffusion models. They + // render with a "downloaded" badge when cached (like GGUF Recommended rows), so + // they must not be hidden once on disk -- otherwise they vanish from the picker + // entirely after the first load. + return models.filter((m) => m.isGguf === false); + }, [models, task]); + // Per-row meta + VRAM badge from the recommended listing's own metadata. const recommendedMeta = useMemo(() => { const map = new Map< @@ -1922,18 +1948,21 @@ export function HubModelPicker({ ), [cachedGguf, downloadedSort, loadTimes, task], ); - // Non-GGUF (safetensors) cached repos aren't single-file diffusion GGUFs, so - // hide them entirely when a task filter is active (the Images picker). In chat, - // drop cached diffusers pipeline repos (a Z-Image / FLUX base) the same way. + // Cached non-GGUF repos. In chat, passesTaskGate drops diffusers image repos. In the + // Images picker (task set) it keeps them, but limit to repos this backend can actually + // load as diffusion: unsloth-hosted ones. Base repos (Qwen/Qwen-Image, FLUX bases) are + // cached as dependencies and fail the diffusion trust gate, so listing them would dead-end. const sortedCachedModels = useMemo( () => - task - ? [] - : sortCachedRepos( - cachedModels.filter((c) => passesTaskGate(c.task, c.repo_id, task)), - downloadedSort, - loadTimes, - ), + sortCachedRepos( + cachedModels.filter( + (c) => + passesTaskGate(c.task, c.repo_id, task) && + (!task || isUnslothRepoId(c.repo_id)), + ), + downloadedSort, + loadTimes, + ), [cachedModels, downloadedSort, loadTimes, task], ); // Each local section's search is scoped to its own models (matched by name). @@ -3426,6 +3455,30 @@ export function HubModelPicker({ {showRecommendedSection ? ( <> + {/* Curated safetensors models (full diffusers pipelines / single-file + fp8). Shown above the GGUF rows; clicking loads directly (no quant + expander), the same path as a non-GGUF Recommended row. */} + {curatedSafetensorsRows.map((m) => { + const optionKey = makeModelOptionKey("curated-safetensors", m.id); + return ( +
+ handleModelClick(m.id)} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + /> +
+ ); + })} {recommendedSearch.isLoading && recommendedRows.length === 0 ? (
@@ -3434,7 +3487,8 @@ export function HubModelPicker({ Loading models…
- ) : recommendedRows.length === 0 ? ( + ) : recommendedRows.length === 0 && + curatedSafetensorsRows.length === 0 ? (
No models found.
diff --git a/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts b/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts index 4e7be09f2c..f2d072fc6c 100644 --- a/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts +++ b/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts @@ -217,9 +217,19 @@ function makeMapModel( ) { return null; } + // A repo cross-tagged "gguf" but that is actually a diffusers pipeline (e.g. + // an unsloth *-bnb-4bit image model) ships no .gguf files, so the GGUF + // variant expander would dead-end at "No GGUF variants found." Trust the bare + // tag only when the repo is not a diffusers pipeline. The "-GGUF" name suffix + // and real gguf metadata (populated via expand=gguf) stay authoritative. + const isDiffusersPipeline = + m.library_name?.toLowerCase() === "diffusers" || + Boolean(m.tags?.some((tag) => tag.toLowerCase().startsWith("diffusers:"))); const isGguf = - Boolean(m.tags?.some((tag) => tag.toLowerCase() === "gguf")) || - isGgufLike(m.name); + isGgufLike(m.name) || + Boolean(m.gguf) || + (Boolean(m.tags?.some((tag) => tag.toLowerCase() === "gguf")) && + !isDiffusersPipeline); if (excludeGguf && isGguf) { return null; } diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 2a3f5ec1a2..3c9c8df062 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -11,7 +11,13 @@ export interface DiffusionStatus { base_repo: string | null; device: string | null; dtype: string | null; + // Resolved load kind: "gguf" | "single_file" | "pipeline". Gates GGUF-only controls + // (the dense transformer_quant fast path only engages on gguf). Null when not loaded. + model_kind?: string | null; cpu_offload: boolean; + // Image workflows the loaded family supports (drives tab gating): txt2img, img2img, + // inpaint. Absent/empty when nothing is loaded or on the native sd.cpp engine. + workflows?: string[]; } export interface DiffusionGenerateProgress { @@ -32,10 +38,33 @@ export interface DiffusionLoadProgress { export interface DiffusionLoadRequest { model_path: string; - gguf_filename: string; + // Optional now: required for the gguf / single_file kinds, omitted for a full + // pipeline (a diffusers repo loaded via from_pretrained). + gguf_filename?: string; + // How to load the model (omit to auto-detect from gguf_filename): "gguf" (single-file + // GGUF transformer), "single_file" (single-file safetensors transformer, e.g. fp8), or + // "pipeline" (a full diffusers repo). Non-GGUF kinds are restricted to unsloth/* repos. + model_kind?: "gguf" | "single_file" | "pipeline"; base_repo?: string; family_override?: string; hf_token?: string; + cpu_offload?: boolean; + // Advanced (load-time) tuning. All optional; omit for the backend's auto defaults. + speed_mode?: "off" | "eager" | "default" | "max"; + transformer_quant?: "auto" | "int8" | "fp8" | "nvfp4" | "mxfp8"; + attention_backend?: + | "auto" + | "native" + | "cudnn" + | "flash" + | "flash2" + | "flash3" + | "flash4" + | "sage" + | "xformers" + | "aiter"; + memory_mode?: "auto" | "fast" | "balanced" | "low_vram"; + transformer_cache?: "off" | "fbcache"; } export interface DiffusionGenerateRequest { @@ -47,6 +76,16 @@ export interface DiffusionGenerateRequest { guidance?: number; seed?: number; batch_size?: number; + // Image-conditioned workflows. init_image alone = img2img; init_image + mask_image = + // inpaint. Base64 or data-URL. strength is the denoise amount (0 keeps source, 1 redraws). + init_image?: string; + mask_image?: string; + strength?: number; + // Upscale (hires fix): factor > 1 with an init_image enlarges the source and re-denoises + // it at low strength. Requires init_image; ignored for txt2img/inpaint/edit. + upscale?: number; + // Additional reference images for the FLUX.2 reference workflow, combined with init_image. + reference_images?: string[]; } // A persisted image's full generation recipe (also embedded in the PNG). diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 0b761aaddd..a563bbd909 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -9,6 +9,8 @@ import { Download01Icon, ImageAdd02Icon, InformationCircleIcon, + LayoutAlignRightIcon, + Settings02Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -28,6 +30,7 @@ import { } from "@/components/ui/select"; import { Slider } from "@/components/ui/slider"; import { Spinner } from "@/components/ui/spinner"; +import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { InfoHint } from "@/components/ui/info-hint"; import { ModelSelector } from "@/components/assistant-ui/model-selector"; @@ -67,6 +70,39 @@ const txt2img = (id: string, name: string): ModelOption => ({ description: "Text-to-image · GGUF", isGguf: true, }); + +// Instruction-editing GGUF (Qwen-Image-Edit). Same single-file GGUF flow as txt2img +// (the picker expands quant variants); the backend resolves it to the edit family, which +// exposes only the "edit" workflow. +const editGguf = (id: string, name: string): ModelOption => ({ + id, + name, + description: "Image editing · GGUF", + isGguf: true, +}); + +// How to load a curated non-GGUF (safetensors) model. "pipeline" = a full diffusers +// repo (from_pretrained, embedded bnb-4bit quant auto-applied); "single_file" = a +// single safetensors transformer (e.g. fp8) assembled onto its base repo. The backend +// gates these to unsloth/* repos. Keyed by repo id so the load handler knows the kind +// (and, for single_file, the exact filename). +type SafetensorsSpec = { kind: "pipeline" | "single_file"; filename?: string }; +const SAFETENSORS_MODELS: Record = { + "unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" }, + "unsloth/Qwen-Image-2512-unsloth-bnb-4bit": { kind: "pipeline" }, + "unsloth/Qwen-Image-2512-FP8": { + kind: "single_file", + filename: "qwen-image-2512-fp8.safetensors", + }, +}; +// Curated non-GGUF picker entries (isGguf:false -> no quant expander, direct load). +const safetensors = (id: string, name: string, label: string): ModelOption => ({ + id, + name, + description: `Text-to-image · ${label}`, + isGguf: false, +}); + const MODELS: ModelOption[] = [ txt2img("unsloth/Z-Image-Turbo-GGUF", "Z-Image-Turbo"), txt2img("unsloth/Z-Image-GGUF", "Z-Image"), @@ -76,6 +112,72 @@ const MODELS: ModelOption[] = [ txt2img("unsloth/FLUX.1-dev-GGUF", "FLUX.1 dev"), txt2img("unsloth/FLUX.2-klein-4B-GGUF", "FLUX.2 klein 4B"), txt2img("unsloth/FLUX.2-klein-9B-GGUF", "FLUX.2 klein 9B"), + editGguf("unsloth/Qwen-Image-Edit-2511-GGUF", "Qwen-Image-Edit 2511"), + editGguf("unsloth/FLUX.1-Kontext-dev-GGUF", "FLUX.1 Kontext dev"), + safetensors( + "unsloth/Z-Image-Turbo-unsloth-bnb-4bit", + "Z-Image-Turbo (bnb-4bit)", + "Safetensors · bnb-4bit", + ), + safetensors( + "unsloth/Qwen-Image-2512-unsloth-bnb-4bit", + "Qwen-Image 2512 (bnb-4bit)", + "Safetensors · bnb-4bit", + ), + safetensors( + "unsloth/Qwen-Image-2512-FP8", + "Qwen-Image 2512 (FP8)", + "Safetensors · fp8", + ), +]; + +// Workflow tabs. `requires` is the backend workflow id (status.workflows) that must +// be supported by the loaded model for the tab to enable; null = always available. +type WorkflowId = "create" | "transform" | "inpaint" | "extend" | "upscale" | "reference" | "edit"; + +const WORKFLOW_TABS: Array<{ + id: WorkflowId; + label: string; + requires: string | null; + hint?: string; +}> = [ + { id: "create", label: "Create", requires: null, hint: "Generate a new image from a prompt" }, + { + id: "transform", + label: "Transform", + requires: "img2img", + hint: "Redraw an uploaded image guided by your prompt (img2img)", + }, + { + id: "inpaint", + label: "Inpaint", + requires: "inpaint", + hint: "Paint over a region to regenerate just that area, keeping the rest", + }, + { + id: "extend", + label: "Extend", + requires: "outpaint", + hint: "Outpaint: grow the canvas and fill the new edges from your prompt", + }, + { + id: "upscale", + label: "Upscale", + requires: "upscale", + hint: "Hires fix: enlarge an uploaded image and re-detail it at higher resolution", + }, + { + id: "reference", + label: "Reference", + requires: "reference", + hint: "Generate a new image guided by a reference image + your prompt (FLUX.2)", + }, + { + id: "edit", + label: "Edit", + requires: "edit", + hint: "Instruction editing: change an image with a prompt (Qwen-Image-Edit)", + }, ]; // Per-model generation defaults (steps + guidance), matched by repo-id substring, @@ -88,8 +190,12 @@ const DEFAULT_GEN = { steps: 9, guidance: 0 }; const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> = [ { match: "z-image-turbo", steps: 9, guidance: 0 }, { match: "flux.1-schnell", steps: 4, guidance: 0 }, + // Kontext (editing) before the generic flux.1: ~28 steps, lower guidance (~2.5). + { match: "kontext", steps: 28, guidance: 2.5 }, { match: "flux.1", steps: 28, guidance: 3.5 }, { match: "flux.2-klein", steps: 4, guidance: 0 }, + // FLUX.2-dev is the full (non-distilled) model: more steps + real guidance, unlike klein. + { match: "flux.2-dev", steps: 28, guidance: 4 }, { match: "qwen-image", steps: 20, guidance: 4 }, { match: "z-image", steps: 20, guidance: 4 }, ]; @@ -288,7 +394,11 @@ function SliderField({ min={min} max={max} step={step} - className="w-12 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [&::-webkit-inner-spin-button]:appearance-none" + // The slider is the primary control, so the native number spinners are + // redundant — and on this narrow field their up/down arrows overlapped and + // covered the value. Remove them on every engine: appearance:textfield for + // Firefox, and zero out the webkit inner/outer spin buttons. + className="w-14 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [appearance:textfield] [&::-webkit-outer-spin-button]:m-0 [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:m-0 [&::-webkit-inner-spin-button]:appearance-none" /> @@ -316,6 +426,360 @@ function Field({ ); } +// A compact labeled Select row for the Advanced Options panel. +function AdvancedSelect({ + label, + hint, + value, + onValueChange, + options, +}: { + label: string; + hint?: ReactNode; + value: string; + onValueChange: (v: string) => void; + options: Array<[string, string]>; +}) { + return ( +
+ + {label} + {hint && {hint}} + + +
+ ); +} + +// Source-image picker for the Transform (img2img) workflow: click or drag-drop an +// image, read it to a data URL the generate request sends as init_image. Shows a +// thumbnail preview with a Clear button once an image is set. +function ImageDropzone({ + value, + onChange, +}: { + value: string | null; + onChange: (dataUrl: string | null) => void; +}) { + const inputRef = useRef(null); + const [dragging, setDragging] = useState(false); + + const readFile = useCallback( + (file: File | undefined | null) => { + if (!file || !file.type.startsWith("image/")) { + if (file) toast.error("Please choose an image file"); + return; + } + const reader = new FileReader(); + reader.onload = () => onChange(typeof reader.result === "string" ? reader.result : null); + reader.onerror = () => toast.error("Could not read the image"); + reader.readAsDataURL(file); + }, + [onChange], + ); + + if (value) { + return ( +
+ Source + +
+ ); + } + + return ( + + ); +} + +// A brush-based mask editor for inpainting. Shows the source image with a paintable +// overlay and exports a grayscale PNG mask at the image's NATIVE resolution, following +// the diffusers inpaint convention (white = repaint, black = keep). Strokes are drawn to +// both a visible tinted overlay (feedback) and an offscreen mask canvas kept in lockstep, +// so the exported mask always matches what the user sees. `brushPct` sizes the brush as a +// fraction of the image's shorter side, so it stays consistent across resolutions. +function MaskCanvas({ + image, + brushPct, + resetKey, + onMaskChange, +}: { + image: string; + brushPct: number; + resetKey: number; + onMaskChange: (dataUrl: string | null) => void; +}) { + const dispRef = useRef(null); + const maskRef = useRef(null); + const dims = useRef<{ w: number; h: number }>({ w: 0, h: 0 }); + const drawing = useRef(false); + const last = useRef<{ x: number; y: number } | null>(null); + const [ready, setReady] = useState(false); + + // (Re)initialise both canvases whenever the image changes or Clear is pressed: + // size them to the image's native pixels and reset the mask to all-black (keep all). + useEffect(() => { + setReady(false); + const img = new Image(); + img.onload = () => { + const w = img.naturalWidth; + const h = img.naturalHeight; + dims.current = { w, h }; + const disp = dispRef.current; + const mask = maskRef.current ?? document.createElement("canvas"); + maskRef.current = mask; + if (!disp) return; + disp.width = w; + disp.height = h; + mask.width = w; + mask.height = h; + const mctx = mask.getContext("2d"); + const dctx = disp.getContext("2d"); + if (!mctx || !dctx) return; + mctx.fillStyle = "#000"; + mctx.fillRect(0, 0, w, h); + dctx.clearRect(0, 0, w, h); + setReady(true); + onMaskChange(null); + }; + img.src = image; + }, [image, resetKey, onMaskChange]); + + const radius = useCallback(() => { + const base = Math.min(dims.current.w, dims.current.h) || 1024; + return Math.max(2, (brushPct / 100) * base); + }, [brushPct]); + + const toNatural = (e: React.PointerEvent) => { + const disp = dispRef.current; + if (!disp) return { x: 0, y: 0 }; + const r = disp.getBoundingClientRect(); + return { + x: ((e.clientX - r.left) / r.width) * dims.current.w, + y: ((e.clientY - r.top) / r.height) * dims.current.h, + }; + }; + + const stroke = (from: { x: number; y: number } | null, to: { x: number; y: number }) => { + const disp = dispRef.current; + const mask = maskRef.current; + if (!disp || !mask) return; + const r = radius(); + const layers: Array<[CanvasRenderingContext2D | null, string]> = [ + [disp.getContext("2d"), "rgba(244,114,114,0.55)"], + [mask.getContext("2d"), "#ffffff"], + ]; + for (const [ctx, style] of layers) { + if (!ctx) continue; + ctx.strokeStyle = style; + ctx.fillStyle = style; + ctx.lineWidth = r * 2; + ctx.lineCap = "round"; + ctx.lineJoin = "round"; + ctx.beginPath(); + ctx.arc(to.x, to.y, r, 0, Math.PI * 2); + ctx.fill(); + if (from) { + ctx.beginPath(); + ctx.moveTo(from.x, from.y); + ctx.lineTo(to.x, to.y); + ctx.stroke(); + } + } + }; + + const onDown = (e: React.PointerEvent) => { + if (!ready) return; + drawing.current = true; + try { + e.currentTarget.setPointerCapture(e.pointerId); + } catch { + // setPointerCapture can throw for synthetic events; safe to ignore. + } + const p = toNatural(e); + last.current = p; + stroke(null, p); + }; + const onMove = (e: React.PointerEvent) => { + if (!drawing.current) return; + const p = toNatural(e); + stroke(last.current, p); + last.current = p; + }; + const onUp = () => { + if (!drawing.current) return; + drawing.current = false; + last.current = null; + const mask = maskRef.current; + if (mask) onMaskChange(mask.toDataURL("image/png")); + }; + + return ( +
+ Inpaint source + +
+ ); +} + +function loadImage(src: string): Promise { + return new Promise((resolve, reject) => { + const img = new Image(); + img.onload = () => resolve(img); + img.onerror = reject; + img.src = src; + }); +} + +// Which sides to grow when outpainting. +type ExtendSides = { left: boolean; right: boolean; top: boolean; bottom: boolean }; + +// Build the (image, mask) pair for outpaint by reusing the inpaint backend: grow the +// canvas by `pct` of each dimension on the selected sides, edge-bleed the original pixels +// into the new bands (so the VAE encodes plausible content), and mask the new bands white +// (= repaint) with a small overlap into the original on each grown side so the seam blends. +async function buildOutpaint( + src: string, + sides: ExtendSides, + pct: number, +): Promise<{ image: string; mask: string }> { + const img = await loadImage(src); + const w = img.naturalWidth; + const h = img.naturalHeight; + const px = Math.round((pct / 100) * w); + const py = Math.round((pct / 100) * h); + const l = sides.left ? px : 0; + const r = sides.right ? px : 0; + const t = sides.top ? py : 0; + const b = sides.bottom ? py : 0; + const nw = w + l + r; + const nh = h + t + b; + + const ic = document.createElement("canvas"); + ic.width = nw; + ic.height = nh; + const ictx = ic.getContext("2d"); + if (!ictx) throw new Error("Could not build the extended canvas"); + ictx.drawImage(img, l, t, w, h); // original, centred by the chosen offsets + // Edge-bleed: stretch the 1px border strips into each new band (and corners). + if (l) ictx.drawImage(img, 0, 0, 1, h, 0, t, l, h); + if (r) ictx.drawImage(img, w - 1, 0, 1, h, l + w, t, r, h); + if (t) ictx.drawImage(img, 0, 0, w, 1, l, 0, w, t); + if (b) ictx.drawImage(img, 0, h - 1, w, 1, l, t + h, w, b); + if (l && t) ictx.drawImage(img, 0, 0, 1, 1, 0, 0, l, t); + if (r && t) ictx.drawImage(img, w - 1, 0, 1, 1, l + w, 0, r, t); + if (l && b) ictx.drawImage(img, 0, h - 1, 1, 1, 0, t + h, l, b); + if (r && b) ictx.drawImage(img, w - 1, h - 1, 1, 1, l + w, t + h, r, b); + + const overlap = Math.round(Math.min(w, h) * 0.02); + const ol = l ? overlap : 0; + const or = r ? overlap : 0; + const ot = t ? overlap : 0; + const ob = b ? overlap : 0; + const mc = document.createElement("canvas"); + mc.width = nw; + mc.height = nh; + const mctx = mc.getContext("2d"); + if (!mctx) throw new Error("Could not build the extend mask"); + mctx.fillStyle = "#ffffff"; // repaint everything... + mctx.fillRect(0, 0, nw, nh); + mctx.fillStyle = "#000000"; // ...except the kept original (inset by the seam overlap). + mctx.fillRect(l + ol, t + ot, w - ol - or, h - ot - ob); + + // The grown canvas can exceed the backend's 4096px-per-side decode limit (e.g. a + // 2048px source at 100% on both sides -> 6144px), which would 400 the load. Scale the + // built pair down proportionally to fit, so Extend still returns an outpaint instead + // of failing. The backend also rounds to /16, so exact dims here are not required. + const MAX_SIDE = 4096; + const longest = Math.max(nw, nh); + if (longest > MAX_SIDE) { + const scale = MAX_SIDE / longest; + const sw = Math.max(1, Math.round(nw * scale)); + const sh = Math.max(1, Math.round(nh * scale)); + const scaleCanvas = (source: HTMLCanvasElement): HTMLCanvasElement => { + const dst = document.createElement("canvas"); + dst.width = sw; + dst.height = sh; + const dctx = dst.getContext("2d"); + if (!dctx) throw new Error("Could not scale the extended canvas"); + dctx.drawImage(source, 0, 0, sw, sh); + return dst; + }; + return { + image: scaleCanvas(ic).toDataURL("image/png"), + mask: scaleCanvas(mc).toDataURL("image/png"), + }; + } + + return { image: ic.toDataURL("image/png"), mask: mc.toDataURL("image/png") }; +} + // One labeled row in the recipe popover. function RecipeRow({ label, @@ -418,6 +882,55 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // Batch size = images per forward pass (VRAM-heavy); count = sequential loops. const [batchSize, setBatchSize] = useState(1); const [count, setCount] = useState(1); + // Active workflow tab. "create" = text-to-image; "transform" = img2img; "inpaint" = + // mask-guided redraw. More tabs (edit/extend/control/enhance) slot in here. + const [workflow, setWorkflow] = useState("create"); + // Transform (img2img) / Inpaint inputs: the uploaded source image as a data URL, and + // the denoise strength (how far to redraw it: low = keep source, high = reimagine). + const [initImage, setInitImage] = useState(null); + const [strength, setStrength] = useState(0.6); + // Inpaint mask (grayscale PNG data URL, white = repaint), the brush size as a percent + // of the image's shorter side, and a key bumped to clear the painted mask. + const [maskImage, setMaskImage] = useState(null); + const [brushPct, setBrushPct] = useState(8); + const [maskResetKey, setMaskResetKey] = useState(0); + // Extend (outpaint): how far to grow each dimension and which sides to grow. Reuses the + // inpaint backend by building a padded image + border mask at generate time. + const [extendPct, setExtendPct] = useState(25); + const [extendSides, setExtendSides] = useState({ + left: true, + right: true, + top: true, + bottom: true, + }); + // Upscale (hires fix): the enlargement factor and the (low) denoise strength used to + // re-detail the enlarged image. The backend caps the factor and rounds the target size. + const [upscaleFactor, setUpscaleFactor] = useState(2); + const [upscaleStrength, setUpscaleStrength] = useState(0.35); + // Reference (FLUX.2): up to 3 ADDITIONAL reference images beyond the primary one, combined + // by the model (subject + style, character + scene). + const [referenceImages, setReferenceImages] = useState([]); + // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by + // default; a single fixed toggle in the top bar opens/closes it (the icon never moves). + const [advancedOpen, setAdvancedOpen] = useState(false); + // Advanced (load-time) options. "auto"/"off"/"none" map to the backend defaults + // (sent through on load). They apply when a model loads; changing them while a model + // is loaded shows a "Reapply" button that reloads the same model with the new values. + const [speedMode, setSpeedMode] = useState<"auto" | "off" | "eager" | "default" | "max">("auto"); + const [transformerQuant, setTransformerQuant] = useState< + "none" | "auto" | "int8" | "fp8" | "nvfp4" | "mxfp8" + >("none"); + const [attentionBackend, setAttentionBackend] = useState<"auto" | "native" | "cudnn" | "flash3" | "sage">( + "auto", + ); + const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto"); + const [transformerCache, setTransformerCache] = useState<"off" | "fbcache">("off"); + const [cpuOffload, setCpuOffload] = useState(false); + // The last load descriptor, so "Reapply" can reload the same model with new advanced + // options without the user re-picking it from the dropdown. + const lastLoad = useRef<{ repoId: string; kind: "gguf" | "single_file" | "pipeline"; filename?: string } | null>( + null, + ); const [busy, setBusy] = useState(null); // {done, total} while a multi-run generation is in flight (for the button). @@ -732,7 +1245,13 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const handleLoad = useCallback( // Resolves true when the background load STARTED (callers may revert // optimistic picker state on false); poll outcomes are handled internally. - async (repoId: string, ggufFilename: string): Promise => { + async ( + repoId: string, + opts: { + kind: "gguf" | "single_file" | "pipeline"; + filename?: string; + }, + ): Promise => { // Cancel any prior poll loop so two can't run at once. if (pollTimer.current) clearTimeout(pollTimer.current); setBusy("loading"); @@ -740,14 +1259,26 @@ export function ImagesPage({ active = true }: { active?: boolean }) { dismissLoadToast(); lastLoadSig.current = null; loadToastId.current = toast(null, loadToastArgs(IDLE_PROGRESS)); + // Remember what was loaded so "Reapply" can reload it with new advanced options. + lastLoad.current = { repoId, kind: opts.kind, filename: opts.filename }; try { // Returns immediately — the load runs in the background; we poll for it. // The backend infers the family + base diffusers repo from the repo id. // Forward the saved HF token so gated bases (FLUX dev/klein) can download. + // A pipeline load carries no filename (the repo IS the pipeline); the + // single-file kinds send the GGUF / safetensors filename. Advanced options map + // sentinels ("auto"/"off"/"none") to omitted so the backend uses its defaults. await loadDiffusionModel({ model_path: repoId, - gguf_filename: ggufFilename, + model_kind: opts.kind, + gguf_filename: opts.filename, hf_token: hfApiToken(getHfToken()), + cpu_offload: cpuOffload, + speed_mode: speedMode === "auto" ? undefined : speedMode, + transformer_quant: transformerQuant === "none" ? undefined : transformerQuant, + attention_backend: attentionBackend === "auto" ? undefined : attentionBackend, + memory_mode: memoryMode === "auto" ? undefined : memoryMode, + transformer_cache: transformerCache === "off" ? undefined : transformerCache, }); } catch (err) { dismissLoadToast(); @@ -759,30 +1290,66 @@ export function ImagesPage({ active = true }: { active?: boolean }) { void pollLoadProgress(); return true; }, - [pollLoadProgress, refreshStatus, dismissLoadToast], + [ + pollLoadProgress, + refreshStatus, + dismissLoadToast, + cpuOffload, + speedMode, + transformerQuant, + attentionBackend, + memoryMode, + transformerCache, + ], ); - // The chat picker emits (modelId, picked quant + its exact filename); load it, - // and seed the inputs with that model's defaults. + // Set (or clear) the Transform/Inpaint source image; always drop any painted mask so it + // can't be applied to a different image (the mask is sized to the previous source). + const handleInitChange = useCallback((dataUrl: string | null) => { + setInitImage(dataUrl); + setMaskImage(null); + setMaskResetKey((k) => k + 1); + }, []); + + // Reload the current model with the current advanced options. + const handleReapply = useCallback(() => { + const l = lastLoad.current; + if (l) void handleLoad(l.repoId, { kind: l.kind, filename: l.filename }); + }, [handleLoad]); + + // The chat picker emits (modelId, picked quant + its exact filename) for a GGUF, + // or just (modelId) for a curated non-GGUF safetensors pick; load it, and seed the + // inputs with that model's defaults. const handleModelSelect = useCallback( (id: string, meta: ModelSelectorChangeMeta) => { // Ignore picks while a load/generation/unload is in flight: starting a // replacement load now would tear down the live poll/toast and reset // busy, while the backend rejects the second load with a 409. if (busy !== null) return; - if (meta.ggufVariant && meta.ggufFilename) { - // Optimistic for instant picker feedback, but revert if the load fails to - // START (400/409/network) or LATER during the poll (download/preflight - // error/eviction) -- in both cases the old pipeline stays loaded, so the - // selector must not advertise the failed quant. The poll owns the after-start - // revert via quantRevert; here we only handle the never-started case. - const prevQuant = quant; - quantRevert.current = { prev: prevQuant }; - setQuant(meta.ggufVariant); + // Curated non-GGUF model: load as a full pipeline or single-file safetensors. + const spec = SAFETENSORS_MODELS[id]; + if (spec) { + setQuant(null); const d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); - void handleLoad(id, meta.ggufFilename).then((started) => { + void handleLoad(id, { kind: spec.kind, filename: spec.filename }); + return; + } + // GGUF quant pick from the variant expander. Optimistic for instant picker + // feedback, but revert if the load fails to START (400/409/network) or LATER + // during the poll (download/preflight error/eviction) -- in both cases the old + // pipeline stays loaded, so the selector must not advertise the failed quant. + // The poll owns the after-start revert via quantRevert; here we only handle + // the never-started case. + if (meta.ggufVariant && meta.ggufFilename) { + const prevQuant = quant; + quantRevert.current = { prev: prevQuant }; + setQuant(meta.ggufVariant); + const dq = defaultsFor(id); + setSteps(dq.steps); + setGuidance(dq.guidance); + void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename }).then((started) => { if (!started) { setQuant(prevQuant); quantRevert.current = null; @@ -806,16 +1373,29 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const prevQuant = quant; quantRevert.current = { prev: prevQuant }; setQuant(filename); - const d = defaultsFor(id); - setSteps(d.steps); - setGuidance(d.guidance); - void handleLoad(dir, filename).then((started) => { + const dq2 = defaultsFor(id); + setSteps(dq2.steps); + setGuidance(dq2.guidance); + void handleLoad(dir, { kind: "gguf", filename }).then((started) => { if (!started) { setQuant(prevQuant); quantRevert.current = null; } }); + return; } + // Otherwise treat it as a full diffusers repo (safetensors / bnb-4bit). The backend + // infers the family + base repo from the id and gates loads to unsloth/* repos or + // on-device paths, so only attempt those; other Hub orgs can't be assembled here. + if (meta.source !== "local" && !id.toLowerCase().startsWith("unsloth/")) { + toast.error("Only unsloth or on-device image models can be loaded here"); + return; + } + setQuant(null); + const d = defaultsFor(id); + setSteps(d.steps); + setGuidance(d.guidance); + void handleLoad(id, { kind: "pipeline" }); }, [busy, handleLoad, quant], ); @@ -846,6 +1426,80 @@ export function ImagesPage({ active = true }: { active?: boolean }) { toast.error("Prompt is empty"); return; } + const isTransform = workflow === "transform"; + const isInpaint = workflow === "inpaint"; + const isExtend = workflow === "extend"; + const isUpscale = workflow === "upscale"; + const isReference = workflow === "reference"; + const isEdit = workflow === "edit"; + const usesInit = isTransform || isInpaint || isExtend || isUpscale || isReference || isEdit; + const tabLabel = isInpaint + ? "Inpaint" + : isExtend + ? "Extend" + : isUpscale + ? "Upscale" + : isReference + ? "Reference" + : isEdit + ? "Edit" + : "Transform"; + if (usesInit && !initImage) { + toast.error(`Upload a source image for ${tabLabel}`); + return; + } + if (isInpaint && !maskImage) { + toast.error("Paint a mask over the region to regenerate"); + return; + } + if (isExtend && !(extendSides.left || extendSides.right || extendSides.top || extendSides.bottom)) { + toast.error("Pick at least one side to extend"); + return; + } + + // Resolve the conditioning image/mask/strength for this workflow up front. Extend + // (outpaint) is built here from the source by padding + masking the new border, then + // sent through the same inpaint path. txt2img leaves all three undefined. + let condInit: string | undefined; + let condMask: string | undefined; + let condStrength: number | undefined; + let condUpscale: number | undefined; + let condRefImages: string[] | undefined; + try { + if (isTransform) { + condInit = initImage ?? undefined; + condStrength = strength; + } else if (isInpaint) { + condInit = initImage ?? undefined; + condMask = maskImage ?? undefined; + condStrength = strength; + } else if (isExtend) { + const built = await buildOutpaint(initImage!, extendSides, extendPct); + condInit = built.image; + condMask = built.mask; + condStrength = 1; // the new border is blank canvas: redraw it fully + } else if (isUpscale) { + // Hires fix: the backend enlarges the source by `upscale` and re-denoises it at + // this low strength so it gains detail without changing the content. + condInit = initImage ?? undefined; + condUpscale = upscaleFactor; + condStrength = upscaleStrength; + } else if (isReference) { + // FLUX.2 reference conditioning: send the primary reference + any extra references + // (combined by the model). The model generates a fresh image at the slider size + // guided by the references + prompt. No mask, no strength (not a denoise blend). + condInit = initImage ?? undefined; + const extras = referenceImages.filter(Boolean); + if (extras.length) condRefImages = extras; + } else if (isEdit) { + // Instruction editing: send the source image; the prompt IS the instruction. + // No mask, no strength (the edit pipeline fully regenerates from the instruction). + condInit = initImage ?? undefined; + } + } catch { + toast.error("Could not prepare the source image"); + return; + } // Resolve a base seed up front. With an explicit seed the run is fully // reproducible; with a random one we still pick a concrete base now so each // sequential image gets a distinct, reproducible seed (base + i). @@ -908,6 +1562,14 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // batch-mates. Unique per image on both engines, reproducible via recipes. seed: baseSeed + i * batchSize, batch_size: batchSize, + // Transform/Inpaint/Extend send the source image (+ mask for inpaint/extend) and + // a denoise strength, resolved above. The backend derives output size from the + // image, so width/height are advisory here. + init_image: condInit, + mask_image: condMask, + strength: condStrength, + upscale: condUpscale, + reference_images: condRefImages, }); if (!isMounted.current) break; // Prepend this run's records (newest first) and load their blobs. @@ -925,14 +1587,128 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setGenDone(null); setGenStep(null); } - }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, ensureSrc]); + }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, ensureSrc]); + + // Keep the active workflow valid for the loaded model: an edit-only model (Qwen-Image- + // Edit) has no Create/Transform tabs, a base model has no Edit tab. Snap to the first + // supported workflow whenever the loaded model's capabilities change. + useEffect(() => { + if (!status?.loaded) return; + const wf = status.workflows ?? []; + const ok = (id: WorkflowId) => { + const t = WORKFLOW_TABS.find((x) => x.id === id); + if (!t) return false; + return t.requires === null ? wf.includes("txt2img") : wf.includes(t.requires); + }; + if (!ok(workflow)) { + const first = WORKFLOW_TABS.find((t) => ok(t.id)); + if (first) setWorkflow(first.id); + } + }, [status?.loaded, status?.workflows, workflow]); + + // The Advanced (load-time) tuning controls, rendered in the right-docked panel below. + const advancedControls = ( + <> + setSpeedMode(v as typeof speedMode)} + options={[ + ["auto", "Auto"], + ["off", "Off (bit-exact)"], + ["eager", "Eager"], + ["default", "Default (compile)"], + ["max", "Max"], + ]} + /> + {/* The dense transformer_quant fast path only engages on the GGUF kind; on a loaded + safetensors pipeline / single-file model it is a silent no-op, so gate the control + to GGUF (or nothing loaded) and otherwise show why it is unavailable. */} + {!status?.loaded || status.model_kind === "gguf" ? ( + setTransformerQuant(v as typeof transformerQuant)} + options={[ + ["none", "Off (run the GGUF)"], + ["auto", "Auto (fastest for GPU)"], + ["fp8", "FP8"], + ["int8", "INT8"], + ["nvfp4", "NVFP4 (Blackwell)"], + ["mxfp8", "MXFP8 (Blackwell)"], + ]} + /> + ) : ( +
+ GGUF speed mode + GGUF models only +
+ )} + setAttentionBackend(v as typeof attentionBackend)} + options={[ + ["auto", "Auto"], + ["native", "Native SDPA"], + ["cudnn", "cuDNN"], + ["flash3", "FlashAttention 3"], + ["sage", "SageAttention (INT8)"], + ]} + /> + setMemoryMode(v as typeof memoryMode)} + options={[ + ["auto", "Auto"], + ["fast", "Fast (resident)"], + ["balanced", "Balanced"], + ["low_vram", "Low VRAM"], + ]} + /> + setTransformerCache(v as typeof transformerCache)} + options={[ + ["off", "Off"], + ["fbcache", "First-Block-Cache"], + ]} + /> +
+ + CPU offload + Offload to CPU to fit low-VRAM cards (slower). Overridden by Memory mode when that is not Auto. + + +
+ {status?.loaded && ( + + )} + + ); return (
{/* ── Top: the model selector, kept at the chat tab's exact position so the shared element matches. The load progress shows in a chat-style toast, not here. ── */} -
+
setSelectorOpen(active && o)} /> + {/* Single fixed toggle for the right-docked Advanced panel (mirrors Chat's settings + toggle, same icon in both states so it never moves). Highlighted when open. */} +
{/* ── Controls rail + preview canvas. Padding mirrors the other tabs @@ -953,8 +1746,284 @@ export function ImagesPage({ active = true }: { active?: boolean }) { {/* The controls rail. Plain card (the gray surface) with no header — the prompt + Generate button make the panel self-explanatory. */}
- -