diff --git a/studio/backend/core/inference/diffusion_inference_info.py b/studio/backend/core/inference/diffusion_inference_info.py new file mode 100644 index 0000000000..e7c0374c90 --- /dev/null +++ b/studio/backend/core/inference/diffusion_inference_info.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pure, torch-free per-family footprint summary for the inference info endpoint. + +The Advanced Dtype selector offers int8/fp8/nvfp4/mxfp8 dense transformer-quant +schemes. This module turns the auto-policy's per-family bf16 component table plus +its per-scheme steady factors into a static "what would this cost resident" summary +the frontend can show BEFORE anything is loaded, so a user can size the tradeoff. + +Estimates here are the STEADY resident footprint (transformer * factor + companions), +not the transient build peak -- the same quantity ``estimate_dense_quant`` reports as +``steady_total`` -- so the numbers match what stays on the card during generation. + +Pure by design: only imports the auto-policy tables (themselves torch-free), does no +GPU probing, so it unit-tests on CPU-only hosts. +""" + +from __future__ import annotations + +from typing import Any + +from .diffusion_auto_policy import _FAMILY_BF16_GB, _QUANT_STEADY_FACTOR + + +def _round1(value: float) -> float: + return round(value, 1) + + +def family_inference_infos() -> list[dict[str, Any]]: + """Per-family bf16 component sizes + estimated resident footprint per quant scheme. + + One dict per family in the auto-policy bf16 table (registry order), each carrying the + bf16-resident component sizes and the estimated resident GB under bf16 and each dense + quant scheme. bf16's estimate is the un-quantised transformer + companions; each + scheme scales the transformer by its steady factor and adds the same companions. + """ + infos: list[dict[str, Any]] = [] + for name, (transformer_gb, text_encoders_gb, vae_gb) in _FAMILY_BF16_GB.items(): + companions_gb = text_encoders_gb + vae_gb + estimated = {"bf16": _round1(transformer_gb + companions_gb)} + for scheme, factor in _QUANT_STEADY_FACTOR.items(): + estimated[scheme] = _round1(transformer_gb * factor + companions_gb) + infos.append( + { + "family": name, + "transformer_bf16_gb": _round1(transformer_gb), + "text_encoders_bf16_gb": _round1(text_encoders_gb), + "vae_bf16_gb": _round1(vae_gb), + "estimated_resident_gb": estimated, + } + ) + return infos diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 93e8dd96a5..698cb1a1f7 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2068,6 +2068,22 @@ class DiffusionLoadProgressResponse(BaseModel): error: Optional[str] = Field(None, description = "Failure message when phase is 'error'") +class DiffusionResolvedControl(BaseModel): + """One Advanced control's engaged value + provenance, for the "Auto: X" badges. + + ``value`` is what actually applied (a scheme string, a mode string, ``null`` when the + control is off, or ``true``/``false`` for cpu_offload), so it is typed ``Any``. + ``source`` is "auto" when this backend decided it or "explicit" when the caller did; + ``reason`` is the short human-readable why the frontend shows as a tooltip. + """ + + value: Any = Field( + None, description = "The engaged value: a string, a boolean (cpu_offload), or null." + ) + source: str = Field(..., description = '"auto" (backend decided) or "explicit" (caller set it)') + reason: str = Field("", description = "Short human-readable reason for the resolved value.") + + class DiffusionStatusResponse(BaseModel): """Current diffusion backend state.""" @@ -2131,14 +2147,44 @@ class DiffusionStatusResponse(BaseModel): "picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False " "for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.", ) - resolved: Optional[dict] = Field( + # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Present + # only on backends that record it; null when nothing is loaded or on older backends. The + # frontend renders an "Auto: X" badge next to each control whose source == "auto". Declared + # explicitly so pydantic's default extra='ignore' does not silently drop the resolved record. + resolved: Optional[Dict[str, DiffusionResolvedControl]] = Field( None, - description = "Per-control auto-policy provenance (value/source/reason for each resolved " - "setting), or null. Declared explicitly so the field is not dropped by the default " - "extra='ignore', which would silently discard the backend's resolved record.", + description = "Per-control resolved value + provenance (source auto|explicit + reason), " + "keyed by Advanced control name; null when unloaded or unavailable.", ) +class DiffusionInferenceInfo(BaseModel): + """One family's bf16 component sizes + estimated resident footprint per quant scheme. + + Mirrors the dicts ``family_inference_infos()`` returns: the bf16-resident transformer / + text-encoder / VAE sizes, and the estimated resident GB under bf16 and each dense + transformer-quant scheme (transformer * factor + companions), rounded to 1 decimal.""" + + family: str = Field(..., description = "Diffusion family name (auto-policy table key).") + transformer_bf16_gb: float = Field(..., description = "bf16-resident transformer size in GB.") + text_encoders_bf16_gb: float = Field( + ..., description = "bf16-resident text encoder(s) size in GB." + ) + vae_bf16_gb: float = Field(..., description = "bf16-resident VAE size in GB.") + estimated_resident_gb: Dict[str, float] = Field( + ..., + description = "Estimated resident GB keyed by scheme: bf16, int8, fp8, mxfp8, nvfp4.", + ) + + +class DiffusionInferenceInfoResponse(BaseModel): + """Static per-family footprint summary for the Advanced Dtype tradeoff (GET + /api/inference/images/info). Hardware-independent: no GPU probing, so it is served + from the pure auto-policy tables and is safe to fetch before anything is loaded.""" + + families: List[DiffusionInferenceInfo] = Field(default_factory = list) + + # ── OpenAI-compatible images API (POST /v1/images/generations) ── # # Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 32f7b48f58..91a15852a5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1080,6 +1080,7 @@ from models.inference import ( DiffusionGenerateResponse, DiffusionGenerateProgressResponse, DiffusionStatusResponse, + DiffusionInferenceInfoResponse, DiffusionLoadProgressResponse, GalleryImage, GalleryListResponse, @@ -11972,6 +11973,16 @@ async def diffusion_status(current_subject: str = Depends(get_current_subject)): return DiffusionStatusResponse(**active_status()) +@studio_router.get("/images/info", response_model = DiffusionInferenceInfoResponse) +async def diffusion_inference_info(current_subject: str = Depends(get_current_subject)): + """Static per-family footprint summary for the Advanced Dtype tradeoff. + + Hardware-independent (served from the pure auto-policy tables, no GPU probing), so it + is cheap and safe to fetch before anything is loaded.""" + from core.inference.diffusion_inference_info import family_inference_infos + return DiffusionInferenceInfoResponse(families = family_inference_infos()) + + @studio_router.get("/images/load-progress", response_model = DiffusionLoadProgressResponse) async def diffusion_load_progress(current_subject: str = Depends(get_current_subject)): from core.inference.diffusion_engine_router import get_active_diffusion_engine diff --git a/studio/backend/tests/test_diffusion_inference_info.py b/studio/backend/tests/test_diffusion_inference_info.py new file mode 100644 index 0000000000..7f2a0ae3f6 --- /dev/null +++ b/studio/backend/tests/test_diffusion_inference_info.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""CPU-only unit tests for the pure per-family inference-info helper. + +Covers ``family_inference_infos()``: every auto-policy family appears, the component sizes +round-trip, and the quant estimates order correctly (quantised < bf16, nvfp4 < int8). No +torch / diffusers / GPU.""" + +from __future__ import annotations + +from core.inference.diffusion_auto_policy import _FAMILY_BF16_GB, _QUANT_STEADY_FACTOR +from core.inference.diffusion_inference_info import family_inference_infos + + +def test_covers_every_auto_policy_family(): + infos = family_inference_infos() + names = {info["family"] for info in infos} + assert names == set(_FAMILY_BF16_GB), "info must list exactly the auto-policy families" + # One entry per family, in registry order. + assert [info["family"] for info in infos] == list(_FAMILY_BF16_GB) + + +def test_each_family_reports_all_schemes(): + for info in family_inference_infos(): + estimated = info["estimated_resident_gb"] + assert set(estimated) == {"bf16", *_QUANT_STEADY_FACTOR} + # Every reported value is a float rounded to one decimal. + for value in estimated.values(): + assert isinstance(value, float) + assert round(value, 1) == value + + +def test_component_sizes_match_the_table(): + infos = {info["family"]: info for info in family_inference_infos()} + for name, (transformer, text_encoders, vae) in _FAMILY_BF16_GB.items(): + info = infos[name] + assert info["transformer_bf16_gb"] == round(transformer, 1) + assert info["text_encoders_bf16_gb"] == round(text_encoders, 1) + assert info["vae_bf16_gb"] == round(vae, 1) + + +def test_quantised_estimate_is_below_bf16(): + # A quantised transformer is smaller than bf16, so its resident estimate must be too + # (the companions are shared, and every steady factor is < 1). + for info in family_inference_infos(): + estimated = info["estimated_resident_gb"] + for scheme in _QUANT_STEADY_FACTOR: + assert estimated[scheme] < estimated["bf16"], f"{info['family']} {scheme}" + + +def test_nvfp4_is_below_int8(): + # nvfp4 packs two params per byte (~0.33x) vs int8's one byte per param (~0.55x), so + # nvfp4's estimate is the smaller of the two on every family. + for info in family_inference_infos(): + estimated = info["estimated_resident_gb"] + assert estimated["nvfp4"] < estimated["int8"], info["family"] diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 567cc39c26..6a4b53961b 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -722,3 +722,46 @@ def test_gpu_native_load_takes_arbiter(client, monkeypatch): ) assert resp.status_code == 200 assert acquired == [gpu_arbiter.DIFFUSION] + + +def test_images_info_lists_every_family(client): + # The pure info endpoint is hardware-independent (no load required): it returns one + # entry per auto-policy family, each with the quant estimates the UI shows. + from core.inference.diffusion_auto_policy import _FAMILY_BF16_GB + + resp = client.get("/api/inference/images/info") + assert resp.status_code == 200 + families = resp.json()["families"] + assert {f["family"] for f in families} == set(_FAMILY_BF16_GB) + sample = families[0] + est = sample["estimated_resident_gb"] + # Quantised estimates undercut bf16, and nvfp4 undercuts int8 (matches the pure helper). + assert est["int8"] < est["bf16"] + assert est["nvfp4"] < est["int8"] + + +def test_status_passes_through_resolved(client, monkeypatch): + # The additive `resolved` provenance record round-trips through the status route so the + # frontend can render the "Auto: X" badges. + backend = diffusion_module.get_diffusion_backend() + resolved = { + "speed_mode": {"value": "eager", "source": "auto", "reason": "per-kind default"}, + "transformer_quant": {"value": "int8", "source": "explicit", "reason": "requested"}, + "cpu_offload": {"value": False, "source": "auto", "reason": "from the memory plan"}, + "transformer_cache": {"value": None, "source": "auto", "reason": "few-step model"}, + } + monkeypatch.setattr( + backend, "status", lambda: {**_unloaded_status(), "loaded": True, "resolved": resolved} + ) + body = client.get("/api/inference/images/status").json() + assert body["resolved"] == resolved + assert body["resolved"]["speed_mode"]["source"] == "auto" + # The cpu_offload value stays a real boolean (not coerced to a string). + assert body["resolved"]["cpu_offload"]["value"] is False + + +def test_status_resolved_defaults_to_null(client): + # A backend status without a `resolved` key leaves the additive field null (older + # backends and the unloaded state). + body = client.get("/api/inference/images/status").json() + assert body["resolved"] is None diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index cde27e83a6..d9cef386ea 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -4,6 +4,16 @@ import { authFetch } from "@/features/auth"; import { readFastApiError } from "@/lib/format-fastapi-error"; +// One Advanced control's resolved value + provenance, for the "Auto: X" badges. `value` +// is the engaged value (a scheme/mode string, null when off, or a boolean for cpu_offload); +// `source` is "auto" (this backend decided) or "explicit" (the caller set it); `reason` is +// the short why shown as a tooltip. +export interface DiffusionResolvedControl { + value: string | boolean | null; + source: "auto" | "explicit"; + reason: string; +} + export interface DiffusionStatus { loaded: boolean; repo_id: string | null; @@ -24,6 +34,11 @@ export interface DiffusionStatus { // Whether the loaded model can apply a ControlNet (drives the ControlNet picker's enabled // state). Diffusers only, for families with a ControlNet pipeline; false otherwise. supports_controlnet?: boolean; + // Per-Advanced-control provenance, keyed by control name (speed_mode, transformer_quant, + // attention_backend, memory_mode, transformer_cache, cpu_offload). Present only when a + // model is loaded on a backend that records it; the "Auto: X" badges read it. Absent on + // older backends. + resolved?: Record | null; } export interface DiffusionGenerateProgress { @@ -174,6 +189,28 @@ export async function getDiffusionStatus(): Promise { return parseJson(await authFetch("/api/inference/images/status")); } +// One family's bf16 component sizes + estimated resident footprint per quant scheme +// (from GET /api/inference/images/info). Hardware-independent, so it can be fetched before +// anything is loaded to size the Advanced Dtype tradeoff. +export interface DiffusionInferenceInfo { + family: string; + transformer_bf16_gb: number; + text_encoders_bf16_gb: number; + vae_bf16_gb: number; + // Estimated resident GB keyed by scheme: bf16, int8, fp8, mxfp8, nvfp4. + estimated_resident_gb: Record; +} + +export interface DiffusionInferenceInfoResponse { + families: DiffusionInferenceInfo[]; +} + +/** Static per-family footprint summary for the Advanced Dtype tradeoff. Hardware-independent + * (served from the pure auto-policy tables), so it is safe to fetch before a load. */ +export async function getDiffusionInferenceInfo(): Promise { + return parseJson(await authFetch("/api/inference/images/info")); +} + export async function getDiffusionLoadProgress(): Promise { return parseJson(await authFetch("/api/inference/images/load-progress")); } diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 87e24d42f4..c7b5abd7f9 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -470,10 +470,44 @@ function Field({ ); } +// The engaged value of a resolved Advanced control, formatted for its "Auto: X" badge. +// Short scheme/mode tokens go uppercase (INT8, FP8, FBCACHE); the attention backend the +// backend reports as `_native_cudnn` shows as cuDNN; cpu_offload's boolean shows On/Off. +function formatResolvedValue(key: string, value: string | boolean | null): string { + if (key === "cpu_offload") return value ? "On" : "Off"; + if (value === null || value === "") return "Off"; + if (typeof value === "boolean") return value ? "On" : "Off"; + if (value === "_native_cudnn" || value.toLowerCase() === "cudnn") return "cuDNN"; + return value.toUpperCase(); +} + +// The "Auto: X" badge for one Advanced control: rendered only when the backend resolved +// that control itself (source === "auto"); an explicit user choice renders nothing. The +// reason is surfaced as a hover tooltip. Muted pill matching the panel's other chips. +function ResolvedBadge({ + status, + controlKey, +}: { + status: DiffusionStatus | null; + controlKey: string; +}) { + const resolved = status?.resolved?.[controlKey]; + if (!resolved || resolved.source !== "auto") return null; + return ( + + Auto: {formatResolvedValue(controlKey, resolved.value)} + + ); +} + // A compact labeled Select row for the Advanced Options panel. function AdvancedSelect({ label, hint, + badge, desc, value, onValueChange, @@ -481,6 +515,8 @@ function AdvancedSelect({ }: { label: string; hint?: ReactNode; + // An optional inline badge next to the label (e.g. the "Auto: X" resolved-value pill). + badge?: ReactNode; // A short always-visible description under the row (the hint tooltip carries the full // detail). Used for controls whose label alone does not convey what they do. desc?: string; @@ -494,6 +530,7 @@ function AdvancedSelect({ {label} {hint && {hint}} + {badge}