diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index e17be82c28..6bbe30b0b0 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -320,6 +320,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: @@ -1131,6 +1138,37 @@ 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 generation was cancelled.") + cn_model = getattr(diffusers, model_cls_name).from_pretrained( + resolved_cn.path, torch_dtype = state.dtype, token = state.hf_token + ).to(state.device) + 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 @@ -1253,6 +1291,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 @@ -1294,6 +1335,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 = None ref_extra: list = [] if getattr(state.family, "edit", False): # Instruction editing: the loaded pipe is the edit pipeline. It always @@ -1357,6 +1400,47 @@ 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 + + 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 # 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. @@ -1403,10 +1487,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: @@ -1417,6 +1501,20 @@ 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 gen = _GenState(total_steps = steps) @@ -1532,6 +1630,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() @@ -1559,8 +1660,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, @@ -1589,6 +1691,15 @@ class DiffusionBackend: model_kind = state.kind, transformer_quant = state.transformer_quant, ), + "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, + ), } diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py new file mode 100644 index 0000000000..22a936daac --- /dev/null +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -0,0 +1,221 @@ +# 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 +import threading +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 _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 + if not (p / "config.json").exists(): + 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, + 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. + """ + entry = _catalog_by_id().get(spec_id) + if entry is not None: + 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). + if "/" in spec_id and " " not in 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" + ) + + +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 a5ef4460c2..5f2c80735f 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -35,6 +35,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 @@ -107,6 +114,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"), @@ -180,6 +189,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 85ff4134e2..5823d88974 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -474,6 +474,9 @@ class SdCppDiffusionBackend: # LoRA adapters as (id, weight) pairs; resolved + materialized into a managed dir # and injected as prompt tags per sd-cli run. None/empty = no LoRA. loras: Optional[list[tuple[str, float]]] = None, + # ControlNet: accepted for interface parity with the diffusers backend. Native sd.cpp + # ControlNet (sd-cli --control-net) is a follow-up; a request is rejected clearly. + controlnet: Optional[tuple[str, str, str, float, float, float]] = None, ) -> dict[str, Any]: import tempfile @@ -486,6 +489,11 @@ class SdCppDiffusionBackend: "img2img / inpaint / reference 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: @@ -661,6 +669,7 @@ class SdCppDiffusionBackend: "transformer_cache": None, "engine": "sd_cpp", "supports_lora": False, + "supports_controlnet": False, } from core.inference import diffusion_lora @@ -692,6 +701,8 @@ class SdCppDiffusionBackend: model_kind = "gguf", transformer_quant = None, ), + # Native ControlNet (sd-cli --control-net) is a follow-up; off for now. + "supports_controlnet": False, } diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 863e673f50..ece1d3ed97 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1816,6 +1816,42 @@ 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" + ) + + class DiffusionGenerateRequest(BaseModel): """Request to generate one image from the loaded diffusion model.""" @@ -1880,6 +1916,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("reference_images") @classmethod @@ -1920,6 +1962,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)") @@ -2010,3 +2055,9 @@ 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.", + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 59b1920c19..1e82bc0394 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10151,6 +10151,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 @@ -10201,6 +10213,12 @@ 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}" + if request.controlnet + else None + ), "created_at": created_at, }, ) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index a83946fb54..b2adf1702b 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2025,6 +2025,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..fb267ca71f --- /dev/null +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -0,0 +1,196 @@ +"""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_local(tmp_path, monkeypatch): + d = tmp_path / "controlnets" + d.mkdir() + cn = d / "my-cn" + cn.mkdir() + (cn / "config.json").write_text("{}") + 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_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/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index b363d3cf67..47ef1ee846 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,27 @@ 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; + control_type: "passthrough" | "canny"; + 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; @@ -125,6 +152,7 @@ export interface GalleryImage { batch_index: number; model: string | null; loras?: string[]; + controlnet?: string | null; created_at: number; } @@ -186,6 +214,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 fe1ec08b5c..f80c64c849 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, @@ -877,6 +880,14 @@ export function ImagesPage() { // 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); + const [controlType, setControlType] = useState<"passthrough" | "canny">("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); @@ -962,6 +973,31 @@ export function ImagesPage() { }; }, [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]); + const selected = useMemo( () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, [images, selectedId], @@ -1386,6 +1422,17 @@ export function ImagesPage() { 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 === "txt2img" + ? { + id: controlnetId, + image: controlImage, + control_type: controlType, + strength: controlStrength, + } + : undefined, }); // Prepend this run's records (newest first) and load their blobs. setImages((prev) => [...res.images, ...prev]); @@ -1402,7 +1449,7 @@ export function ImagesPage() { 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 @@ -1921,6 +1968,58 @@ export function ImagesPage() { )} + {/* 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 === "txt2img" && ( + +
+ + {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 && (