ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call cannot load a ControlNet built for another family through the wrong pipeline. - Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError to ValueError) instead of a generic 500. - strength 0 disables ControlNet entirely, so a no-op selection never pays the download / VRAM cost; the control image is decoded and validated BEFORE the ControlNet is resolved or built, so a malformed image fails fast for the same reason. - ControlNet loads use the base compute dtype (state.dtype is a display string, not a torch.dtype, so it silently fell back to float32) and honor the base offload policy via group offloading instead of forcing the module resident. - Empty/malformed HF token coerced to anonymous access. - Flux Union ControlNet control_mode mapped from the selected control type. - resolve_controlnet drops the unused hf_token/cancel_event params. - ControlNetSpec validates guidance_start <= guidance_end (clean 422). - Images UI ControlNet Select shows its placeholder when nothing is selected. Adds regression tests for family enforcement and the union control-mode map.
This commit is contained in:
Daniel Han
parent
060fac0a9d
commit
cdfc6f0f56
5 changed files with 173 additions and 42 deletions
|
|
@ -1216,11 +1216,30 @@ class DiffusionBackend:
|
|||
if cn_model is None:
|
||||
if cancel.is_set():
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
cn_model = (
|
||||
getattr(diffusers, model_cls_name)
|
||||
.from_pretrained(resolved_cn.path, torch_dtype = state.dtype, token = state.hf_token)
|
||||
.to(state.device)
|
||||
import torch
|
||||
|
||||
# state.dtype is the display string saved at load ("bfloat16"), NOT a
|
||||
# torch.dtype; pass the real dtype so diffusers loads the ControlNet at the
|
||||
# base compute dtype instead of silently defaulting to float32 (extra VRAM).
|
||||
cn_dtype = getattr(torch, str(state.dtype).replace("torch.", ""), None)
|
||||
cn_model = getattr(diffusers, model_cls_name).from_pretrained(
|
||||
resolved_cn.path,
|
||||
torch_dtype = cn_dtype,
|
||||
# An empty / malformed token means anonymous access; the HF client can
|
||||
# raise on a blank credential instead of falling back, so coerce to None.
|
||||
token = state.hf_token or None,
|
||||
)
|
||||
# Placement must follow the base model's offload policy. A resident base moves
|
||||
# the ControlNet resident too; an offloaded (low-VRAM) base streams it through
|
||||
# the device with group offloading instead of forcing the whole module onto the
|
||||
# GPU, which would defeat the offload and risk an OOM. Best-effort: any failure
|
||||
# falls back to the resident placement (the prior behaviour).
|
||||
if getattr(state, "offload_policy", OFFLOAD_NONE) != OFFLOAD_NONE and (
|
||||
_offload_controlnet_module(cn_model, state.device, logger)
|
||||
):
|
||||
pass
|
||||
else:
|
||||
cn_model = cn_model.to(state.device)
|
||||
if cancel.is_set():
|
||||
# An unload raced the blocking download above and already cleared the
|
||||
# ControlNet caches; caching now would pin the module past the unload.
|
||||
|
|
@ -1401,7 +1420,7 @@ class DiffusionBackend:
|
|||
pipe = state.pipe
|
||||
init_pil = mask_pil = None
|
||||
control_pil = None
|
||||
cn_scale = cn_gstart = cn_gend = None
|
||||
cn_scale = cn_gstart = cn_gend = cn_mode = None
|
||||
ref_extra: list = []
|
||||
if getattr(state.family, "edit", False):
|
||||
# Instruction editing: the loaded pipe is the edit pipeline. It always
|
||||
|
|
@ -1472,40 +1491,55 @@ class DiffusionBackend:
|
|||
if controlnet is not None:
|
||||
from core.inference import diffusion_controlnet
|
||||
|
||||
if workflow != "txt2img":
|
||||
raise ValueError(
|
||||
"ControlNet currently combines with plain text-to-image only, not the "
|
||||
f"{workflow} workflow."
|
||||
)
|
||||
if not diffusion_controlnet.supports_controlnet(
|
||||
engine = "diffusers",
|
||||
family = state.family.name,
|
||||
has_controlnet_pipeline = bool(
|
||||
getattr(state.family, "controlnet_pipeline_class", None)
|
||||
),
|
||||
model_kind = state.kind,
|
||||
transformer_quant = state.transformer_quant,
|
||||
):
|
||||
raise ValueError(
|
||||
"ControlNet is not supported for this model/quantisation on the "
|
||||
"diffusers engine (needs a bf16 or bnb-4bit load of a family with a "
|
||||
"ControlNet pipeline; not GGUF-via-diffusers or torchao fp8/int8)."
|
||||
)
|
||||
cn_id, cn_image_b64, cn_type, cn_strength, cn_gs, cn_ge = controlnet
|
||||
resolved_cn = diffusion_controlnet.resolve_controlnet(
|
||||
cn_id,
|
||||
family = state.family.name,
|
||||
hf_token = state.hf_token,
|
||||
cancel_event = cancel,
|
||||
)
|
||||
pipe = self._controlnet_pipe(state, resolved_cn, cancel)
|
||||
workflow = "controlnet"
|
||||
src = _decode_b64_image(cn_image_b64, mode = "RGB")
|
||||
# Control map at the OUTPUT size so it aligns with the generated latents.
|
||||
control_pil = diffusion_controlnet.preprocess_control(src, cn_type).resize(
|
||||
(width, height), Image.LANCZOS
|
||||
)
|
||||
cn_scale, cn_gstart, cn_gend = cn_strength, cn_gs, cn_ge
|
||||
# strength 0 disables ControlNet (documented on the request model, and the
|
||||
# frontend slider allows it): skip the whole path so a no-op selection never
|
||||
# pays the multi-GB ControlNet download / VRAM cost.
|
||||
if cn_strength in (None, 0, 0.0):
|
||||
controlnet = None
|
||||
else:
|
||||
if workflow != "txt2img":
|
||||
raise ValueError(
|
||||
"ControlNet currently combines with plain text-to-image only, not "
|
||||
f"the {workflow} workflow."
|
||||
)
|
||||
if not diffusion_controlnet.supports_controlnet(
|
||||
engine = "diffusers",
|
||||
family = state.family.name,
|
||||
has_controlnet_pipeline = bool(
|
||||
getattr(state.family, "controlnet_pipeline_class", None)
|
||||
),
|
||||
model_kind = state.kind,
|
||||
transformer_quant = state.transformer_quant,
|
||||
):
|
||||
raise ValueError(
|
||||
"ControlNet is not supported for this model/quantisation on the "
|
||||
"diffusers engine (needs a bf16 or bnb-4bit load of a family with a "
|
||||
"ControlNet pipeline; not GGUF-via-diffusers or torchao fp8/int8)."
|
||||
)
|
||||
# Decode + preprocess the control image FIRST so a malformed / unsupported
|
||||
# image fails as a clean 400 BEFORE any ControlNet download or pipe build,
|
||||
# rather than after paying that cost. Control map at the OUTPUT size so it
|
||||
# aligns with the generated latents.
|
||||
src = _decode_b64_image(cn_image_b64, mode = "RGB")
|
||||
control_pil = diffusion_controlnet.preprocess_control(src, cn_type).resize(
|
||||
(width, height), Image.LANCZOS
|
||||
)
|
||||
try:
|
||||
resolved_cn = diffusion_controlnet.resolve_controlnet(
|
||||
cn_id, family = state.family.name
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
# An unknown / missing ControlNet id is a bad selection -> 400, not a
|
||||
# generic 500 (the route maps ValueError, not FileNotFoundError).
|
||||
raise ValueError(str(exc)) from exc
|
||||
pipe = self._controlnet_pipe(state, resolved_cn, cancel)
|
||||
workflow = "controlnet"
|
||||
cn_scale, cn_gstart, cn_gend = cn_strength, cn_gs, cn_ge
|
||||
# Flux Union ControlNet selects the active mode by an integer
|
||||
# ``control_mode`` (canny/depth/pose/...); map the chosen control type so
|
||||
# the union model applies the right head instead of a default/wrong one.
|
||||
cn_mode = diffusion_controlnet.union_control_mode(cn_id, cn_type)
|
||||
# Auto-resize odd-sized inputs to a multiple of 16 for the workflows whose
|
||||
# OUTPUT size is taken from the input image (img2img / inpaint / extend / edit),
|
||||
# so an upload like 186px tall no longer fails the pipeline's divisibility check.
|
||||
|
|
@ -1580,6 +1614,10 @@ class DiffusionBackend:
|
|||
kwargs["control_guidance_start"] = cn_gstart
|
||||
if "control_guidance_end" in call_params and cn_gend is not None:
|
||||
kwargs["control_guidance_end"] = cn_gend
|
||||
# Union ControlNet mode index (Flux); only when the pipe accepts it and the
|
||||
# selected control type maps to a known mode.
|
||||
if "control_mode" in call_params and cn_mode is not None:
|
||||
kwargs["control_mode"] = cn_mode
|
||||
|
||||
gen = _GenState(total_steps = steps)
|
||||
|
||||
|
|
@ -1821,6 +1859,34 @@ def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]:
|
|||
return base if isinstance(base, str) and base.strip() else None
|
||||
|
||||
|
||||
def _offload_controlnet_module(cn_model: Any, device: str, logger: Any) -> bool:
|
||||
"""Stream a ControlNet module through ``device`` via diffusers group offloading.
|
||||
|
||||
Used when the base model was loaded with an offload policy: forcing the ControlNet
|
||||
fully resident with ``.to(device)`` would defeat that low-VRAM placement and can OOM.
|
||||
Group offloading is applied to this single module (it does not touch the base pipe's
|
||||
existing hooks), so it is isolated and reversible. Returns True on success; on any
|
||||
failure the caller falls back to a resident placement, so this never blocks a load."""
|
||||
try:
|
||||
import torch
|
||||
from diffusers.hooks import apply_group_offloading
|
||||
|
||||
onload = torch.device(device)
|
||||
apply_group_offloading(
|
||||
cn_model,
|
||||
onload_device = onload,
|
||||
offload_device = torch.device("cpu"),
|
||||
offload_type = "block_level",
|
||||
num_blocks_per_group = 1,
|
||||
use_stream = onload.type == "cuda",
|
||||
)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — offload is best-effort; resident is the fallback
|
||||
if logger is not None:
|
||||
logger.warning("diffusion.controlnet: group offload failed (%s); loading resident", exc)
|
||||
return False
|
||||
|
||||
|
||||
def _base_file_downloaded(rfilename: str) -> bool:
|
||||
"""True for base-repo files ``from_pretrained`` actually fetches.
|
||||
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@ backend read an arbitrary location.
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
|
@ -142,17 +141,28 @@ def resolve_controlnet(
|
|||
spec_id: str,
|
||||
*,
|
||||
family: Optional[str] = None,
|
||||
hf_token: Optional[str] = None,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
) -> ResolvedControlNet:
|
||||
"""Resolve a ControlNet id to a loadable repo id / local dir.
|
||||
|
||||
Accepts a catalog/local id, or a bare public HF repo id (``owner/name``). The backend
|
||||
loads the result with ``ControlNetModelClass.from_pretrained(path)`` (download + cache
|
||||
handled there, like the base pipeline). Raises on an unknown id -> the caller maps to 400.
|
||||
|
||||
``family`` (the loaded base family) enforces catalog compatibility: a ControlNet is
|
||||
architecture-specific, so a catalog entry tagged for another family is rejected here
|
||||
with a clear error rather than being loaded through the wrong pipeline class later.
|
||||
"""
|
||||
entry = _catalog_by_id().get(spec_id)
|
||||
if entry is not None:
|
||||
# A curated/local entry may declare the families it is built for. A client that
|
||||
# bypasses the UI filter (direct API call) could send an entry for another family;
|
||||
# reject it before any download so it never reaches the wrong ControlNet pipeline.
|
||||
fam = (family or "").strip().lower()
|
||||
if entry.families and fam and fam not in {f.lower() for f in entry.families}:
|
||||
raise ValueError(
|
||||
f"ControlNet '{spec_id}' is for {', '.join(entry.families)}, not the loaded "
|
||||
f"'{family}' model; pick a ControlNet built for this family."
|
||||
)
|
||||
if entry.source == "local":
|
||||
path = entry.local_path or ""
|
||||
if not path or not Path(path).is_dir():
|
||||
|
|
@ -174,6 +184,33 @@ def resolve_controlnet(
|
|||
)
|
||||
|
||||
|
||||
# Union ControlNet mode indices. A single "union" model covers several control modes and
|
||||
# selects the active one via an integer ``control_mode`` argument; these are the standard
|
||||
# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" (an already-made
|
||||
# map) carries no intrinsic mode, so it maps to nothing (the caller omits control_mode).
|
||||
_UNION_CONTROL_MODES: dict[str, int] = {
|
||||
"canny": 0,
|
||||
"tile": 1,
|
||||
"depth": 2,
|
||||
"blur": 3,
|
||||
"pose": 4,
|
||||
"gray": 5,
|
||||
"lq": 6,
|
||||
}
|
||||
|
||||
|
||||
def union_control_mode(spec_id: str, control_type: str) -> Optional[int]:
|
||||
"""The integer ``control_mode`` for a union ControlNet, or None.
|
||||
|
||||
Returns a mode only for a curated *union* catalog entry AND a control type that maps to a
|
||||
known index; otherwise None so the caller omits the kwarg (a non-union ControlNet has a
|
||||
single fixed mode, and 'passthrough' does not name one). Pure lookup, no network."""
|
||||
entry = _catalog_by_id().get(spec_id)
|
||||
if entry is None or not entry.is_union:
|
||||
return None
|
||||
return _UNION_CONTROL_MODES.get((control_type or "").strip().lower())
|
||||
|
||||
|
||||
def preprocess_control(image: Any, control_type: str) -> Any:
|
||||
"""Turn a source image into a control map.
|
||||
|
||||
|
|
|
|||
|
|
@ -1856,6 +1856,14 @@ class ControlNetSpec(BaseModel):
|
|||
1.0, ge = 0.0, le = 1.0, description = "Fraction of steps at which ControlNet ends"
|
||||
)
|
||||
|
||||
@model_validator(mode = "after")
|
||||
def _check_guidance_range(self) -> "ControlNetSpec":
|
||||
# An inverted range (start > end) means "act over no steps"; reject it as a clean
|
||||
# 422 instead of letting the diffusers pipeline raise a 500 deep in the denoise.
|
||||
if self.guidance_start > self.guidance_end:
|
||||
raise ValueError("guidance_start must be <= guidance_end")
|
||||
return self
|
||||
|
||||
|
||||
class DiffusionGenerateRequest(BaseModel):
|
||||
"""Request to generate one image from the loaded diffusion model."""
|
||||
|
|
|
|||
|
|
@ -45,6 +45,26 @@ def test_resolve_controlnet_rejects_filesystem_like_ids():
|
|||
dc.resolve_controlnet(bad)
|
||||
|
||||
|
||||
def test_resolve_controlnet_enforces_family_match():
|
||||
# A curated entry tagged for another family must be rejected before download so it
|
||||
# never reaches the wrong ControlNet pipeline class.
|
||||
with pytest.raises(ValueError, match = "not the"):
|
||||
dc.resolve_controlnet("qwen-union", family = "flux.1")
|
||||
# The matching family resolves fine, and no family (unfiltered) is permissive.
|
||||
assert dc.resolve_controlnet("qwen-union", family = "qwen-image").path
|
||||
assert dc.resolve_controlnet("qwen-union").path
|
||||
|
||||
|
||||
def test_union_control_mode_maps_only_union_entries():
|
||||
# Union entries map a known control type to its integer mode; passthrough / unknown
|
||||
# types and non-union ids return None so the caller omits control_mode.
|
||||
assert dc.union_control_mode("flux-union-pro", "canny") == 0
|
||||
assert dc.union_control_mode("flux-union-pro", "depth") == 2
|
||||
assert dc.union_control_mode("flux-union-pro", "pose") == 4
|
||||
assert dc.union_control_mode("flux-union-pro", "passthrough") is None
|
||||
assert dc.union_control_mode("some/bare-repo", "canny") is None
|
||||
|
||||
|
||||
def test_resolve_controlnet_local(tmp_path, monkeypatch):
|
||||
d = tmp_path / "controlnets"
|
||||
d.mkdir()
|
||||
|
|
|
|||
|
|
@ -2095,7 +2095,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
hint="Condition the image on a control map (edges / depth / pose). Union models cover many types. Use 'Canny' to trace edges from your image, or 'Passthrough' if it is already a control map."
|
||||
>
|
||||
<div className="space-y-2 rounded-lg border border-border bg-muted/30 p-2">
|
||||
<Select value={controlnetId} onValueChange={setControlnetId}>
|
||||
<Select value={controlnetId || undefined} onValueChange={setControlnetId}>
|
||||
<SelectTrigger className="h-8 w-full text-xs">
|
||||
<SelectValue placeholder="Select a ControlNet" />
|
||||
</SelectTrigger>
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue