diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index c06e3cc359..17ae1c5687 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -329,6 +329,13 @@ class DiffusionBackend: # resident modules (no extra VRAM, no reload), so we build each once per load and # reuse it. Keyed by pipeline class name; cleared on unload with the base pipe. self._aux_pipes: dict[str, Any] = {} + # Cache of loaded ControlNet models (id -> module) and the ControlNet workflow + # pipelines built around them ((pipeline_class, cn_id) -> pipe). ControlNet models + # are a small extra module loaded via from_pretrained; the pipeline is assembled via + # Pipeline.from_pipe(base, controlnet=model), reusing the resident base modules (no + # reload). Both are cleared on unload with the base pipe. + self._cn_models: dict[str, Any] = {} + self._cn_pipes: dict[tuple[str, str], Any] = {} @property def is_loaded(self) -> bool: @@ -1381,6 +1388,70 @@ class DiffusionBackend: self._aux_pipes[class_name] = pipe return pipe + def _controlnet_pipe(self, state: _LoadState, resolved_cn: Any, cancel: threading.Event) -> Any: + """Build (once, cached) the family's diffusers ControlNet pipeline around the requested + ControlNet model. The ControlNet model is a small extra module loaded via from_pretrained + and cached by id; the pipeline is assembled with ``Pipeline.from_pipe(base, + controlnet=model)`` -- reusing the resident base modules at their loaded dtype (no reload, + no recast; torch_dtype=None for the same reason as _workflow_pipe). Raises a clear + ValueError when the family declares no ControlNet classes.""" + fam = state.family + pipe_cls_name = getattr(fam, "controlnet_pipeline_class", None) + model_cls_name = getattr(fam, "controlnet_model_class", None) + if not pipe_cls_name or not model_cls_name: + raise ValueError(f"ControlNet is not supported for the '{fam.name}' model family.") + import diffusers + + cn_model = self._cn_models.get(resolved_cn.id) + if cn_model is None: + if cancel.is_set(): + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + 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, + ) + if cancel.is_set(): + # An unload/eviction raced the blocking download above and may have already + # cleared the load. Bail BEFORE any device placement so we don't allocate + # several GB onto the GPU after _unload_locked() freed it (which would OOM + # or make the unload appear to free memory only to repopulate it). + del cn_model + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + # 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. + del cn_model + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + self._cn_models[resolved_cn.id] = cn_model + key = (pipe_cls_name, resolved_cn.id) + pipe = self._cn_pipes.get(key) + if pipe is None: + pipe = getattr(diffusers, pipe_cls_name).from_pipe( + state.pipe, controlnet = cn_model, torch_dtype = None + ) + self._cn_pipes[key] = pipe + return pipe + @staticmethod def _align_vae_dtype(pipe: Any) -> None: """Cast the VAE to the transformer's compute dtype before an image-conditioned @@ -1532,6 +1603,9 @@ class DiffusionBackend: # LoRA adapters as (id, weight) pairs; loaded onto the pipe (non-fused) and activated # with set_adapters for this generation. None/empty = no LoRA (adapters cleared). loras: Optional[list[tuple[str, float]]] = None, + # ControlNet as (id, control_image_b64, control_type, strength, guidance_start, + # guidance_end); conditions the text-to-image path on a spatial control map. None = off. + controlnet: Optional[tuple[str, str, str, float, float, float]] = None, ) -> dict[str, Any]: import torch from PIL import Image @@ -1573,6 +1647,8 @@ class DiffusionBackend: # an edit model's OWN loaded pipe is already the edit pipeline. pipe = state.pipe init_pil = mask_pil = None + control_pil = None + cn_scale = cn_gstart = cn_gend = cn_mode = None ref_extra: list = [] # Validate parameter dependencies up front: mask / upscale / reference all # need an input image, and reference conditioning needs a family that @@ -1661,6 +1737,61 @@ class DiffusionBackend: init_pil = _decode_b64_image(init_image, mode = "RGB") else: workflow = "txt2img" + + # ControlNet conditioning (diffusers): applies to the plain text-to-image path. + # Builds the family's ControlNet pipeline around the resident modules (no reload) + # and passes a control map. v1 conditions txt2img only (not img2img/inpaint/edit). + if controlnet is not None: + from core.inference import diffusion_controlnet + cn_id, cn_image_b64, cn_type, cn_strength, cn_gs, cn_ge = controlnet + # 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. @@ -1707,10 +1838,10 @@ class DiffusionBackend: # outpaint vs a 1024 slider -> "tensor a (128) must match tensor b (192)"). # Many img2img/inpaint pipelines drop width/height entirely; pass them only # when accepted, derived from the image so they are always consistent. - if workflow in ("txt2img", "reference"): - # txt2img and FLUX.2 reference both generate at the REQUESTED size; the - # reference pipe resizes the conditioning image itself, so it must not be - # pinned to the input image's size like img2img/inpaint/upscale are. + if workflow in ("txt2img", "reference", "controlnet"): + # txt2img, FLUX.2 reference, and ControlNet all generate at the REQUESTED + # size; the reference/control image is resized to match, so it must not be + # pinned to an input image's size like img2img/inpaint/upscale are. kwargs["width"] = width kwargs["height"] = height elif init_pil is not None: @@ -1721,6 +1852,24 @@ class DiffusionBackend: kwargs["height"] = ih if negative_prompt and "negative_prompt" in call_params: kwargs["negative_prompt"] = negative_prompt + if workflow == "controlnet" and control_pil is not None: + # The ControlNet pipeline takes the control map + its conditioning scale; + # guidance start/end bound the step range it acts over. Every kwarg is gated + # on the pipe signature so a family whose CN pipe omits one still runs. + if "control_image" in call_params: + kwargs["control_image"] = control_pil + elif "image" in call_params: # some CN pipelines name it "image" + kwargs["image"] = control_pil + if "controlnet_conditioning_scale" in call_params and cn_scale is not None: + kwargs["controlnet_conditioning_scale"] = cn_scale + if "control_guidance_start" in call_params and cn_gstart is not None: + 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) @@ -1847,6 +1996,9 @@ class DiffusionBackend: # freed pipeline (they only re-wire its components, but holding the wrappers # would keep the modules alive past unload). self._aux_pipes.clear() + # Drop any ControlNet models + pipelines so the freed load carries no extra modules. + self._cn_pipes.clear() + self._cn_models.clear() self._state = None del state clear_gpu_cache() @@ -1874,8 +2026,9 @@ class DiffusionBackend: "transformer_cache": None, "workflows": [], "supports_lora": False, + "supports_controlnet": False, } - from core.inference import diffusion_lora + from core.inference import diffusion_controlnet, diffusion_lora return { "loaded": True, @@ -1905,6 +2058,15 @@ class DiffusionBackend: transformer_quant = state.transformer_quant, compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()), ), + "supports_controlnet": 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, + ), } @@ -1961,6 +2123,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, *, include_transformer: bool = False) -> bool: """True for base-repo files ``from_pretrained`` actually fetches. diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py new file mode 100644 index 0000000000..b318b040c8 --- /dev/null +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Diffusion ControlNet support: family-gated discovery of ControlNet models, resolution to +a loadable diffusers repo/dir, control-image preprocessing, and a capability gate. + +Mirrors ``diffusion_lora.py``. Two differences from LoRA: (1) a ControlNet is a full diffusers +repo (loaded via ``from_pretrained``), not a single-file adapter, so resolution yields a repo id +or local directory rather than a file path; (2) ControlNet needs a spatial *control image*, which +is either supplied already-preprocessed ("passthrough", as in ComfyUI where preprocessing is a +separate step) or derived here ("canny", a dependency-free edge map). + +ControlNet models are architecture-specific (a FLUX ControlNet cannot drive a Qwen base), so +discovery is family-gated exactly like the LoRA picker. The request never carries a filesystem +path -- only a discovery id or a public ``owner/name`` repo id -- so a client cannot make the +backend read an arbitrary location. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +from utils.paths.storage_roots import studio_root + +# Control map types. "passthrough": the supplied image IS the control map (a depth/pose/etc. +# map produced elsewhere). "canny": derive an edge map here (no heavy detector dependency). +CONTROL_TYPES = ("passthrough", "canny") + +# Families whose diffusers pipeline supports ControlNet (declared in diffusion_families via +# controlnet_pipeline_class). Native sd.cpp ControlNet is a follow-up. Torchao fp8/int8 dense +# and GGUF-via-diffusers are gated off, same rule as LoRA. +_DIFFUSERS_BLOCKED_QUANT = ("int8", "fp8", "nvfp4", "mxfp8") + + +@dataclass(frozen = True) +class ControlNetCatalogEntry: + """One discoverable ControlNet model.""" + + id: str + display_name: str + source: str # "local" | "hub" + families: tuple[str, ...] = () # compatible family names (empty = shown, not gated) + repo_id: Optional[str] = None # for source == "hub" + local_path: Optional[str] = None # for source == "local" + control_types: tuple[str, ...] = ("passthrough",) # recommended control types + is_union: bool = False # a single model covering many control modes + + +@dataclass(frozen = True) +class ResolvedControlNet: + """A ControlNet resolved to something ``from_pretrained`` can load.""" + + id: str + path: str # repo id (hub) or local directory + is_local: bool + + +# Curated, family-tagged catalog. Union models (one model, many control modes) dominate real +# usage, so they are the default picks. Extend as more are curated; local dirs + a bare public +# ``owner/name`` repo id also work. +_CURATED: tuple[ControlNetCatalogEntry, ...] = ( + ControlNetCatalogEntry( + id = "flux-union-pro", + display_name = "FLUX.1 ControlNet Union Pro (Shakker-Labs)", + source = "hub", + families = ("flux.1",), + repo_id = "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", + control_types = ("canny", "depth", "pose", "passthrough"), + is_union = True, + ), + ControlNetCatalogEntry( + id = "qwen-union", + display_name = "Qwen-Image ControlNet Union (InstantX)", + source = "hub", + families = ("qwen-image",), + repo_id = "InstantX/Qwen-Image-ControlNet-Union", + control_types = ("canny", "depth", "pose", "passthrough"), + is_union = True, + ), +) + + +def controlnets_dir() -> Path: + """Local directory Studio scans for user-provided ControlNet model folders.""" + d = studio_root() / "controlnets" / "diffusion" + d.mkdir(parents = True, exist_ok = True) + return d + + +def sanitize_id(raw: str) -> str: + """Filesystem-safe id from a repo id / folder name.""" + stem = raw.rsplit("/", 1)[-1] + stem = re.sub(r"[^A-Za-z0-9._-]+", "_", stem).strip("._-") + return stem or "controlnet" + + +def _has_controlnet_weights(p: Path) -> bool: + """True when ``p`` holds a loadable diffusers ControlNet weight (or shard index). + + A config-only folder (interrupted copy/download) would otherwise be advertised and + then fail deep inside ``from_pretrained`` as a generic 500. Accept the standard + single-file weights, a sharded weight index, or any ``.safetensors`` shard.""" + names = ( + "diffusion_pytorch_model.safetensors", + "diffusion_pytorch_model.bin", + "diffusion_pytorch_model.safetensors.index.json", + "diffusion_pytorch_model.bin.index.json", + ) + if any((p / n).exists() for n in names): + return True + try: + return any(child.suffix == ".safetensors" for child in p.iterdir()) + except OSError: + return False + + +def _scan_local() -> list[ControlNetCatalogEntry]: + """A local ControlNet is a directory containing a diffusers config + weights.""" + entries: list[ControlNetCatalogEntry] = [] + root = controlnets_dir() + try: + children = sorted(root.iterdir()) + except OSError: + return entries + for p in children: + if not p.is_dir(): + continue + # Require BOTH the config and a loadable weight/index: a config-only folder is an + # incomplete copy/download, and advertising it would fail later in from_pretrained. + if not (p / "config.json").exists() or not _has_controlnet_weights(p): + continue + entries.append( + ControlNetCatalogEntry( + id = p.name, + display_name = p.name, + source = "local", + local_path = str(p), + control_types = CONTROL_TYPES, + ) + ) + return entries + + +def list_controlnets(*, family: Optional[str] = None) -> list[ControlNetCatalogEntry]: + """Merged catalog (curated + local), optionally family-filtered. Cheap: one dir scan + plus the in-memory curated list. Network is only touched on resolve().""" + merged = list(_CURATED) + _scan_local() + if family: + fam = family.strip().lower() + merged = [e for e in merged if not e.families or fam in {f.lower() for f in e.families}] + merged.sort(key = lambda e: (e.source != "local", e.display_name.lower())) + return merged + + +def _catalog_by_id() -> dict[str, ControlNetCatalogEntry]: + return {e.id: e for e in (list(_CURATED) + _scan_local())} + + +def resolve_controlnet(spec_id: str, *, family: Optional[str] = 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(): + raise FileNotFoundError(f"ControlNet '{spec_id}' is no longer present on disk") + return ResolvedControlNet(spec_id, path, is_local = True) + if not entry.repo_id: + raise ValueError(f"ControlNet '{spec_id}' has no repo") + return ResolvedControlNet(spec_id, entry.repo_id, is_local = False) + + # A bare public HF repo id (owner/name). STRICT shape -- exactly one slash and + # alphanumeric-leading segments -- so a filesystem-looking id (/tmp/x, ../x, ~/x, + # C:\x) can never reach from_pretrained, which would happily treat it as a local + # directory and bypass the controlnets_dir() no-raw-path contract. + if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*", spec_id): + return ResolvedControlNet(spec_id, spec_id, is_local = False) + + raise FileNotFoundError( + f"unknown ControlNet '{spec_id}': not a local model, catalog entry, or HF repo id" + ) + + +# 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. + + ``passthrough`` returns the image unchanged (it is already a depth/pose/edge map made + elsewhere). ``canny`` derives a dependency-free gradient edge map (a rough stand-in for a + true Canny; a cv2/kornia detector and depth/pose detectors are a follow-up). Unknown types + pass through so a new type never hard-fails generation. + """ + ct = (control_type or "passthrough").strip().lower() + if ct != "canny": + return image + import numpy as np + from PIL import Image + + gray = np.asarray(image.convert("L"), dtype = np.float32) + gy, gx = np.gradient(gray) + mag = np.hypot(gx, gy) + peak = float(mag.max()) + if peak <= 1e-6: + return image # flat image -> nothing to trace; don't emit a black map + mag = mag / peak * 255.0 + edges = (mag > 40.0).astype(np.uint8) * 255 # white edges on black, the ControlNet convention + return Image.fromarray(edges).convert("RGB") + + +def supports_controlnet( + *, + engine: str, + family: Optional[str], + has_controlnet_pipeline: bool, + model_kind: Optional[str], + transformer_quant: Optional[str], +) -> bool: + """Whether the loaded model can apply a ControlNet. + + diffusers only for now (native sd.cpp is a follow-up). Requires the family to declare a + ControlNet pipeline. Blocked for the diffusers GGUF path and torchao fp8/int8 dense + (same constraints as LoRA): those transformers cannot host the extra conditioning cleanly. + """ + if not family or not has_controlnet_pipeline: + return False + if engine != "diffusers": + return False + if model_kind == "gguf": + return False + if transformer_quant and str(transformer_quant).strip().lower() in _DIFFUSERS_BLOCKED_QUANT: + return False + return True diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 2f9ba12c62..86493a3070 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -44,6 +44,13 @@ class DiffusionFamily: # 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 + # ControlNet: the diffusers ControlNet pipeline + model classes for this family. The backend + # loads the (small) ControlNet model via from_pretrained and builds the pipeline via + # ``Pipeline.from_pipe(base, controlnet=model)`` around the resident modules (no reload), + # then passes the control image + conditioning scale at generate time. None on both = the + # family has no diffusers ControlNet support and the UI gates the workflow off. + controlnet_pipeline_class: Optional[str] = None + controlnet_model_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 @@ -116,6 +123,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( aliases = ("flux1", "flux-1"), img2img_pipeline_class = "FluxImg2ImgPipeline", inpaint_pipeline_class = "FluxInpaintPipeline", + controlnet_pipeline_class = "FluxControlNetPipeline", + controlnet_model_class = "FluxControlNetModel", sd_cpp_vae = ("black-forest-labs/FLUX.1-schnell", "ae.safetensors"), sd_cpp_text_encoders = ( ("comfyanonymous/flux_text_encoders", "clip_l.safetensors", "clip_l"), @@ -211,6 +220,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( aliases = ("qwen_image", "qwenimage"), img2img_pipeline_class = "QwenImageImg2ImgPipeline", inpaint_pipeline_class = "QwenImageInpaintPipeline", + controlnet_pipeline_class = "QwenImageControlNetPipeline", + controlnet_model_class = "QwenImageControlNetModel", 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. diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 059b85974c..7632551661 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -713,6 +713,10 @@ class SdCppDiffusionBackend: # path: prompt tags for one-shot sd-cli, structured `lora` entries # for the resident sd-server. None/empty = no LoRA. loras: Optional[list[tuple[str, float]]] = None, + # Accepted for the uniform engine interface; the guard below rejects it on the native + # engine (ControlNet is diffusers-only) like img2img/inpaint, so a direct API call with + # ControlNet set fails clearly instead of TypeError'ing on an unexpected kwarg. + controlnet: Optional[tuple[str, str, str, float, float, float]] = None, ) -> dict[str, Any]: import tempfile @@ -733,6 +737,11 @@ class SdCppDiffusionBackend: "img2img / inpaint / reference / upscale are not yet supported on the native " "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows." ) + if controlnet is not None: + raise ValueError( + "ControlNet is not yet supported on the native sd.cpp engine; run on a GPU " + "(diffusers) for ControlNet conditioning." + ) cancel = threading.Event() with self._generate_lock: @@ -1108,6 +1117,7 @@ class SdCppDiffusionBackend: "engine": "sd_cpp", "native_mode": None, "supports_lora": False, + "supports_controlnet": False, "workflows": [], } from core.inference import diffusion_lora @@ -1140,6 +1150,8 @@ class SdCppDiffusionBackend: model_kind = "gguf", transformer_quant = None, ), + # ControlNet is diffusers-only; the native engine's generate() rejects it. + "supports_controlnet": False, # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. "native_mode": state.mode, # The native engine supports plain text-to-image only (generate() rejects diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d3feff32f7..58701b5828 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1829,6 +1829,50 @@ class LoraSpec(BaseModel): ) +class ControlNetSpec(BaseModel): + """A ControlNet to condition this generation on: a discovery id plus a control image. + + The id resolves against the backend's ControlNet catalog + local scan (see + core/inference/diffusion_controlnet.py); the client never supplies a raw filesystem path. + ``image`` is either an already-made control map (``control_type='passthrough'``) or a source + image the backend turns into a map (``control_type='canny'``). strength 0 disables it. + """ + + id: str = Field( + ..., + min_length = 1, + max_length = 512, + description = "ControlNet discovery id (repo id or local name)", + ) + image: str = Field( + ..., + min_length = 1, + max_length = 32 * 1024 * 1024, + description = "Base64/data-URL control image (a source image or a preprocessed map)", + ) + control_type: str = Field( + "passthrough", + description = "How to derive the control map: 'passthrough' (already a map) or 'canny'", + ) + strength: float = Field( + 1.0, ge = 0.0, le = 2.0, description = "ControlNet conditioning scale; 0 disables" + ) + guidance_start: float = Field( + 0.0, ge = 0.0, le = 1.0, description = "Fraction of steps at which ControlNet begins" + ) + guidance_end: float = Field( + 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.""" @@ -1896,6 +1940,12 @@ class DiffusionGenerateRequest(BaseModel): "Omitted/empty applies none and behaves exactly as before. Rejected with a clear " "message when the loaded model or its quantisation can't apply LoRA.", ) + controlnet: Optional[ControlNetSpec] = Field( + None, + description = "ControlNet conditioning for this generation (id + control image + strength). " + "Omitted applies none and behaves exactly as before. Rejected with a clear message when " + "the loaded model or its quantisation can't apply ControlNet.", + ) @field_validator("loras") @classmethod @@ -1956,6 +2006,9 @@ class GalleryImage(BaseModel): loras: list[str] = Field( default_factory = list, description = "LoRA adapters applied, formatted as 'id:weight'" ) + controlnet: Optional[str] = Field( + None, description = "ControlNet applied, formatted as 'id:control_type:strength'" + ) created_at: float = Field(..., description = "Creation time (epoch seconds)") @@ -2051,6 +2104,12 @@ class DiffusionStatusResponse(BaseModel): "LoRA picker's enabled state). False on unsupported families/quant (e.g. torchao fp8/int8 " "dense, GGUF-via-diffusers, or Qwen-Image on the native engine).", ) + supports_controlnet: bool = Field( + False, + description = "Whether the loaded model can apply a ControlNet (drives the ControlNet " + "picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False " + "for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.", + ) # ── OpenAI-compatible images API (POST /v1/images/generations) ── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 929aa19afa..b8a970a106 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11162,6 +11162,18 @@ async def generate_diffusion_image( upscale = request.upscale, reference_images = request.reference_images, loras = [(l.id, l.weight) for l in request.loras] if request.loras else None, + controlnet = ( + ( + request.controlnet.id, + request.controlnet.image, + request.controlnet.control_type, + request.controlnet.strength, + request.controlnet.guidance_start, + request.controlnet.guidance_end, + ) + if request.controlnet + else None + ), ) except ValueError as exc: # Bad client input (undecodable image/mask, or a workflow the loaded family @@ -11224,6 +11236,15 @@ async def generate_diffusion_image( "loras": ( [f"{l.id}:{l.weight:g}" for l in request.loras] if request.loras else [] ), + "controlnet": ( + f"{request.controlnet.id}:{request.controlnet.control_type}:" + f"{request.controlnet.strength:g}" + # strength 0 is treated as disabled and skipped before loading / + # conditioning, so the image is unconditioned; don't claim a + # ControlNet was applied in the recipe/metadata. + if request.controlnet and request.controlnet.strength > 0 + else None + ), "created_at": created_at, }, ) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index b22d47c545..9ea12644eb 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2065,6 +2065,39 @@ async def scan_diffusion_loras( } +@router.get("/diffusion-controlnets") +async def scan_diffusion_controlnets( + family: Optional[str] = Query( + default = None, description = "Filter to ControlNets compatible with this diffusion family" + ), + current_subject: str = Depends(get_current_subject), +): + """List diffusion ControlNet models for the Images workflow. + + Merges the curated, family-tagged catalog with local model folders in + ``/controlnets/diffusion``, optionally filtered to the loaded model's family. + Cheap: one directory scan, no network (a hub model is only downloaded when selected). + """ + from core.inference import diffusion_controlnet + + entries = diffusion_controlnet.list_controlnets(family = family) + return { + "controlnets": [ + { + "id": e.id, + "display_name": e.display_name, + "source": e.source, + "families": list(e.families), + "control_types": list(e.control_types), + "is_union": e.is_union, + } + for e in entries + ], + "control_types": list(diffusion_controlnet.CONTROL_TYPES), + "controlnets_dir": str(diffusion_controlnet.controlnets_dir()), + } + + def _is_path_under(path: Path, root: Path) -> bool: try: path.resolve().relative_to(root.resolve()) diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py new file mode 100644 index 0000000000..5708c3a3ae --- /dev/null +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -0,0 +1,261 @@ +"""Tests for diffusion ControlNet support: discovery/resolve/preprocess/gate helpers, the +request-model validation, the family wiring, and the diffusers ControlNet pipe manager.""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference import diffusion_controlnet as dc + + +# ── Pure helpers ──────────────────────────────────────────────────────────── + + +def test_sanitize_id(): + assert dc.sanitize_id("owner/My ControlNet") == "My_ControlNet" + assert dc.sanitize_id("weird:<>name") == "weird_name" + assert dc.sanitize_id("") == "controlnet" + + +def test_list_controlnets_family_filter(): + flux = {e.id for e in dc.list_controlnets(family = "flux.1")} + qwen = {e.id for e in dc.list_controlnets(family = "qwen-image")} + assert "flux-union-pro" in flux and "qwen-union" not in flux + assert "qwen-union" in qwen and "flux-union-pro" not in qwen + + +def test_resolve_controlnet_catalog_bare_repo_and_unknown(): + r = dc.resolve_controlnet("flux-union-pro", family = "flux.1") + assert r.path == "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro" and not r.is_local + # A bare public repo id passes through. + r2 = dc.resolve_controlnet("owner/some-controlnet") + assert r2.path == "owner/some-controlnet" and not r2.is_local + with pytest.raises(FileNotFoundError): + dc.resolve_controlnet("not-a-known-id") + + +def test_resolve_controlnet_rejects_filesystem_like_ids(): + # The bare-repo fallback must never accept a path-shaped id: from_pretrained + # would treat it as a local directory, bypassing the controlnets_dir() contract. + for bad in ("/tmp/model", "../some/model", "./x/y", "~/x/y", "a/b/c", "C:\\x/y", ".hidden/x"): + with pytest.raises(FileNotFoundError): + 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() + cn = d / "my-cn" + cn.mkdir() + (cn / "config.json").write_text("{}") + (cn / "diffusion_pytorch_model.safetensors").write_bytes(b"x") # a loadable weight + monkeypatch.setattr(dc, "controlnets_dir", lambda: d) + entries = {e.id for e in dc.list_controlnets()} + assert "my-cn" in entries + r = dc.resolve_controlnet("my-cn") + assert r.is_local and r.path == str(cn) + + +def test_scan_local_skips_config_only_folder(tmp_path, monkeypatch): + # A folder with config.json but no weight/index (interrupted copy) must NOT be + # advertised: it would otherwise fail deep in from_pretrained as a generic 500. + d = tmp_path / "controlnets" + d.mkdir() + incomplete = d / "incomplete-cn" + incomplete.mkdir() + (incomplete / "config.json").write_text("{}") + monkeypatch.setattr(dc, "controlnets_dir", lambda: d) + assert "incomplete-cn" not in {e.id for e in dc.list_controlnets()} + # A sharded weight index counts as a loadable weight. + (incomplete / "diffusion_pytorch_model.safetensors.index.json").write_text("{}") + assert "incomplete-cn" in {e.id for e in dc.list_controlnets()} + + +def test_preprocess_control_passthrough_and_canny(): + from PIL import Image + + img = Image.new("RGB", (32, 24), (10, 20, 30)) + # passthrough returns the same object. + assert dc.preprocess_control(img, "passthrough") is img + # a flat image has no edges -> canny falls back to passthrough (no black map). + assert dc.preprocess_control(img, "canny") is img + # an image with structure yields an edge map: RGB, same size, some white pixels. + import numpy as np + + arr = np.zeros((24, 32, 3), np.uint8) + arr[:, 16:, :] = 255 # a hard vertical edge + edged = dc.preprocess_control(Image.fromarray(arr), "canny") + assert edged.mode == "RGB" and edged.size == (32, 24) + assert np.asarray(edged).max() == 255 # traced the edge + + +def test_supports_controlnet_matrix(): + ok = dict(engine = "diffusers", family = "flux.1", has_controlnet_pipeline = True) + assert dc.supports_controlnet(**ok, model_kind = "pipeline", transformer_quant = None) + assert dc.supports_controlnet(**ok, model_kind = "single_file", transformer_quant = None) + # GGUF-via-diffusers and fp8/int8 dense are gated off, like LoRA. + assert not dc.supports_controlnet(**ok, model_kind = "gguf", transformer_quant = None) + assert not dc.supports_controlnet(**ok, model_kind = "single_file", transformer_quant = "fp8") + assert not dc.supports_controlnet(**ok, model_kind = "single_file", transformer_quant = "int8") + # native engine + a family without a CN pipeline are off. + assert not dc.supports_controlnet( + engine = "sd_cpp", + family = "flux.1", + has_controlnet_pipeline = True, + model_kind = "gguf", + transformer_quant = None, + ) + assert not dc.supports_controlnet( + engine = "diffusers", + family = "z-image", + has_controlnet_pipeline = False, + model_kind = "pipeline", + transformer_quant = None, + ) + + +# ── Request-model validation ──────────────────────────────────────────────── + + +def test_controlnet_spec_and_request_validation(): + from models.inference import ControlNetSpec, DiffusionGenerateRequest + + assert DiffusionGenerateRequest(prompt = "x").controlnet is None + req = DiffusionGenerateRequest( + prompt = "x", + controlnet = { + "id": "flux-union-pro", + "image": "data", + "control_type": "canny", + "strength": 0.6, + }, + ) + assert req.controlnet.id == "flux-union-pro" and req.controlnet.strength == 0.6 + # defaults + s = ControlNetSpec(id = "a", image = "b") + assert s.control_type == "passthrough" and s.strength == 1.0 + assert s.guidance_start == 0.0 and s.guidance_end == 1.0 + # bounds + with pytest.raises(Exception): + ControlNetSpec(id = "a", image = "b", strength = 3.0) + with pytest.raises(Exception): + ControlNetSpec(id = "a", image = "b", guidance_end = 1.5) + + +# ── Family wiring ─────────────────────────────────────────────────────────── + + +def test_families_declare_controlnet_classes(): + from core.inference.diffusion_families import _FAMILIES + + by_name = {f.name: f for f in _FAMILIES} + assert by_name["flux.1"].controlnet_pipeline_class == "FluxControlNetPipeline" + assert by_name["flux.1"].controlnet_model_class == "FluxControlNetModel" + assert by_name["qwen-image"].controlnet_pipeline_class == "QwenImageControlNetPipeline" + # z-image has no diffusers ControlNet pipeline -> gated off. + assert by_name["z-image"].controlnet_pipeline_class is None + + +# ── Diffusers ControlNet pipe manager ─────────────────────────────────────── + + +class _FakeCNModel: + @classmethod + def from_pretrained( + cls, + path, + torch_dtype = None, + token = None, + ): + m = cls() + m.path = path + return m + + def to(self, device): + self.device = device + return self + + +class _FakeCNPipe: + @classmethod + def from_pipe( + cls, + base, + controlnet = None, + torch_dtype = None, + ): + p = cls() + p.base = base + p.controlnet = controlnet + return p + + +def _fake_diffusers(): + mod = types.ModuleType("diffusers") + mod.FluxControlNetModel = _FakeCNModel + mod.FluxControlNetPipeline = _FakeCNPipe + return mod + + +def _state(): + fam = types.SimpleNamespace( + name = "flux.1", + controlnet_pipeline_class = "FluxControlNetPipeline", + controlnet_model_class = "FluxControlNetModel", + ) + return types.SimpleNamespace( + family = fam, dtype = "bf16", device = "cpu", hf_token = None, pipe = object() + ) + + +def test_controlnet_pipe_loads_once_and_caches(monkeypatch): + import threading + + from core.inference.diffusion import DiffusionBackend + + monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers()) + b = DiffusionBackend() + st = _state() + resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False) + p1 = b._controlnet_pipe(st, resolved, threading.Event()) + assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel) + assert p1.controlnet.path == "repo/id" and p1.controlnet.device == "cpu" + # cached: same id -> same model + same pipe, no reload. + p2 = b._controlnet_pipe(st, resolved, threading.Event()) + assert p2 is p1 + assert b._cn_models["flux-union-pro"] is p1.controlnet + + +def test_controlnet_pipe_rejects_family_without_classes(): + import threading + + from core.inference.diffusion import DiffusionBackend + + b = DiffusionBackend() + st = _state() + st.family.controlnet_pipeline_class = None + with pytest.raises(ValueError, match = "not supported"): + b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event()) diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 109724ca0b..0b52ea6e14 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -752,3 +752,25 @@ def test_generate_zero_weight_loras_are_noop(monkeypatch): b.generate(prompt = "x", steps = 4, seed = 1, loras = [("id1", 0.0)]) _, params, _, _ = eng.calls[0] assert params.lora_dir is None # nothing applied + + +def test_generate_rejects_controlnet_on_native_engine(): + # ControlNet is diffusers-only. The route passes `controlnet` to whichever engine is + # active, so the native backend must reject it with a clean ValueError (-> 400) rather + # than TypeError on an unexpected kwarg (-> opaque 500). + b = _loaded_backend(engine = _FakeEngine()) + with pytest.raises(ValueError, match = "ControlNet is not yet supported on the native"): + b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0)) + + +def test_generate_rejects_image_conditioned_on_native_engine(): + # img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call + # with an init image on the native engine gets a clean ValueError, not a silent txt2img. + b = _loaded_backend(engine = _FakeEngine()) + with pytest.raises(ValueError, match = "not yet supported on the native"): + b.generate(prompt = "x", steps = 4, seed = 1, init_image = "data:image/png;base64,AAAA") + + +def test_status_native_reports_supports_controlnet_false(): + b = _loaded_backend() + assert b.status()["supports_controlnet"] is False diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index aa50c4c16e..ad38261705 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -21,6 +21,9 @@ export interface DiffusionStatus { // Whether the loaded model + quantisation can apply LoRA adapters (drives the LoRA // picker's enabled state). False on unsupported families/quant. supports_lora?: boolean; + // Whether the loaded model can apply a ControlNet (drives the ControlNet picker's enabled + // state). Diffusers only, for families with a ControlNet pipeline; false otherwise. + supports_controlnet?: boolean; } export interface DiffusionGenerateProgress { @@ -92,6 +95,9 @@ export interface DiffusionGenerateRequest { // LoRA adapters to apply for this generation (by discovery id + weight, 0..2). Omitted // or empty applies none. Rejected (400) when the loaded model/quant can't apply LoRA. loras?: LoraSpecInput[]; + // ControlNet conditioning for this generation. Omitted applies none. Rejected (400) when + // the loaded model/quant can't apply ControlNet. + controlnet?: ControlNetSpecInput; } // One LoRA selection sent with a generation. @@ -100,6 +106,29 @@ export interface LoraSpecInput { weight: number; } +// A ControlNet selection sent with a generation. +export interface ControlNetSpecInput { + id: string; + // Base64/data-URL control image (a source image or an already-made control map). + image: string; + // "canny" preprocesses edges from a source image; any other type (passthrough, or a + // union type like depth/pose) is an already-made map the backend maps to a control mode. + control_type: string; + strength: number; + guidance_start?: number; + guidance_end?: number; +} + +// A discoverable ControlNet model (from GET /api/models/diffusion-controlnets). +export interface DiffusionControlNetInfo { + id: string; + display_name: string; + source: "local" | "hub"; + families: string[]; + control_types: string[]; + is_union: boolean; +} + // A discoverable diffusion LoRA adapter (from GET /api/models/diffusion-loras). export interface DiffusionLoraInfo { id: string; @@ -126,6 +155,7 @@ export interface GalleryImage { batch_size: number; model: string | null; loras?: string[]; + controlnet?: string | null; created_at: number; } @@ -187,6 +217,17 @@ export async function listDiffusionLoras(family?: string): Promise { + const qs = family ? `?family=${encodeURIComponent(family)}` : ""; + const data = await parseJson<{ controlnets: DiffusionControlNetInfo[] }>( + await authFetch(`/api/models/diffusion-controlnets${qs}`), + ); + return data.controlnets ?? []; +} + export interface GalleryPage { images: GalleryImage[]; has_more: boolean; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 58340003c9..d9f9004a27 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -46,6 +46,8 @@ import { cn } from "@/lib/utils"; import { toast } from "@/lib/toast"; import { + type ControlNetSpecInput, + type DiffusionControlNetInfo, type DiffusionGenerateProgress, type DiffusionLoadProgress, type DiffusionLoraInfo, @@ -59,6 +61,7 @@ import { getDiffusionStatus, getGallery, getGenerateProgress, + listDiffusionControlNets, listDiffusionLoras, loadDiffusionModel, unloadDiffusionModel, @@ -221,6 +224,16 @@ const ASPECT_RATIOS: Record = { }; const ASPECT_OPTIONS = ["custom", ...Object.keys(ASPECT_RATIOS)]; +// Friendly labels for ControlNet control types. "canny" traces edges from a source image; +// every other type is an already-made map (passthrough/depth/pose/...). Unknown types fall +// back to a capitalized "(map)" label so a new backend type still renders. +const CONTROL_TYPE_LABELS: Record = { + passthrough: "Passthrough (already a map)", + canny: "Canny (trace edges)", + depth: "Depth (map)", + pose: "Pose (map)", +}; + // Z-Image accepts 256–2048, in multiples of 16. Snap any value into range. const MIN_DIM = 256; const MAX_DIM = 2048; @@ -917,6 +930,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // offers. Applied at generate time; available adapters are refreshed per loaded family. const [loras, setLoras] = useState([]); const [availableLoras, setAvailableLoras] = useState([]); + // ControlNet for the next generation: the chosen model id, a control image (data URL), + // how to derive the control map, and the conditioning strength. Available models refresh + // per loaded family; applied at generate time only when a model + control image are set. + const [controlnetId, setControlnetId] = useState(""); + const [controlImage, setControlImage] = useState(null); + // Free-form: a union ControlNet advertises depth/pose/etc alongside the preprocessing + // "canny", and the backend maps the exact control_type to the union control_mode. The + // picker is built from the selected model's control_types, so it isn't limited to two. + const [controlType, setControlType] = useState("passthrough"); + const [controlStrength, setControlStrength] = useState(0.7); + const [availableControlNets, setAvailableControlNets] = useState([]); // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by // default; a single fixed toggle in the top bar opens/closes it (the icon never moves). const [advancedOpen, setAdvancedOpen] = useState(false); @@ -1022,6 +1046,50 @@ export function ImagesPage({ active = true }: { active?: boolean }) { }; }, [loraCapable, status?.family]); + // Refresh the ControlNet picker's options when the loaded model (family) changes, and clear + // a stale selection the new model can't use so an incompatible ControlNet is never sent. + const controlnetCapable = Boolean(status?.loaded && status?.supports_controlnet); + useEffect(() => { + if (!controlnetCapable) { + setAvailableControlNets([]); + setControlnetId(""); + setControlImage(null); + return; + } + let cancelled = false; + listDiffusionControlNets(status?.family ?? undefined) + .then((list) => { + if (cancelled) return; + setAvailableControlNets(list); + setControlnetId((prev) => (list.some((c) => c.id === prev) ? prev : "")); + }) + .catch(() => { + if (!cancelled) setAvailableControlNets([]); + }); + return () => { + cancelled = true; + }; + }, [controlnetCapable, status?.family]); + + // The control types offered for the selected ControlNet. A union model advertises + // several (canny/depth/pose/passthrough); a plain model advertises its own. Fall back + // to the preprocessing pair when nothing is selected. + const controlTypeOptions = useMemo(() => { + const cn = availableControlNets.find((c) => c.id === controlnetId); + const types = cn?.control_types?.length ? cn.control_types : ["passthrough", "canny"]; + return types; + }, [availableControlNets, controlnetId]); + + // Keep controlType valid for the selected model: if the current choice isn't among the + // model's advertised types, snap to the first (prefer passthrough when offered). + useEffect(() => { + if (!controlTypeOptions.includes(controlType)) { + setControlType( + controlTypeOptions.includes("passthrough") ? "passthrough" : controlTypeOptions[0], + ); + } + }, [controlTypeOptions, controlType]); + const selected = useMemo( () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, [images, selectedId], @@ -1622,6 +1690,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) { reference_images: condRefImages, // Drop zero-weight rows so the recipe records only adapters that actually applied. loras: loras.length ? loras.filter((l) => l.weight > 0) : undefined, + // ControlNet: sent only when a model + control image are chosen; v1 conditions plain + // text-to-image only, so skip it for image-conditioned workflows. + controlnet: + controlnetCapable && controlnetId && controlImage && workflow === "create" + ? { + id: controlnetId, + image: controlImage, + control_type: controlType, + strength: controlStrength, + } + : undefined, }); if (!isMounted.current) break; // Prepend this run's records (newest first) and load their blobs. @@ -1639,7 +1718,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setGenDone(null); setGenStep(null); } - }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, ensureSrc]); + }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, controlnetCapable, controlnetId, controlImage, controlType, controlStrength, ensureSrc]); // Keep the active workflow valid for the loaded model: an edit-only model (Qwen-Image- // Edit) has no Create/Transform tabs, a base model has no Edit tab. Snap to the first @@ -2160,6 +2239,59 @@ export function ImagesPage({ active = true }: { active?: boolean }) { )} + {/* ControlNet: shown when the loaded model supports it, a model is discoverable, and + the plain text-to-image workflow is active (v1 conditions txt2img only). Pick a + model, add a control image, choose how to derive the map, and set the strength. */} + {controlnetCapable && availableControlNets.length > 0 && workflow === "create" && ( + +
+ + {controlnetId && ( + <> + +
+ Control type + +
+ + + )} +
+ + )} {/* A negative prompt only does anything with guidance on, so hide it at guidance 0 (Z-Image-Turbo's default) instead of showing a dead field. */} {guidance > 0 && (