From a1423e05d23962a084a7ba0dd476b609bda77742 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Tue, 23 Jun 2026 00:39:36 -0300 Subject: [PATCH 01/35] Studio: local diffusion image generation page --- studio/backend/core/inference/diffusion.py | 403 +++++++++++++++ .../core/inference/diffusion_families.py | 120 +++++ studio/backend/core/inference/gpu_arbiter.py | 78 +++ studio/backend/models/inference.py | 76 +++ studio/backend/requirements/studio.txt | 4 + studio/backend/routes/inference.py | 105 ++++ .../backend/tests/test_diffusion_backend.py | 270 ++++++++++ studio/backend/tests/test_diffusion_routes.py | 156 ++++++ studio/backend/tests/test_gpu_arbiter.py | 70 +++ studio/frontend/src/app/router.tsx | 2 + studio/frontend/src/app/routes/images.tsx | 21 + .../frontend/src/components/app-sidebar.tsx | 13 + studio/frontend/src/features/images/api.ts | 90 ++++ .../src/features/images/images-page.tsx | 469 ++++++++++++++++++ studio/frontend/src/features/images/index.ts | 4 + studio/frontend/src/i18n/locales/en.ts | 1 + studio/frontend/src/i18n/locales/zh-CN.ts | 1 + 17 files changed, 1883 insertions(+) create mode 100644 studio/backend/core/inference/diffusion.py create mode 100644 studio/backend/core/inference/diffusion_families.py create mode 100644 studio/backend/core/inference/gpu_arbiter.py create mode 100644 studio/backend/tests/test_diffusion_backend.py create mode 100644 studio/backend/tests/test_diffusion_routes.py create mode 100644 studio/backend/tests/test_gpu_arbiter.py create mode 100644 studio/frontend/src/app/routes/images.tsx create mode 100644 studio/frontend/src/features/images/api.ts create mode 100644 studio/frontend/src/features/images/images-page.tsx create mode 100644 studio/frontend/src/features/images/index.ts diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py new file mode 100644 index 0000000000..b6aeaf4c88 --- /dev/null +++ b/studio/backend/core/inference/diffusion.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local diffusion (text-to-image) backend. + +A small torch-only singleton that loads a diffusers pipeline and generates +images. Single-file GGUF checkpoints are dequantised on-device via +``diffusers.GGUFQuantizationConfig``; the rest of the pipeline (VAE, text +encoders, scheduler) comes from the matching base repo (see +``diffusion_families``). torch/diffusers are imported lazily inside the methods +so this module imports cleanly in a no-torch runtime. + +Loading runs on a background thread (``begin_load``) so the route returns +immediately and the frontend can poll ``load_progress`` for a download bar. + +This backend owns no GPU-handoff policy: it assumes it is the heavy GPU user +while loaded. Coordinating with the chat backend lives in the GPU arbiter the +routes call, not here. +""" + +from __future__ import annotations + +import base64 +import io +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +from loggers import get_logger +from utils.hardware import clear_gpu_cache + +from .diffusion_families import ( + detect_family, + resolve_base_repo, + resolve_local_gguf_child, +) + +logger = get_logger(__name__) + + +def encode_png_base64(image: Any) -> str: + """PIL image -> base64-encoded PNG string.""" + buf = io.BytesIO() + image.save(buf, format = "PNG") + return base64.b64encode(buf.getvalue()).decode("ascii") + + +@dataclass(frozen = True) +class _LoadState: + """Everything about the currently-loaded pipeline, swapped as one unit.""" + + pipe: Any + family: Any + repo_id: str + base_repo: str + device: str + dtype: str + cpu_offload: bool + + +@dataclass +class _LoadingState: + """An in-flight background load, polled for download progress.""" + + repo_id: str + base_repo: str + expected_bytes: int = 0 + error: Optional[str] = None + + +class DiffusionBackend: + """Holds at most one loaded diffusers pipeline. All mutations are serialised.""" + + def __init__(self) -> None: + # One lock serialises load / generate / unload — a generate must not run + # while the pipeline is being swapped out from under it. status() and + # load_progress() read the single state references without the lock, so + # polling never blocks a slow load. + self._lock = threading.Lock() + self._state: Optional[_LoadState] = None + self._loading: Optional[_LoadingState] = None + + @property + def is_loaded(self) -> bool: + return self._state is not None + + def _pick_device_and_dtype(self) -> tuple[str, Any]: + import torch + + if torch.cuda.is_available(): + return "cuda", torch.bfloat16 + mps = getattr(torch.backends, "mps", None) + if mps is not None and mps.is_available(): + return "mps", torch.float16 + return "cpu", torch.float32 + + def _resolve_gguf_path(self, repo_id: str, gguf_filename: str, hf_token: Optional[str]) -> str: + local_root = Path(repo_id).expanduser() + if local_root.exists(): + return str(resolve_local_gguf_child(local_root, gguf_filename)) + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id, gguf_filename, token = hf_token) + + # ── Background load + progress ───────────────────────────────────────── + + def begin_load( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + hf_token: Optional[str] = None, + cpu_offload: bool = False, + ) -> dict[str, Any]: + """Validate, then run the (slow) load on a daemon thread. Returns at once.""" + fam = detect_family(repo_id, family_override) + if fam is None: + raise ValueError( + f"Could not infer a diffusion family for '{repo_id}'. " + "Pass family_override (one of: flux.2-klein, flux.2, flux.1, qwen-image)." + ) + base = resolve_base_repo(fam, base_repo) + + with self._lock: + # Allow starting over a previously-failed load, but not over a live one. + if self._loading is not None and self._loading.error is None: + raise RuntimeError("A diffusion load is already in progress.") + self._loading = _LoadingState(repo_id = repo_id, base_repo = base) + + threading.Thread( + target = self._run_load, + kwargs = dict( + repo_id = repo_id, + gguf_filename = gguf_filename, + base_repo = base, + family_override = family_override, + hf_token = hf_token, + cpu_offload = cpu_offload, + ), + daemon = True, + ).start() + return self.status() + + def _run_load(self, **kwargs: Any) -> None: + try: + # Estimate sizes on this thread (a network call) so begin_load returns + # instantly; the bar shows raw bytes until the total lands. This is the + # only writer of _loading's fields, so no lock is needed here. + loading = self._loading + if loading is not None: + loading.expected_bytes = self._estimate_download_bytes( + kwargs["repo_id"], + kwargs.get("gguf_filename"), + kwargs["base_repo"], + kwargs.get("hf_token"), + ) + self.load_pipeline(**kwargs) + with self._lock: + self._loading = None + except Exception as exc: # noqa: BLE001 — surfaced to the client via load_progress + logger.error("diffusion.load_failed: %s", exc) + with self._lock: + if self._loading is not None: + self._loading.error = str(exc) + + def load_progress(self) -> dict[str, Any]: + """Phase + downloaded/total bytes for the in-flight load (cache-scan based).""" + loading = self._loading + if loading is not None and loading.error: + return _progress("error", error = loading.error) + 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) + expected = loading.expected_bytes + # Downloads done but pipeline still dequantising / moving to GPU. + if expected > 0 and downloaded >= expected * 0.999: + return _progress("finalizing", downloaded, expected, 1.0) + fraction = downloaded / expected if expected > 0 else 0.0 + return _progress("downloading", downloaded, expected, fraction) + + @staticmethod + def _estimate_download_bytes( + repo_id: str, gguf_filename: Optional[str], base_repo: str, hf_token: Optional[str] + ) -> int: + from huggingface_hub import HfApi + + api = HfApi() + total = 0 + try: + if gguf_filename: + info = api.model_info(repo_id, files_metadata = True, token = hf_token) + total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) + base_info = api.model_info(base_repo, files_metadata = True, token = hf_token) + # The transformer subfolder is supplied by the GGUF, so from_pretrained + # never downloads it — exclude it from the expected total. + total += sum( + s.size or 0 + for s in base_info.siblings + if not s.rfilename.startswith("transformer/") + ) + except Exception as exc: # noqa: BLE001 — estimate is best-effort + logger.warning("diffusion.size_estimate_failed: %s", exc) + return total + + @staticmethod + def _cache_bytes(repo_id: str) -> int: + from huggingface_hub import constants + + blobs = Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" / "blobs" + total = 0 + try: + for entry in blobs.iterdir(): + try: + total += entry.stat().st_size + except OSError: + continue # broken symlink / unreadable + except OSError: + return 0 # repo not in cache yet + return total + + # ── Synchronous load / generate / unload ─────────────────────────────── + + def load_pipeline( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + hf_token: Optional[str] = None, + cpu_offload: bool = False, + ) -> dict[str, Any]: + import diffusers + + fam = detect_family(repo_id, family_override) + if fam is None: + raise ValueError( + f"Could not infer a diffusion family for '{repo_id}'. " + "Pass family_override (one of: flux.2-klein, flux.2, flux.1, qwen-image)." + ) + base = resolve_base_repo(fam, base_repo) + device, dtype = self._pick_device_and_dtype() + + with self._lock: + # Free the old pipeline before allocating the new one so two FLUX + # checkpoints never sit in VRAM at once. + self._unload_locked() + + transformer = None + if gguf_filename: + gguf_path = self._resolve_gguf_path(repo_id, gguf_filename, hf_token) + transformer_cls = getattr(diffusers, fam.transformer_class) + transformer = transformer_cls.from_single_file( + gguf_path, + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), + torch_dtype = dtype, + config = base, + subfolder = "transformer", + ) + + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} + if transformer is not None: + pipe_kwargs["transformer"] = transformer + if hf_token: + pipe_kwargs["token"] = hf_token + pipeline_cls = getattr(diffusers, fam.pipeline_class) + pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + + use_offload = bool(cpu_offload and device == "cuda") + if use_offload: + pipe.enable_model_cpu_offload() + else: + pipe.to(device) + + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = use_offload, + ) + + logger.info("diffusion.loaded: repo=%s base=%s device=%s", repo_id, base, device) + return self.status() + + def generate( + self, + *, + prompt: str, + negative_prompt: Optional[str] = None, + width: int = 1024, + height: int = 1024, + steps: int = 24, + guidance: float = 3.5, + seed: Optional[int] = None, + ) -> dict[str, Any]: + import torch + + with self._lock: + state = self._state + if state is None: + raise RuntimeError("No diffusion model is loaded.") + + generator = torch.Generator(device = state.device) + if seed is None: + # Generator.seed() seeds the generator with a fresh random value + # AND returns it, so the reported seed reproduces the image. + seed = generator.seed() + else: + seed = int(seed) + generator.manual_seed(seed) + + kwargs: dict[str, Any] = { + "prompt": prompt, + "width": width, + "height": height, + "num_inference_steps": steps, + "guidance_scale": guidance, + "generator": generator, + } + if negative_prompt: + kwargs["negative_prompt"] = negative_prompt + + image = state.pipe(**kwargs).images[0] + return { + "image_b64": encode_png_base64(image), + "mime": "image/png", + "seed": int(seed), + } + + def unload(self) -> dict[str, Any]: + with self._lock: + self._unload_locked() + # Drop a finished/failed load marker so the next load starts clean. + if self._loading is not None and self._loading.error is not None: + self._loading = None + return self.status() + + def _unload_locked(self) -> None: + state = self._state + if state is None: + return + self._state = None + del state + clear_gpu_cache() + + def status(self) -> dict[str, Any]: + state = self._state + loading = self._loading is not None and (self._loading.error is None) + if state is None: + return { + "loaded": False, + "loading": loading, + "repo_id": None, + "family": None, + "base_repo": None, + "device": None, + "dtype": None, + "cpu_offload": False, + } + return { + "loaded": True, + "loading": loading, + "repo_id": state.repo_id, + "family": state.family.name, + "base_repo": state.base_repo, + "device": state.device, + "dtype": state.dtype, + "cpu_offload": state.cpu_offload, + } + + +def _progress( + phase: Optional[str], + bytes_downloaded: int = 0, + bytes_total: int = 0, + fraction: float = 0.0, + *, + error: Optional[str] = None, +) -> dict[str, Any]: + return { + "phase": phase, + "bytes_downloaded": bytes_downloaded, + "bytes_total": bytes_total, + "fraction": fraction, + "error": error, + } + + +_diffusion_backend: Optional[DiffusionBackend] = None + + +def get_diffusion_backend() -> DiffusionBackend: + global _diffusion_backend + if _diffusion_backend is None: + _diffusion_backend = DiffusionBackend() + return _diffusion_backend diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py new file mode 100644 index 0000000000..523c8b29d9 --- /dev/null +++ b/studio/backend/core/inference/diffusion_families.py @@ -0,0 +1,120 @@ +# 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 helpers for diffusion model identification. + +No torch/diffusers imports here: everything in this module is a pure function of +its string/path arguments so it can be unit-tested without the heavy runtime. + +A diffusion checkpoint published as a single-file GGUF only carries the +transformer weights; the matching VAE / text encoders / scheduler come from a +companion ``diffusers`` base repo. ``DiffusionFamily`` maps a checkpoint to the +diffusers classes and base repo needed to assemble the full pipeline. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path, PurePosixPath +from typing import Optional + + +@dataclass(frozen = True) +class DiffusionFamily: + name: str + pipeline_class: str + transformer_class: str + base_repo: str + # Extra lowercased substrings (besides ``name``) that map a repo id here. + aliases: tuple[str, ...] = field(default_factory = tuple) + + +# Ordered most-specific first so "flux.2-klein" matches klein, not flux.2. +_FAMILIES: tuple[DiffusionFamily, ...] = ( + DiffusionFamily( + name = "flux.2-klein", + pipeline_class = "Flux2KleinPipeline", + transformer_class = "Flux2Transformer2DModel", + base_repo = "black-forest-labs/FLUX.2-klein-base-4B", + aliases = ("flux2-klein", "flux-2-klein"), + ), + DiffusionFamily( + name = "flux.2", + pipeline_class = "Flux2Pipeline", + transformer_class = "Flux2Transformer2DModel", + base_repo = "black-forest-labs/FLUX.2-dev", + aliases = ("flux2-dev", "flux-2-dev"), + ), + DiffusionFamily( + name = "flux.1", + pipeline_class = "FluxPipeline", + transformer_class = "FluxTransformer2DModel", + base_repo = "black-forest-labs/FLUX.1-dev", + aliases = ("flux1-dev", "flux-1-dev", "flux-dev"), + ), + DiffusionFamily( + name = "qwen-image", + pipeline_class = "QwenImagePipeline", + transformer_class = "QwenImageTransformer2DModel", + base_repo = "Qwen/Qwen-Image", + aliases = ("qwenimage", "qwen_image"), + ), +) + + +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. Returns + ``None`` when nothing matches so the caller can raise a clean error. + """ + if override: + key = override.strip().lower() + for fam in _FAMILIES: + if key == fam.name or key in fam.aliases: + return fam + return None + needle = repo_id.lower() + for fam in _FAMILIES: + if fam.name in needle or any(alias in needle for alias in fam.aliases): + return fam + return None + + +def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: + """The companion diffusers repo: caller-supplied if given, else the family default.""" + base = (base_repo or "").strip() + return base or fam.base_repo + + +def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: + """Resolve ``gguf_filename`` to a file under ``repo_root``, rejecting escapes. + + ``gguf_filename`` is user-supplied, so an absolute path or a ``..`` segment + could otherwise read outside the local repo directory. + """ + if ( + Path(gguf_filename).is_absolute() + or PurePosixPath(gguf_filename).is_absolute() + or gguf_filename.startswith(("/", "\\")) + or "\\" in gguf_filename + ): + raise ValueError("gguf_filename must be a relative path inside the repo.") + rel = PurePosixPath(gguf_filename) + if any(part in ("", ".", "..") for part in rel.parts): + raise ValueError("gguf_filename must not contain '', '.', or '..' segments.") + # Resolve symlinks before the containment check: the guards above stop + # lexical escapes, but a symlink inside the repo could still point outside it. + repo_real = repo_root.resolve() + child = repo_root.joinpath(*rel.parts).resolve() + if child != repo_real and repo_real not in child.parents: + raise ValueError("gguf_filename must resolve to a file inside the repo.") + if not child.exists(): + raise FileNotFoundError(f"'{gguf_filename}' not found under {repo_root}.") + return child + + +def supported_families() -> list[dict]: + """Name + base repo for each known family (for status / UI listing).""" + return [{"name": fam.name, "base_repo": fam.base_repo} for fam in _FAMILIES] diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py new file mode 100644 index 0000000000..70335e7ea6 --- /dev/null +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Single-GPU arbiter for Studio's two heavy GPU consumers. + +The chat backends (llama-server + the Unsloth subprocess) and the diffusion +backend share one GPU. Before either takes the GPU it calls ``acquire_for(owner)``, +which evicts the current *other* owner so two large models never sit in VRAM at +once. The arbiter only sequences ownership; the actual freeing is delegated to +each backend's existing teardown. + +Eviction runs under the arbiter lock, so an ownership transfer is atomic with +respect to other acquires. +""" + +from __future__ import annotations + +import threading +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +CHAT = "chat" +DIFFUSION = "diffusion" + +_lock = threading.Lock() +_owner: Optional[str] = None + + +def _evict_chat() -> None: + from core.inference import get_inference_backend + from routes.inference import get_llama_cpp_backend + + llama = get_llama_cpp_backend() + if llama.is_loaded: + llama.unload_model() + orchestrator = get_inference_backend() + if orchestrator.active_model_name: + orchestrator.unload_model(orchestrator.active_model_name) + # Kill the subprocess too, not just the model: its base CUDA context holds + # VRAM the diffusion pipeline needs. + orchestrator._shutdown_subprocess(timeout = 5.0) + + +def _evict_diffusion() -> None: + from core.inference.diffusion import get_diffusion_backend + + get_diffusion_backend().unload() + + +# Patchable in tests via monkeypatch.setitem. +_EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion} + + +def acquire_for(owner: str) -> None: + """Make ``owner`` the sole GPU owner, evicting the other if it holds it.""" + global _owner + if owner not in _EVICTORS: + raise ValueError(f"unknown GPU owner: {owner!r}") + with _lock: + if _owner is not None and _owner != owner: + logger.info("gpu_arbiter: evicting %s for %s", _owner, owner) + _EVICTORS[_owner]() + _owner = owner + + +def release(owner: str) -> None: + """Drop ``owner``'s claim (no-op if it isn't the current owner).""" + global _owner + with _lock: + if _owner == owner: + _owner = None + + +def current_owner() -> Optional[str]: + return _owner diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b8432f588c..5c2c92109a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1677,3 +1677,79 @@ class AnthropicMessagesResponse(BaseModel): stop_reason: Optional[str] = None stop_sequence: Optional[str] = None usage: AnthropicUsage = Field(default_factory = AnthropicUsage) + + +# ── Diffusion (local text-to-image) ── + + +class DiffusionLoadRequest(BaseModel): + """Request to load a local diffusion (text-to-image) checkpoint.""" + + model_path: str = Field(..., description = "Diffusion repo id or local path") + gguf_filename: Optional[str] = Field( + None, description = "Single-file GGUF inside model_path (omit for a full diffusers repo)" + ) + base_repo: Optional[str] = Field( + None, description = "Companion diffusers repo for VAE/text-encoders (default: family base)" + ) + family_override: Optional[str] = Field( + None, description = "Force a family when it can't be inferred from the repo id" + ) + hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated repos") + cpu_offload: bool = Field(False, description = "Enable model CPU offload to fit low-VRAM cards") + + +class DiffusionGenerateRequest(BaseModel): + """Request to generate one image from the loaded diffusion model.""" + + prompt: str = Field(..., min_length = 1, description = "Text prompt") + negative_prompt: Optional[str] = Field(None, description = "What to avoid (if the model supports it)") + width: int = Field(1024, ge = 256, le = 2048, description = "Image width in pixels (multiple of 8)") + height: int = Field(1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 8)") + steps: int = Field(24, ge = 1, le = 100, description = "Number of denoising steps") + guidance: float = Field(3.5, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") + seed: Optional[int] = Field( + None, ge = 0, le = 2**64 - 1, description = "Seed for reproducibility (random if omitted)" + ) + + @field_validator("width", "height") + @classmethod + def _multiple_of_8(cls, value: int) -> int: + # VAEs downsample by 8; non-multiples crash deep in the pipeline, so + # reject them here for a clean 422 instead of a cryptic 500. + if value % 8 != 0: + raise ValueError("must be a multiple of 8") + return value + + +class DiffusionGenerateResponse(BaseModel): + """A generated image plus the seed actually used.""" + + image_b64: str = Field(..., description = "Base64-encoded PNG") + mime: str = Field("image/png", description = "MIME type of image_b64") + seed: int = Field(..., description = "Seed used (echoes the request, or the random one chosen)") + + +class DiffusionLoadProgressResponse(BaseModel): + """Download/finalize progress for an in-flight diffusion load.""" + + phase: Optional[Literal["downloading", "finalizing", "ready", "error"]] = Field( + None, description = "Load phase; null when idle" + ) + bytes_downloaded: int = Field(0, description = "Bytes present in the HF cache so far") + bytes_total: int = Field(0, description = "Estimated total bytes to download (0 = unknown)") + fraction: float = Field(0.0, description = "bytes_downloaded / bytes_total, clamped to [0,1]") + error: Optional[str] = Field(None, description = "Failure message when phase is 'error'") + + +class DiffusionStatusResponse(BaseModel): + """Current diffusion backend state.""" + + loaded: bool = Field(False, description = "Whether a diffusion model is loaded") + loading: bool = Field(False, description = "Whether a load is currently in progress") + repo_id: Optional[str] = Field(None, description = "Loaded repo id or local path") + family: Optional[str] = Field(None, description = "Detected diffusion family") + 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") + cpu_offload: bool = Field(False, description = "Whether CPU offload is engaged") diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 96fef60471..0173082560 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -20,6 +20,10 @@ cryptography>=42.0.0 boto3>=1.34.0 # optional: S3 dataset loading httpx>=0.27.0 fastmcp>=3.0.2 +# Local image generation (Images page): diffusers' GGUFQuantizationConfig needs +# the gguf reader. diffusers itself is torch-bound and comes via the ML stack; +# gguf is torch-free so it lives here, where it's installed in every mode. +gguf # RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in # extras-no-deps.txt; these add the lexical+dense store and document parsing. sqlite-vec==0.1.9 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 53d81961c3..677cc15e0b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -905,6 +905,11 @@ from models.inference import ( LoadRequest, UnloadRequest, GenerateRequest, + DiffusionLoadRequest, + DiffusionGenerateRequest, + DiffusionGenerateResponse, + DiffusionStatusResponse, + DiffusionLoadProgressResponse, LoadResponse, LoadProgressResponse, UnloadResponse, @@ -2346,6 +2351,13 @@ async def load_model( user_override = request.chat_template_override, ) + # Reclaim the GPU from the diffusion (Images) backend before any chat + # load path — including the already-loaded fast path below, so chat + # ownership is asserted without depending on chat/diffusion exclusivity + # holding. No-op when diffusion isn't loaded. + from core.inference.gpu_arbiter import acquire_for, CHAT + await asyncio.to_thread(acquire_for, CHAT) + # ── Already-loaded check: skip reload if the exact model is active ── backend = get_inference_backend() llama_backend = get_llama_cpp_backend() @@ -9424,3 +9436,96 @@ async def _openai_passthrough_non_streaming( # redundant parse + re-serialize round-trip. return Response(content = resp.content, media_type = "application/json") return JSONResponse(content = data) + + +# ────────────────────────────────────────────────────────────────────────── +# Diffusion (local text-to-image) +# +# Studio-only routes (studio_router is not mounted under /v1). The diffusion +# backend runs in-process and is synchronous, so the blocking load/generate/ +# unload calls are offloaded with asyncio.to_thread to keep the event loop free. +# This is the single error boundary: backend methods raise, we map to HTTP here. +# ────────────────────────────────────────────────────────────────────────── + + +@studio_router.post("/images/load", response_model = DiffusionStatusResponse) +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.gpu_arbiter import acquire_for, DIFFUSION + from utils.native_path_leases import redact_native_paths + + backend = get_diffusion_backend() + try: + # Take the GPU from the chat backend, then kick the (slow) load onto a + # background thread and return at once — the client polls images/load-progress. + await asyncio.to_thread(acquire_for, DIFFUSION) + status_dict = await asyncio.to_thread( + backend.begin_load, + request.model_path, + gguf_filename = request.gguf_filename, + base_repo = request.base_repo, + family_override = request.family_override, + hf_token = request.hf_token, + cpu_offload = request.cpu_offload, + ) + return DiffusionStatusResponse(**status_dict) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) + except RuntimeError as exc: + # A load is already in progress. + raise HTTPException(status_code = 409, detail = str(exc)) + + +@studio_router.post("/images/generate", response_model = DiffusionGenerateResponse) +async def generate_diffusion_image( + request: DiffusionGenerateRequest, + current_subject: str = Depends(get_current_subject), +): + from core.inference.diffusion import get_diffusion_backend + + backend = get_diffusion_backend() + try: + result = await asyncio.to_thread( + backend.generate, + prompt = request.prompt, + negative_prompt = request.negative_prompt, + width = request.width, + height = request.height, + steps = request.steps, + guidance = request.guidance, + seed = request.seed, + ) + return DiffusionGenerateResponse(**result) + except RuntimeError as exc: + # No model loaded (or unloaded mid-flight) — a client-state problem. + raise HTTPException(status_code = 409, detail = str(exc)) + except Exception as exc: + logger.error("diffusion.generate_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Image generation failed.") + + +@studio_router.post("/images/unload", response_model = DiffusionStatusResponse) +async def unload_diffusion_model(current_subject: str = Depends(get_current_subject)): + from core.inference.diffusion import get_diffusion_backend + from core.inference.gpu_arbiter import release, DIFFUSION + + status_dict = await asyncio.to_thread(get_diffusion_backend().unload) + release(DIFFUSION) + return DiffusionStatusResponse(**status_dict) + + +@studio_router.get("/images/status", response_model = DiffusionStatusResponse) +async def diffusion_status(current_subject: str = Depends(get_current_subject)): + from core.inference.diffusion import get_diffusion_backend + + return DiffusionStatusResponse(**get_diffusion_backend().status()) + + +@studio_router.get("/images/load-progress", response_model = DiffusionLoadProgressResponse) +async def diffusion_load_progress(current_subject: str = Depends(get_current_subject)): + from core.inference.diffusion import get_diffusion_backend + + return DiffusionLoadProgressResponse(**get_diffusion_backend().load_progress()) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py new file mode 100644 index 0000000000..a6956d1125 --- /dev/null +++ b/studio/backend/tests/test_diffusion_backend.py @@ -0,0 +1,270 @@ +# 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 diffusion backend. + +The family helpers are pure functions, tested directly. The backend lifecycle is +exercised with ``torch`` / ``diffusers`` stubbed via ``sys.modules`` so no real +GPU, weights, or network access is needed (sub-second, CI-friendly). +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.diffusion import DiffusionBackend, encode_png_base64 +from core.inference.diffusion_families import ( + detect_family, + resolve_base_repo, + resolve_local_gguf_child, +) + + +# Pure family helpers + + +def test_detect_family_from_repo_id(): + assert detect_family("unsloth/FLUX.2-klein-4B-GGUF").name == "flux.2-klein" + assert detect_family("unsloth/FLUX.2-dev-GGUF").name == "flux.2" + assert detect_family("city96/FLUX.1-dev-gguf").name == "flux.1" + assert detect_family("unsloth/Qwen-Image-GGUF").name == "qwen-image" + assert detect_family("meta-llama/Llama-3-8B") is None + + +def test_detect_family_override(): + assert detect_family("local/path", override = "flux.1").name == "flux.1" + assert detect_family("local/path", override = "qwen_image").name == "qwen-image" + assert detect_family("local/path", override = "not-a-family") is None + + +def test_resolve_base_repo(): + fam = detect_family("x", override = "flux.1") + assert resolve_base_repo(fam, None) == fam.base_repo + assert resolve_base_repo(fam, " ") == fam.base_repo + assert resolve_base_repo(fam, "custom/base") == "custom/base" + + +def test_resolve_local_gguf_child(tmp_path): + (tmp_path / "model.gguf").write_bytes(b"x") + assert resolve_local_gguf_child(tmp_path, "model.gguf") == (tmp_path / "model.gguf").resolve() + with pytest.raises(ValueError): + resolve_local_gguf_child(tmp_path, "/etc/passwd") + with pytest.raises(ValueError): + resolve_local_gguf_child(tmp_path, "../secret.gguf") + with pytest.raises(ValueError): + resolve_local_gguf_child(tmp_path, "..\\secret.gguf") + with pytest.raises(FileNotFoundError): + resolve_local_gguf_child(tmp_path, "missing.gguf") + + +def test_resolve_local_gguf_child_blocks_symlink_escape(tmp_path): + outside = tmp_path / "outside.gguf" + outside.write_bytes(b"secret") + repo = tmp_path / "repo" + repo.mkdir() + try: + (repo / "model.gguf").symlink_to(outside) + except (OSError, NotImplementedError): + pytest.skip("symlinks not supported on this platform") + with pytest.raises(ValueError): + resolve_local_gguf_child(repo, "model.gguf") + + +# Stubbed runtime for backend lifecycle + + +class _FakeDtype: + def __init__(self, name: str) -> None: + self._name = name + + def __repr__(self) -> str: + return f"torch.{self._name}" + + __str__ = __repr__ + + +class _FakeGenerator: + def __init__(self, device = None) -> None: + self.device = device + self.manual = None + + def seed(self) -> int: + return 4242 + + def manual_seed(self, value: int): + self.manual = value + return self + + +class _FakeImage: + def save(self, buf, format = None) -> None: + buf.write(b"\x89PNG\r\n fake image bytes") + + +class _FakePipe: + def __init__(self) -> None: + self.moved_to = None + self.offloaded = False + self.last_kwargs = None + + def to(self, device): + self.moved_to = device + return self + + def enable_model_cpu_offload(self) -> None: + self.offloaded = True + + def __call__(self, **kwargs): + self.last_kwargs = kwargs + return types.SimpleNamespace(images = [_FakeImage()]) + + +class _FakePipeline: + last: dict = {} + + @classmethod + def from_pretrained(cls, base, **kwargs): + _FakePipeline.last = {"base": base, **kwargs} + return _FakePipe() + + +class _FakeTransformer: + last: dict = {} + + @classmethod + def from_single_file(cls, path, **kwargs): + _FakeTransformer.last = {"path": path, **kwargs} + return object() + + +@pytest.fixture +def fake_runtime(monkeypatch): + torch = types.ModuleType("torch") + torch.bfloat16 = _FakeDtype("bfloat16") + torch.float16 = _FakeDtype("float16") + torch.float32 = _FakeDtype("float32") + torch.Generator = _FakeGenerator + torch.cuda = types.SimpleNamespace(is_available = lambda: False) + torch.backends = types.SimpleNamespace(mps = None) + + diffusers = types.ModuleType("diffusers") + diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype) + diffusers.Flux2KleinPipeline = _FakePipeline + diffusers.Flux2Transformer2DModel = _FakeTransformer + + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "diffusers", diffusers) + # The backend imports clear_gpu_cache by reference; no-op it so unload doesn't + # run real hardware detection against the stubbed torch. + monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: None) + _FakePipeline.last = {} + _FakeTransformer.last = {} + yield + + +def test_load_generate_unload_gguf(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "flux.2-klein", + ) + assert status["loaded"] is True + assert status["family"] == "flux.2-klein" + assert status["base_repo"] == "base/repo" + assert status["device"] == "cpu" + assert status["dtype"] == "float32" + assert status["cpu_offload"] is False + # Transformer built from the local GGUF, pipeline assembled from the base repo. + assert _FakeTransformer.last["path"] == str((tmp_path / "model.gguf").resolve()) + assert _FakeTransformer.last["subfolder"] == "transformer" + assert _FakePipeline.last["base"] == "base/repo" + assert "transformer" in _FakePipeline.last + + gen = backend.generate(prompt = "a sloth", width = 512, height = 512, steps = 4, guidance = 3.0) + assert gen["mime"] == "image/png" + assert gen["seed"] == 4242 # random seed reported back + assert isinstance(gen["image_b64"], str) and gen["image_b64"] + + gen2 = backend.generate(prompt = "again", seed = 99) + assert gen2["seed"] == 99 + + assert backend.unload()["loaded"] is False + assert backend.is_loaded is False + + +def test_cpu_offload_ignored_off_cuda(fake_runtime, tmp_path): + backend = DiffusionBackend() + status = backend.load_pipeline( + str(tmp_path), family_override = "flux.2-klein", base_repo = "base/repo", cpu_offload = True + ) + # No CUDA in the stub, so offload is not engaged. + assert status["cpu_offload"] is False + + +def test_generate_without_load_raises(fake_runtime): + backend = DiffusionBackend() + with pytest.raises(RuntimeError): + backend.generate(prompt = "x") + + +def test_load_unknown_family_raises(): + backend = DiffusionBackend() + with pytest.raises(ValueError): + backend.load_pipeline("some/unrecognised-repo") + + +def test_encode_png_base64_roundtrips(): + encoded = encode_png_base64(_FakeImage()) + import base64 + + assert base64.b64decode(encoded).startswith(b"\x89PNG") + + +# load_progress state machine (no threads / network / real cache) + +from core.inference.diffusion import _LoadingState, _LoadState # noqa: E402 + + +def test_load_progress_idle_and_ready(): + backend = DiffusionBackend() + assert backend.load_progress()["phase"] is None + backend._state = _LoadState(object(), None, "r", "b", "cpu", "float32", False) + assert backend.load_progress()["phase"] == "ready" + + +def test_load_progress_error(): + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "r", base_repo = "b", error = "boom") + p = backend.load_progress() + assert p["phase"] == "error" and p["error"] == "boom" + + +def test_load_progress_downloading_then_finalizing(monkeypatch): + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "r", base_repo = "b", expected_bytes = 1000) + + monkeypatch.setattr(DiffusionBackend, "_cache_bytes", staticmethod(lambda repo: 150)) + p = backend.load_progress() + assert p["phase"] == "downloading" + assert p["bytes_downloaded"] == 300 # summed across repo + base + assert abs(p["fraction"] - 0.3) < 1e-9 + + monkeypatch.setattr(DiffusionBackend, "_cache_bytes", staticmethod(lambda repo: 500)) + assert backend.load_progress()["phase"] == "finalizing" # 1000/1000 + + +def test_begin_load_rejects_concurrent(monkeypatch): + backend = DiffusionBackend() + monkeypatch.setattr(DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0)) + # Block the spawned worker so the load stays "in progress". + monkeypatch.setattr(DiffusionBackend, "load_pipeline", lambda self, **k: __import__("time").sleep(0.2)) + backend.begin_load("unsloth/FLUX.2-klein-4B-GGUF") + with pytest.raises(RuntimeError): + backend.begin_load("unsloth/FLUX.2-klein-4B-GGUF") diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py new file mode 100644 index 0000000000..b0dae9d656 --- /dev/null +++ b/studio/backend/tests/test_diffusion_routes.py @@ -0,0 +1,156 @@ +# 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 diffusion image routes. + +The diffusion backend is replaced with a lightweight fake, so these exercise the +route wiring, validation (422), error mapping, and response shapes without torch, +diffusers, weights, or a GPU. +""" + +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.gpu_arbiter as gpu_arbiter +from auth.authentication import get_current_subject +from routes.inference import studio_router + + +class _FakeBackend: + def __init__(self) -> None: + self.loaded = False + + @property + def is_loaded(self) -> bool: + return self.loaded + + def begin_load(self, model_path, **kwargs): + # The real backend loads on a thread; the fake completes instantly. + self.loaded = True + return { + "loaded": True, + "loading": False, + "repo_id": model_path, + "family": "flux.2-klein", + "base_repo": kwargs.get("base_repo") or "base/repo", + "device": "cpu", + "dtype": "float32", + "cpu_offload": False, + } + + def load_progress(self): + return { + "phase": "ready" if self.loaded else None, + "bytes_downloaded": 0, + "bytes_total": 0, + "fraction": 1.0 if self.loaded else 0.0, + "error": None, + } + + def generate(self, *, seed = None, **kwargs): + if not self.loaded: + raise RuntimeError("No diffusion model is loaded.") + return {"image_b64": "QUJD", "mime": "image/png", "seed": seed if seed is not None else 4242} + + def unload(self): + self.loaded = False + return _unloaded_status() + + def status(self): + return {**_unloaded_status(), "loaded": self.loaded} + + +def _unloaded_status(): + return { + "loaded": False, + "loading": False, + "repo_id": None, + "family": None, + "base_repo": None, + "device": None, + "dtype": None, + "cpu_offload": False, + } + + +@pytest.fixture +def client(monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + # Isolate from the real GPU arbiter: reset ownership and stub the evictors so + # the load route's acquire_for() never touches live backend singletons. + monkeypatch.setattr(gpu_arbiter, "_owner", None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None) + + app = FastAPI() + app.include_router(studio_router, prefix = "/api/inference") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + return TestClient(app) + + +def test_load_generate_status_unload_roundtrip(client): + loaded = client.post("/api/inference/images/load", json = { + "model_path": "unsloth/FLUX.2-klein-4B-GGUF", + "gguf_filename": "model.gguf", + "base_repo": "base/repo", + }) + assert loaded.status_code == 200 + body = loaded.json() + assert body["loaded"] is True and body["family"] == "flux.2-klein" + + assert client.get("/api/inference/images/status").json()["loaded"] is True + + gen = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7}) + assert gen.status_code == 200 + gbody = gen.json() + assert gbody["mime"] == "image/png" and gbody["seed"] == 7 and gbody["image_b64"] + + unloaded = client.post("/api/inference/images/unload") + assert unloaded.status_code == 200 and unloaded.json()["loaded"] is False + assert client.get("/api/inference/images/status").json()["loaded"] is False + + +def test_generate_rejects_non_multiple_of_8(client): + client.post("/api/inference/images/load", json = {"model_path": "x/flux.2-klein"}) + resp = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": 1001}) + assert resp.status_code == 422 + + +def test_generate_without_load_returns_409(client): + resp = client.post("/api/inference/images/generate", json = {"prompt": "p"}) + assert resp.status_code == 409 + + +def test_load_unknown_family_returns_400(client, monkeypatch): + def _raise(*a, **k): + raise ValueError("Could not infer a diffusion family for 'x/y'.") + + backend = _FakeBackend() + backend.begin_load = _raise + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post("/api/inference/images/load", json = {"model_path": "x/y"}) + assert resp.status_code == 400 + assert "family" in resp.json()["detail"] + + +def test_load_progress_route(client): + # Before load: idle. + idle = client.get("/api/inference/images/load-progress") + assert idle.status_code == 200 and idle.json()["phase"] is None + # After load: the fake reports ready. + client.post("/api/inference/images/load", json = {"model_path": "x/flux.2-klein"}) + ready = client.get("/api/inference/images/load-progress") + assert ready.json()["phase"] == "ready" + + +def test_routes_require_auth(): + # No dependency override: the auth dependency must reject the request. + app = FastAPI() + app.include_router(studio_router, prefix = "/api/inference") + unauth = TestClient(app) + assert unauth.get("/api/inference/images/status").status_code in (401, 403) diff --git a/studio/backend/tests/test_gpu_arbiter.py b/studio/backend/tests/test_gpu_arbiter.py new file mode 100644 index 0000000000..d6c7d808bf --- /dev/null +++ b/studio/backend/tests/test_gpu_arbiter.py @@ -0,0 +1,70 @@ +# 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 single-GPU arbiter. + +The real evictors (which tear down live backends) are replaced with recorders, so +these verify only the ownership/eviction sequencing — no torch, GPU, or subprocess. +""" + +from __future__ import annotations + +import pytest + +import core.inference.gpu_arbiter as arb + + +@pytest.fixture +def calls(monkeypatch): + recorded: list[str] = [] + monkeypatch.setattr(arb, "_owner", None) + monkeypatch.setitem(arb._EVICTORS, arb.CHAT, lambda: recorded.append("evict-chat")) + monkeypatch.setitem(arb._EVICTORS, arb.DIFFUSION, lambda: recorded.append("evict-diffusion")) + return recorded + + +def test_first_acquire_evicts_nothing(calls): + arb.acquire_for(arb.CHAT) + assert calls == [] + assert arb.current_owner() == arb.CHAT + + +def test_diffusion_load_evicts_chat(calls): + arb.acquire_for(arb.CHAT) + arb.acquire_for(arb.DIFFUSION) + assert calls == ["evict-chat"] + assert arb.current_owner() == arb.DIFFUSION + + +def test_chat_load_evicts_diffusion(calls): + arb.acquire_for(arb.DIFFUSION) + arb.acquire_for(arb.CHAT) + assert calls == ["evict-diffusion"] + assert arb.current_owner() == arb.CHAT + + +def test_reacquiring_same_owner_does_not_evict(calls): + arb.acquire_for(arb.CHAT) + arb.acquire_for(arb.CHAT) + assert calls == [] + assert arb.current_owner() == arb.CHAT + + +def test_release_clears_owner(calls): + arb.acquire_for(arb.DIFFUSION) + arb.release(arb.DIFFUSION) + assert arb.current_owner() is None + # A subsequent chat load then has nothing to evict. + arb.acquire_for(arb.CHAT) + assert calls == [] + + +def test_release_by_non_owner_is_noop(calls): + arb.acquire_for(arb.CHAT) + arb.release(arb.DIFFUSION) + assert arb.current_owner() == arb.CHAT + + +def test_unknown_owner_raises(calls): + with pytest.raises(ValueError): + arb.acquire_for("gpu") diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 5c18e637e2..9272dfab2b 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -10,6 +10,7 @@ import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; import { Route as chatRoute } from "./routes/chat"; import { Route as exportRoute } from "./routes/export"; +import { Route as imagesRoute } from "./routes/images"; import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; @@ -32,6 +33,7 @@ const routeTree = rootRoute.addChildren([ chatRoute, projectsRoute, exportRoute, + imagesRoute, dataRecipesRoute, dataRecipeRoute, ]); diff --git a/studio/frontend/src/app/routes/images.tsx b/studio/frontend/src/app/routes/images.tsx new file mode 100644 index 0000000000..1761612140 --- /dev/null +++ b/studio/frontend/src/app/routes/images.tsx @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createRoute } from "@tanstack/react-router"; +import { lazy } from "react"; +import { requireAuth } from "../auth-guards"; +import { Route as rootRoute } from "./__root"; + +const ImagesPage = lazy(() => + import("@/features/images").then((m) => ({ + default: m.ImagesPage, + })), +); + +export const Route = createRoute({ + getParentRoute: () => rootRoute, + path: "/images", + staticData: { title: "Images" }, + beforeLoad: () => requireAuth(), + component: ImagesPage, +}); diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 734860138a..390a4aa302 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -66,6 +66,7 @@ import { PinOffIcon, PlusSignIcon, PowerIcon, + PaintBrush02Icon, PencilEdit02Icon, LayoutAlignLeftIcon, Settings02Icon, @@ -1182,6 +1183,18 @@ export function AppSidebar() { closeMobileIfOpen(); }} /> + { + if (chatOnly) return; + navigate({ to: "/images" }); + closeMobileIfOpen(); + }} + /> (response: Response): Promise { + if (!response.ok) { + throw new Error(await readFastApiError(response)); + } + return (await response.json()) as T; +} + +export async function getDiffusionStatus(): Promise { + return parseJson(await authFetch("/api/inference/images/status")); +} + +export async function getDiffusionLoadProgress(): Promise { + return parseJson(await authFetch("/api/inference/images/load-progress")); +} + +export async function loadDiffusionModel(body: DiffusionLoadRequest): Promise { + return parseJson( + await authFetch("/api/inference/images/load", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }), + ); +} + +export async function generateDiffusionImage( + body: DiffusionGenerateRequest, +): Promise { + return parseJson( + await authFetch("/api/inference/images/generate", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }), + ); +} + +export async function unloadDiffusionModel(): Promise { + return parseJson(await authFetch("/api/inference/images/unload", { method: "POST" })); +} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx new file mode 100644 index 0000000000..43a0bdba37 --- /dev/null +++ b/studio/frontend/src/features/images/images-page.tsx @@ -0,0 +1,469 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react"; +import { Download01Icon, ImageAdd02Icon, PaintBrush02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { Separator } from "@/components/ui/separator"; +import { Slider } from "@/components/ui/slider"; +import { Spinner } from "@/components/ui/spinner"; +import { Textarea } from "@/components/ui/textarea"; +import { SectionCard } from "@/components/section-card"; +import { ModelLoadDescription } from "@/features/chat/components/model-load-status"; +import { formatBytes } from "@/features/hub/lib/format"; +import { cn } from "@/lib/utils"; +import { toast } from "@/lib/toast"; + +import { + type DiffusionLoadProgress, + type DiffusionStatus, + generateDiffusionImage, + getDiffusionLoadProgress, + getDiffusionStatus, + loadDiffusionModel, + unloadDiffusionModel, +} from "./api"; + +interface CuratedModel { + label: string; + repo_id: string; + gguf_filename: string; + base_repo: string; + family: string; + // The base repo (VAE/text-encoders) is gated on HF, so loading needs a token — + // even though the GGUF repo itself is an open redistribution. + gated: boolean; +} + +// Repo ids and GGUF filenames are the Hub-canonical (lowercase) names; verified +// to exist on HF. The base_repo provides the VAE/text-encoders the GGUF omits. +const CURATED_MODELS: CuratedModel[] = [ + { + label: "FLUX.2 klein base 4B", + repo_id: "unsloth/FLUX.2-klein-base-4B-GGUF", + gguf_filename: "flux-2-klein-base-4b-Q4_K_S.gguf", + base_repo: "black-forest-labs/FLUX.2-klein-base-4B", + family: "flux.2-klein", + gated: false, + }, + { + label: "FLUX.2 dev", + repo_id: "unsloth/FLUX.2-dev-GGUF", + gguf_filename: "flux2-dev-Q4_K_S.gguf", + base_repo: "black-forest-labs/FLUX.2-dev", + family: "flux.2", + gated: true, + }, + { + label: "FLUX.1 dev", + repo_id: "city96/FLUX.1-dev-gguf", + gguf_filename: "flux1-dev-Q4_K_S.gguf", + base_repo: "black-forest-labs/FLUX.1-dev", + family: "flux.1", + gated: true, + }, +]; + +const RESOLUTIONS: Array<{ label: string; w: number; h: number }> = [ + { label: "Square 1024", w: 1024, h: 1024 }, + { label: "Square 768", w: 768, h: 768 }, + { label: "Portrait 832×1216", w: 832, h: 1216 }, + { label: "Landscape 1216×832", w: 1216, h: 832 }, +]; + +interface ResultItem { + id: number; + src: string; + prompt: string; + width: number; + height: number; + steps: number; + guidance: number; + seed: number; +} + +function downloadImage(src: string, seed: number) { + const link = document.createElement("a"); + link.href = src; + link.download = `unsloth-${seed}.png`; + link.click(); +} + +// Title + percent/label for the chat progress component, from a progress poll. +function progressView(p: DiffusionLoadProgress): { + title: string; + progressPercent: number | null; + progressLabel: string | null; +} { + const title = p.phase === "finalizing" ? "Loading to GPU…" : "Downloading model…"; + if (p.bytes_total > 0) { + return { + title, + progressPercent: p.fraction * 100, + progressLabel: `${formatBytes(p.bytes_downloaded)} of ${formatBytes(p.bytes_total)}`, + }; + } + // Unknown total: show bytes counting up, no bar. + return { + title, + progressPercent: null, + progressLabel: p.bytes_downloaded > 0 ? `${formatBytes(p.bytes_downloaded)} downloaded` : null, + }; +} + +// Matches the field-label style used across Studio (export/chat settings). +function Field({ label, children }: { label: string; children: ReactNode }) { + return ( +
+ + {children} +
+ ); +} + +type Busy = "loading" | "unloading" | "generating" | null; + +export function ImagesPage() { + const [modelIdx, setModelIdx] = useState(0); + const [hfToken, setHfToken] = useState(""); + const [prompt, setPrompt] = useState( + "a tiny ginger sloth coding in a sunlit treehouse, photorealistic", + ); + const [negativePrompt, setNegativePrompt] = useState(""); + const [resolutionIdx, setResolutionIdx] = useState(0); + const [steps, setSteps] = useState(24); + const [guidance, setGuidance] = useState(3.5); + const [seed, setSeed] = useState(""); + + const [busy, setBusy] = useState(null); + const [status, setStatus] = useState(null); + const [loadProgress, setLoadProgress] = useState(null); + const [results, setResults] = useState([]); + const [selectedId, setSelectedId] = useState(null); + // Stable, ever-increasing ids: results are prepended, so an array index would + // re-key every existing image on each new generation. + const nextResultId = useRef(0); + const pollTimer = useRef | null>(null); + + const model = CURATED_MODELS[modelIdx]; + const resolution = RESOLUTIONS[resolutionIdx]; + const selected = useMemo( + () => results.find((r) => r.id === selectedId) ?? results[0] ?? null, + [results, selectedId], + ); + + const refreshStatus = useCallback(async () => { + try { + setStatus(await getDiffusionStatus()); + } catch { + // Status is best-effort; a failed poll shouldn't surface an error toast. + } + }, []); + + useEffect(() => { + void refreshStatus(); + // Stop polling if the page unmounts mid-load. + return () => { + if (pollTimer.current) clearTimeout(pollTimer.current); + }; + }, [refreshStatus]); + + // Poll load-progress until the background load reaches "ready" or "error". + const pollLoadProgress = useCallback(async () => { + try { + const p = await getDiffusionLoadProgress(); + setLoadProgress(p); + if (p.phase === "ready") { + setStatus(await getDiffusionStatus()); + toast.success("Model loaded"); + setBusy(null); + setLoadProgress(null); + return; + } + if (p.phase === "error") { + toast.error(p.error || "Failed to load model"); + setBusy(null); + setLoadProgress(null); + return; + } + } catch { + // Transient poll failure: keep trying. + } + pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000); + }, []); + + const handleLoad = useCallback(async () => { + // Cancel any prior poll loop so two can't run at once. + if (pollTimer.current) clearTimeout(pollTimer.current); + setBusy("loading"); + setLoadProgress(null); + try { + // Returns immediately — the load runs in the background; we poll for it. + await loadDiffusionModel({ + model_path: model.repo_id, + gguf_filename: model.gguf_filename, + base_repo: model.base_repo, + family_override: model.family, + hf_token: hfToken.trim() || undefined, + }); + } catch (err) { + toast.error(err instanceof Error ? err.message : "Failed to start load"); + setBusy(null); + void refreshStatus(); + return; + } + void pollLoadProgress(); + }, [model, hfToken, refreshStatus, pollLoadProgress]); + + const handleUnload = useCallback(async () => { + setBusy("unloading"); + try { + setStatus(await unloadDiffusionModel()); + } catch (err) { + toast.error(err instanceof Error ? err.message : "Failed to unload model"); + void refreshStatus(); + } finally { + setBusy(null); + } + }, [refreshStatus]); + + const handleGenerate = useCallback(async () => { + if (!prompt.trim()) { + toast.error("Prompt is empty"); + return; + } + let parsedSeed: number | undefined; + if (seed.trim()) { + const n = Number(seed); + if (!Number.isInteger(n) || n < 0) { + toast.error("Seed must be a non-negative integer"); + return; + } + parsedSeed = n; + } + + setBusy("generating"); + try { + const res = await generateDiffusionImage({ + prompt: prompt.trim(), + negative_prompt: negativePrompt.trim() || undefined, + width: resolution.w, + height: resolution.h, + steps, + guidance, + seed: parsedSeed, + }); + const id = nextResultId.current++; + setResults((prev) => [ + { + id, + src: `data:${res.mime};base64,${res.image_b64}`, + prompt: prompt.trim(), + width: resolution.w, + height: resolution.h, + steps, + guidance, + seed: res.seed, + }, + ...prev, + ]); + setSelectedId(id); + } catch (err) { + toast.error(err instanceof Error ? err.message : "Image generation failed"); + } finally { + setBusy(null); + } + }, [prompt, negativePrompt, resolution, steps, guidance, seed]); + + const statusLabel = useMemo(() => { + if (busy === "loading") return "Loading…"; + if (status?.loaded) return `Loaded · ${status.repo_id}`; + return "No model loaded"; + }, [busy, status]); + + return ( +
+ {/* ── Left control rail ─────────────────────────────── */} + } + title="Image generation" + description="Generate from local diffusion GGUFs" + accent="indigo" + className="w-[360px] shrink-0 gap-4 overflow-y-auto" + > + + + + + {model.gated && ( + + setHfToken(e.target.value)} + /> + + )} + +
+ + +
+ {busy === "loading" && loadProgress ? ( + + ) : ( +

+ {statusLabel} +

+ )} + + + + +