489 lines
18 KiB
Python
489 lines
18 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Diffusion image generation backend.
|
|
|
|
Loads Hugging Face diffusion checkpoints in either the standard
|
|
``diffusers`` layout or the single-file GGUF layout published under
|
|
``unsloth/*-GGUF`` (Flux 2, Flux 2 Klein, Qwen-Image, SD3, SDXL, ...).
|
|
GGUF files are dynamically dequantised on-device via
|
|
``diffusers.GGUFQuantizationConfig``, then the rest of the pipeline
|
|
(VAE, text encoders, scheduler) is pulled from the matching ``diffusers``
|
|
repo so end users only ever need one local file plus the metadata repo.
|
|
|
|
The module is intentionally torch-only: it never spawns a subprocess and
|
|
shares the active CUDA / MPS device with the rest of Studio. The cost of
|
|
not having a separate process is that loading a diffusion model and a
|
|
GGUF chat model at the same time can OOM on consumer GPUs; the routes
|
|
layer must therefore swap between the two as needed (the orchestrator
|
|
unloads llama-server before any diffusion load on hosts with < 24 GB).
|
|
|
|
The class deliberately exposes a small, llama-cpp-style surface:
|
|
|
|
load_model(repo_id, ...)
|
|
generate_image(prompt, ...) -> PIL.Image
|
|
unload_model()
|
|
status() -> dict
|
|
|
|
so the route layer at ``studio/backend/routes/inference.py`` can mirror
|
|
the existing llama-server lifecycle (probe + load + generate + unload)
|
|
without learning a second API.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import gc
|
|
import io
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
from loggers import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
# ─── Pipeline registry ────────────────────────────────────────────────
|
|
#
|
|
# Keep this list narrow on purpose: only ship the small text-to-image
|
|
# families with first-class GGUF coverage on the Hub. Anything else is
|
|
# either video (LTX*, Wan) or research-grade (Sana, SD3.5) and can be
|
|
# added once it has a working GGUF release plus a smoke test.
|
|
#
|
|
# Each entry maps a substring of the loaded repo id (case-insensitive)
|
|
# to the (pipeline_class_name, transformer_class_name, default base
|
|
# repo for missing pieces). ``base_repo`` is what we pass to
|
|
# ``Pipeline.from_pretrained`` to pick up the VAE + text encoders when
|
|
# the user gave us a GGUF-only repo. The base_repo is documented to the
|
|
# user via ``status()`` so they understand why a second download fires.
|
|
|
|
|
|
@dataclass(frozen = True)
|
|
class DiffusionFamily:
|
|
name: str
|
|
pipeline_class: str
|
|
transformer_class: str
|
|
base_repo: str
|
|
# Optional: list of HF "trigger" substrings besides ``name`` that map
|
|
# to this family (e.g. "flux1-dev" plus "flux.1-dev"). Lowercased.
|
|
aliases: tuple[str, ...] = field(default_factory = tuple)
|
|
|
|
|
|
_FAMILIES: tuple[DiffusionFamily, ...] = (
|
|
DiffusionFamily(
|
|
name = "flux.2-klein",
|
|
pipeline_class = "Flux2KleinPipeline",
|
|
transformer_class = "Flux2Transformer2DModel",
|
|
base_repo = "black-forest-labs/FLUX.2-klein",
|
|
aliases = ("flux2-klein", "flux-2-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", "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.1.dev", "flux-dev"),
|
|
),
|
|
DiffusionFamily(
|
|
name = "qwen-image",
|
|
pipeline_class = "QwenImagePipeline",
|
|
transformer_class = "QwenImageTransformer2DModel",
|
|
base_repo = "Qwen/Qwen-Image",
|
|
aliases = ("qwenimage", "qwen_image"),
|
|
),
|
|
DiffusionFamily(
|
|
name = "stable-diffusion-3",
|
|
pipeline_class = "StableDiffusion3Pipeline",
|
|
transformer_class = "SD3Transformer2DModel",
|
|
base_repo = "stabilityai/stable-diffusion-3-medium-diffusers",
|
|
aliases = ("sd3-medium", "stable-diffusion-3-medium", "sd3.5"),
|
|
),
|
|
DiffusionFamily(
|
|
name = "stable-diffusion-xl",
|
|
pipeline_class = "StableDiffusionXLPipeline",
|
|
transformer_class = "", # SDXL uses a UNet, not a transformer
|
|
base_repo = "stabilityai/stable-diffusion-xl-base-1.0",
|
|
aliases = ("sdxl",),
|
|
),
|
|
)
|
|
|
|
|
|
def detect_family(
|
|
repo_id: str, *, override_family: Optional[str] = None
|
|
) -> Optional[DiffusionFamily]:
|
|
"""Return the diffusion family matching ``repo_id``.
|
|
|
|
Matching is substring-based and case-insensitive. ``override_family``
|
|
bypasses substring matching and looks up by ``DiffusionFamily.name``.
|
|
Returns ``None`` when no family applies so callers can surface a clear
|
|
"unsupported model" error rather than guessing wrong.
|
|
"""
|
|
if override_family:
|
|
wanted = override_family.strip().lower()
|
|
for fam in _FAMILIES:
|
|
if fam.name == wanted:
|
|
return fam
|
|
return None
|
|
needle = (repo_id or "").lower()
|
|
if not needle:
|
|
return None
|
|
for fam in _FAMILIES:
|
|
if fam.name in needle:
|
|
return fam
|
|
for alias in fam.aliases:
|
|
if alias and alias in needle:
|
|
return fam
|
|
return None
|
|
|
|
|
|
def supported_families() -> list[dict[str, str]]:
|
|
"""Public-facing list of families for ``/api/inference/images/status``."""
|
|
return [
|
|
{
|
|
"name": fam.name,
|
|
"pipeline_class": fam.pipeline_class,
|
|
"base_repo": fam.base_repo,
|
|
}
|
|
for fam in _FAMILIES
|
|
]
|
|
|
|
|
|
# ─── Backend ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class DiffusionBackend:
|
|
"""Singleton-style diffusion backend.
|
|
|
|
One pipeline at a time; ``load_model`` swaps the previous one out.
|
|
Generation is mutex'd so concurrent requests serialise rather than
|
|
racing GPU memory.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._pipe: Any = None
|
|
self._lock = threading.Lock()
|
|
self._family: Optional[DiffusionFamily] = None
|
|
self._repo_id: Optional[str] = None
|
|
self._gguf_path: Optional[str] = None
|
|
self._base_repo: Optional[str] = None
|
|
self._device: Optional[str] = None
|
|
self._dtype: Optional[str] = None
|
|
self._loaded_at: Optional[float] = None
|
|
self._loading: bool = False
|
|
self._last_error: Optional[str] = None
|
|
|
|
# ── lifecycle ─────────────────────────────────────────────────
|
|
|
|
@property
|
|
def is_loaded(self) -> bool:
|
|
return self._pipe is not None
|
|
|
|
@property
|
|
def repo_id(self) -> Optional[str]:
|
|
return self._repo_id
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
return {
|
|
"is_loaded": self.is_loaded,
|
|
"is_loading": self._loading,
|
|
"repo_id": self._repo_id,
|
|
"family": self._family.name if self._family else None,
|
|
"pipeline_class": self._family.pipeline_class if self._family else None,
|
|
"base_repo": self._base_repo,
|
|
"gguf_path": self._gguf_path,
|
|
"device": self._device,
|
|
"dtype": self._dtype,
|
|
"loaded_at": self._loaded_at,
|
|
"last_error": self._last_error,
|
|
"supported_families": supported_families(),
|
|
}
|
|
|
|
def _pick_device_and_dtype(self) -> tuple[str, "Any"]:
|
|
"""Pick (device, dtype) for the current host.
|
|
|
|
CUDA-first because that is the only path our diffusion GGUFs are
|
|
validated on. On macOS we use MPS in float16 to keep the pipeline
|
|
on the Metal GPU. CPU is allowed only as a last resort because
|
|
running FLUX on CPU is unusably slow (> 10 minutes per image).
|
|
"""
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
return "cuda", torch.bfloat16
|
|
if (
|
|
hasattr(torch, "backends")
|
|
and getattr(torch.backends, "mps", None)
|
|
and torch.backends.mps.is_available()
|
|
):
|
|
return "mps", torch.float16
|
|
return "cpu", torch.float32
|
|
|
|
def load_model(
|
|
self,
|
|
repo_id: str,
|
|
*,
|
|
gguf_filename: Optional[str] = None,
|
|
base_repo: Optional[str] = None,
|
|
hf_token: Optional[str] = None,
|
|
family_override: Optional[str] = None,
|
|
enable_model_cpu_offload: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""Load a diffusion model.
|
|
|
|
``repo_id`` is the Hugging Face repo id of either a GGUF-only
|
|
repo (e.g. ``unsloth/FLUX.2-klein-4B-GGUF``) or a full diffusers
|
|
repo (e.g. ``black-forest-labs/FLUX.2-klein``). When the repo
|
|
contains a GGUF, ``gguf_filename`` picks which quant to load;
|
|
otherwise diffusers' standard config-driven load runs.
|
|
|
|
``base_repo`` overrides the auto-detected diffusers base used
|
|
for VAE / text encoders. ``family_override`` short-circuits the
|
|
substring matcher when an exotic repo name confuses it.
|
|
|
|
Raises ``RuntimeError`` on failure with a user-facing message;
|
|
the previous pipeline (if any) stays loaded so a failed swap
|
|
does not leave Studio in an unusable state.
|
|
"""
|
|
from huggingface_hub import hf_hub_download
|
|
import diffusers
|
|
import torch
|
|
|
|
fam = detect_family(repo_id, override_family = family_override)
|
|
if fam is None:
|
|
raise RuntimeError(
|
|
f"Could not infer a diffusion family for '{repo_id}'. "
|
|
"Pass family_override = 'flux.2-klein' / 'flux.2' / "
|
|
"'flux.1' / 'qwen-image' / 'stable-diffusion-3' / "
|
|
"'stable-diffusion-xl' to disambiguate."
|
|
)
|
|
|
|
device, dtype = self._pick_device_and_dtype()
|
|
|
|
with self._lock:
|
|
self._loading = True
|
|
self._last_error = None
|
|
try:
|
|
pipeline_cls = getattr(diffusers, fam.pipeline_class, None)
|
|
if pipeline_cls is None:
|
|
raise RuntimeError(
|
|
f"diffusers {diffusers.__version__} has no "
|
|
f"{fam.pipeline_class}; upgrade diffusers and retry."
|
|
)
|
|
transformer_cls = (
|
|
getattr(diffusers, fam.transformer_class, None)
|
|
if fam.transformer_class
|
|
else None
|
|
)
|
|
|
|
effective_base = base_repo or fam.base_repo
|
|
logger.info(
|
|
"Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)",
|
|
repo_id,
|
|
fam.name,
|
|
device,
|
|
dtype,
|
|
effective_base,
|
|
)
|
|
|
|
transformer = None
|
|
local_gguf_path: Optional[str] = None
|
|
if gguf_filename:
|
|
if transformer_cls is None:
|
|
raise RuntimeError(
|
|
f"Family {fam.name} does not have a GGUF transformer "
|
|
"path; load the full repo instead."
|
|
)
|
|
local_gguf_path = hf_hub_download(
|
|
repo_id = repo_id,
|
|
filename = gguf_filename,
|
|
token = hf_token,
|
|
)
|
|
quant_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype)
|
|
transformer = transformer_cls.from_single_file(
|
|
local_gguf_path,
|
|
quantization_config = quant_config,
|
|
torch_dtype = dtype,
|
|
)
|
|
|
|
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
|
|
|
|
pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs)
|
|
if enable_model_cpu_offload and device == "cuda":
|
|
pipe.enable_model_cpu_offload()
|
|
else:
|
|
pipe.to(device)
|
|
|
|
# Drop the old pipeline only after the new one is in place.
|
|
old = self._pipe
|
|
with self._lock:
|
|
self._pipe = pipe
|
|
self._family = fam
|
|
self._repo_id = repo_id
|
|
self._gguf_path = local_gguf_path
|
|
self._base_repo = effective_base
|
|
self._device = device
|
|
self._dtype = str(dtype).replace("torch.", "")
|
|
self._loaded_at = time.time()
|
|
_release(old)
|
|
|
|
return self.status()
|
|
except Exception as exc:
|
|
with self._lock:
|
|
self._last_error = str(exc)
|
|
logger.exception("Diffusion load failed for %s", repo_id)
|
|
raise RuntimeError(f"Failed to load diffusion model: {exc}") from exc
|
|
finally:
|
|
with self._lock:
|
|
self._loading = False
|
|
|
|
def unload_model(self) -> dict[str, Any]:
|
|
with self._lock:
|
|
old = self._pipe
|
|
self._pipe = None
|
|
self._family = None
|
|
self._repo_id = None
|
|
self._gguf_path = None
|
|
self._base_repo = None
|
|
self._device = None
|
|
self._dtype = None
|
|
self._loaded_at = None
|
|
_release(old)
|
|
return {"is_loaded": False}
|
|
|
|
# ── generation ────────────────────────────────────────────────
|
|
|
|
def generate_image(
|
|
self,
|
|
*,
|
|
prompt: str,
|
|
negative_prompt: Optional[str] = None,
|
|
num_inference_steps: int = 24,
|
|
guidance_scale: float = 3.5,
|
|
width: int = 1024,
|
|
height: int = 1024,
|
|
seed: Optional[int] = None,
|
|
) -> "Any":
|
|
"""Generate a single PIL image and return it.
|
|
|
|
The mutex is held for the entire call: diffusion pipelines are
|
|
not thread-safe, and overlapping ``__call__``s on a shared
|
|
pipeline frequently corrupt their internal scheduler state.
|
|
"""
|
|
if not prompt or not prompt.strip():
|
|
raise ValueError("prompt is empty")
|
|
if num_inference_steps < 1 or num_inference_steps > 200:
|
|
raise ValueError("num_inference_steps must be in [1, 200]")
|
|
if width <= 0 or height <= 0 or width > 2048 or height > 2048:
|
|
raise ValueError("width and height must be in (0, 2048]")
|
|
# Snap to a multiple of 8: Flux / SD pipelines require it and a
|
|
# silent crash deep in the VAE is much worse than a clear error
|
|
# message up front.
|
|
if width % 8 or height % 8:
|
|
raise ValueError("width and height must be multiples of 8")
|
|
|
|
import torch
|
|
|
|
with self._lock:
|
|
if self._pipe is None:
|
|
raise RuntimeError("No diffusion model is loaded.")
|
|
pipe = self._pipe
|
|
device = self._device or "cpu"
|
|
|
|
generator = None
|
|
if seed is not None:
|
|
# Match the device of the pipeline so determinism holds
|
|
# across reload cycles. For CPU offload, the noise still
|
|
# has to live on the device the diffusion forward runs on.
|
|
gen_device = (
|
|
"cuda" if device == "cuda" and torch.cuda.is_available() else "cpu"
|
|
)
|
|
generator = torch.Generator(device = gen_device).manual_seed(int(seed))
|
|
|
|
call_kwargs: dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"num_inference_steps": int(num_inference_steps),
|
|
"guidance_scale": float(guidance_scale),
|
|
"width": int(width),
|
|
"height": int(height),
|
|
}
|
|
if negative_prompt is not None and negative_prompt.strip():
|
|
call_kwargs["negative_prompt"] = negative_prompt
|
|
if generator is not None:
|
|
call_kwargs["generator"] = generator
|
|
|
|
out = pipe(**call_kwargs)
|
|
images = getattr(out, "images", None) or []
|
|
if not images:
|
|
raise RuntimeError("Diffusion pipeline returned no images.")
|
|
return images[0]
|
|
|
|
|
|
def encode_png_base64(pil_image: "Any") -> str:
|
|
"""Encode a PIL image to base64-encoded PNG."""
|
|
import base64
|
|
|
|
buf = io.BytesIO()
|
|
pil_image.save(buf, format = "PNG", optimize = True)
|
|
return base64.b64encode(buf.getvalue()).decode("ascii")
|
|
|
|
|
|
# ─── Helpers ──────────────────────────────────────────────────────────
|
|
|
|
|
|
def _release(obj: Any) -> None:
|
|
"""Best-effort GPU-memory release for a pipeline being swapped out."""
|
|
if obj is None:
|
|
return
|
|
try:
|
|
del obj
|
|
except Exception:
|
|
pass
|
|
gc.collect()
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ─── Module-level singleton ───────────────────────────────────────────
|
|
|
|
|
|
_singleton: Optional[DiffusionBackend] = None
|
|
_singleton_lock = threading.Lock()
|
|
|
|
|
|
def get_diffusion_backend() -> DiffusionBackend:
|
|
"""Return the process-wide diffusion backend (lazy-instantiated)."""
|
|
global _singleton
|
|
if _singleton is None:
|
|
with _singleton_lock:
|
|
if _singleton is None:
|
|
_singleton = DiffusionBackend()
|
|
return _singleton
|
|
|
|
|
|
async def async_generate(
|
|
backend: DiffusionBackend,
|
|
**kwargs: Any,
|
|
) -> "Any":
|
|
"""Run ``generate_image`` in the default executor so route handlers
|
|
do not block the event loop for the 5-30 s a diffusion step takes."""
|
|
loop = asyncio.get_event_loop()
|
|
return await loop.run_in_executor(None, lambda: backend.generate_image(**kwargs))
|