Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI
Backend: - Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file fp8 transformers, gated to the unsloth org plus a curated allowlist. - Image-conditioned workflows built with Pipeline.from_pipe so they reuse the loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint, outpaint, and a hires-fix upscale pass. - Instruction editing as its own family kind (Qwen-Image-Edit-2511, FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and multi-reference) plus klein inpaint. - Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the decoded image size and cap upscale output to avoid OOM on large inputs. - Fixes: from_pipe defaulting to a float32 recast that crashed torchao quantized transformers; image-conditioned calls forcing the slider size onto the input image. Native sd.cpp engine rejects image-conditioned and reference requests it cannot serve. Frontend: - Redesigned Images page with capability-gated workflow tabs (Create, Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor, client-side outpaint, and a multi-reference picker. - Advanced options moved to a right-docked panel mirroring Chat: closed by default, toggled by a single fixed top-bar button that stays in place. sd.cpp installer: pin the release, verify each download's sha256, add a download timeout, and make the source repo configurable for a future mirror.
This commit is contained in:
Daniel Han
parent
f24384b4e9
commit
b14e2f9be7
14 changed files with 2920 additions and 183 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -99,7 +99,9 @@ def _activate(name: str, reason: Optional[str]) -> Any:
|
|||
return get_active_diffusion_engine()
|
||||
|
||||
|
||||
def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] = None) -> Any:
|
||||
def select_and_activate_engine(
|
||||
fam: DiffusionFamily, *, hf_token: Optional[str] = None, model_kind: Optional[str] = None
|
||||
) -> Any:
|
||||
"""Pick + activate the engine for loading ``fam`` on this host; return the engine.
|
||||
|
||||
Falls back to diffusers (recording a reason) whenever the native route is
|
||||
|
|
@ -107,6 +109,12 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str]
|
|||
native asset mapping, or the sd-cli binary is unavailable -- always BEFORE the
|
||||
slow load begins, so a fallback never strands a half-native load.
|
||||
"""
|
||||
# Non-GGUF loads (a single-file safetensors transformer, or a full diffusers
|
||||
# pipeline) only run on diffusers: the native sd.cpp engine consumes single-file
|
||||
# GGUF checkpoints only, so force diffusers before the device/native checks below.
|
||||
if model_kind and model_kind != "gguf":
|
||||
return _activate(ENGINE_DIFFUSERS, f"non-GGUF load ({model_kind}) requires diffusers")
|
||||
|
||||
forced, sd_cpp_pref, mps_enabled = _engine_config()
|
||||
|
||||
if forced == ENGINE_DIFFUSERS:
|
||||
|
|
|
|||
|
|
@ -28,6 +28,32 @@ class DiffusionFamily:
|
|||
# Pipeline kwarg carrying the guidance value. Most use "guidance_scale";
|
||||
# Qwen-Image's distilled guidance is off, so its real CFG is "true_cfg_scale".
|
||||
cfg_kwarg: str = "guidance_scale"
|
||||
# Optional diffusers pipeline classes for image-conditioned workflows. The backend
|
||||
# builds these around the ALREADY-loaded transformer/VAE/text-encoder via
|
||||
# ``Pipeline.from_pipe`` (no extra weights, no reload), so a family only needs the
|
||||
# class name here to gain the workflow. None = the family does not support it (the
|
||||
# UI gates the workflow off). The base text-to-image pipeline is ``pipeline_class``.
|
||||
img2img_pipeline_class: Optional[str] = None
|
||||
inpaint_pipeline_class: Optional[str] = None
|
||||
# True when the inpaint pipeline keeps the input canvas size, so it can also drive
|
||||
# outpaint (extend), where the padded canvas is LARGER than the original. False for
|
||||
# FLUX.2 (its pipelines scale any >1MP input down to ~1MP, which shrinks an outpaint
|
||||
# canvas back and defeats the extend). Such families get Inpaint but not Extend.
|
||||
inpaint_preserves_size: bool = True
|
||||
# True for instruction-editing families (Qwen-Image-Edit / FLUX Kontext): the model's
|
||||
# OWN pipeline (``pipeline_class``) is the edit pipeline -- it takes an input image plus
|
||||
# a text instruction and has no plain text-to-image mode. So these expose only the
|
||||
# "edit" workflow, require an input image at generate time, and the loaded pipe is used
|
||||
# directly (no from_pipe). ``base_repo`` here is the matching diffusers repo that
|
||||
# supplies the VAE / text-encoder / processor / scheduler for the GGUF transformer.
|
||||
edit: bool = False
|
||||
# True for families whose OWN text-to-image pipeline ALSO accepts reference image(s)
|
||||
# (FLUX.2: Flux2KleinPipeline takes an optional ``image`` arg). Unlike ``edit`` these
|
||||
# families still do plain text-to-image (no image), and unlike img2img the conditioning
|
||||
# is reference-based, not a denoise blend: there is no ``strength`` and the output size
|
||||
# comes from the requested width/height, not the reference's size. The loaded pipe is
|
||||
# used directly (no from_pipe). Exposes a "reference" workflow alongside "txt2img".
|
||||
reference: bool = False
|
||||
# Extra lowercased substrings (besides ``name``) that map a repo id here.
|
||||
aliases: tuple[str, ...] = field(default_factory = tuple)
|
||||
# True for families whose activations overflow float16's finite range
|
||||
|
|
@ -79,6 +105,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
|
|||
transformer_class = "FluxTransformer2DModel",
|
||||
base_repo = "black-forest-labs/FLUX.1-schnell",
|
||||
aliases = ("flux1", "flux-1"),
|
||||
img2img_pipeline_class = "FluxImg2ImgPipeline",
|
||||
inpaint_pipeline_class = "FluxInpaintPipeline",
|
||||
sd_cpp_vae = ("black-forest-labs/FLUX.1-schnell", "ae.safetensors"),
|
||||
sd_cpp_text_encoders = (
|
||||
("comfyanonymous/flux_text_encoders", "clip_l.safetensors", "clip_l"),
|
||||
|
|
@ -94,6 +122,13 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
|
|||
transformer_class = "Flux2Transformer2DModel",
|
||||
base_repo = "black-forest-labs/FLUX.2-klein-4B",
|
||||
aliases = ("flux2-klein",),
|
||||
# Flux2KleinPipeline natively accepts reference image(s) via its `image` arg, so it
|
||||
# exposes a "reference" workflow on top of plain text-to-image. It has a dedicated
|
||||
# inpaint pipeline too (no img2img one), so it also gets inpaint + extend (outpaint).
|
||||
reference = True,
|
||||
inpaint_pipeline_class = "Flux2KleinInpaintPipeline",
|
||||
# FLUX.2 scales >1MP inputs down to ~1MP, so outpaint (a larger canvas) can't grow.
|
||||
inpaint_preserves_size = False,
|
||||
# FLUX.2 uses a distinct 32-channel autoencoder; sd-cli needs the latent
|
||||
# format override. The single-file VAE ships in Comfy-Org/flux2-dev (the
|
||||
# klein-4B repo only has a sharded diffusers VAE). Shares Qwen3-4B with z-image.
|
||||
|
|
@ -103,6 +138,39 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
|
|||
("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
|
||||
),
|
||||
),
|
||||
DiffusionFamily(
|
||||
# Instruction editing with FLUX. FluxKontextPipeline takes an input image + an edit
|
||||
# instruction; the GGUF transformer is the standard FluxTransformer2DModel, with the
|
||||
# T5/CLIP text encoders + VAE from the base diffusers repo. cfg defaults to
|
||||
# guidance_scale (FLUX). Most-specific aliases first so detect_family prefers this
|
||||
# over the plain "flux.1" family and un-rejects the "kontext" keyword for it.
|
||||
name = "flux.1-kontext",
|
||||
pipeline_class = "FluxKontextPipeline",
|
||||
transformer_class = "FluxTransformer2DModel",
|
||||
base_repo = "black-forest-labs/FLUX.1-Kontext-dev",
|
||||
aliases = ("flux.1-kontext-dev", "flux1-kontext", "flux-kontext", "kontext"),
|
||||
edit = True,
|
||||
),
|
||||
DiffusionFamily(
|
||||
# Instruction editing (image-in + text-instruction-out). The 2511 checkpoint ships
|
||||
# as QwenImageEditPlusPipeline (multi-image-capable); the GGUF transformer is the
|
||||
# standard QwenImageTransformer2DModel, with the VAE / Qwen2.5-VL text-encoder /
|
||||
# image processor / scheduler coming from the base diffusers repo. Most-specific
|
||||
# aliases first so detect_family prefers this over the plain "qwen-image" family.
|
||||
name = "qwen-image-edit",
|
||||
pipeline_class = "QwenImageEditPlusPipeline",
|
||||
transformer_class = "QwenImageTransformer2DModel",
|
||||
base_repo = "Qwen/Qwen-Image-Edit-2511",
|
||||
cfg_kwarg = "true_cfg_scale",
|
||||
aliases = (
|
||||
"qwen-image-edit-2511",
|
||||
"qwen-image-edit-2509",
|
||||
"qwen-image-edit",
|
||||
"qwen_image_edit",
|
||||
"qwenimageedit",
|
||||
),
|
||||
edit = True,
|
||||
),
|
||||
DiffusionFamily(
|
||||
name = "qwen-image",
|
||||
pipeline_class = "QwenImagePipeline",
|
||||
|
|
@ -110,6 +178,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
|
|||
base_repo = "Qwen/Qwen-Image",
|
||||
cfg_kwarg = "true_cfg_scale",
|
||||
aliases = ("qwen_image", "qwenimage"),
|
||||
img2img_pipeline_class = "QwenImageImg2ImgPipeline",
|
||||
inpaint_pipeline_class = "QwenImageInpaintPipeline",
|
||||
sd_cpp_vae = ("Comfy-Org/Qwen-Image_ComfyUI", "split_files/vae/qwen_image_vae.safetensors"),
|
||||
# The Qwen2.5-VL text encoder as a Q4_K_M GGUF keeps the CPU RAM win (the
|
||||
# bf16 safetensors encoder is ~15 GB). sd-cli's --qwen2vl is an alias of --llm.
|
||||
|
|
@ -130,6 +200,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
|
|||
transformer_class = "ZImageTransformer2DModel",
|
||||
base_repo = "Tongyi-MAI/Z-Image-Turbo",
|
||||
aliases = ("zimage", "z_image"),
|
||||
img2img_pipeline_class = "ZImageImg2ImgPipeline",
|
||||
inpaint_pipeline_class = "ZImageInpaintPipeline",
|
||||
# Z-Image's MLP down-projections peak near 9e5, which overflows float16.
|
||||
fp16_incompatible = True,
|
||||
sd_cpp_vae = ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"),
|
||||
|
|
@ -141,15 +213,33 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
|
|||
|
||||
# Editing / inpaint checkpoints share an arch keyword but need a different
|
||||
# pipeline and an input image, which this text-to-image backend doesn't drive.
|
||||
_EDIT_KEYWORDS = ("edit", "kontext", "inpaint")
|
||||
# "layered" rejects Qwen-Image-Layered: its transformer sets additional_t_cond=True
|
||||
# and expects an extra addition_t_cond input that the standard QwenImagePipeline
|
||||
# never supplies, so it loads but crashes at the first denoise step. Rejecting it
|
||||
# here fails the load fast with a clear message and hides it from the picker.
|
||||
_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "layered")
|
||||
|
||||
|
||||
def _best_family_match(needle: str) -> Optional[DiffusionFamily]:
|
||||
"""The family whose name/alias is the LONGEST substring of ``needle``. Longest =
|
||||
most specific, so an edit checkpoint ('...qwen-image-edit-2511...') matches the
|
||||
'qwen-image-edit' family rather than the generic 'qwen-image' one."""
|
||||
best: Optional[tuple[DiffusionFamily, int]] = None
|
||||
for fam in _FAMILIES:
|
||||
for token in (fam.name, *fam.aliases):
|
||||
if token in needle and (best is None or len(token) > best[1]):
|
||||
best = (fam, len(token))
|
||||
return best[0] if best else None
|
||||
|
||||
|
||||
def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[DiffusionFamily]:
|
||||
"""Resolve a ``DiffusionFamily`` from a repo id, or an explicit override.
|
||||
|
||||
``override`` matches a family ``name`` or alias exactly; otherwise the repo
|
||||
id is scanned for the first family whose name/alias appears in it. Image
|
||||
editing checkpoints are rejected (None) since this backend is text-to-image.
|
||||
``override`` matches a family ``name`` or alias exactly. Otherwise the most-specific
|
||||
family whose name/alias is a substring of the repo id wins. Supported editing families
|
||||
(Qwen-Image-Edit) match here; unsupported editing/inpaint/layered checkpoints that only
|
||||
share a base family's arch keyword are still rejected (None), because they need a
|
||||
different pipeline + input this backend's base text-to-image path doesn't drive.
|
||||
"""
|
||||
if override:
|
||||
key = override.strip().lower()
|
||||
|
|
@ -158,11 +248,18 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
|
|||
return fam
|
||||
return None
|
||||
needle = repo_id.lower()
|
||||
if any(kw in needle for kw in _EDIT_KEYWORDS):
|
||||
return None
|
||||
for fam in _FAMILIES:
|
||||
if fam.name in needle or any(alias in needle for alias in fam.aliases):
|
||||
return fam
|
||||
match = _best_family_match(needle)
|
||||
if match is not None:
|
||||
# Don't let a generic base family (e.g. qwen-image) swallow a variant it can't run
|
||||
# (qwen-image-LAYERED, ...-Inpaint): if the id still carries a reject keyword the
|
||||
# matched family does not itself declare, reject so the load fails fast + clearly.
|
||||
matched_tokens = (match.name, *match.aliases)
|
||||
if any(
|
||||
kw in needle and not any(kw in tok for tok in matched_tokens)
|
||||
for kw in _EDIT_KEYWORDS
|
||||
):
|
||||
return None
|
||||
return match
|
||||
return None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -262,6 +262,17 @@ def estimate_gguf_dense_mib(storage_mib: Optional[int], quant: Optional[str]) ->
|
|||
return int(storage_mib * 4.0) # unknown: assume 4-bit-ish
|
||||
|
||||
|
||||
def estimate_safetensors_dense_mib(storage_mib: Optional[int]) -> Optional[int]:
|
||||
"""Resident size of a safetensors checkpoint, in MiB.
|
||||
|
||||
Unlike a GGUF (which is dequantised to bf16/fp16 on load, so a 4-bit file
|
||||
expands ~4x), a safetensors checkpoint loads near its on-disk size: a dense
|
||||
bf16 file is already bf16, and a bnb-4bit / fp8 file stays compressed in VRAM.
|
||||
So the on-disk size is the estimate, returned unchanged (None passes through).
|
||||
"""
|
||||
return storage_mib
|
||||
|
||||
|
||||
def estimate_image_runtime_mib(
|
||||
*,
|
||||
width: Optional[int],
|
||||
|
|
|
|||
|
|
@ -232,6 +232,9 @@ class SdCppDiffusionBackend:
|
|||
attention_backend: Optional[str] = None,
|
||||
transformer_cache: Optional[str] = None,
|
||||
transformer_cache_threshold: Optional[float] = None,
|
||||
# Accepted for a uniform engine interface; the native engine is GGUF-only, so a
|
||||
# non-GGUF kind never routes here (the router forces diffusers for those).
|
||||
model_kind: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Validate, then fetch assets on a daemon thread. Returns at once."""
|
||||
# An empty / whitespace token is "no token": passing "" verbatim to HfApi /
|
||||
|
|
@ -454,11 +457,28 @@ class SdCppDiffusionBackend:
|
|||
guidance: float = 0.0,
|
||||
seed: Optional[int] = None,
|
||||
batch_size: int = 1,
|
||||
# Accepted for a uniform engine interface. The native engine is text-to-image
|
||||
# only for now (sd-cli's init-img/mask plumbing is not wired), so an image-
|
||||
# conditioned request is rejected clearly rather than silently dropping the input.
|
||||
init_image: Optional[str] = None,
|
||||
mask_image: Optional[str] = None,
|
||||
strength: Optional[float] = None,
|
||||
# Accepted for the uniform engine interface; upscale needs an init image, so the
|
||||
# init_image guard below rejects it on the native engine like img2img/inpaint.
|
||||
upscale: Optional[float] = None,
|
||||
# Reference workflow is GPU/diffusers-only (FLUX.2); accepted for interface parity.
|
||||
reference_images: Optional[list[str]] = None,
|
||||
) -> dict[str, Any]:
|
||||
import tempfile
|
||||
|
||||
from PIL import Image
|
||||
|
||||
if init_image is not None or mask_image is not None or reference_images:
|
||||
raise ValueError(
|
||||
"img2img / inpaint / reference are not yet supported on the native sd.cpp "
|
||||
"engine; run on a GPU (diffusers) for image-conditioned workflows."
|
||||
)
|
||||
|
||||
cancel = threading.Event()
|
||||
with self._generate_lock:
|
||||
with self._lock:
|
||||
|
|
|
|||
|
|
@ -1689,9 +1689,19 @@ class AnthropicMessagesResponse(BaseModel):
|
|||
class DiffusionLoadRequest(BaseModel):
|
||||
"""Request to load a local diffusion (text-to-image) checkpoint."""
|
||||
|
||||
model_path: str = Field(..., description = "Diffusion GGUF repo id or local path")
|
||||
gguf_filename: str = Field(
|
||||
..., description = "The chosen single-file GGUF quant inside model_path"
|
||||
model_path: str = Field(..., description = "Diffusion repo id or local path")
|
||||
gguf_filename: Optional[str] = Field(
|
||||
None,
|
||||
description = "The chosen single-file checkpoint (GGUF or safetensors) inside "
|
||||
"model_path. Required for the gguf / single_file kinds; omit for a full pipeline.",
|
||||
)
|
||||
model_kind: Optional[Literal["gguf", "single_file", "pipeline"]] = Field(
|
||||
None,
|
||||
description = "How to load the model (null = auto-detect from gguf_filename): gguf "
|
||||
"(single-file GGUF transformer, dequantised on-device), single_file (single-file "
|
||||
"safetensors transformer, e.g. fp8), or pipeline (a full diffusers repo via "
|
||||
"from_pretrained, embedded quant auto-applied). Non-GGUF kinds are restricted to "
|
||||
"unsloth/* repos (or a local path).",
|
||||
)
|
||||
base_repo: Optional[str] = Field(
|
||||
None, description = "Companion diffusers repo for VAE/text-encoders (default: family base)"
|
||||
|
|
@ -1708,10 +1718,12 @@ class DiffusionLoadRequest(BaseModel):
|
|||
"cut), low_vram (offload every component, lowest VRAM, slower). "
|
||||
"Overrides cpu_offload when set.",
|
||||
)
|
||||
speed_mode: Optional[Literal["off", "default", "max"]] = Field(
|
||||
speed_mode: Optional[Literal["off", "eager", "default", "max"]] = Field(
|
||||
None,
|
||||
description = "Opt-in speed optims (default off -> bit-identical output): "
|
||||
"default (channels_last + regional torch.compile where eligible), "
|
||||
"eager (channels_last + cudnn + attention + fused RMSNorm/AdaLayerNorm patches, "
|
||||
"NO torch.compile -> fast first image, no compile tax), "
|
||||
"default (also regional torch.compile where eligible), "
|
||||
"max (also TF32 + fused QKV).",
|
||||
)
|
||||
text_encoder_quant: Optional[Literal["fp8", "nvfp4"]] = Field(
|
||||
|
|
@ -1807,6 +1819,55 @@ class DiffusionGenerateRequest(BaseModel):
|
|||
batch_size: int = Field(
|
||||
1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)"
|
||||
)
|
||||
# Image-conditioned workflows (base64 or data-URL). An init_image alone runs img2img;
|
||||
# init_image + mask_image runs inpaint. Both require a model family with the matching
|
||||
# pipeline (img2img/inpaint) or the load is rejected with a clear message.
|
||||
# Cap each base64 image string so a single request can't buffer a multi-GB payload (the
|
||||
# decoded dimensions are bounded separately in the backend). ~32 MiB comfortably fits a
|
||||
# full 4096px image yet rejects abuse.
|
||||
init_image: Optional[str] = Field(
|
||||
None,
|
||||
max_length = 32 * 1024 * 1024,
|
||||
description = "Base64/data-URL source image for img2img or inpaint (omit for txt2img)",
|
||||
)
|
||||
mask_image: Optional[str] = Field(
|
||||
None,
|
||||
max_length = 32 * 1024 * 1024,
|
||||
description = "Base64/data-URL mask for inpaint (white = repaint, black = keep). "
|
||||
"Requires init_image.",
|
||||
)
|
||||
strength: Optional[float] = Field(
|
||||
None,
|
||||
ge = 0.0,
|
||||
le = 1.0,
|
||||
description = "img2img/inpaint denoise strength: 0 keeps the source, 1 fully "
|
||||
"redraws it. Ignored for txt2img.",
|
||||
)
|
||||
upscale: Optional[float] = Field(
|
||||
None,
|
||||
ge = 1.0,
|
||||
le = 4.0,
|
||||
description = "Upscale (hires fix) factor for an init_image: enlarges the source "
|
||||
"by this multiple and re-denoises at low strength. Requires init_image; "
|
||||
"ignored for txt2img/inpaint/edit.",
|
||||
)
|
||||
reference_images: Optional[list[str]] = Field(
|
||||
None,
|
||||
max_length = 3,
|
||||
description = "Additional reference images (base64/data-URL) for the FLUX.2 reference "
|
||||
"workflow, combined with init_image. Up to 3; ignored by other workflows.",
|
||||
)
|
||||
|
||||
@field_validator("reference_images")
|
||||
@classmethod
|
||||
def _bounded_reference_items(cls, value: Optional[list[str]]) -> Optional[list[str]]:
|
||||
# Each reference is a base64 image; bound its length like init_image/mask_image so a
|
||||
# request carrying several references can't buffer a multi-GB payload.
|
||||
if value is not None:
|
||||
for item in value:
|
||||
if len(item) > 32 * 1024 * 1024:
|
||||
raise ValueError("each reference image must be at most 32 MiB (base64)")
|
||||
return value
|
||||
|
||||
@field_validator("width", "height")
|
||||
@classmethod
|
||||
|
|
@ -1880,6 +1941,9 @@ class DiffusionStatusResponse(BaseModel):
|
|||
base_repo: Optional[str] = Field(None, description = "Companion diffusers base repo")
|
||||
device: Optional[str] = Field(None, description = "Device the pipeline is on")
|
||||
dtype: Optional[str] = Field(None, description = "Compute dtype")
|
||||
model_kind: Optional[str] = Field(
|
||||
None, description = "Resolved load kind: gguf | single_file | pipeline (gates GGUF-only UI)"
|
||||
)
|
||||
cpu_offload: bool = Field(False, description = "Whether CPU offload is engaged")
|
||||
offload_policy: Optional[str] = Field(
|
||||
None, description = "Resolved offload policy: none | group | model | sequential"
|
||||
|
|
@ -1904,6 +1968,11 @@ class DiffusionStatusResponse(BaseModel):
|
|||
"_native_cudnn), or null for the default SDPA",
|
||||
)
|
||||
transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null")
|
||||
workflows: list[str] = Field(
|
||||
default_factory = list,
|
||||
description = "Image workflows the loaded family supports (drives UI tab gating): "
|
||||
"txt2img, img2img, inpaint. Empty when nothing is loaded or on the native engine.",
|
||||
)
|
||||
engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp")
|
||||
fallback_reason: Optional[str] = Field(
|
||||
None,
|
||||
|
|
|
|||
|
|
@ -10057,7 +10057,7 @@ async def _openai_passthrough_non_streaming(
|
|||
async def load_diffusion_model(
|
||||
request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
|
||||
):
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
from core.inference.diffusion import get_diffusion_backend, resolve_model_kind
|
||||
from core.inference.diffusion_device import resolve_diffusion_device_target
|
||||
from core.inference.diffusion_engine_router import (
|
||||
active_engine_name,
|
||||
|
|
@ -10070,19 +10070,25 @@ async def load_diffusion_model(
|
|||
|
||||
backend = get_diffusion_backend()
|
||||
try:
|
||||
# Resolve the load kind once (gguf / single_file / pipeline) so validation,
|
||||
# engine selection, and the load all agree. A bad explicit kind raises here -> 400.
|
||||
kind = resolve_model_kind(request.gguf_filename, request.model_kind)
|
||||
# Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family,
|
||||
# missing local GGUF) must not evict a working chat model and then 400. The
|
||||
# validated family also drives engine selection below.
|
||||
# missing local GGUF, a non-unsloth non-GGUF repo) must not evict a working chat
|
||||
# model and then 400. The validated family also drives engine selection below.
|
||||
fam = await asyncio.to_thread(
|
||||
backend.validate_load_request,
|
||||
request.model_path,
|
||||
gguf_filename = request.gguf_filename,
|
||||
family_override = request.family_override,
|
||||
model_kind = kind,
|
||||
)
|
||||
# Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU),
|
||||
# installing the sd-cli binary if needed -- all BEFORE evicting chat, so a
|
||||
# native fallback never strands a half-loaded state.
|
||||
engine = await asyncio.to_thread(select_and_activate_engine, fam, hf_token = request.hf_token)
|
||||
# native fallback never strands a half-loaded state. Non-GGUF kinds force diffusers.
|
||||
engine = await asyncio.to_thread(
|
||||
select_and_activate_engine, fam, hf_token = request.hf_token, model_kind = kind
|
||||
)
|
||||
# Take the GPU from the chat backend only when this load will actually use it.
|
||||
# diffusers always does; a *force-native* sd.cpp load on a CUDA/XPU/MPS box does
|
||||
# too. But a native sd.cpp load on a pure-CPU host never touches the GPU, so
|
||||
|
|
@ -10110,6 +10116,7 @@ async def load_diffusion_model(
|
|||
attention_backend = request.attention_backend,
|
||||
transformer_cache = request.transformer_cache,
|
||||
transformer_cache_threshold = request.transformer_cache_threshold,
|
||||
model_kind = kind,
|
||||
)
|
||||
return DiffusionStatusResponse(**annotate_status(status_dict))
|
||||
except (ValueError, FileNotFoundError) as exc:
|
||||
|
|
@ -10138,7 +10145,16 @@ async def generate_diffusion_image(
|
|||
guidance = request.guidance,
|
||||
seed = request.seed,
|
||||
batch_size = request.batch_size,
|
||||
init_image = request.init_image,
|
||||
mask_image = request.mask_image,
|
||||
strength = request.strength,
|
||||
upscale = request.upscale,
|
||||
reference_images = request.reference_images,
|
||||
)
|
||||
except ValueError as exc:
|
||||
# Bad client input (undecodable image/mask, or a workflow the loaded family
|
||||
# doesn't support) — a 400 with the reason, not a generic 500.
|
||||
raise HTTPException(status_code = 400, detail = str(exc))
|
||||
except RuntimeError as exc:
|
||||
# Only "no model loaded" / cancelled are client-state (409). The native
|
||||
# sd.cpp engine also raises RuntimeError for execution failures (nonzero
|
||||
|
|
@ -10146,10 +10162,10 @@ async def generate_diffusion_image(
|
|||
msg = str(exc)
|
||||
if "No diffusion model is loaded" in msg or "cancelled" in msg.lower():
|
||||
raise HTTPException(status_code = 409, detail = msg)
|
||||
logger.error("diffusion.generate_failed: %s", exc)
|
||||
logger.error("diffusion.generate_failed: %s", exc, exc_info = True)
|
||||
raise HTTPException(status_code = 500, detail = "Image generation failed.")
|
||||
except Exception as exc:
|
||||
logger.error("diffusion.generate_failed: %s", exc)
|
||||
logger.error("diffusion.generate_failed: %s", exc, exc_info = True)
|
||||
raise HTTPException(status_code = 500, detail = "Image generation failed.")
|
||||
|
||||
# Persist each image with its full recipe embedded. The diffusers batch shares
|
||||
|
|
|
|||
|
|
@ -50,9 +50,24 @@ def test_detect_family_from_repo_id():
|
|||
# Qwen-Image guides via true_cfg_scale, not guidance_scale.
|
||||
assert detect_family("unsloth/Qwen-Image-2512-GGUF").cfg_kwarg == "true_cfg_scale"
|
||||
assert detect_family("unsloth/Z-Image-GGUF").cfg_kwarg == "guidance_scale"
|
||||
# Image-editing checkpoints are rejected (text-to-image backend only).
|
||||
assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None
|
||||
assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") is None
|
||||
# Qwen-Image-Edit is a SUPPORTED instruction-editing family (its own edit pipeline);
|
||||
# the most-specific match wins so it doesn't fall back to the generic qwen-image.
|
||||
edit = detect_family("unsloth/Qwen-Image-Edit-2511-GGUF")
|
||||
assert edit.name == "qwen-image-edit"
|
||||
assert edit.pipeline_class == "QwenImageEditPlusPipeline"
|
||||
assert edit.edit is True
|
||||
assert detect_family("unsloth/Qwen-Image-Edit-2509-GGUF").name == "qwen-image-edit"
|
||||
# FLUX Kontext is a SUPPORTED editing family (FluxKontextPipeline); the "kontext"
|
||||
# keyword is un-rejected for it, and it must win over the generic "flux.1" match.
|
||||
kontext = detect_family("unsloth/FLUX.1-Kontext-dev-GGUF")
|
||||
assert kontext.name == "flux.1-kontext"
|
||||
assert kontext.pipeline_class == "FluxKontextPipeline"
|
||||
assert kontext.edit is True
|
||||
assert kontext.cfg_kwarg == "guidance_scale"
|
||||
# A plain FLUX.1 checkpoint must still resolve to the base flux.1 family, not kontext.
|
||||
assert detect_family("unsloth/FLUX.1-dev-GGUF").name == "flux.1"
|
||||
# A plain Qwen-Image checkpoint must still resolve to the base family, not edit.
|
||||
assert detect_family("unsloth/Qwen-Image-2512-GGUF").name == "qwen-image"
|
||||
assert detect_family("meta-llama/Llama-3-8B") is None
|
||||
|
||||
|
||||
|
|
@ -194,6 +209,86 @@ class _FakeTransformer:
|
|||
return object()
|
||||
|
||||
|
||||
class _FakeImg2ImgPipe:
|
||||
"""An img2img pipeline call: records the image-conditioned kwargs. Its signature
|
||||
declares image/strength but NOT width/height, mirroring real img2img pipelines
|
||||
(which derive the output size from the input image)."""
|
||||
|
||||
last_kwargs: dict = {}
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
prompt = None,
|
||||
image = None,
|
||||
strength = None,
|
||||
negative_prompt = None,
|
||||
callback_on_step_end = None,
|
||||
guidance_scale = None,
|
||||
true_cfg_scale = None,
|
||||
**kwargs,
|
||||
):
|
||||
_FakeImg2ImgPipe.last_kwargs = {
|
||||
"prompt": prompt,
|
||||
"image": image,
|
||||
"strength": strength,
|
||||
**kwargs,
|
||||
}
|
||||
n = kwargs.get("num_images_per_prompt", 1)
|
||||
return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)])
|
||||
|
||||
|
||||
class _FakeImg2ImgPipeline:
|
||||
built_from: object = None
|
||||
from_pipe_kwargs: dict = {}
|
||||
|
||||
@classmethod
|
||||
def from_pipe(cls, base_pipe, **kwargs):
|
||||
_FakeImg2ImgPipeline.built_from = base_pipe
|
||||
_FakeImg2ImgPipeline.from_pipe_kwargs = kwargs
|
||||
return _FakeImg2ImgPipe()
|
||||
|
||||
|
||||
class _FakeInpaintPipe:
|
||||
"""An inpaint pipeline call: records image + mask_image + strength. Real inpaint
|
||||
pipelines take both an init image and a grayscale mask and derive output size from
|
||||
the input, so width/height are not in its signature."""
|
||||
|
||||
last_kwargs: dict = {}
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
prompt = None,
|
||||
image = None,
|
||||
mask_image = None,
|
||||
strength = None,
|
||||
negative_prompt = None,
|
||||
callback_on_step_end = None,
|
||||
guidance_scale = None,
|
||||
true_cfg_scale = None,
|
||||
**kwargs,
|
||||
):
|
||||
_FakeInpaintPipe.last_kwargs = {
|
||||
"prompt": prompt,
|
||||
"image": image,
|
||||
"mask_image": mask_image,
|
||||
"strength": strength,
|
||||
**kwargs,
|
||||
}
|
||||
n = kwargs.get("num_images_per_prompt", 1)
|
||||
return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)])
|
||||
|
||||
|
||||
class _FakeInpaintPipeline:
|
||||
built_from: object = None
|
||||
|
||||
@classmethod
|
||||
def from_pipe(cls, base_pipe, **kwargs):
|
||||
_FakeInpaintPipeline.built_from = base_pipe
|
||||
return _FakeInpaintPipe()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_runtime(monkeypatch):
|
||||
torch = types.ModuleType("torch")
|
||||
|
|
@ -210,9 +305,15 @@ def fake_runtime(monkeypatch):
|
|||
diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype)
|
||||
diffusers.ZImagePipeline = _FakePipeline
|
||||
diffusers.ZImageTransformer2DModel = _FakeTransformer
|
||||
diffusers.ZImageImg2ImgPipeline = _FakeImg2ImgPipeline
|
||||
diffusers.ZImageInpaintPipeline = _FakeInpaintPipeline
|
||||
# Qwen-Image too, so the true_cfg_scale cfg-kwarg path is exercisable.
|
||||
diffusers.QwenImagePipeline = _FakePipeline
|
||||
diffusers.QwenImageTransformer2DModel = _FakeTransformer
|
||||
diffusers.QwenImageImg2ImgPipeline = _FakeImg2ImgPipeline
|
||||
diffusers.QwenImageInpaintPipeline = _FakeInpaintPipeline
|
||||
# Instruction-editing pipeline (Qwen-Image-Edit): its own pipeline IS the loaded one.
|
||||
diffusers.QwenImageEditPlusPipeline = _FakePipeline
|
||||
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
|
||||
|
|
@ -221,6 +322,10 @@ def fake_runtime(monkeypatch):
|
|||
monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: None)
|
||||
_FakePipeline.last = {}
|
||||
_FakeTransformer.last = {}
|
||||
_FakeImg2ImgPipeline.built_from = None
|
||||
_FakeImg2ImgPipe.last_kwargs = {}
|
||||
_FakeInpaintPipeline.built_from = None
|
||||
_FakeInpaintPipe.last_kwargs = {}
|
||||
yield
|
||||
|
||||
|
||||
|
|
@ -273,6 +378,453 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path):
|
|||
assert backend.is_loaded is False
|
||||
|
||||
|
||||
def _tiny_png_b64() -> str:
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image
|
||||
|
||||
buf = io.BytesIO()
|
||||
Image.new("RGB", (64, 64), (120, 30, 30)).save(buf, format = "PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def test_generate_img2img_uses_from_pipe(fake_runtime, tmp_path):
|
||||
"""An init_image routes generate() through the family's img2img pipeline, built via
|
||||
Pipeline.from_pipe around the loaded pipe (no reload), with image + strength passed
|
||||
and width/height dropped (the img2img pipe derives size from the input image)."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
# The loaded family advertises the image-conditioned workflows for UI gating
|
||||
# (upscale rides the img2img pipeline, so it appears whenever img2img does).
|
||||
assert backend.status()["workflows"] == ["txt2img", "img2img", "upscale", "inpaint", "outpaint"]
|
||||
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a car at sunset", steps = 4, guidance = 0.0, seed = 3,
|
||||
init_image = _tiny_png_b64(), strength = 0.5,
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
# from_pipe was handed the loaded text-to-image pipe (component reuse, no reload).
|
||||
assert _FakeImg2ImgPipeline.built_from is loaded_pipe
|
||||
# ...and with torch_dtype=None so from_pipe SKIPS its default float32 recast, which
|
||||
# both upcasts the reused bf16 modules and crashes on torchao-quantized weights.
|
||||
assert _FakeImg2ImgPipeline.from_pipe_kwargs.get("torch_dtype", "MISSING") is None
|
||||
call = _FakeImg2ImgPipe.last_kwargs
|
||||
assert call["image"] is not None # decoded source image passed through
|
||||
assert call["strength"] == 0.5
|
||||
assert "width" not in call and "height" not in call # img2img derives size from image
|
||||
|
||||
# A txt2img call after it still uses the base pipe (no image kwarg).
|
||||
backend.generate(prompt = "plain", steps = 4, seed = 1)
|
||||
assert backend._state.pipe.last_kwargs.get("image") is None
|
||||
|
||||
|
||||
def test_generate_img2img_unsupported_family_raises(fake_runtime, tmp_path, monkeypatch):
|
||||
"""A family with no image-conditioning at all (no img2img/inpaint/edit/reference) rejects
|
||||
an init_image with a clear error rather than failing deep in the pipeline."""
|
||||
from core.inference.diffusion_families import DiffusionFamily
|
||||
|
||||
# A synthetic txt2img-only family: no img2img/inpaint pipeline, not edit, not reference.
|
||||
# (Every shipped family now supports some image workflow, so build one for this case.)
|
||||
plain = DiffusionFamily(
|
||||
name = "plain-test",
|
||||
pipeline_class = "ZImagePipeline",
|
||||
transformer_class = "ZImageTransformer2DModel",
|
||||
base_repo = "base/repo",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion.detect_family", lambda repo_id, override = None: plain
|
||||
)
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo")
|
||||
assert backend.status()["workflows"] == ["txt2img"]
|
||||
with pytest.raises(ValueError, match = "img2img"):
|
||||
backend.generate(prompt = "x", steps = 4, init_image = _tiny_png_b64())
|
||||
|
||||
|
||||
def test_generate_upscale_enlarges_and_low_strength(fake_runtime, tmp_path):
|
||||
"""An init_image + upscale factor routes generate() through the family's img2img
|
||||
pipeline (hires fix): the source is enlarged to size*factor (rounded to /16) before the
|
||||
denoise, the strength defaults low, and the factor is capped so a huge value can't OOM."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
# Upscale rides the img2img pipeline, so it is advertised alongside img2img.
|
||||
assert "upscale" in backend.status()["workflows"]
|
||||
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a crisp photo", steps = 4, guidance = 0.0, seed = 3,
|
||||
init_image = _tiny_png_b64(), upscale = 2.0, # 64 -> 128, no explicit strength
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
# Reuses the resident modules via from_pipe (no reload, no extra VRAM).
|
||||
assert _FakeImg2ImgPipeline.built_from is loaded_pipe
|
||||
call = _FakeImg2ImgPipe.last_kwargs
|
||||
# The image handed to the pipe is the ENLARGED source (64 * 2 = 128, already /16).
|
||||
assert call["image"].size == (128, 128)
|
||||
# Strength defaults to the hires-fix value when the caller sends none.
|
||||
assert call["strength"] == 0.35
|
||||
|
||||
# The factor is capped at 4x so a large request can't blow up the VAE/transformer.
|
||||
backend.generate(
|
||||
prompt = "x", steps = 4, seed = 1, init_image = _tiny_png_b64(), upscale = 99.0,
|
||||
)
|
||||
assert _FakeImg2ImgPipe.last_kwargs["image"].size == (256, 256) # 64 * 4 (capped)
|
||||
|
||||
# An explicit strength overrides the hires-fix default.
|
||||
backend.generate(
|
||||
prompt = "x", steps = 4, seed = 1, init_image = _tiny_png_b64(),
|
||||
upscale = 1.5, strength = 0.2,
|
||||
)
|
||||
assert _FakeImg2ImgPipe.last_kwargs["strength"] == 0.2
|
||||
# 64 * 1.5 = 96, already a multiple of 16.
|
||||
assert _FakeImg2ImgPipe.last_kwargs["image"].size == (96, 96)
|
||||
|
||||
|
||||
def _png_b64(side: int) -> str:
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image
|
||||
|
||||
buf = io.BytesIO()
|
||||
Image.new("RGB", (side, side), (10, 20, 30)).save(buf, format = "PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def test_decode_image_rejects_oversized(fake_runtime, tmp_path):
|
||||
"""An input image larger than the per-side cap is rejected with a clear error (protects
|
||||
img2img / inpaint / reference from decompression-bomb / OOM inputs), not a 500."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
with pytest.raises(ValueError, match = "too large"):
|
||||
backend.generate(prompt = "x", steps = 4, init_image = _png_b64(4112)) # > 4096/side
|
||||
|
||||
|
||||
def test_upscale_output_is_capped(fake_runtime, tmp_path):
|
||||
"""Upscale bounds the absolute output side to 2048 even when input*factor exceeds it, so a
|
||||
large upload at 4x can't OOM the VAE/transformer."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
backend.generate(prompt = "x", steps = 4, seed = 1, init_image = _png_b64(1024), upscale = 4.0)
|
||||
# 1024 * 4 = 4096 -> clamped to 2048 (longest side), still a multiple of 16.
|
||||
assert _FakeImg2ImgPipe.last_kwargs["image"].size == (2048, 2048)
|
||||
|
||||
|
||||
def _mask_b64(side: int) -> str:
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image
|
||||
|
||||
buf = io.BytesIO()
|
||||
img = Image.new("L", (side, side), 0)
|
||||
for y in range(side // 4, 3 * side // 4):
|
||||
for x in range(side // 4, 3 * side // 4):
|
||||
img.putpixel((x, y), 255)
|
||||
img.save(buf, format = "PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def test_img2img_snaps_non_multiple_of_16(fake_runtime, tmp_path):
|
||||
"""An odd-sized img2img upload (not divisible by 16) is auto-resized to the nearest
|
||||
multiple of 16 so the pipeline's divisibility check passes instead of erroring."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
backend.generate(prompt = "x", steps = 4, seed = 1, init_image = _png_b64(186), strength = 0.5)
|
||||
# 186 / 16 = 11.625 -> round to 12 -> 192.
|
||||
assert _FakeImg2ImgPipe.last_kwargs["image"].size == (192, 192)
|
||||
|
||||
|
||||
def test_inpaint_snaps_image_and_mask_together(fake_runtime, tmp_path):
|
||||
"""Inpaint snaps the odd-sized input to /16 AND resizes the mask to match, so the image
|
||||
and mask stay aligned (a mismatch would crash the inpaint pipeline)."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
backend.generate(
|
||||
prompt = "x", steps = 4, seed = 1,
|
||||
init_image = _png_b64(186), mask_image = _mask_b64(186), strength = 0.5,
|
||||
)
|
||||
assert _FakeInpaintPipe.last_kwargs["image"].size == (192, 192)
|
||||
assert _FakeInpaintPipe.last_kwargs["mask_image"].size == (192, 192)
|
||||
|
||||
|
||||
def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_path):
|
||||
"""A reference family (FLUX.2-klein) advertises txt2img + reference, and a generate with
|
||||
an init_image passes it as the loaded pipe's `image` arg (no from_pipe, no strength) while
|
||||
the output size stays the REQUESTED slider size (the pipe resizes the reference itself)."""
|
||||
import diffusers
|
||||
|
||||
diffusers.Flux2KleinPipeline = _FakePipeline
|
||||
diffusers.Flux2KleinInpaintPipeline = _FakeInpaintPipeline
|
||||
diffusers.Flux2Transformer2DModel = _FakeTransformer
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo",
|
||||
family_override = "flux.2-klein",
|
||||
)
|
||||
# FLUX.2-klein: txt2img + reference (own pipe) + inpaint (dedicated pipe). No img2img class,
|
||||
# so no img2img/upscale.
|
||||
assert backend.status()["workflows"] == ["txt2img", "reference", "inpaint"]
|
||||
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a portrait in this style", steps = 6, guidance = 4.0, seed = 5,
|
||||
width = 768, height = 512, init_image = _tiny_png_b64(), strength = 0.5,
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
call = loaded_pipe.last_kwargs
|
||||
assert call["image"] is not None # reference handed to the loaded pipe
|
||||
assert call["width"] == 768 and call["height"] == 512 # OUTPUT size = sliders, not input
|
||||
assert "strength" not in call # reference conditioning has no strength
|
||||
assert "mask_image" not in call
|
||||
# Guidance flows via guidance_scale (FLUX.2 default behaviour).
|
||||
assert call["guidance_scale"] == 4.0
|
||||
|
||||
# Multi-reference: extra reference_images are combined with init_image into a LIST so the
|
||||
# model can blend several references (subject + style).
|
||||
backend.generate(
|
||||
prompt = "combine these", steps = 6, seed = 9, width = 1024, height = 1024,
|
||||
init_image = _tiny_png_b64(), reference_images = [_tiny_png_b64(), _tiny_png_b64()],
|
||||
)
|
||||
img_arg = loaded_pipe.last_kwargs["image"]
|
||||
assert isinstance(img_arg, list) and len(img_arg) == 3 # primary + 2 extras
|
||||
|
||||
# Branch ordering: an init image + MASK on a reference family must route to inpaint (the
|
||||
# dedicated pipeline), NOT be swallowed by the reference branch (which ignores the mask).
|
||||
backend.generate(
|
||||
prompt = "repaint here", steps = 6, seed = 2,
|
||||
init_image = _tiny_png_b64(), mask_image = _tiny_mask_b64(), strength = 0.8,
|
||||
)
|
||||
assert _FakeInpaintPipeline.built_from is loaded_pipe # built via from_pipe off the load
|
||||
assert _FakeInpaintPipe.last_kwargs["mask_image"] is not None
|
||||
assert _FakeInpaintPipe.last_kwargs["strength"] == 0.8
|
||||
|
||||
# Without an init image the same family does plain txt2img (no image arg).
|
||||
backend.generate(prompt = "just text", steps = 6, seed = 1)
|
||||
assert backend._state.pipe.last_kwargs.get("image") is None
|
||||
|
||||
|
||||
def _tiny_mask_b64() -> str:
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image
|
||||
|
||||
buf = io.BytesIO()
|
||||
# A grayscale mask: white square (repaint) on black (keep).
|
||||
img = Image.new("L", (64, 64), 0)
|
||||
for y in range(16, 48):
|
||||
for x in range(16, 48):
|
||||
img.putpixel((x, y), 255)
|
||||
img.save(buf, format = "PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def test_generate_inpaint_uses_from_pipe(fake_runtime, tmp_path):
|
||||
"""An init_image + mask_image routes generate() through the family's inpaint pipeline,
|
||||
built via Pipeline.from_pipe around the loaded pipe (no reload), with the decoded image
|
||||
+ mask + strength passed through and width/height dropped (size derives from the input)."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
loaded_pipe = backend._state.pipe
|
||||
out = backend.generate(
|
||||
prompt = "a red door", steps = 4, guidance = 0.0, seed = 5,
|
||||
init_image = _tiny_png_b64(), mask_image = _tiny_mask_b64(), strength = 0.7,
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
# The inpaint pipe (not img2img) was selected and built from the loaded pipe.
|
||||
assert _FakeInpaintPipeline.built_from is loaded_pipe
|
||||
assert _FakeImg2ImgPipeline.built_from is None
|
||||
call = _FakeInpaintPipe.last_kwargs
|
||||
assert call["image"] is not None and call["mask_image"] is not None
|
||||
assert call["strength"] == 0.7
|
||||
assert "width" not in call and "height" not in call # inpaint derives size from image
|
||||
|
||||
|
||||
def test_image_conditioned_passes_image_size_not_slider(fake_runtime, tmp_path):
|
||||
"""When the workflow pipe DOES accept width/height, an image-conditioned call must pass
|
||||
the INPUT IMAGE's size, never the txt2img slider size -- otherwise a non-slider-sized
|
||||
input (e.g. a 1536px outpaint canvas with a 1024 slider) mismatches the latents
|
||||
("tensor a (128) must match tensor b (192)"). Covers Transform + Extend with any size."""
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image
|
||||
|
||||
class _SizePipe:
|
||||
last: dict = {}
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
prompt = None,
|
||||
image = None,
|
||||
strength = None,
|
||||
width = None,
|
||||
height = None,
|
||||
negative_prompt = None,
|
||||
callback_on_step_end = None,
|
||||
guidance_scale = None,
|
||||
true_cfg_scale = None,
|
||||
**kwargs,
|
||||
):
|
||||
_SizePipe.last = {"width": width, "height": height}
|
||||
n = kwargs.get("num_images_per_prompt", 1)
|
||||
return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)])
|
||||
|
||||
class _SizePipeline:
|
||||
@classmethod
|
||||
def from_pipe(cls, base_pipe, **kwargs):
|
||||
return _SizePipe()
|
||||
|
||||
import diffusers
|
||||
|
||||
diffusers.ZImageImg2ImgPipeline = _SizePipeline
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
|
||||
)
|
||||
buf = io.BytesIO()
|
||||
Image.new("RGB", (96, 64), (10, 20, 30)).save(buf, format = "PNG") # non-square, non-slider
|
||||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||||
backend.generate(prompt = "x", steps = 4, width = 1024, height = 1024, init_image = b64, strength = 0.5)
|
||||
# The pipe got the IMAGE's 96x64, not the 1024x1024 slider.
|
||||
assert _SizePipe.last == {"width": 96, "height": 64}
|
||||
|
||||
|
||||
def test_edit_family_uses_own_pipeline_and_requires_image(fake_runtime, tmp_path):
|
||||
"""An instruction-editing family (Qwen-Image-Edit) exposes only the 'edit' workflow,
|
||||
runs the image through its OWN loaded pipeline (no from_pipe), and rejects a call with
|
||||
no input image."""
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
str(tmp_path), gguf_filename = "model.gguf", base_repo = "Qwen/Qwen-Image-Edit-2511",
|
||||
family_override = "qwen-image-edit",
|
||||
)
|
||||
# Edit families advertise only the edit workflow (no txt2img / img2img / inpaint).
|
||||
assert backend.status()["workflows"] == ["edit"]
|
||||
loaded_pipe = backend._state.pipe
|
||||
|
||||
out = backend.generate(
|
||||
prompt = "make it night", steps = 8, guidance = 4.0, seed = 1,
|
||||
init_image = _tiny_png_b64(),
|
||||
)
|
||||
assert len(out["images"]) == 1
|
||||
# The loaded pipe handled it directly -- no from_pipe img2img/inpaint was built.
|
||||
assert backend._state.pipe is loaded_pipe
|
||||
assert _FakeImg2ImgPipeline.built_from is None and _FakeInpaintPipeline.built_from is None
|
||||
assert loaded_pipe.last_kwargs.get("image") is not None
|
||||
|
||||
# An edit model with no input image fails fast with a clear message.
|
||||
with pytest.raises(ValueError, match = "image"):
|
||||
backend.generate(prompt = "make it night", steps = 8)
|
||||
|
||||
|
||||
def test_load_pipeline_kind_uses_from_pretrained(fake_runtime):
|
||||
"""A full-pipeline (no single-file) load on an unsloth/* repo builds the pipe with
|
||||
pipeline_cls.from_pretrained(repo_id) -- NO single-file transformer build, NO GGUF
|
||||
quant config -- so an embedded bnb-4bit config is reloaded by diffusers itself."""
|
||||
backend = DiffusionBackend()
|
||||
status = backend.load_pipeline(
|
||||
"unsloth/Z-Image-Turbo-unsloth-bnb-4bit", family_override = "z-image"
|
||||
)
|
||||
assert status["loaded"] is True
|
||||
assert status["family"] == "z-image"
|
||||
# from_pretrained pointed at the repo itself (it IS its own base), with no transformer.
|
||||
assert _FakePipeline.last["base"] == "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"
|
||||
assert "transformer" not in _FakePipeline.last
|
||||
# The GGUF single-file build path was never taken.
|
||||
assert _FakeTransformer.last == {}
|
||||
|
||||
|
||||
def test_load_single_file_safetensors_no_gguf_config(fake_runtime, tmp_path):
|
||||
"""A single-file *.safetensors transformer is built with from_single_file WITHOUT the
|
||||
GGUF dequant config (it carries its own dtype), then assembled from the base repo."""
|
||||
(tmp_path / "model.safetensors").write_bytes(b"weights")
|
||||
backend = DiffusionBackend()
|
||||
status = backend.load_pipeline(
|
||||
str(tmp_path),
|
||||
gguf_filename = "model.safetensors",
|
||||
base_repo = "base/repo",
|
||||
family_override = "qwen-image",
|
||||
)
|
||||
assert status["loaded"] is True
|
||||
assert _FakeTransformer.last["path"] == str((tmp_path / "model.safetensors").resolve())
|
||||
assert _FakeTransformer.last["subfolder"] == "transformer"
|
||||
# No GGUF quant config on the safetensors path (the GGUF path sets one).
|
||||
assert "quantization_config" not in _FakeTransformer.last
|
||||
assert _FakePipeline.last["base"] == "base/repo"
|
||||
assert "transformer" in _FakePipeline.last
|
||||
|
||||
|
||||
def test_load_pipeline_rejects_non_unsloth_repo(fake_runtime):
|
||||
backend = DiffusionBackend()
|
||||
with pytest.raises(ValueError, match = "unsloth"):
|
||||
backend.load_pipeline("randomorg/Z-Image-bnb-4bit", family_override = "z-image")
|
||||
|
||||
|
||||
def test_detect_family_rejects_layered():
|
||||
# Qwen-Image-Layered needs a dedicated pipeline (additional_t_cond); it must be
|
||||
# rejected so it fails fast at load instead of crashing at the first denoise step.
|
||||
assert detect_family("unsloth/Qwen-Image-Layered-GGUF") is None
|
||||
assert detect_family("unsloth/qwen_image_layered") is None
|
||||
|
||||
|
||||
def test_failed_load_rolls_back_eager_patches(fake_runtime, tmp_path, monkeypatch):
|
||||
"""A load failure AFTER the eager patches install but BEFORE the _LoadState commit must
|
||||
roll the process-wide patches back, so the next bit-identical `off` load is not
|
||||
contaminated (the asymmetric-cleanup bug the reviewers flagged)."""
|
||||
from core.inference import diffusion as diff_mod
|
||||
from core.inference import diffusion_eager_patches as ep
|
||||
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
ep.uninstall_patches() # clean slate
|
||||
|
||||
def _boom(*_a, **_k):
|
||||
raise RuntimeError("placement boom")
|
||||
|
||||
# apply_memory_plan runs AFTER the patches are installed, before _LoadState commits.
|
||||
monkeypatch.setattr(diff_mod, "apply_memory_plan", _boom)
|
||||
backend = DiffusionBackend()
|
||||
with pytest.raises(RuntimeError):
|
||||
backend.load_pipeline(
|
||||
str(tmp_path),
|
||||
gguf_filename = "model.gguf",
|
||||
family_override = "z-image",
|
||||
base_repo = "base/repo",
|
||||
speed_mode = "eager", # != off -> installs the shared patches
|
||||
)
|
||||
assert ep.is_installed() is False # rolled back by the load-failure finally
|
||||
assert backend.is_loaded is False
|
||||
|
||||
|
||||
def test_cpu_offload_ignored_off_cuda(fake_runtime, tmp_path):
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
|
|
@ -319,8 +871,10 @@ def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch)
|
|||
|
||||
def test_load_without_gguf_raises():
|
||||
backend = DiffusionBackend()
|
||||
with pytest.raises(ValueError):
|
||||
backend.load_pipeline("unsloth/Z-Image-Turbo-GGUF") # no gguf_filename
|
||||
# No gguf_filename -> a full-pipeline load, gated to unsloth/*; a non-unsloth repo
|
||||
# is rejected before any GPU/network work.
|
||||
with pytest.raises(ValueError, match = "unsloth"):
|
||||
backend.load_pipeline("some-org/Z-Image-bnb-4bit")
|
||||
|
||||
|
||||
def test_load_unknown_family_raises():
|
||||
|
|
@ -685,8 +1239,24 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime):
|
|||
|
||||
def test_validate_load_request(tmp_path):
|
||||
backend = DiffusionBackend()
|
||||
with pytest.raises(ValueError, match = "gguf_filename"):
|
||||
backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF")
|
||||
# No filename + unsloth repo -> a full-pipeline load (allowed for unsloth/*).
|
||||
assert (
|
||||
backend.validate_load_request("unsloth/Z-Image-Turbo-unsloth-bnb-4bit").name == "z-image"
|
||||
)
|
||||
# No filename + non-unsloth repo -> a pipeline load, gated to unsloth/* -> rejected.
|
||||
with pytest.raises(ValueError, match = "unsloth"):
|
||||
backend.validate_load_request("some-org/Z-Image-bnb-4bit")
|
||||
# An explicit gguf/single_file kind still requires a single-file name.
|
||||
with pytest.raises(ValueError, match = "single-file"):
|
||||
backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", model_kind = "gguf")
|
||||
# A pipeline kind must NOT carry a single-file name.
|
||||
with pytest.raises(ValueError, match = "pipeline"):
|
||||
backend.validate_load_request(
|
||||
"unsloth/Z-Image-Turbo-bnb-4bit", gguf_filename = "q.gguf", model_kind = "pipeline"
|
||||
)
|
||||
# A single-file safetensors load is also gated to unsloth/* repos.
|
||||
with pytest.raises(ValueError, match = "unsloth"):
|
||||
backend.validate_load_request("some-org/Z-Image", gguf_filename = "model.safetensors")
|
||||
with pytest.raises(ValueError, match = "family"):
|
||||
backend.validate_load_request("meta/Llama-3", gguf_filename = "q.gguf")
|
||||
assert (
|
||||
|
|
|
|||
|
|
@ -35,13 +35,21 @@ class _FakeBackend:
|
|||
*,
|
||||
gguf_filename = None,
|
||||
family_override = None,
|
||||
model_kind = None,
|
||||
):
|
||||
# Mirror the real backend's cheap validation so the route's
|
||||
# validate-before-evict ordering is exercised.
|
||||
from core.inference.diffusion import resolve_model_kind
|
||||
from core.inference.diffusion_families import detect_family
|
||||
|
||||
if not gguf_filename:
|
||||
raise ValueError("gguf_filename is required.")
|
||||
kind = resolve_model_kind(gguf_filename, model_kind)
|
||||
if kind in ("gguf", "single_file") and not gguf_filename:
|
||||
raise ValueError("a single-file checkpoint name is required.")
|
||||
# Non-GGUF loads are gated to unsloth/* (or a local path), like the real backend.
|
||||
if kind != "gguf" and not model_path.lower().startswith("unsloth/"):
|
||||
raise ValueError(
|
||||
f"Non-GGUF diffusion loads are restricted to unsloth/* repos; got '{model_path}'."
|
||||
)
|
||||
fam = detect_family(model_path, family_override)
|
||||
if fam is None:
|
||||
raise ValueError(f"Could not infer a diffusion family for '{model_path}'.")
|
||||
|
|
@ -254,10 +262,24 @@ def test_generate_rejects_non_multiple_of_16(client):
|
|||
assert ok.status_code == 200
|
||||
|
||||
|
||||
def test_load_requires_gguf_filename(client):
|
||||
# gguf_filename is now mandatory — a load without it is a 422.
|
||||
def test_non_gguf_load_restricted_to_unsloth(client):
|
||||
# gguf_filename is optional now; with none, the load is a full-pipeline kind, which
|
||||
# is gated to unsloth/* repos. A non-unsloth repo (no filename) is rejected -> 400.
|
||||
resp = client.post("/api/inference/images/load", json = {"model_path": "x/z-image"})
|
||||
assert resp.status_code == 422
|
||||
assert resp.status_code == 400
|
||||
assert "unsloth" in resp.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_pipeline_load_allowed_for_unsloth_repo(client):
|
||||
# An unsloth/* repo with no filename loads as a full diffusers pipeline (kind auto
|
||||
# = pipeline); the route forwards model_kind="pipeline" to begin_load.
|
||||
resp = client.post(
|
||||
"/api/inference/images/load", json = {"model_path": "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"}
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
backend = diffusion_module.get_diffusion_backend()
|
||||
assert backend.last_load_kwargs["model_kind"] == "pipeline"
|
||||
assert backend.last_load_kwargs.get("gguf_filename") is None
|
||||
|
||||
|
||||
def test_generate_without_load_returns_409(client):
|
||||
|
|
|
|||
|
|
@ -17,7 +17,26 @@ _STUDIO = Path(__file__).resolve().parents[2]
|
|||
if str(_STUDIO) not in sys.path:
|
||||
sys.path.insert(0, str(_STUDIO))
|
||||
|
||||
from install_sd_cpp_prebuilt import default_install_dir, resolve_release_asset # noqa: E402
|
||||
import hashlib # noqa: E402
|
||||
import io # noqa: E402
|
||||
import json # noqa: E402
|
||||
import urllib.error # noqa: E402
|
||||
import zipfile # noqa: E402
|
||||
|
||||
import pytest # noqa: E402
|
||||
|
||||
import install_sd_cpp_prebuilt as sdmod # noqa: E402
|
||||
from install_sd_cpp_prebuilt import ( # noqa: E402
|
||||
DEFAULT_REPO,
|
||||
DEFAULT_TAG,
|
||||
_fetch_release,
|
||||
_pinned_tag,
|
||||
_repo,
|
||||
_verify_sha256,
|
||||
default_install_dir,
|
||||
install,
|
||||
resolve_release_asset,
|
||||
)
|
||||
|
||||
# A real stable-diffusion.cpp latest-release asset list.
|
||||
_ASSETS = [
|
||||
|
|
@ -114,3 +133,127 @@ def test_default_install_dir_is_sibling_of_llama(monkeypatch):
|
|||
d = default_install_dir()
|
||||
assert d.name == "stable-diffusion.cpp"
|
||||
assert d.parent.name == ".unsloth"
|
||||
|
||||
|
||||
# ── version pin + source repo (reproducibility) ─────────────────────────────
|
||||
|
||||
|
||||
def test_pinned_tag_default_and_override(monkeypatch):
|
||||
monkeypatch.delenv("UNSLOTH_SD_CPP_TAG", raising = False)
|
||||
assert _pinned_tag() == DEFAULT_TAG # pinned, not "latest"
|
||||
monkeypatch.setenv("UNSLOTH_SD_CPP_TAG", "master-999-deadbee")
|
||||
assert _pinned_tag() == "master-999-deadbee"
|
||||
monkeypatch.setenv("UNSLOTH_SD_CPP_TAG", "") # explicit empty -> track latest
|
||||
assert _pinned_tag() is None
|
||||
|
||||
|
||||
def test_repo_default_and_mirror_override(monkeypatch):
|
||||
monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False)
|
||||
assert _repo() == DEFAULT_REPO == "leejet/stable-diffusion.cpp"
|
||||
monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", "unslothai/stable-diffusion.cpp")
|
||||
assert _repo() == "unslothai/stable-diffusion.cpp"
|
||||
|
||||
|
||||
# ── sha256 integrity check ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_verify_sha256_accepts_matching_digest(tmp_path):
|
||||
f = tmp_path / "asset.zip"
|
||||
f.write_bytes(b"hello sd-cli")
|
||||
digest = "sha256:" + hashlib.sha256(b"hello sd-cli").hexdigest()
|
||||
_verify_sha256(f, digest) # no raise
|
||||
|
||||
|
||||
def test_verify_sha256_rejects_mismatch(tmp_path):
|
||||
f = tmp_path / "asset.zip"
|
||||
f.write_bytes(b"tampered")
|
||||
bad = "sha256:" + hashlib.sha256(b"original").hexdigest()
|
||||
with pytest.raises(RuntimeError, match = "sha256 mismatch"):
|
||||
_verify_sha256(f, bad)
|
||||
|
||||
|
||||
def test_verify_sha256_skips_when_absent_or_unknown(tmp_path):
|
||||
f = tmp_path / "asset.zip"
|
||||
f.write_bytes(b"x")
|
||||
_verify_sha256(f, None) # no digest published -> warn + proceed (no raise)
|
||||
_verify_sha256(f, "md5:abc") # unrecognised algo -> skip (no raise)
|
||||
|
||||
|
||||
# ── _fetch_release: pinned-tag 404 -> latest fallback ───────────────────────
|
||||
|
||||
|
||||
def test_fetch_release_falls_back_to_latest_on_404(monkeypatch):
|
||||
calls: list[str] = []
|
||||
|
||||
class _Resp:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def read(self):
|
||||
return json.dumps({"tag_name": "latest-xyz", "assets": []}).encode()
|
||||
|
||||
def fake_urlopen(req, timeout = 30.0):
|
||||
url = getattr(req, "full_url", req)
|
||||
calls.append(url)
|
||||
if "/tags/" in url:
|
||||
raise urllib.error.HTTPError(url, 404, "not found", None, None)
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(sdmod.urllib.request, "urlopen", fake_urlopen)
|
||||
rel = _fetch_release("gone-tag", repo = "leejet/stable-diffusion.cpp")
|
||||
assert rel["tag_name"] == "latest-xyz"
|
||||
assert any("/tags/gone-tag" in c for c in calls) and any(c.endswith("/latest") for c in calls)
|
||||
|
||||
|
||||
def test_fetch_release_propagates_non_404(monkeypatch):
|
||||
def fake_urlopen(req, timeout = 30.0):
|
||||
url = getattr(req, "full_url", req)
|
||||
raise urllib.error.HTTPError(url, 403, "rate limited", None, None)
|
||||
|
||||
monkeypatch.setattr(sdmod.urllib.request, "urlopen", fake_urlopen)
|
||||
with pytest.raises(urllib.error.HTTPError):
|
||||
_fetch_release("any-tag")
|
||||
|
||||
|
||||
# ── install(): download -> verify -> extract -> locate (offline) ────────────
|
||||
|
||||
|
||||
def _zip_with_sd_cli() -> bytes:
|
||||
buf = io.BytesIO()
|
||||
with zipfile.ZipFile(buf, "w") as zf:
|
||||
zf.writestr("build/bin/sd-cli", b"#!/bin/sh\necho sd-cli\n")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _stub_release(monkeypatch, *, zip_bytes: bytes, digest: str):
|
||||
name = "sd-master-deadbee-bin-Linux-Ubuntu-24.04-x86_64.zip"
|
||||
release = {
|
||||
"tag_name": "master-1-deadbee",
|
||||
"assets": [
|
||||
{"name": name, "browser_download_url": f"https://example.invalid/{name}", "digest": digest}
|
||||
],
|
||||
}
|
||||
monkeypatch.setattr(sdmod, "_fetch_release", lambda *a, **k: release)
|
||||
monkeypatch.setattr(sdmod, "_download", lambda url, dest, **k: dest.write_bytes(zip_bytes))
|
||||
monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux")
|
||||
monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64")
|
||||
return name
|
||||
|
||||
|
||||
def test_install_downloads_verifies_extracts(tmp_path, monkeypatch):
|
||||
zb = _zip_with_sd_cli()
|
||||
name = _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest())
|
||||
sd_cli = install(install_dir = tmp_path)
|
||||
assert sd_cli.name == "sd-cli" and sd_cli.is_file()
|
||||
assert not (tmp_path / name).exists() # archive cleaned up after extract
|
||||
|
||||
|
||||
def test_install_sha256_mismatch_raises_and_cleans_up(tmp_path, monkeypatch):
|
||||
zb = _zip_with_sd_cli()
|
||||
name = _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + "0" * 64)
|
||||
with pytest.raises(RuntimeError, match = "sha256 mismatch"):
|
||||
install(install_dir = tmp_path)
|
||||
assert not (tmp_path / name).exists() # the finally: drops the bad archive
|
||||
|
|
|
|||
|
|
@ -989,11 +989,17 @@ export const IMAGE_GEN_TASKS = [
|
|||
// which the text-to-image backend rejects (mirrors its _EDIT_KEYWORDS). Hidden by
|
||||
// id so they don't show in the Images picker only to 400 on load. Keeping the
|
||||
// image-to-image task itself is required: some supported models (FLUX.2-klein)
|
||||
// carry that tag too.
|
||||
const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint"] as const;
|
||||
// carry that tag too. "layered" hides Qwen-Image-Layered, which needs a dedicated
|
||||
// pipeline (additional_t_cond) the standard text-to-image path can't drive.
|
||||
const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint", "layered"] as const;
|
||||
// Editing families the backend now SUPPORTS (their own Edit workflow) -- must not be
|
||||
// hidden even though their id contains an edit keyword. Mirrors the backend's
|
||||
// qwen-image-edit family in diffusion_families.py.
|
||||
const SUPPORTED_EDIT_KEYWORDS = ["qwen-image-edit", "kontext"] as const;
|
||||
function isImageEditModel(repoId: string | null | undefined): boolean {
|
||||
if (!repoId) return false;
|
||||
const id = repoId.toLowerCase();
|
||||
if (SUPPORTED_EDIT_KEYWORDS.some((kw) => id.includes(kw))) return false;
|
||||
return IMAGE_EDIT_KEYWORDS.some((kw) => id.includes(kw));
|
||||
}
|
||||
|
||||
|
|
@ -1676,6 +1682,19 @@ export function HubModelPicker({
|
|||
isChatSupported,
|
||||
]);
|
||||
|
||||
// Curated non-GGUF (safetensors) models for the Images picker. The HF listing +
|
||||
// Recommended gate only surface GGUF on a GPU host (isRecommendableFormat), so a
|
||||
// bnb-4bit / fp8 safetensors model would never appear there. These curated entries
|
||||
// (the non-GGUF ModelOptions passed in) are shown explicitly above the GGUF rows so
|
||||
// the user can pick a full diffusers pipeline. Only the Images picker (task set)
|
||||
// curates them; already-downloaded ones show under Downloaded instead.
|
||||
const curatedSafetensorsRows = useMemo(() => {
|
||||
if (!task) return [];
|
||||
return models.filter(
|
||||
(m) => m.isGguf === false && !downloadedSet.has(m.id.toLowerCase()),
|
||||
);
|
||||
}, [models, task, downloadedSet]);
|
||||
|
||||
// Per-row meta + VRAM badge from the recommended listing's own metadata.
|
||||
const recommendedMeta = useMemo(() => {
|
||||
const map = new Map<
|
||||
|
|
@ -3244,6 +3263,30 @@ export function HubModelPicker({
|
|||
|
||||
{showRecommendedSection ? (
|
||||
<>
|
||||
{/* Curated safetensors models (full diffusers pipelines / single-file
|
||||
fp8). Shown above the GGUF rows; clicking loads directly (no quant
|
||||
expander), the same path as a non-GGUF Recommended row. */}
|
||||
{curatedSafetensorsRows.map((m) => {
|
||||
const optionKey = makeModelOptionKey("curated-safetensors", m.id);
|
||||
return (
|
||||
<div key={m.id}>
|
||||
<ModelRow
|
||||
label={m.id}
|
||||
hideOwner={true}
|
||||
downloaded={downloadedSet.has(m.id.toLowerCase())}
|
||||
capabilities={capsById.get(m.id)}
|
||||
meta={m.description ?? "Safetensors"}
|
||||
selected={value === m.id}
|
||||
optionProps={hubModelList.getOptionProps(
|
||||
optionKey,
|
||||
value === m.id,
|
||||
)}
|
||||
onClick={() => handleModelClick(m.id)}
|
||||
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
{recommendedSearch.isLoading &&
|
||||
recommendedRows.length === 0 ? (
|
||||
<div className="flex items-center gap-2 px-5 py-3">
|
||||
|
|
@ -3252,7 +3295,8 @@ export function HubModelPicker({
|
|||
Loading models…
|
||||
</span>
|
||||
</div>
|
||||
) : recommendedRows.length === 0 ? (
|
||||
) : recommendedRows.length === 0 &&
|
||||
curatedSafetensorsRows.length === 0 ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">
|
||||
No models found.
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -11,7 +11,13 @@ export interface DiffusionStatus {
|
|||
base_repo: string | null;
|
||||
device: string | null;
|
||||
dtype: string | null;
|
||||
// Resolved load kind: "gguf" | "single_file" | "pipeline". Gates GGUF-only controls
|
||||
// (the dense transformer_quant fast path only engages on gguf). Null when not loaded.
|
||||
model_kind?: string | null;
|
||||
cpu_offload: boolean;
|
||||
// Image workflows the loaded family supports (drives tab gating): txt2img, img2img,
|
||||
// inpaint. Absent/empty when nothing is loaded or on the native sd.cpp engine.
|
||||
workflows?: string[];
|
||||
}
|
||||
|
||||
export interface DiffusionGenerateProgress {
|
||||
|
|
@ -32,11 +38,33 @@ export interface DiffusionLoadProgress {
|
|||
|
||||
export interface DiffusionLoadRequest {
|
||||
model_path: string;
|
||||
gguf_filename: string;
|
||||
// Optional now: required for the gguf / single_file kinds, omitted for a full
|
||||
// pipeline (a diffusers repo loaded via from_pretrained).
|
||||
gguf_filename?: string;
|
||||
// How to load the model (omit to auto-detect from gguf_filename): "gguf" (single-file
|
||||
// GGUF transformer), "single_file" (single-file safetensors transformer, e.g. fp8), or
|
||||
// "pipeline" (a full diffusers repo). Non-GGUF kinds are restricted to unsloth/* repos.
|
||||
model_kind?: "gguf" | "single_file" | "pipeline";
|
||||
base_repo?: string;
|
||||
family_override?: string;
|
||||
hf_token?: string;
|
||||
cpu_offload?: boolean;
|
||||
// Advanced (load-time) tuning. All optional; omit for the backend's auto defaults.
|
||||
speed_mode?: "off" | "eager" | "default" | "max";
|
||||
transformer_quant?: "auto" | "int8" | "fp8" | "nvfp4" | "mxfp8";
|
||||
attention_backend?:
|
||||
| "auto"
|
||||
| "native"
|
||||
| "cudnn"
|
||||
| "flash"
|
||||
| "flash2"
|
||||
| "flash3"
|
||||
| "flash4"
|
||||
| "sage"
|
||||
| "xformers"
|
||||
| "aiter";
|
||||
memory_mode?: "auto" | "fast" | "balanced" | "low_vram";
|
||||
transformer_cache?: "off" | "fbcache";
|
||||
}
|
||||
|
||||
export interface DiffusionGenerateRequest {
|
||||
|
|
@ -48,6 +76,16 @@ export interface DiffusionGenerateRequest {
|
|||
guidance?: number;
|
||||
seed?: number;
|
||||
batch_size?: number;
|
||||
// Image-conditioned workflows. init_image alone = img2img; init_image + mask_image =
|
||||
// inpaint. Base64 or data-URL. strength is the denoise amount (0 keeps source, 1 redraws).
|
||||
init_image?: string;
|
||||
mask_image?: string;
|
||||
strength?: number;
|
||||
// Upscale (hires fix): factor > 1 with an init_image enlarges the source and re-denoises
|
||||
// it at low strength. Requires init_image; ignored for txt2img/inpaint/edit.
|
||||
upscale?: number;
|
||||
// Additional reference images for the FLUX.2 reference workflow, combined with init_image.
|
||||
reference_images?: string[];
|
||||
}
|
||||
|
||||
// A persisted image's full generation recipe (also embedded in the PNG).
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -25,18 +25,40 @@ Usage:
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import stat
|
||||
import sys
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Optional, Sequence
|
||||
|
||||
REPO = "leejet/stable-diffusion.cpp"
|
||||
RELEASES_API = f"https://api.github.com/repos/{REPO}/releases/latest"
|
||||
# Default upstream source. Overridable with UNSLOTH_SD_CPP_REPO so a pinned unslothai
|
||||
# mirror (built the same way as unslothai/llama.cpp's prebuilts) can be used without a
|
||||
# code change once it exists; otherwise this falls back to leejet upstream.
|
||||
DEFAULT_REPO = "leejet/stable-diffusion.cpp"
|
||||
# Pinned release tag for REPRODUCIBILITY: "releases/latest" silently swaps the binary
|
||||
# under users on every upstream push. Override with UNSLOTH_SD_CPP_TAG; set it empty to
|
||||
# track latest. If the pinned tag is gone upstream, install falls back to latest.
|
||||
DEFAULT_TAG = "master-737-3b6c9ca"
|
||||
|
||||
# Back-compat alias (some callers/tests import REPO).
|
||||
REPO = DEFAULT_REPO
|
||||
|
||||
|
||||
def _repo() -> str:
|
||||
return (os.environ.get("UNSLOTH_SD_CPP_REPO") or DEFAULT_REPO).strip() or DEFAULT_REPO
|
||||
|
||||
|
||||
def _pinned_tag() -> Optional[str]:
|
||||
"""The release tag to install: env override, else the pinned default; '' = latest."""
|
||||
val = os.environ.get("UNSLOTH_SD_CPP_TAG", DEFAULT_TAG).strip()
|
||||
return val or None
|
||||
|
||||
# accelerator -> the token that must appear in a Linux/Windows asset name.
|
||||
_LINUX_ACCEL_TOKEN = {"rocm": "rocm", "vulkan": "vulkan"}
|
||||
|
|
@ -108,14 +130,59 @@ def resolve_release_asset(
|
|||
return sel[0] if sel else None
|
||||
|
||||
|
||||
def _fetch_latest_release(*, token: Optional[str] = None, timeout: float = 30.0) -> dict:
|
||||
"""GET the latest-release JSON from GitHub (token optional, lifts rate limit)."""
|
||||
req = urllib.request.Request(RELEASES_API, headers = {"Accept": "application/vnd.github+json"})
|
||||
def _fetch_release(
|
||||
tag: Optional[str] = None, *, repo: Optional[str] = None,
|
||||
token: Optional[str] = None, timeout: float = 30.0,
|
||||
) -> dict:
|
||||
"""GET a release JSON from GitHub. With ``tag`` set, fetch that exact release (and fall
|
||||
back to latest if the tag is gone upstream); otherwise fetch latest. ``token`` is
|
||||
optional and lifts the API rate limit."""
|
||||
repo = repo or _repo()
|
||||
token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
|
||||
if token:
|
||||
req.add_header("Authorization", f"Bearer {token}")
|
||||
with urllib.request.urlopen(req, timeout = timeout) as resp: # noqa: S310 (fixed https host)
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
|
||||
def _get(url: str) -> dict:
|
||||
req = urllib.request.Request(url, headers = {"Accept": "application/vnd.github+json"})
|
||||
if token:
|
||||
req.add_header("Authorization", f"Bearer {token}")
|
||||
with urllib.request.urlopen(req, timeout = timeout) as resp: # noqa: S310 (fixed https host)
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
|
||||
base = f"https://api.github.com/repos/{repo}/releases"
|
||||
if tag:
|
||||
try:
|
||||
return _get(f"{base}/tags/{tag}")
|
||||
except urllib.error.HTTPError as exc: # pinned tag removed upstream -> latest
|
||||
if exc.code != 404:
|
||||
raise
|
||||
print(f"sd-cli: pinned tag {tag} not found on {repo}; falling back to latest", flush = True)
|
||||
return _get(f"{base}/latest")
|
||||
|
||||
|
||||
# Back-compat alias: the old name fetched latest.
|
||||
def _fetch_latest_release(*, token: Optional[str] = None, timeout: float = 30.0) -> dict:
|
||||
return _fetch_release(None, token = token, timeout = timeout)
|
||||
|
||||
|
||||
def _verify_sha256(path: Path, expected_digest: Optional[str]) -> None:
|
||||
"""Verify ``path`` against a GitHub asset ``digest`` ('sha256:<hex>'). Integrity check
|
||||
against a corrupted/tampered download before we extract + execute the binary. When the
|
||||
release publishes no digest (older releases), warn and proceed rather than hard-fail."""
|
||||
if not expected_digest:
|
||||
print(f"sd-cli: WARNING no digest for {path.name}; cannot verify integrity", flush = True)
|
||||
return
|
||||
algo, _, want = expected_digest.partition(":")
|
||||
if algo.lower() != "sha256" or not want:
|
||||
print(f"sd-cli: WARNING unrecognised digest {expected_digest!r}; skipping check", flush = True)
|
||||
return
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1 << 20), b""):
|
||||
h.update(chunk)
|
||||
got = h.hexdigest()
|
||||
if got != want.lower():
|
||||
raise RuntimeError(
|
||||
f"sha256 mismatch for {path.name}: expected {want.lower()}, got {got}"
|
||||
)
|
||||
|
||||
|
||||
def default_install_dir() -> Path:
|
||||
|
|
@ -132,6 +199,14 @@ def _make_executable(path: Path) -> None:
|
|||
path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
|
||||
|
||||
|
||||
def _download(url: str, dest: Path, *, timeout: float = 300.0) -> None:
|
||||
"""Stream a release asset to ``dest`` with a timeout. ``urlretrieve`` has no timeout,
|
||||
so a stalled connection would hang the lazy first-load (ensure_sd_cpp_binary) forever.
|
||||
Anonymous, matching the public release URL -- the API fetch carries any token."""
|
||||
with urllib.request.urlopen(url, timeout = timeout) as resp, open(dest, "wb") as f: # noqa: S310
|
||||
shutil.copyfileobj(resp, f)
|
||||
|
||||
|
||||
def _locate_sd_cli(root: Path) -> Optional[Path]:
|
||||
name = "sd-cli.exe" if sys.platform == "win32" else "sd-cli"
|
||||
for p in root.rglob(name):
|
||||
|
|
@ -152,7 +227,8 @@ def install(
|
|||
build from source) or the archive has no ``sd-cli``.
|
||||
"""
|
||||
target = install_dir or default_install_dir()
|
||||
release = _fetch_latest_release(token = token)
|
||||
release = _fetch_release(_pinned_tag(), token = token)
|
||||
print(f"sd-cli: source {_repo()} release {release.get('tag_name', '?')}", flush = True)
|
||||
names = [a["name"] for a in release.get("assets", [])]
|
||||
chosen = resolve_release_asset(
|
||||
names,
|
||||
|
|
@ -164,17 +240,24 @@ def install(
|
|||
raise RuntimeError(
|
||||
f"No prebuilt sd-cli for {platform.system()}/{platform.machine()} "
|
||||
f"(accelerator={accelerator}). Build from source: "
|
||||
f"https://github.com/{REPO}"
|
||||
f"https://github.com/{_repo()}"
|
||||
)
|
||||
url = next(a["browser_download_url"] for a in release["assets"] if a["name"] == chosen)
|
||||
asset = next(a for a in release["assets"] if a["name"] == chosen)
|
||||
url = asset["browser_download_url"]
|
||||
target.mkdir(parents = True, exist_ok = True)
|
||||
archive = target / chosen
|
||||
print(f"downloading {chosen} -> {archive}", flush = True)
|
||||
urllib.request.urlretrieve(url, archive) # noqa: S310 (github release URL)
|
||||
print("extracting ...", flush = True)
|
||||
with zipfile.ZipFile(archive) as zf:
|
||||
zf.extractall(target)
|
||||
archive.unlink(missing_ok = True)
|
||||
try:
|
||||
_download(url, archive)
|
||||
# Verify integrity BEFORE extracting + executing.
|
||||
_verify_sha256(archive, asset.get("digest"))
|
||||
print("extracting ...", flush = True)
|
||||
with zipfile.ZipFile(archive) as zf:
|
||||
zf.extractall(target)
|
||||
finally:
|
||||
# Always drop the archive: on a sha256 mismatch / corrupt zip / network error it
|
||||
# must not linger (and a stale partial would defeat a later retry).
|
||||
archive.unlink(missing_ok = True)
|
||||
sd_cli = _locate_sd_cli(target)
|
||||
if not sd_cli:
|
||||
raise RuntimeError(f"archive {chosen} contained no sd-cli binary")
|
||||
|
|
@ -196,7 +279,7 @@ def main(argv: Optional[list[str]] = None) -> int:
|
|||
args = p.parse_args(argv)
|
||||
|
||||
if args.print_asset:
|
||||
release = _fetch_latest_release()
|
||||
release = _fetch_release(_pinned_tag())
|
||||
names = [a["name"] for a in release.get("assets", [])]
|
||||
chosen = resolve_release_asset(
|
||||
names,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue