chore: import upstream snapshot with attribution

invokeai/backend/__init__.py

@@ -0,0 +1,3 @@
+"""
+Initialization file for invokeai.backend
+"""

invokeai/backend/anima/__init__.py

@@ -0,0 +1,6 @@
+"""Anima model backend module.
+
+Anima is a 2B-parameter anime-focused text-to-image model built on NVIDIA's
+Cosmos Predict2 DiT architecture with a custom LLM Adapter that bridges Qwen3
+0.6B text encoder outputs to the DiT backbone.
+"""

invokeai/backend/anima/anima_transformer_patch.py

@@ -0,0 +1,106 @@
+"""Utilities for patching the AnimaTransformer to support regional cross-attention masks."""
+
+from contextlib import contextmanager
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from invokeai.backend.anima.regional_prompting import AnimaRegionalPromptingExtension
+
+
+def _patched_cross_attn_forward(
+    original_forward,
+    attn_mask: torch.Tensor,
+):
+    """Create a patched forward for CosmosAttention that injects a cross-attention mask.
+
+    Args:
+        original_forward: The original CosmosAttention.forward method (bound to self).
+        attn_mask: Cross-attention mask of shape (img_seq_len, context_seq_len).
+    """
+
+    def forward(x, context=None, rope_emb=None):
+        # If the context sequence length doesn't match the mask (e.g. negative conditioning
+        # has a different number of tokens than positive regional conditioning), skip masking
+        # and use the original unmasked forward.
+        actual_context = x if context is None else context
+        if actual_context.shape[-2] != attn_mask.shape[1]:
+            return original_forward(x, context, rope_emb=rope_emb)
+
+        self = original_forward.__self__
+
+        q = self.q_proj(x)
+        context = x if context is None else context
+        k = self.k_proj(context)
+        v = self.v_proj(context)
+        q, k, v = (rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim) for t in (q, k, v))
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        v = self.v_norm(v)
+
+        if self.is_selfattn and rope_emb is not None:
+            from invokeai.backend.anima.anima_transformer import apply_rotary_pos_emb_cosmos
+
+            q = apply_rotary_pos_emb_cosmos(q, rope_emb)
+            k = apply_rotary_pos_emb_cosmos(k, rope_emb)
+
+        in_q_shape = q.shape
+        in_k_shape = k.shape
+        q = rearrange(q, "b ... h d -> b h ... d").reshape(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1])
+        k = rearrange(k, "b ... h d -> b h ... d").reshape(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
+        v = rearrange(v, "b ... h d -> b h ... d").reshape(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
+
+        # Convert boolean mask to float additive mask for SDPA
+        # True (attend) -> 0.0, False (block) -> -inf
+        # Shape: (img_seq_len, context_seq_len) -> (1, 1, img_seq_len, context_seq_len)
+        float_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+        float_mask[~attn_mask] = float("-inf")
+        expanded_mask = float_mask.unsqueeze(0).unsqueeze(0)
+
+        result = F.scaled_dot_product_attention(q, k, v, attn_mask=expanded_mask)
+        result = rearrange(result, "b h s d -> b s (h d)")
+        return self.output_dropout(self.output_proj(result))
+
+    return forward
+
+
+@contextmanager
+def patch_anima_for_regional_prompting(
+    transformer,
+    regional_extension: Optional[AnimaRegionalPromptingExtension],
+):
+    """Context manager to temporarily patch the Anima transformer for regional prompting.
+
+    Patches the cross-attention in each DiT block to use a regional attention mask.
+    Uses alternating pattern: masked on even blocks, unmasked on odd blocks for
+    global coherence.
+
+    Args:
+        transformer: The AnimaTransformer instance.
+        regional_extension: The regional prompting extension. If None or no mask, no patching.
+
+    Yields:
+        The (possibly patched) transformer.
+    """
+    if regional_extension is None or regional_extension.cross_attn_mask is None:
+        yield transformer
+        return
+
+    # Store original forwards
+    original_forwards = []
+    for block_idx, block in enumerate(transformer.blocks):
+        original_forwards.append(block.cross_attn.forward)
+
+        mask = regional_extension.get_cross_attn_mask(block_idx)
+        if mask is not None:
+            block.cross_attn.forward = _patched_cross_attn_forward(block.cross_attn.forward, mask)
+
+    try:
+        yield transformer
+    finally:
+        # Restore original forwards
+        for block_idx, block in enumerate(transformer.blocks):
+            block.cross_attn.forward = original_forwards[block_idx]

invokeai/backend/anima/conditioning_data.py

@@ -0,0 +1,64 @@
+"""Anima text conditioning data structures.
+
+Anima uses a dual-conditioning scheme:
+- Qwen3 0.6B hidden states (continuous embeddings)
+- T5-XXL token IDs (discrete IDs, embedded by the LLM Adapter inside the transformer)
+
+Both are produced by the text encoder invocation and stored together.
+
+For regional prompting, multiple conditionings (each with an optional spatial mask)
+are concatenated and processed together. The LLM Adapter runs on each region's
+conditioning separately, producing per-region context vectors that are concatenated
+for the DiT's cross-attention layers. An attention mask restricts which image tokens
+attend to which regional context tokens.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range
+
+
+@dataclass
+class AnimaTextConditioning:
+    """Anima text conditioning with Qwen3 hidden states, T5-XXL token IDs, and optional mask.
+
+    Attributes:
+        qwen3_embeds: Text embeddings from Qwen3 0.6B encoder.
+            Shape: (seq_len, hidden_size) where hidden_size=1024.
+        t5xxl_ids: T5-XXL token IDs for the same prompt.
+            Shape: (seq_len,).
+        t5xxl_weights: Per-token weights for prompt weighting.
+            Shape: (seq_len,). Defaults to all ones if not provided.
+        mask: Optional binary mask for regional prompting. If None, the prompt is global.
+              Shape: (1, 1, img_seq_len) where img_seq_len = (H // patch_size) * (W // patch_size).
+    """
+
+    qwen3_embeds: torch.Tensor
+    t5xxl_ids: torch.Tensor
+    t5xxl_weights: torch.Tensor | None = None
+    mask: torch.Tensor | None = None
+
+
+@dataclass
+class AnimaRegionalTextConditioning:
+    """Container for multiple regional text conditionings processed by the LLM Adapter.
+
+    After the LLM Adapter processes each region's conditioning, the outputs are concatenated.
+    The DiT cross-attention then uses an attention mask to restrict which image tokens
+    attend to which region's context tokens.
+
+    Attributes:
+        context_embeds: Concatenated LLM Adapter outputs from all regional prompts.
+                        Shape: (total_context_len, 1024).
+        image_masks: List of binary masks for each regional prompt.
+                     If None, the prompt is global (applies to entire image).
+                     Shape: (1, 1, img_seq_len).
+        context_ranges: List of ranges indicating which portion of context_embeds
+                       corresponds to each regional prompt.
+    """
+
+    context_embeds: torch.Tensor
+    image_masks: list[torch.Tensor | None]
+    context_ranges: list[Range]

invokeai/backend/anima/control_net_lllite.py

@@ -0,0 +1,553 @@
+"""ControlNet-LLLite adapter for Anima (DiT), v2 weight format.
+
+A LLLite adapter is a shared conv trunk (``conditioning1``) that encodes a
+conditioning image into per-token embeddings, plus one tiny zero-init MLP per
+target Linear in the DiT. Each module perturbs its Linear's *input*:
+``y = org_forward(x + up(film_mlp(x, cond + depth_embed)) * multiplier)``, so
+multiplier 0 or a missing cond image is an exact passthrough.
+
+On-disk format (v2, named-key): shared trunk under ``lllite_conditioning1.*``,
+per-module weights under ``lllite_dit_blocks_{i}_{target}.{down,mid,cond_to_film,up}.*``
+plus a per-module ``.depth_embed``; hyperparams in safetensors metadata
+(``lllite.*`` keys) with state-dict-shape fallbacks.
+
+Unlike the reference implementation, the model is constructed from
+``(state_dict, metadata)`` alone — no transformer instance is needed until
+:meth:`AnimaControlNetLLLite.apply_to` binds the modules to the target Linears
+by their saved names.
+
+Original source code:
+- kohya-ss ControlNet-LLLite for Anima: ComfyUI-Anima-LLLite port
+  (``control_net_lllite_anima.py``, ``nodes.py``) of kohya-ss/sd-scripts
+  ``networks/control_net_lllite_anima.py``.
+  SPDX-License-Identifier: Apache-2.0
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Callable, Sequence
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+ASPP_DEFAULT_DILATIONS: tuple[int, ...] = (1, 2, 4, 8)
+
+_SAVED_COND_PREFIX = "lllite_conditioning1."
+_INTERNAL_COND_PREFIX = "conditioning1."
+_INTERNAL_MODULES_PREFIX = "lllite_modules."
+_LEGACY_MODULES_PREFIX = "lllite_modules."
+
+MODULE_NAME_PATTERN = re.compile(
+    r"^lllite_dit_blocks_(\d+)_(self_attn_q_proj|self_attn_k_proj|self_attn_v_proj|cross_attn_q_proj|mlp_layer1)$"
+)
+
+# Saved module name suffix -> attribute path under transformer.blocks[i].
+_SUFFIX_TO_ATTR_PATH: dict[str, tuple[str, ...]] = {
+    "self_attn_q_proj": ("self_attn", "q_proj"),
+    "self_attn_k_proj": ("self_attn", "k_proj"),
+    "self_attn_v_proj": ("self_attn", "v_proj"),
+    "cross_attn_q_proj": ("cross_attn", "q_proj"),
+    "mlp_layer1": ("mlp", "layer1"),
+}
+_SUFFIX_ORDER = list(_SUFFIX_TO_ATTR_PATH)
+
+
+# ----------------------------------------------------------------------------
+# Conditioning image preprocessing (torch-only, PIL-free)
+# ----------------------------------------------------------------------------
+
+
+def target_cond_hw(latent_h: int, latent_w: int, patch_spatial: int = 2) -> tuple[int, int]:
+    """Return the (H, W) the cond image / mask must be resized to.
+
+    The LLLite ``conditioning1`` trunk has total conv stride 16, so the cond
+    image must be sized to ``latent_HW * 8`` in input pixel space (= token_HW
+    * 16 after DiT patchify with patch_spatial=2). The DiT internally pads the
+    latent up to a multiple of ``patch_spatial`` before patchify, so the same
+    rounding is mirrored here — otherwise odd latent dims yield a token-count
+    mismatch that silently bypasses every LLLite module.
+    """
+    padded_h = ((latent_h + patch_spatial - 1) // patch_spatial) * patch_spatial
+    padded_w = ((latent_w + patch_spatial - 1) // patch_spatial) * patch_spatial
+    return padded_h * 8, padded_w * 8
+
+
+def prepare_cond_image(rgb_bchw_01: torch.Tensor, latent_h: int, latent_w: int, patch_spatial: int = 2) -> torch.Tensor:
+    """RGB image (B, 3, H, W) in [0, 1] -> (1, 3, H_t, W_t) in [-1, 1]."""
+    if rgb_bchw_01.ndim != 4 or rgb_bchw_01.shape[1] != 3:
+        raise ValueError(f"Unexpected cond image shape: {tuple(rgb_bchw_01.shape)} (expected B,3,H,W)")
+    img = rgb_bchw_01[:1]
+    target_h, target_w = target_cond_hw(latent_h, latent_w, patch_spatial)
+    if img.shape[-2] != target_h or img.shape[-1] != target_w:
+        img = F.interpolate(img, size=(target_h, target_w), mode="bicubic", align_corners=False)
+        img = img.clamp(0.0, 1.0)
+    return img * 2.0 - 1.0
+
+
+def prepare_mask(mask_b1hw_01: torch.Tensor, latent_h: int, latent_w: int, patch_spatial: int = 2) -> torch.Tensor:
+    """Mask (B, 1, H, W) or (B, H, W) in [0, 1] -> (1, 1, H_t, W_t) in {0.0, 1.0}.
+
+    1 = inpaint area, 0 = keep. The caller is responsible for the ``*2-1``
+    rescale before concat with RGB (see :func:`build_inpaint_cond_image`).
+    """
+    if mask_b1hw_01.ndim == 3:
+        m = mask_b1hw_01.unsqueeze(1)
+    elif mask_b1hw_01.ndim == 4 and mask_b1hw_01.shape[1] == 1:
+        m = mask_b1hw_01
+    else:
+        raise ValueError(f"Unexpected mask shape: {tuple(mask_b1hw_01.shape)} (expected B,H,W or B,1,H,W)")
+    m = m[:1].float()
+    target_h, target_w = target_cond_hw(latent_h, latent_w, patch_spatial)
+    if m.shape[-2] != target_h or m.shape[-1] != target_w:
+        m = F.interpolate(m, size=(target_h, target_w), mode="nearest")
+    return (m >= 0.5).float()
+
+
+def build_inpaint_cond_image(rgb_pm1: torch.Tensor, mask01: torch.Tensor, masked_input: bool) -> torch.Tensor:
+    """rgb_pm1: (1, 3, H, W) in [-1, 1], mask01: (1, 1, H, W) in {0, 1}. Returns (1, 4, H, W).
+
+    The mask channel is rescaled to [-1, +1] (matching the RGB range), and if
+    ``masked_input`` is set the RGB is zeroed where ``mask >= 0.5``.
+    """
+    if masked_input:
+        keep = (mask01 < 0.5).to(rgb_pm1.dtype)
+        rgb_pm1 = rgb_pm1 * keep
+    mask_pm1 = mask01.to(rgb_pm1.dtype) * 2.0 - 1.0
+    return torch.cat([rgb_pm1, mask_pm1], dim=1)
+
+
+# ----------------------------------------------------------------------------
+# Conditioning1 trunk (v2)
+# ----------------------------------------------------------------------------
+
+
+def _gn(channels: int) -> nn.GroupNorm:
+    g = 8
+    while g > 1 and channels % g != 0:
+        g //= 2
+    return nn.GroupNorm(g, channels)
+
+
+class _ResBlock(nn.Module):
+    def __init__(self, ch: int):
+        super().__init__()
+        self.norm1 = _gn(ch)
+        self.conv1 = nn.Conv2d(ch, ch, kernel_size=3, padding=1)
+        self.norm2 = _gn(ch)
+        self.conv2 = nn.Conv2d(ch, ch, kernel_size=3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.conv1(F.silu(self.norm1(x)))
+        h = self.conv2(F.silu(self.norm2(h)))
+        return x + h
+
+
+class _ASPP(nn.Module):
+    def __init__(self, ch: int, dilations: tuple[int, ...] = ASPP_DEFAULT_DILATIONS):
+        super().__init__()
+        assert len(dilations) >= 1, "ASPP needs at least one dilation"
+        branches = []
+        for d in dilations:
+            if d == 1:
+                conv = nn.Conv2d(ch, ch, kernel_size=1)
+            else:
+                conv = nn.Conv2d(ch, ch, kernel_size=3, padding=d, dilation=d)
+            branches.append(nn.Sequential(conv, _gn(ch), nn.SiLU()))
+        self.branches = nn.ModuleList(branches)
+
+        self.global_pool = nn.AdaptiveAvgPool2d(1)
+        self.global_conv = nn.Sequential(nn.Conv2d(ch, ch, kernel_size=1), _gn(ch), nn.SiLU())
+
+        n_branches = len(dilations) + 1
+        self.proj = nn.Sequential(nn.Conv2d(ch * n_branches, ch, kernel_size=1), _gn(ch), nn.SiLU())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h, w = x.shape[-2:]
+        outs = [b(x) for b in self.branches]
+        g = self.global_conv(self.global_pool(x))
+        g = F.interpolate(g, size=(h, w), mode="bilinear", align_corners=False)
+        outs.append(g)
+        return self.proj(torch.cat(outs, dim=1))
+
+
+class _Conditioning1(nn.Module):
+    def __init__(
+        self,
+        cond_dim: int,
+        cond_emb_dim: int,
+        n_resblocks: int,
+        use_aspp: bool = False,
+        aspp_dilations: tuple[int, ...] = ASPP_DEFAULT_DILATIONS,
+        cond_in_channels: int = 3,
+    ):
+        super().__init__()
+        assert cond_dim % 2 == 0, f"cond_dim must be even, got {cond_dim}"
+        assert cond_in_channels >= 1, f"cond_in_channels must be >= 1, got {cond_in_channels}"
+        ch_half = cond_dim // 2
+
+        self.cond_in_channels = cond_in_channels
+        self.conv1 = nn.Conv2d(cond_in_channels, ch_half, kernel_size=4, stride=4, padding=0)
+        self.norm1 = _gn(ch_half)
+        self.conv2 = nn.Conv2d(ch_half, ch_half, kernel_size=3, stride=1, padding=1)
+        self.norm2 = _gn(ch_half)
+        self.conv3 = nn.Conv2d(ch_half, cond_dim, kernel_size=4, stride=4, padding=0)
+        self.norm3 = _gn(cond_dim)
+
+        self.resblocks = nn.ModuleList([_ResBlock(cond_dim) for _ in range(n_resblocks)])
+        self.aspp = _ASPP(cond_dim, aspp_dilations) if use_aspp else None
+
+        self.proj = nn.Conv2d(cond_dim, cond_emb_dim, kernel_size=1)
+        self.out_norm = nn.LayerNorm(cond_emb_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = F.silu(self.norm1(self.conv1(x)))
+        h = F.silu(self.norm2(self.conv2(h)))
+        h = F.silu(self.norm3(self.conv3(h)))
+        for rb in self.resblocks:
+            h = rb(h)
+        if self.aspp is not None:
+            h = self.aspp(h)
+        h = self.proj(h)
+        b, c, hh, ww = h.shape
+        h = h.view(b, c, hh * ww).permute(0, 2, 1).contiguous()
+        h = self.out_norm(h)
+        return h
+
+
+# ----------------------------------------------------------------------------
+# LLLite module (v2: FiLM + SiLU + 5D path + per-module depth embedding)
+# ----------------------------------------------------------------------------
+
+
+class LLLiteModuleDiT(nn.Module):
+    def __init__(
+        self,
+        name: str,
+        in_dim: int,
+        cond_emb_dim: int,
+        mlp_dim: int,
+        dropout: float | None = None,
+        multiplier: float = 1.0,
+    ):
+        super().__init__()
+        self.lllite_name = name
+        self.in_dim = in_dim
+        self.cond_emb_dim = cond_emb_dim
+        self.mlp_dim = mlp_dim
+        self.dropout = dropout
+        self.multiplier = multiplier
+
+        self.down = nn.Linear(in_dim, mlp_dim)
+        self.mid = nn.Linear(mlp_dim + cond_emb_dim, mlp_dim)
+
+        # FiLM: cond_local -> (gamma, beta), zero-init for identity at start.
+        self.cond_to_film = nn.Linear(cond_emb_dim, 2 * mlp_dim)
+        nn.init.zeros_(self.cond_to_film.weight)
+        nn.init.zeros_(self.cond_to_film.bias)
+
+        self.up = nn.Linear(mlp_dim, in_dim)
+        nn.init.zeros_(self.up.weight)
+        nn.init.zeros_(self.up.bias)
+
+        self.depth_embed = nn.Parameter(torch.zeros(cond_emb_dim))
+
+        self.cond_emb: torch.Tensor | None = None
+        # Wrapped in a list so the original Linear is not registered as a
+        # submodule and its weights stay out of state_dict.
+        self._org_module: list[nn.Linear] = []
+        self._org_forward: Callable[[torch.Tensor], torch.Tensor] | None = None
+        self._org_forward_was_instance_attr = False
+
+    def bind(self, org_module: nn.Linear) -> None:
+        self.unbind()
+        self._org_module = [org_module]
+        self._org_forward = org_module.forward
+        self._org_forward_was_instance_attr = "forward" in org_module.__dict__
+        org_module.forward = self.forward  # type: ignore[method-assign]
+
+    def unbind(self) -> None:
+        if self._org_forward is not None:
+            org_module = self._org_module[0]
+            if self._org_forward_was_instance_attr:
+                org_module.forward = self._org_forward  # type: ignore[method-assign]
+            else:
+                # Restoring by assignment would pin a frozen bound method in the
+                # instance __dict__, which silently bypasses later class-level
+                # forward swaps that share the module __dict__ (see
+                # wrap_custom_layer notes in model_manager/load/load_default.py).
+                del org_module.__dict__["forward"]
+            self._org_forward = None
+            self._org_module = []
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Input layouts:
+        #   self/cross attention q/k/v: (B, S, D) — already flattened in the Anima block
+        #   mlp.layer1:                 (B, T, H, W, D) — passed un-flattened
+        # Flatten the 5D case to 3D for the LLLite path and reshape on exit.
+        assert self._org_forward is not None
+        if self.multiplier == 0.0 or self.cond_emb is None:
+            return self._org_forward(x)
+
+        orig_shape = x.shape
+        is_5d = x.dim() == 5
+        if is_5d:
+            b, t, hh, ww, d = orig_shape
+            x = x.reshape(b, t * hh * ww, d)
+
+        cx = self.cond_emb  # (B_c, S, cond_emb_dim)
+
+        # Broadcast cond_emb to the runtime batch (CFG cond+uncond, multi-cond).
+        if x.shape[0] != cx.shape[0]:
+            if x.shape[0] % cx.shape[0] != 0:
+                return self._org_forward(x.reshape(orig_shape) if is_5d else x)
+            cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1)
+
+        if x.shape[1] != cx.shape[1]:
+            return self._org_forward(x.reshape(orig_shape) if is_5d else x)
+
+        # Run the LLLite mini-MLP in its own parameter dtype, then cast the
+        # correction back to ``x``'s dtype before adding. Robust to autocast
+        # flows where x and LLLite weights have different dtypes.
+        param_dtype = self.down.weight.dtype
+        x_proc = x if x.dtype == param_dtype else x.to(param_dtype)
+        if cx.dtype != param_dtype or cx.device != x.device:
+            cx = cx.to(device=x.device, dtype=param_dtype)
+
+        depth_e = self.depth_embed
+        if depth_e.dtype != param_dtype or depth_e.device != x.device:
+            depth_e = depth_e.to(device=x.device, dtype=param_dtype)
+        cond_local = cx + depth_e
+
+        h = F.silu(self.down(x_proc))
+
+        gb = self.cond_to_film(cond_local)
+        gamma, beta = gb.chunk(2, dim=-1)
+
+        m = self.mid(torch.cat([cond_local, h], dim=-1))
+        m = m * (1 + gamma) + beta
+        m = F.silu(m)
+
+        if self.dropout is not None and self.training:
+            m = F.dropout(m, p=self.dropout)
+
+        out = self.up(m) * self.multiplier
+        if out.dtype != x.dtype:
+            out = out.to(x.dtype)
+
+        y = self._org_forward(x + out)
+
+        if is_5d:
+            # org Linear out_features may differ from in_features — recover with -1.
+            y = y.reshape(orig_shape[0], orig_shape[1], orig_shape[2], orig_shape[3], -1)
+        return y
+
+
+# ----------------------------------------------------------------------------
+# AnimaControlNetLLLite
+# ----------------------------------------------------------------------------
+
+
+def _meta_int(metadata: dict[str, str], key: str, fallback: int) -> int:
+    value = metadata.get(key)
+    return int(value) if value is not None else fallback
+
+
+def _meta_bool(metadata: dict[str, str], key: str, fallback: bool) -> bool:
+    value = metadata.get(key)
+    return str(value).lower() == "true" if value is not None else fallback
+
+
+class AnimaControlNetLLLite(nn.Module):
+    """Self-contained, cacheable LLLite adapter for the Anima transformer.
+
+    Construct via :meth:`from_state_dict`; bind to a transformer with
+    :meth:`apply_to` and undo with :meth:`restore`.
+    """
+
+    def __init__(
+        self,
+        module_specs: Sequence[tuple[str, int]],
+        cond_emb_dim: int,
+        mlp_dim: int,
+        cond_dim: int,
+        cond_resblocks: int,
+        use_aspp: bool = False,
+        aspp_dilations: tuple[int, ...] = ASPP_DEFAULT_DILATIONS,
+        cond_in_channels: int = 3,
+        inpaint_masked_input: bool = False,
+        multiplier: float = 1.0,
+    ):
+        super().__init__()
+        self.cond_emb_dim = cond_emb_dim
+        self.mlp_dim = mlp_dim
+        self.cond_dim = cond_dim
+        self.cond_resblocks = cond_resblocks
+        self.use_aspp = use_aspp
+        self.cond_in_channels = cond_in_channels
+        # Training-time RGB-masking policy for cond image preparation; does not
+        # alter the forward pass.
+        self.inpaint_masked_input = inpaint_masked_input
+        self.multiplier = multiplier
+
+        self.conditioning1 = _Conditioning1(
+            cond_dim,
+            cond_emb_dim,
+            cond_resblocks,
+            use_aspp=use_aspp,
+            aspp_dilations=aspp_dilations,
+            cond_in_channels=cond_in_channels,
+        )
+
+        modules = []
+        for name, in_dim in module_specs:
+            if MODULE_NAME_PATTERN.match(name) is None:
+                raise ValueError(f"Unrecognized LLLite module name: '{name}'")
+            modules.append(LLLiteModuleDiT(name, in_dim, cond_emb_dim, mlp_dim, multiplier=multiplier))
+        self.lllite_modules = nn.ModuleList(modules)
+
+    @classmethod
+    def from_state_dict(
+        cls, state_dict: dict[str, torch.Tensor], metadata: dict[str, str] | None
+    ) -> AnimaControlNetLLLite:
+        """Build the adapter from a saved v2 named-key state dict.
+
+        Hyperparams come from ``lllite.*`` metadata when present, with
+        state-dict-shape fallbacks. ``inpaint_masked_input`` is metadata-only
+        (not derivable from shapes; defaults to False).
+        """
+        meta = metadata or {}
+
+        if any(k.startswith(_LEGACY_MODULES_PREFIX) for k in state_dict):
+            raise ValueError(
+                f"State dict appears to be in a legacy ControlNet-LLLite weight format (keys starting "
+                f"with '{_LEGACY_MODULES_PREFIX}'). Only the v2 named-key format is supported."
+            )
+
+        module_names: set[str] = set()
+        for key in state_dict:
+            head, dot, _tail = key.partition(".")
+            if dot and MODULE_NAME_PATTERN.match(head):
+                module_names.add(head)
+        if not module_names:
+            raise ValueError("State dict contains no LLLite modules (no 'lllite_dit_blocks_*' keys).")
+
+        def sort_key(name: str) -> tuple[int, int]:
+            match = MODULE_NAME_PATTERN.match(name)
+            assert match is not None
+            return int(match.group(1)), _SUFFIX_ORDER.index(match.group(2))
+
+        sorted_names = sorted(module_names, key=sort_key)
+        module_specs: list[tuple[str, int]] = []
+        for name in sorted_names:
+            down_key = f"{name}.down.weight"
+            if down_key not in state_dict:
+                raise ValueError(f"LLLite module '{name}' is missing key '{down_key}'")
+            module_specs.append((name, state_dict[down_key].shape[1]))
+
+        conv1_weight = state_dict[f"{_SAVED_COND_PREFIX}conv1.weight"]
+        conv3_weight = state_dict[f"{_SAVED_COND_PREFIX}conv3.weight"]
+        proj_weight = state_dict[f"{_SAVED_COND_PREFIX}proj.weight"]
+        resblock_indices = {
+            m.group(1) for m in (re.match(rf"^{_SAVED_COND_PREFIX}resblocks\.(\d+)\.", k) for k in state_dict) if m
+        }
+        has_aspp_keys = any(k.startswith(f"{_SAVED_COND_PREFIX}aspp.") for k in state_dict)
+
+        use_aspp = _meta_bool(meta, "lllite.use_aspp", has_aspp_keys)
+        aspp_dilations_meta = meta.get("lllite.aspp_dilations")
+        if use_aspp and aspp_dilations_meta:
+            aspp_dilations = tuple(int(d) for d in aspp_dilations_meta.split(",") if d.strip())
+        else:
+            aspp_dilations = ASPP_DEFAULT_DILATIONS
+
+        model = cls(
+            module_specs=module_specs,
+            cond_emb_dim=_meta_int(meta, "lllite.cond_emb_dim", proj_weight.shape[0]),
+            mlp_dim=_meta_int(meta, "lllite.mlp_dim", state_dict[f"{sorted_names[0]}.down.weight"].shape[0]),
+            cond_dim=_meta_int(meta, "lllite.cond_dim", conv3_weight.shape[0]),
+            cond_resblocks=_meta_int(meta, "lllite.cond_resblocks", len(resblock_indices)),
+            use_aspp=use_aspp,
+            aspp_dilations=aspp_dilations,
+            cond_in_channels=_meta_int(meta, "lllite.cond_in_channels", conv1_weight.shape[1]),
+            inpaint_masked_input=_meta_bool(meta, "lllite.inpaint_masked_input", False),
+        )
+
+        name_to_idx = {name: i for i, name in enumerate(sorted_names)}
+        remapped: dict[str, torch.Tensor] = {}
+        for key, value in state_dict.items():
+            if key.startswith(_SAVED_COND_PREFIX):
+                remapped[_INTERNAL_COND_PREFIX + key[len(_SAVED_COND_PREFIX) :]] = value
+                continue
+            head, dot, tail = key.partition(".")
+            if dot and head in name_to_idx:
+                remapped[f"{_INTERNAL_MODULES_PREFIX}{name_to_idx[head]}.{tail}"] = value
+            else:
+                # Unknown keys are passed through so strict loading reports them.
+                remapped[key] = value
+
+        model.load_state_dict(remapped, strict=True)
+        model.eval().requires_grad_(False)
+        return model
+
+    def set_cond_image(self, cond: torch.Tensor | None) -> None:
+        """cond: (B, cond_in_channels, H_t, W_t) in [-1, 1]; ``None`` clears."""
+        if cond is None:
+            for m in self.lllite_modules:
+                m.cond_emb = None
+            return
+        trunk_weight = self.conditioning1.conv1.weight
+        cond = cond.to(device=trunk_weight.device, dtype=trunk_weight.dtype)
+        cx = self.conditioning1(cond)  # (B, S, cond_emb_dim)
+        for m in self.lllite_modules:
+            m.cond_emb = cx
+
+    def clear_cond_image(self) -> None:
+        self.set_cond_image(None)
+
+    def set_multiplier(self, multiplier: float) -> None:
+        self.multiplier = multiplier
+        for m in self.lllite_modules:
+            m.multiplier = multiplier
+
+    def apply_to(self, transformer: nn.Module) -> None:
+        """Swap the forward of each target Linear in ``transformer``. Idempotent."""
+        self.restore()
+        for m in self.lllite_modules:
+            target = self._resolve_target(transformer, m.lllite_name)
+            if not isinstance(target, nn.Linear):
+                raise TypeError(f"LLLite target for '{m.lllite_name}' is {type(target).__name__}, expected nn.Linear")
+            if target.in_features != m.in_dim:
+                raise ValueError(
+                    f"LLLite module '{m.lllite_name}' was trained for in_features={m.in_dim}, but the "
+                    f"target Linear has in_features={target.in_features}"
+                )
+            m.bind(target)
+
+    def restore(self) -> None:
+        """Undo :meth:`apply_to`. Safe to call when not applied.
+
+        LIFO contract: each bind saves the forward that was CURRENT at bind
+        time, so when multiple adapters are stacked on one transformer they
+        must be restored in reverse apply order. Restoring an earlier adapter
+        first would delete a later adapter's wrapper and re-pin the earlier
+        one's saved forward.
+        """
+        for m in self.lllite_modules:
+            m.unbind()
+
+    @staticmethod
+    def _resolve_target(transformer: nn.Module, name: str) -> nn.Module:
+        match = MODULE_NAME_PATTERN.match(name)
+        if match is None:
+            raise ValueError(f"Unrecognized LLLite module name: '{name}'")
+        block_idx = int(match.group(1))
+        blocks = transformer.blocks
+        if block_idx >= len(blocks):
+            raise ValueError(
+                f"LLLite module '{name}' targets block {block_idx}, but the transformer has only {len(blocks)} blocks"
+            )
+        target: nn.Module = blocks[block_idx]
+        for attr in _SUFFIX_TO_ATTR_PATH[match.group(2)]:
+            target = getattr(target, attr)
+        return target

invokeai/backend/anima/regional_prompting.py

@@ -0,0 +1,173 @@
+"""Regional prompting extension for Anima.
+
+Anima's architecture uses separate cross-attention in each DiT block: image tokens
+(in 5D spatial layout) cross-attend to context tokens (LLM Adapter output). This is
+different from Z-Image's unified [img, txt] sequence with self-attention.
+
+For regional prompting, we:
+1. Run the LLM Adapter separately for each regional prompt
+2. Concatenate the resulting context vectors
+3. Build a cross-attention mask that restricts each image region to attend only to
+   its corresponding context tokens
+4. Patch the DiT's cross-attention to use this mask
+
+The mask alternation strategy (masked on even blocks, full on odd blocks) helps
+maintain global coherence across regions.
+"""
+
+from typing import Optional
+
+import torch
+import torchvision
+
+from invokeai.backend.anima.conditioning_data import AnimaRegionalTextConditioning
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.mask import to_standard_float_mask
+
+
+class AnimaRegionalPromptingExtension:
+    """Manages regional prompting for Anima's cross-attention.
+
+    Unlike Z-Image which uses a unified [img, txt] sequence, Anima has separate
+    cross-attention where image tokens (query) attend to context tokens (key/value).
+    The cross-attention mask shape is (img_seq_len, context_seq_len).
+    """
+
+    def __init__(
+        self,
+        regional_text_conditioning: AnimaRegionalTextConditioning,
+        cross_attn_mask: torch.Tensor | None = None,
+    ):
+        self.regional_text_conditioning = regional_text_conditioning
+        self.cross_attn_mask = cross_attn_mask
+
+    def get_cross_attn_mask(self, block_index: int) -> torch.Tensor | None:
+        """Get the cross-attention mask for a given block index.
+
+        Uses alternating pattern: apply mask on even blocks, no mask on odd blocks.
+        This helps balance regional control with global coherence.
+        """
+        if block_index % 2 == 0:
+            return self.cross_attn_mask
+        return None
+
+    @classmethod
+    def from_regional_conditioning(
+        cls,
+        regional_text_conditioning: AnimaRegionalTextConditioning,
+        img_seq_len: int,
+    ) -> "AnimaRegionalPromptingExtension":
+        """Create extension from pre-processed regional conditioning.
+
+        Args:
+            regional_text_conditioning: Regional conditioning with concatenated context and masks.
+            img_seq_len: Number of image tokens (H_patches * W_patches).
+        """
+        cross_attn_mask = cls._prepare_cross_attn_mask(regional_text_conditioning, img_seq_len)
+        return cls(
+            regional_text_conditioning=regional_text_conditioning,
+            cross_attn_mask=cross_attn_mask,
+        )
+
+    @classmethod
+    def _prepare_cross_attn_mask(
+        cls,
+        regional_text_conditioning: AnimaRegionalTextConditioning,
+        img_seq_len: int,
+    ) -> torch.Tensor | None:
+        """Prepare a cross-attention mask for regional prompting.
+
+        The mask shape is (img_seq_len, context_seq_len) where:
+        - Each image token can attend to context tokens from its assigned region
+        - Global prompts (mask=None) attend to background regions
+
+        Args:
+            regional_text_conditioning: The regional text conditioning data.
+            img_seq_len: Number of image tokens.
+
+        Returns:
+            Cross-attention mask of shape (img_seq_len, context_seq_len), or None
+            if no regional masks are present.
+        """
+        has_regional_masks = any(mask is not None for mask in regional_text_conditioning.image_masks)
+        if not has_regional_masks:
+            return None
+
+        # Identify background region (area not covered by any mask)
+        background_region_mask: torch.Tensor | None = None
+        for image_mask in regional_text_conditioning.image_masks:
+            if image_mask is not None:
+                mask_flat = image_mask.view(-1)
+                if background_region_mask is None:
+                    background_region_mask = torch.ones_like(mask_flat)
+                background_region_mask = background_region_mask * (1 - mask_flat)
+
+        device = TorchDevice.choose_torch_device()
+        context_seq_len = regional_text_conditioning.context_embeds.shape[0]
+
+        # Cross-attention mask: (img_seq_len, context_seq_len)
+        # img tokens are queries, context tokens are keys/values
+        cross_attn_mask = torch.zeros((img_seq_len, context_seq_len), device=device, dtype=torch.float16)
+
+        for image_mask, context_range in zip(
+            regional_text_conditioning.image_masks,
+            regional_text_conditioning.context_ranges,
+            strict=True,
+        ):
+            ctx_start = context_range.start
+            ctx_end = context_range.end
+
+            if image_mask is not None:
+                # Regional prompt: only masked image tokens attend to this region's context
+                mask_flat = image_mask.view(img_seq_len)
+                cross_attn_mask[:, ctx_start:ctx_end] = mask_flat.view(img_seq_len, 1)
+            else:
+                # Global prompt: background image tokens attend to this context
+                if background_region_mask is not None:
+                    cross_attn_mask[:, ctx_start:ctx_end] = background_region_mask.view(img_seq_len, 1)
+                else:
+                    cross_attn_mask[:, ctx_start:ctx_end] = 1.0
+
+        # Convert to boolean
+        cross_attn_mask = cross_attn_mask > 0.5
+        return cross_attn_mask
+
+    @staticmethod
+    def preprocess_regional_prompt_mask(
+        mask: Optional[torch.Tensor],
+        target_height: int,
+        target_width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Preprocess a regional prompt mask to match the target image token grid.
+
+        Args:
+            mask: Input mask tensor. If None, returns a mask of all ones.
+            target_height: Height of the image token grid (H // patch_size).
+            target_width: Width of the image token grid (W // patch_size).
+            dtype: Target dtype for the mask.
+            device: Target device for the mask.
+
+        Returns:
+            Processed mask of shape (1, 1, target_height * target_width).
+        """
+        img_seq_len = target_height * target_width
+
+        if mask is None:
+            return torch.ones((1, 1, img_seq_len), dtype=dtype, device=device)
+
+        mask = to_standard_float_mask(mask, out_dtype=dtype)
+
+        tf = torchvision.transforms.Resize(
+            (target_height, target_width),
+            interpolation=torchvision.transforms.InterpolationMode.NEAREST,
+        )
+
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(0)
+
+        resized_mask = tf(mask)
+        return resized_mask.flatten(start_dim=2).to(device=device)

invokeai/backend/anima/scheduler_driver.py

@@ -0,0 +1,150 @@
+"""Anima scheduler driver.
+
+Encapsulates the per-scheduler API quirks that ``anima_denoise._run_diffusion``
+would otherwise have to know about:
+
+* Schedulers that accept ``set_timesteps(sigmas=...)`` get the pre-shifted
+  Anima schedule passed directly.
+* Schedulers that don't accept ``sigmas=`` use ``set_begin_index()`` over their
+  own internal flow-shifted schedule. For Heun, the doubled-array index
+  translation (logical step ``k`` → doubled index ``2k``) is handled here.
+* SDE-style schedulers receive a seeded ``torch.Generator`` on every step.
+
+The denoise loop iterates :meth:`AnimaSchedulerDriver.iterations` and calls
+:meth:`AnimaSchedulerDriver.step` per iteration; the driver yields the
+``sigma_prev`` and ``completes_user_step`` flags the caller needs for inpaint
+mixing and progress reporting.
+"""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Iterator
+
+import torch
+from diffusers import FlowMatchHeunDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from invokeai.backend.flux.schedulers import ANIMA_SCHEDULER_MAP
+
+
+@dataclass(frozen=True)
+class AnimaSchedulerIteration:
+    """Per-iteration metadata yielded by :meth:`AnimaSchedulerDriver.iterations`.
+
+    ``sigma_prev`` is the noise level the latents will be at after this iteration's
+    :meth:`AnimaSchedulerDriver.step` call. ``completes_user_step`` is True when
+    this iteration finishes a user-visible step — for Heun, the second-order
+    half of each pair plus the unpaired terminal first-order step; for every
+    other scheduler, always True.
+    """
+
+    sched_timestep: torch.Tensor
+    sigma_curr: float
+    sigma_prev: float
+    completes_user_step: bool
+    order: int
+
+
+class AnimaSchedulerDriver:
+    """Drives a diffusers scheduler over Anima's pre-shifted sigma schedule."""
+
+    def __init__(
+        self,
+        scheduler_name: str,
+        sigmas: list[float],
+        steps: int,
+        denoising_start: float,
+        denoising_end: float,
+        device: torch.device,
+        seed: int,
+    ):
+        scheduler_class, scheduler_kwargs = ANIMA_SCHEDULER_MAP[scheduler_name]
+        self.scheduler: SchedulerMixin = scheduler_class(num_train_timesteps=1000, **scheduler_kwargs)
+        # Heun toggles state_in_first_order during step(); detect by class so we
+        # can read it before set_timesteps has run.
+        self.is_heun: bool = isinstance(self.scheduler, FlowMatchHeunDiscreteScheduler)
+        self._begin_index: int = 0
+        self._step_generator = torch.Generator(device=device).manual_seed(seed)
+
+        is_lcm = scheduler_name == "lcm"
+        accepts_sigmas = "sigmas" in inspect.signature(self.scheduler.set_timesteps).parameters
+        clipped = denoising_start > 0 or denoising_end < 1
+
+        if not is_lcm and accepts_sigmas:
+            self.scheduler.set_timesteps(sigmas=sigmas, device=device)
+            self._num_iterations = len(self.scheduler.timesteps)
+        elif not is_lcm and clipped and hasattr(self.scheduler, "set_begin_index"):
+            k_start = int(denoising_start * steps)
+            k_end = int(denoising_end * steps)
+            self.scheduler.set_timesteps(num_inference_steps=steps, device=device)
+            if self.is_heun:
+                # Heun's timesteps array is 2N-1 entries; logical step k maps to
+                # doubled index 2k. min() clamps denoising_end=1.0 to the
+                # unpaired terminal first-order step.
+                self._begin_index = 2 * k_start
+                self._num_iterations = min(
+                    2 * (k_end - k_start),
+                    len(self.scheduler.timesteps) - self._begin_index,
+                )
+            else:
+                self._begin_index = k_start
+                self._num_iterations = k_end - self._begin_index
+            self.scheduler.set_begin_index(self._begin_index)
+        else:
+            self.scheduler.set_timesteps(num_inference_steps=len(sigmas) - 1, device=device)
+            self._num_iterations = len(self.scheduler.timesteps)
+
+    @property
+    def num_iterations(self) -> int:
+        """Total :meth:`step` calls. For Heun this is roughly 2× the user-visible step count."""
+        return self._num_iterations
+
+    @property
+    def begin_index(self) -> int:
+        return self._begin_index
+
+    def iterations(self) -> Iterator[AnimaSchedulerIteration]:
+        for i in range(self._num_iterations):
+            sched_idx = i + self._begin_index
+            sched_timestep = self.scheduler.timesteps[sched_idx]
+            sigma_curr = sched_timestep.item() / self.scheduler.config.num_train_timesteps
+
+            # Read state_in_first_order before step (Heun toggles it inside step()).
+            in_first_order = self.scheduler.state_in_first_order if self.is_heun else True
+
+            next_idx = sched_idx + 1
+            sigma_prev = self.scheduler.sigmas[next_idx].item() if next_idx < len(self.scheduler.sigmas) else 0.0
+
+            # For Heun, a user step completes on the second-order half of each
+            # pair AND on the unpaired terminal first-order step (sigma_prev==0).
+            is_terminal = sigma_prev == 0.0
+            completes_user_step = (not self.is_heun) or (not in_first_order) or is_terminal
+            order = 2 if self.is_heun else 1
+
+            yield AnimaSchedulerIteration(
+                sched_timestep=sched_timestep,
+                sigma_curr=sigma_curr,
+                sigma_prev=sigma_prev,
+                completes_user_step=completes_user_step,
+                order=order,
+            )
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: torch.Tensor,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        step_output = self.scheduler.step(
+            model_output=model_output,
+            timestep=timestep,
+            sample=sample,
+            generator=self._step_generator,
+        )
+        return step_output.prev_sample
+
+    @property
+    def step_generator(self) -> torch.Generator:
+        return self._step_generator

invokeai/backend/anima/t5_tokenizer.py

@@ -0,0 +1,25 @@
+"""Bundled T5-XXL tokenizer for Anima.
+
+Anima tokenizes the prompt with the T5-XXL tokenizer to produce token IDs that
+index the LLM Adapter's learned embedding table. Only the tokenizer is needed —
+never the 9GB T5-XXL weights — so the tokenizer is vendored in the package as a
+self-contained fast tokenizer (tokenizer.json), avoiding both the large download
+and the sentencepiece runtime path.
+"""
+
+from functools import lru_cache
+from pathlib import Path
+
+from transformers import T5TokenizerFast
+
+# Size of the LLM Adapter's token embedding table (T5 v1.1 vocab incl. 100 sentinel
+# extra_id tokens). Token IDs must stay within this range.
+ANIMA_T5_VOCAB_SIZE = 32128
+
+_TOKENIZER_DIR = Path(__file__).parent / "tokenizer"
+
+
+@lru_cache(maxsize=1)
+def load_bundled_t5_tokenizer() -> T5TokenizerFast:
+    """Load the vendored T5-XXL fast tokenizer. Result is cached for the process."""
+    return T5TokenizerFast.from_pretrained(_TOKENIZER_DIR)

invokeai/backend/anima/tokenizer/special_tokens_map.json

@@ -0,0 +1,125 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

invokeai/backend/anima/tokenizer/tokenizer_config.json

@@ -0,0 +1,941 @@
+{
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<extra_id_99>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<extra_id_98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<extra_id_97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<extra_id_96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<extra_id_95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<extra_id_94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<extra_id_93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<extra_id_92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<extra_id_91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32009": {
+      "content": "<extra_id_90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32010": {
+      "content": "<extra_id_89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32011": {
+      "content": "<extra_id_88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32012": {
+      "content": "<extra_id_87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32013": {
+      "content": "<extra_id_86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32014": {
+      "content": "<extra_id_85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32015": {
+      "content": "<extra_id_84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32016": {
+      "content": "<extra_id_83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32017": {
+      "content": "<extra_id_82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32018": {
+      "content": "<extra_id_81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32019": {
+      "content": "<extra_id_80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32020": {
+      "content": "<extra_id_79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32021": {
+      "content": "<extra_id_78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32022": {
+      "content": "<extra_id_77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32023": {
+      "content": "<extra_id_76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32024": {
+      "content": "<extra_id_75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32025": {
+      "content": "<extra_id_74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32026": {
+      "content": "<extra_id_73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32027": {
+      "content": "<extra_id_72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32028": {
+      "content": "<extra_id_71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32029": {
+      "content": "<extra_id_70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32030": {
+      "content": "<extra_id_69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32031": {
+      "content": "<extra_id_68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32032": {
+      "content": "<extra_id_67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32033": {
+      "content": "<extra_id_66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32034": {
+      "content": "<extra_id_65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32035": {
+      "content": "<extra_id_64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32036": {
+      "content": "<extra_id_63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32037": {
+      "content": "<extra_id_62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32038": {
+      "content": "<extra_id_61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32039": {
+      "content": "<extra_id_60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32040": {
+      "content": "<extra_id_59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32041": {
+      "content": "<extra_id_58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32042": {
+      "content": "<extra_id_57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32043": {
+      "content": "<extra_id_56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32044": {
+      "content": "<extra_id_55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32045": {
+      "content": "<extra_id_54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32046": {
+      "content": "<extra_id_53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32047": {
+      "content": "<extra_id_52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32048": {
+      "content": "<extra_id_51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32049": {
+      "content": "<extra_id_50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32050": {
+      "content": "<extra_id_49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32051": {
+      "content": "<extra_id_48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32052": {
+      "content": "<extra_id_47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32053": {
+      "content": "<extra_id_46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32054": {
+      "content": "<extra_id_45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32055": {
+      "content": "<extra_id_44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32056": {
+      "content": "<extra_id_43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32057": {
+      "content": "<extra_id_42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32058": {
+      "content": "<extra_id_41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32059": {
+      "content": "<extra_id_40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32060": {
+      "content": "<extra_id_39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32061": {
+      "content": "<extra_id_38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32062": {
+      "content": "<extra_id_37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32063": {
+      "content": "<extra_id_36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32064": {
+      "content": "<extra_id_35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32065": {
+      "content": "<extra_id_34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32066": {
+      "content": "<extra_id_33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32067": {
+      "content": "<extra_id_32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32068": {
+      "content": "<extra_id_31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32069": {
+      "content": "<extra_id_30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32070": {
+      "content": "<extra_id_29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32071": {
+      "content": "<extra_id_28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32072": {
+      "content": "<extra_id_27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32073": {
+      "content": "<extra_id_26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32074": {
+      "content": "<extra_id_25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32075": {
+      "content": "<extra_id_24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32076": {
+      "content": "<extra_id_23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32077": {
+      "content": "<extra_id_22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32078": {
+      "content": "<extra_id_21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32079": {
+      "content": "<extra_id_20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32080": {
+      "content": "<extra_id_19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32081": {
+      "content": "<extra_id_18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32082": {
+      "content": "<extra_id_17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32083": {
+      "content": "<extra_id_16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32084": {
+      "content": "<extra_id_15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32085": {
+      "content": "<extra_id_14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32086": {
+      "content": "<extra_id_13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32087": {
+      "content": "<extra_id_12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32088": {
+      "content": "<extra_id_11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32089": {
+      "content": "<extra_id_10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32090": {
+      "content": "<extra_id_9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32091": {
+      "content": "<extra_id_8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32092": {
+      "content": "<extra_id_7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32093": {
+      "content": "<extra_id_6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32094": {
+      "content": "<extra_id_5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32095": {
+      "content": "<extra_id_4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32096": {
+      "content": "<extra_id_3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32097": {
+      "content": "<extra_id_2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32098": {
+      "content": "<extra_id_1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32099": {
+      "content": "<extra_id_0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

invokeai/backend/flux/controlnet/controlnet_flux_output.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class ControlNetFluxOutput:
+    single_block_residuals: list[torch.Tensor] | None
+    double_block_residuals: list[torch.Tensor] | None
+
+    def apply_weight(self, weight: float):
+        if self.single_block_residuals is not None:
+            for i in range(len(self.single_block_residuals)):
+                self.single_block_residuals[i] = self.single_block_residuals[i] * weight
+        if self.double_block_residuals is not None:
+            for i in range(len(self.double_block_residuals)):
+                self.double_block_residuals[i] = self.double_block_residuals[i] * weight
+
+
+def add_tensor_lists_elementwise(
+    list1: list[torch.Tensor] | None, list2: list[torch.Tensor] | None
+) -> list[torch.Tensor] | None:
+    """Add two tensor lists elementwise that could be None."""
+    if list1 is None and list2 is None:
+        return None
+    if list1 is None:
+        return list2
+    if list2 is None:
+        return list1
+
+    new_list: list[torch.Tensor] = []
+    for list1_tensor, list2_tensor in zip(list1, list2, strict=True):
+        new_list.append(list1_tensor + list2_tensor)
+    return new_list
+
+
+def add_controlnet_flux_outputs(
+    controlnet_output_1: ControlNetFluxOutput, controlnet_output_2: ControlNetFluxOutput
+) -> ControlNetFluxOutput:
+    return ControlNetFluxOutput(
+        single_block_residuals=add_tensor_lists_elementwise(
+            controlnet_output_1.single_block_residuals, controlnet_output_2.single_block_residuals
+        ),
+        double_block_residuals=add_tensor_lists_elementwise(
+            controlnet_output_1.double_block_residuals, controlnet_output_2.double_block_residuals
+        ),
+    )
+
+
+def sum_controlnet_flux_outputs(
+    controlnet_outputs: list[ControlNetFluxOutput],
+) -> ControlNetFluxOutput:
+    controlnet_output_sum = ControlNetFluxOutput(single_block_residuals=None, double_block_residuals=None)
+
+    for controlnet_output in controlnet_outputs:
+        controlnet_output_sum = add_controlnet_flux_outputs(controlnet_output_sum, controlnet_output)
+
+    return controlnet_output_sum

invokeai/backend/flux/controlnet/instantx_controlnet_flux.py

@@ -0,0 +1,180 @@
+# This file was initially copied from:
+# https://github.com/huggingface/diffusers/blob/99f608218caa069a2f16dcf9efab46959b15aec0/src/diffusers/models/controlnet_flux.py
+
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from invokeai.backend.flux.controlnet.zero_module import zero_module
+from invokeai.backend.flux.model import FluxParams
+from invokeai.backend.flux.modules.layers import (
+    DoubleStreamBlock,
+    EmbedND,
+    MLPEmbedder,
+    SingleStreamBlock,
+    timestep_embedding,
+)
+
+
+@dataclass
+class InstantXControlNetFluxOutput:
+    controlnet_block_samples: list[torch.Tensor] | None
+    controlnet_single_block_samples: list[torch.Tensor] | None
+
+
+# NOTE(ryand): Mapping between diffusers FLUX transformer params and BFL FLUX transformer params:
+# - Diffusers: BFL
+# - in_channels: in_channels
+# - num_layers: depth
+# - num_single_layers: depth_single_blocks
+# - attention_head_dim: hidden_size // num_heads
+# - num_attention_heads: num_heads
+# - joint_attention_dim: context_in_dim
+# - pooled_projection_dim: vec_in_dim
+# - guidance_embeds: guidance_embed
+# - axes_dims_rope: axes_dim
+
+
+class InstantXControlNetFlux(torch.nn.Module):
+    def __init__(self, params: FluxParams, num_control_modes: int | None = None):
+        """
+        Args:
+            params (FluxParams): The parameters for the FLUX model.
+            num_control_modes (int | None, optional): The number of controlnet modes. If non-None, then the model is a
+                'union controlnet' model and expects a mode conditioning input at runtime.
+        """
+        super().__init__()
+
+        # The following modules mirror the base FLUX transformer model.
+        # -------------------------------------------------------------
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}")
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
+
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        # The following modules are specific to the ControlNet model.
+        # -----------------------------------------------------------
+        self.controlnet_blocks = nn.ModuleList([])
+        for _ in range(len(self.double_blocks)):
+            self.controlnet_blocks.append(zero_module(nn.Linear(self.hidden_size, self.hidden_size)))
+
+        self.controlnet_single_blocks = nn.ModuleList([])
+        for _ in range(len(self.single_blocks)):
+            self.controlnet_single_blocks.append(zero_module(nn.Linear(self.hidden_size, self.hidden_size)))
+
+        self.is_union = False
+        if num_control_modes is not None:
+            self.is_union = True
+            self.controlnet_mode_embedder = nn.Embedding(num_control_modes, self.hidden_size)
+
+        self.controlnet_x_embedder = zero_module(torch.nn.Linear(self.in_channels, self.hidden_size))
+
+    def forward(
+        self,
+        controlnet_cond: torch.Tensor,
+        controlnet_mode: torch.Tensor | None,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt: torch.Tensor,
+        txt_ids: torch.Tensor,
+        timesteps: torch.Tensor,
+        y: torch.Tensor,
+        guidance: torch.Tensor | None = None,
+    ) -> InstantXControlNetFluxOutput:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        img = self.img_in(img)
+
+        # Add controlnet_cond embedding.
+        img = img + self.controlnet_x_embedder(controlnet_cond)
+
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        # If this is a union ControlNet, then concat the control mode embedding to the T5 text embedding.
+        if self.is_union:
+            if controlnet_mode is None:
+                # We allow users to enter 'None' as the controlnet_mode if they don't want to worry about this input.
+                # We've chosen to use a zero-embedding in this case.
+                zero_index = torch.zeros([1, 1], dtype=torch.long, device=txt.device)
+                controlnet_mode_emb = torch.zeros_like(self.controlnet_mode_embedder(zero_index))
+            else:
+                controlnet_mode_emb = self.controlnet_mode_embedder(controlnet_mode)
+            txt = torch.cat([controlnet_mode_emb, txt], dim=1)
+            txt_ids = torch.cat([txt_ids[:, :1, :], txt_ids], dim=1)
+        else:
+            assert controlnet_mode is None
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        double_block_samples: list[torch.Tensor] = []
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+            double_block_samples.append(img)
+
+        img = torch.cat((txt, img), 1)
+
+        single_block_samples: list[torch.Tensor] = []
+        for block in self.single_blocks:
+            img = block(img, vec=vec, pe=pe)
+            single_block_samples.append(img[:, txt.shape[1] :])
+
+        # ControlNet Block
+        controlnet_double_block_samples: list[torch.Tensor] = []
+        for double_block_sample, controlnet_block in zip(double_block_samples, self.controlnet_blocks, strict=True):
+            double_block_sample = controlnet_block(double_block_sample)
+            controlnet_double_block_samples.append(double_block_sample)
+
+        controlnet_single_block_samples: list[torch.Tensor] = []
+        for single_block_sample, controlnet_block in zip(
+            single_block_samples, self.controlnet_single_blocks, strict=True
+        ):
+            single_block_sample = controlnet_block(single_block_sample)
+            controlnet_single_block_samples.append(single_block_sample)
+
+        return InstantXControlNetFluxOutput(
+            controlnet_block_samples=controlnet_double_block_samples or None,
+            controlnet_single_block_samples=controlnet_single_block_samples or None,
+        )

invokeai/backend/flux/controlnet/state_dict_utils.py

@@ -0,0 +1,295 @@
+from typing import Any, Dict
+
+import torch
+
+from invokeai.backend.flux.model import FluxParams
+
+
+def is_state_dict_xlabs_controlnet(sd: dict[str | int, Any]) -> bool:
+    """Is the state dict for an XLabs ControlNet model?
+
+    This is intended to be a reasonably high-precision detector, but it is not guaranteed to have perfect precision.
+    """
+    # If all of the expected keys are present, then this is very likely an XLabs ControlNet model.
+    expected_keys = {
+        "controlnet_blocks.0.bias",
+        "controlnet_blocks.0.weight",
+        "input_hint_block.0.bias",
+        "input_hint_block.0.weight",
+        "pos_embed_input.bias",
+        "pos_embed_input.weight",
+    }
+
+    if expected_keys.issubset(sd.keys()):
+        return True
+    return False
+
+
+def is_state_dict_instantx_controlnet(sd: dict[str | int, Any]) -> bool:
+    """Is the state dict for an InstantX ControlNet model?
+
+    This is intended to be a reasonably high-precision detector, but it is not guaranteed to have perfect precision.
+    """
+    # If all of the expected keys are present, then this is very likely an InstantX ControlNet model.
+    expected_keys = {
+        "controlnet_blocks.0.bias",
+        "controlnet_blocks.0.weight",
+        "controlnet_x_embedder.bias",
+        "controlnet_x_embedder.weight",
+    }
+
+    if expected_keys.issubset(sd.keys()):
+        return True
+    return False
+
+
+def _fuse_weights(*t: torch.Tensor) -> torch.Tensor:
+    """Fuse weights along dimension 0.
+
+    Used to fuse q, k, v attention weights into a single qkv tensor when converting from diffusers to BFL format.
+    """
+    # TODO(ryand): Double check dim=0 is correct.
+    return torch.cat(t, dim=0)
+
+
+def _convert_flux_double_block_sd_from_diffusers_to_bfl_format(
+    sd: Dict[str, torch.Tensor], double_block_index: int
+) -> Dict[str, torch.Tensor]:
+    """Convert the state dict for a double block from diffusers format to BFL format."""
+    to_prefix = f"double_blocks.{double_block_index}"
+    from_prefix = f"transformer_blocks.{double_block_index}"
+
+    new_sd: dict[str, torch.Tensor] = {}
+
+    # Check one key to determine if this block exists.
+    if f"{from_prefix}.attn.add_q_proj.bias" not in sd:
+        return new_sd
+
+    # txt_attn.qkv
+    new_sd[f"{to_prefix}.txt_attn.qkv.bias"] = _fuse_weights(
+        sd.pop(f"{from_prefix}.attn.add_q_proj.bias"),
+        sd.pop(f"{from_prefix}.attn.add_k_proj.bias"),
+        sd.pop(f"{from_prefix}.attn.add_v_proj.bias"),
+    )
+    new_sd[f"{to_prefix}.txt_attn.qkv.weight"] = _fuse_weights(
+        sd.pop(f"{from_prefix}.attn.add_q_proj.weight"),
+        sd.pop(f"{from_prefix}.attn.add_k_proj.weight"),
+        sd.pop(f"{from_prefix}.attn.add_v_proj.weight"),
+    )
+
+    # img_attn.qkv
+    new_sd[f"{to_prefix}.img_attn.qkv.bias"] = _fuse_weights(
+        sd.pop(f"{from_prefix}.attn.to_q.bias"),
+        sd.pop(f"{from_prefix}.attn.to_k.bias"),
+        sd.pop(f"{from_prefix}.attn.to_v.bias"),
+    )
+    new_sd[f"{to_prefix}.img_attn.qkv.weight"] = _fuse_weights(
+        sd.pop(f"{from_prefix}.attn.to_q.weight"),
+        sd.pop(f"{from_prefix}.attn.to_k.weight"),
+        sd.pop(f"{from_prefix}.attn.to_v.weight"),
+    )
+
+    # Handle basic 1-to-1 key conversions.
+    key_map = {
+        # img_attn
+        "attn.norm_k.weight": "img_attn.norm.key_norm.scale",
+        "attn.norm_q.weight": "img_attn.norm.query_norm.scale",
+        "attn.to_out.0.weight": "img_attn.proj.weight",
+        "attn.to_out.0.bias": "img_attn.proj.bias",
+        # img_mlp
+        "ff.net.0.proj.weight": "img_mlp.0.weight",
+        "ff.net.0.proj.bias": "img_mlp.0.bias",
+        "ff.net.2.weight": "img_mlp.2.weight",
+        "ff.net.2.bias": "img_mlp.2.bias",
+        # img_mod
+        "norm1.linear.weight": "img_mod.lin.weight",
+        "norm1.linear.bias": "img_mod.lin.bias",
+        # txt_attn
+        "attn.norm_added_q.weight": "txt_attn.norm.query_norm.scale",
+        "attn.norm_added_k.weight": "txt_attn.norm.key_norm.scale",
+        "attn.to_add_out.weight": "txt_attn.proj.weight",
+        "attn.to_add_out.bias": "txt_attn.proj.bias",
+        # txt_mlp
+        "ff_context.net.0.proj.weight": "txt_mlp.0.weight",
+        "ff_context.net.0.proj.bias": "txt_mlp.0.bias",
+        "ff_context.net.2.weight": "txt_mlp.2.weight",
+        "ff_context.net.2.bias": "txt_mlp.2.bias",
+        # txt_mod
+        "norm1_context.linear.weight": "txt_mod.lin.weight",
+        "norm1_context.linear.bias": "txt_mod.lin.bias",
+    }
+    for from_key, to_key in key_map.items():
+        new_sd[f"{to_prefix}.{to_key}"] = sd.pop(f"{from_prefix}.{from_key}")
+
+    return new_sd
+
+
+def _convert_flux_single_block_sd_from_diffusers_to_bfl_format(
+    sd: Dict[str, torch.Tensor], single_block_index: int
+) -> Dict[str, torch.Tensor]:
+    """Convert the state dict for a single block from diffusers format to BFL format."""
+    to_prefix = f"single_blocks.{single_block_index}"
+    from_prefix = f"single_transformer_blocks.{single_block_index}"
+
+    new_sd: dict[str, torch.Tensor] = {}
+
+    # Check one key to determine if this block exists.
+    if f"{from_prefix}.attn.to_q.bias" not in sd:
+        return new_sd
+
+    # linear1 (qkv)
+    new_sd[f"{to_prefix}.linear1.bias"] = _fuse_weights(
+        sd.pop(f"{from_prefix}.attn.to_q.bias"),
+        sd.pop(f"{from_prefix}.attn.to_k.bias"),
+        sd.pop(f"{from_prefix}.attn.to_v.bias"),
+        sd.pop(f"{from_prefix}.proj_mlp.bias"),
+    )
+    new_sd[f"{to_prefix}.linear1.weight"] = _fuse_weights(
+        sd.pop(f"{from_prefix}.attn.to_q.weight"),
+        sd.pop(f"{from_prefix}.attn.to_k.weight"),
+        sd.pop(f"{from_prefix}.attn.to_v.weight"),
+        sd.pop(f"{from_prefix}.proj_mlp.weight"),
+    )
+
+    # Handle basic 1-to-1 key conversions.
+    key_map = {
+        # linear2
+        "proj_out.weight": "linear2.weight",
+        "proj_out.bias": "linear2.bias",
+        # modulation
+        "norm.linear.weight": "modulation.lin.weight",
+        "norm.linear.bias": "modulation.lin.bias",
+        # norm
+        "attn.norm_k.weight": "norm.key_norm.scale",
+        "attn.norm_q.weight": "norm.query_norm.scale",
+    }
+    for from_key, to_key in key_map.items():
+        new_sd[f"{to_prefix}.{to_key}"] = sd.pop(f"{from_prefix}.{from_key}")
+
+    return new_sd
+
+
+def convert_diffusers_instantx_state_dict_to_bfl_format(sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Convert an InstantX ControlNet state dict to the format that can be loaded by our internal
+    InstantXControlNetFlux model.
+
+    The original InstantX ControlNet model was developed to be used in diffusers. We have ported the original
+    implementation to InstantXControlNetFlux to make it compatible with BFL-style models. This function converts the
+    original state dict to the format expected by InstantXControlNetFlux.
+    """
+    # Shallow copy sd so that we can pop keys from it without modifying the original.
+    sd = sd.copy()
+
+    new_sd: dict[str, torch.Tensor] = {}
+
+    # Handle basic 1-to-1 key conversions.
+    basic_key_map = {
+        # Base model keys.
+        # ----------------
+        # txt_in keys.
+        "context_embedder.bias": "txt_in.bias",
+        "context_embedder.weight": "txt_in.weight",
+        # guidance_in MLPEmbedder keys.
+        "time_text_embed.guidance_embedder.linear_1.bias": "guidance_in.in_layer.bias",
+        "time_text_embed.guidance_embedder.linear_1.weight": "guidance_in.in_layer.weight",
+        "time_text_embed.guidance_embedder.linear_2.bias": "guidance_in.out_layer.bias",
+        "time_text_embed.guidance_embedder.linear_2.weight": "guidance_in.out_layer.weight",
+        # vector_in MLPEmbedder keys.
+        "time_text_embed.text_embedder.linear_1.bias": "vector_in.in_layer.bias",
+        "time_text_embed.text_embedder.linear_1.weight": "vector_in.in_layer.weight",
+        "time_text_embed.text_embedder.linear_2.bias": "vector_in.out_layer.bias",
+        "time_text_embed.text_embedder.linear_2.weight": "vector_in.out_layer.weight",
+        # time_in MLPEmbedder keys.
+        "time_text_embed.timestep_embedder.linear_1.bias": "time_in.in_layer.bias",
+        "time_text_embed.timestep_embedder.linear_1.weight": "time_in.in_layer.weight",
+        "time_text_embed.timestep_embedder.linear_2.bias": "time_in.out_layer.bias",
+        "time_text_embed.timestep_embedder.linear_2.weight": "time_in.out_layer.weight",
+        # img_in keys.
+        "x_embedder.bias": "img_in.bias",
+        "x_embedder.weight": "img_in.weight",
+    }
+    for old_key, new_key in basic_key_map.items():
+        v = sd.pop(old_key, None)
+        if v is not None:
+            new_sd[new_key] = v
+
+    # Handle the double_blocks.
+    block_index = 0
+    while True:
+        converted_double_block_sd = _convert_flux_double_block_sd_from_diffusers_to_bfl_format(sd, block_index)
+        if len(converted_double_block_sd) == 0:
+            break
+        new_sd.update(converted_double_block_sd)
+        block_index += 1
+
+    # Handle the single_blocks.
+    block_index = 0
+    while True:
+        converted_singe_block_sd = _convert_flux_single_block_sd_from_diffusers_to_bfl_format(sd, block_index)
+        if len(converted_singe_block_sd) == 0:
+            break
+        new_sd.update(converted_singe_block_sd)
+        block_index += 1
+
+    # Transfer controlnet keys as-is.
+    for k in list(sd.keys()):
+        if k.startswith("controlnet_"):
+            new_sd[k] = sd.pop(k)
+
+    # Assert that all keys have been handled.
+    assert len(sd) == 0
+    return new_sd
+
+
+def infer_flux_params_from_state_dict(sd: Dict[str, torch.Tensor]) -> FluxParams:
+    """Infer the FluxParams from the shape of a FLUX state dict. When a model is distributed in diffusers format, this
+    information is all contained in the config.json file that accompanies the model. However, being apple to infer the
+    params from the state dict enables us to load models (e.g. an InstantX ControlNet) from a single weight file.
+    """
+    hidden_size = sd["img_in.weight"].shape[0]
+    mlp_hidden_dim = sd["double_blocks.0.img_mlp.0.weight"].shape[0]
+    # mlp_ratio is a float, but we treat it as an int here to avoid having to think about possible float precision
+    # issues. In practice, mlp_ratio is usually 4.
+    mlp_ratio = mlp_hidden_dim // hidden_size
+
+    head_dim = sd["double_blocks.0.img_attn.norm.query_norm.scale"].shape[0]
+    num_heads = hidden_size // head_dim
+
+    # Count the number of double blocks.
+    double_block_index = 0
+    while f"double_blocks.{double_block_index}.img_attn.qkv.weight" in sd:
+        double_block_index += 1
+
+    # Count the number of single blocks.
+    single_block_index = 0
+    while f"single_blocks.{single_block_index}.linear1.weight" in sd:
+        single_block_index += 1
+
+    return FluxParams(
+        in_channels=sd["img_in.weight"].shape[1],
+        vec_in_dim=sd["vector_in.in_layer.weight"].shape[1],
+        context_in_dim=sd["txt_in.weight"].shape[1],
+        hidden_size=hidden_size,
+        mlp_ratio=mlp_ratio,
+        num_heads=num_heads,
+        depth=double_block_index,
+        depth_single_blocks=single_block_index,
+        # axes_dim cannot be inferred from the state dict. The hard-coded value is correct for dev/schnell models.
+        axes_dim=[16, 56, 56],
+        # theta cannot be inferred from the state dict. The hard-coded value is correct for dev/schnell models.
+        theta=10_000,
+        qkv_bias="double_blocks.0.img_attn.qkv.bias" in sd,
+        guidance_embed="guidance_in.in_layer.weight" in sd,
+    )
+
+
+def infer_instantx_num_control_modes_from_state_dict(sd: Dict[str, torch.Tensor]) -> int | None:
+    """Infer the number of ControlNet Union modes from the shape of a InstantX ControlNet state dict.
+
+    Returns None if the model is not a ControlNet Union model. Otherwise returns the number of modes.
+    """
+    mode_embedder_key = "controlnet_mode_embedder.weight"
+    if mode_embedder_key not in sd:
+        return None
+
+    return sd[mode_embedder_key].shape[0]

invokeai/backend/flux/controlnet/xlabs_controlnet_flux.py

@@ -0,0 +1,130 @@
+# This file was initially based on:
+# https://github.com/XLabs-AI/x-flux/blob/47495425dbed499be1e8e5a6e52628b07349cba2/src/flux/controlnet.py
+
+
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+
+from invokeai.backend.flux.controlnet.zero_module import zero_module
+from invokeai.backend.flux.model import FluxParams
+from invokeai.backend.flux.modules.layers import DoubleStreamBlock, EmbedND, MLPEmbedder, timestep_embedding
+
+
+@dataclass
+class XLabsControlNetFluxOutput:
+    controlnet_double_block_residuals: list[torch.Tensor] | None
+
+
+class XLabsControlNetFlux(torch.nn.Module):
+    """A ControlNet model for FLUX.
+
+    The architecture is very similar to the base FLUX model, with the following differences:
+    - A `controlnet_depth` parameter is passed to control the number of double_blocks that the ControlNet is applied to.
+      In order to keep the ControlNet small, this is typically much less than the depth of the base FLUX model.
+    - There is a set of `controlnet_blocks` that are applied to the output of each double_block.
+    """
+
+    def __init__(self, params: FluxParams, controlnet_depth: int = 2):
+        super().__init__()
+
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}")
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = torch.nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else torch.nn.Identity()
+        )
+        self.txt_in = torch.nn.Linear(params.context_in_dim, self.hidden_size)
+
+        self.double_blocks = torch.nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(controlnet_depth)
+            ]
+        )
+
+        # Add ControlNet blocks.
+        self.controlnet_blocks = torch.nn.ModuleList([])
+        for _ in range(controlnet_depth):
+            controlnet_block = torch.nn.Linear(self.hidden_size, self.hidden_size)
+            controlnet_block = zero_module(controlnet_block)
+            self.controlnet_blocks.append(controlnet_block)
+        self.pos_embed_input = torch.nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.input_hint_block = torch.nn.Sequential(
+            torch.nn.Conv2d(3, 16, 3, padding=1),
+            torch.nn.SiLU(),
+            torch.nn.Conv2d(16, 16, 3, padding=1),
+            torch.nn.SiLU(),
+            torch.nn.Conv2d(16, 16, 3, padding=1, stride=2),
+            torch.nn.SiLU(),
+            torch.nn.Conv2d(16, 16, 3, padding=1),
+            torch.nn.SiLU(),
+            torch.nn.Conv2d(16, 16, 3, padding=1, stride=2),
+            torch.nn.SiLU(),
+            torch.nn.Conv2d(16, 16, 3, padding=1),
+            torch.nn.SiLU(),
+            torch.nn.Conv2d(16, 16, 3, padding=1, stride=2),
+            torch.nn.SiLU(),
+            zero_module(torch.nn.Conv2d(16, 16, 3, padding=1)),
+        )
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        controlnet_cond: torch.Tensor,
+        txt: torch.Tensor,
+        txt_ids: torch.Tensor,
+        timesteps: torch.Tensor,
+        y: torch.Tensor,
+        guidance: torch.Tensor | None = None,
+    ) -> XLabsControlNetFluxOutput:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        controlnet_cond = self.input_hint_block(controlnet_cond)
+        controlnet_cond = rearrange(controlnet_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        controlnet_cond = self.pos_embed_input(controlnet_cond)
+        img = img + controlnet_cond
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        block_res_samples: list[torch.Tensor] = []
+
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+            block_res_samples.append(img)
+
+        controlnet_block_res_samples: list[torch.Tensor] = []
+        for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks, strict=True):
+            block_res_sample = controlnet_block(block_res_sample)
+            controlnet_block_res_samples.append(block_res_sample)
+
+        return XLabsControlNetFluxOutput(controlnet_double_block_residuals=controlnet_block_res_samples)

invokeai/backend/flux/controlnet/zero_module.py

@@ -0,0 +1,12 @@
+from typing import TypeVar
+
+import torch
+
+T = TypeVar("T", bound=torch.nn.Module)
+
+
+def zero_module(module: T) -> T:
+    """Initialize the parameters of a module to zero."""
+    for p in module.parameters():
+        torch.nn.init.zeros_(p)
+    return module

invokeai/backend/flux/custom_block_processor.py

@@ -0,0 +1,138 @@
+import einops
+import torch
+
+from invokeai.backend.flux.extensions.regional_prompting_extension import RegionalPromptingExtension
+from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
+from invokeai.backend.flux.math import attention
+from invokeai.backend.flux.modules.layers import DoubleStreamBlock, SingleStreamBlock
+
+
+class CustomDoubleStreamBlockProcessor:
+    """A class containing a custom implementation of DoubleStreamBlock.forward() with additional features
+    (IP-Adapter, etc.).
+    """
+
+    @staticmethod
+    def _double_stream_block_forward(
+        block: DoubleStreamBlock,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """This function is a direct copy of DoubleStreamBlock.forward(), but it returns some of the intermediate
+        values.
+        """
+        img_mod1, img_mod2 = block.img_mod(vec)
+        txt_mod1, txt_mod2 = block.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = block.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = block.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = einops.rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=block.num_heads)
+        img_q, img_k = block.img_attn.norm(img_q, img_k, img_v)
+
+        # prepare txt for attention
+        txt_modulated = block.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = block.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = einops.rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=block.num_heads)
+        txt_q, txt_k = block.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+
+        attn = attention(q, k, v, pe=pe, attn_mask=attn_mask)
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * block.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * block.img_mlp((1 + img_mod2.scale) * block.img_norm2(img) + img_mod2.shift)
+
+        # calculate the txt bloks
+        txt = txt + txt_mod1.gate * block.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * block.txt_mlp((1 + txt_mod2.scale) * block.txt_norm2(txt) + txt_mod2.shift)
+        return img, txt, img_q
+
+    @staticmethod
+    def custom_double_block_forward(
+        timestep_index: int,
+        total_num_timesteps: int,
+        block_index: int,
+        block: DoubleStreamBlock,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        ip_adapter_extensions: list[XLabsIPAdapterExtension],
+        regional_prompting_extension: RegionalPromptingExtension,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """A custom implementation of DoubleStreamBlock.forward() with additional features:
+        - IP-Adapter support
+        """
+        attn_mask = regional_prompting_extension.get_double_stream_attn_mask(block_index)
+        img, txt, img_q = CustomDoubleStreamBlockProcessor._double_stream_block_forward(
+            block, img, txt, vec, pe, attn_mask=attn_mask
+        )
+
+        # Apply IP-Adapter conditioning.
+        for ip_adapter_extension in ip_adapter_extensions:
+            img = ip_adapter_extension.run_ip_adapter(
+                timestep_index=timestep_index,
+                total_num_timesteps=total_num_timesteps,
+                block_index=block_index,
+                block=block,
+                img_q=img_q,
+                img=img,
+            )
+
+        return img, txt
+
+
+class CustomSingleStreamBlockProcessor:
+    """A class containing a custom implementation of SingleStreamBlock.forward() with additional features (masking,
+    etc.)
+    """
+
+    @staticmethod
+    def _single_stream_block_forward(
+        block: SingleStreamBlock,
+        x: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """This function is a direct copy of SingleStreamBlock.forward()."""
+        mod, _ = block.modulation(vec)
+        x_mod = (1 + mod.scale) * block.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(block.linear1(x_mod), [3 * block.hidden_size, block.mlp_hidden_dim], dim=-1)
+
+        q, k, v = einops.rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=block.num_heads)
+        q, k = block.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe, attn_mask=attn_mask)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = block.linear2(torch.cat((attn, block.mlp_act(mlp)), 2))
+        return x + mod.gate * output
+
+    @staticmethod
+    def custom_single_block_forward(
+        timestep_index: int,
+        total_num_timesteps: int,
+        block_index: int,
+        block: SingleStreamBlock,
+        img: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        regional_prompting_extension: RegionalPromptingExtension,
+    ) -> torch.Tensor:
+        """A custom implementation of SingleStreamBlock.forward() with additional features:
+        - Masking
+        """
+        attn_mask = regional_prompting_extension.get_single_stream_attn_mask(block_index)
+        return CustomSingleStreamBlockProcessor._single_stream_block_forward(block, img, vec, pe, attn_mask=attn_mask)

invokeai/backend/flux/denoise.py

@@ -0,0 +1,408 @@
+import inspect
+import math
+from typing import Callable
+
+import torch
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from tqdm import tqdm
+
+from invokeai.backend.flux.controlnet.controlnet_flux_output import ControlNetFluxOutput, sum_controlnet_flux_outputs
+from invokeai.backend.flux.extensions.dype_extension import DyPEExtension
+from invokeai.backend.flux.extensions.instantx_controlnet_extension import InstantXControlNetExtension
+from invokeai.backend.flux.extensions.regional_prompting_extension import RegionalPromptingExtension
+from invokeai.backend.flux.extensions.xlabs_controlnet_extension import XLabsControlNetExtension
+from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
+from invokeai.backend.flux.model import Flux
+from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
+from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
+
+
+def denoise(
+    model: Flux,
+    # model input
+    img: torch.Tensor,
+    img_ids: torch.Tensor,
+    pos_regional_prompting_extension: RegionalPromptingExtension,
+    neg_regional_prompting_extension: RegionalPromptingExtension | None,
+    # sampling parameters
+    timesteps: list[float],
+    step_callback: Callable[[PipelineIntermediateState], None],
+    guidance: float,
+    cfg_scale: list[float],
+    inpaint_extension: RectifiedFlowInpaintExtension | None,
+    controlnet_extensions: list[XLabsControlNetExtension | InstantXControlNetExtension],
+    pos_ip_adapter_extensions: list[XLabsIPAdapterExtension],
+    neg_ip_adapter_extensions: list[XLabsIPAdapterExtension],
+    # extra img tokens (channel-wise)
+    img_cond: torch.Tensor | None,
+    # extra img tokens (sequence-wise) - for Kontext conditioning
+    img_cond_seq: torch.Tensor | None = None,
+    img_cond_seq_ids: torch.Tensor | None = None,
+    # DyPE extension for high-resolution generation
+    dype_extension: DyPEExtension | None = None,
+    # Optional scheduler for alternative sampling methods
+    scheduler: SchedulerMixin | None = None,
+):
+    # Determine if we're using a diffusers scheduler or the built-in Euler method
+    use_scheduler = scheduler is not None
+
+    if use_scheduler:
+        # Initialize scheduler with timesteps
+        # The timesteps list contains values in [0, 1] range (sigmas)
+        # LCM should use num_inference_steps (it has its own sigma schedule),
+        # while other schedulers can use custom sigmas if supported
+        is_lcm = scheduler.__class__.__name__ == "FlowMatchLCMScheduler"
+        set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
+        if not is_lcm and "sigmas" in set_timesteps_sig.parameters:
+            # Scheduler supports custom sigmas - use InvokeAI's time-shifted schedule
+            scheduler.set_timesteps(sigmas=timesteps, device=img.device)
+        else:
+            # LCM or scheduler doesn't support custom sigmas - use num_inference_steps
+            # The schedule will be computed by the scheduler itself.
+            #
+            # Important for img2img callers: if the initial latent/noise blend was
+            # computed from a separate pre-scheduler schedule, that preblend may not
+            # match this scheduler's true first step exactly.
+            num_inference_steps = len(timesteps) - 1
+            scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=img.device)
+
+        # For schedulers like Heun, the number of actual steps may differ
+        # (Heun doubles timesteps internally)
+        num_scheduler_steps = len(scheduler.timesteps)
+        # For user-facing step count, use the original number of denoising steps
+        total_steps = len(timesteps) - 1
+    else:
+        total_steps = len(timesteps) - 1
+        num_scheduler_steps = total_steps
+
+    # guidance_vec is ignored for schnell.
+    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+
+    # Store original sequence length for slicing predictions
+    original_seq_len = img.shape[1]
+
+    # DyPE: Patch model with DyPE-aware position embedder
+    dype_embedder = None
+    original_pe_embedder = None
+    if dype_extension is not None:
+        dype_embedder, original_pe_embedder = dype_extension.patch_model(model)
+
+    try:
+        # Track the actual step for user-facing progress (accounts for Heun's double steps)
+        user_step = 0
+
+        if use_scheduler:
+            # Use diffusers scheduler for stepping
+            # Use tqdm with total_steps (user-facing steps) not num_scheduler_steps (internal steps)
+            # This ensures progress bar shows 1/8, 2/8, etc. even when scheduler uses more internal steps
+            pbar = tqdm(total=total_steps, desc="Denoising")
+            for step_index in range(num_scheduler_steps):
+                timestep = scheduler.timesteps[step_index]
+                # Convert scheduler timestep (0-1000) to normalized (0-1) for the model
+                t_curr = timestep.item() / scheduler.config.num_train_timesteps
+                dype_sigma = DyPEExtension.resolve_step_sigma(
+                    fallback_sigma=t_curr,
+                    step_index=step_index,
+                    scheduler_sigmas=getattr(scheduler, "sigmas", None),
+                )
+                t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+                # DyPE: Update step state for timestep-dependent scaling
+                if dype_extension is not None and dype_embedder is not None:
+                    dype_extension.update_step_state(
+                        embedder=dype_embedder,
+                        sigma=dype_sigma,
+                    )
+
+                # For Heun scheduler, track if we're in first or second order step
+                is_heun = hasattr(scheduler, "state_in_first_order")
+                in_first_order = scheduler.state_in_first_order if is_heun else True
+
+                # Run ControlNet models
+                controlnet_residuals: list[ControlNetFluxOutput] = []
+                for controlnet_extension in controlnet_extensions:
+                    controlnet_residuals.append(
+                        controlnet_extension.run_controlnet(
+                            timestep_index=user_step,
+                            total_num_timesteps=total_steps,
+                            img=img,
+                            img_ids=img_ids,
+                            txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                            txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                            y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                            timesteps=t_vec,
+                            guidance=guidance_vec,
+                        )
+                    )
+
+                merged_controlnet_residuals = sum_controlnet_flux_outputs(controlnet_residuals)
+
+                # Prepare input for model
+                img_input = img
+                img_input_ids = img_ids
+
+                if img_cond is not None:
+                    img_input = torch.cat((img_input, img_cond), dim=-1)
+
+                if img_cond_seq is not None:
+                    assert img_cond_seq_ids is not None
+                    img_input = torch.cat((img_input, img_cond_seq), dim=1)
+                    img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids), dim=1)
+
+                pred = model(
+                    img=img_input,
+                    img_ids=img_input_ids,
+                    txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                    txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                    y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                    timesteps=t_vec,
+                    guidance=guidance_vec,
+                    timestep_index=user_step,
+                    total_num_timesteps=total_steps,
+                    controlnet_double_block_residuals=merged_controlnet_residuals.double_block_residuals,
+                    controlnet_single_block_residuals=merged_controlnet_residuals.single_block_residuals,
+                    ip_adapter_extensions=pos_ip_adapter_extensions,
+                    regional_prompting_extension=pos_regional_prompting_extension,
+                )
+
+                if img_cond_seq is not None:
+                    pred = pred[:, :original_seq_len]
+
+                # Get CFG scale for current user step
+                step_cfg_scale = cfg_scale[min(user_step, len(cfg_scale) - 1)]
+
+                if not math.isclose(step_cfg_scale, 1.0):
+                    if neg_regional_prompting_extension is None:
+                        raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
+
+                    neg_img_input = img
+                    neg_img_input_ids = img_ids
+
+                    if img_cond is not None:
+                        neg_img_input = torch.cat((neg_img_input, img_cond), dim=-1)
+
+                    if img_cond_seq is not None:
+                        neg_img_input = torch.cat((neg_img_input, img_cond_seq), dim=1)
+                        neg_img_input_ids = torch.cat((neg_img_input_ids, img_cond_seq_ids), dim=1)
+
+                    neg_pred = model(
+                        img=neg_img_input,
+                        img_ids=neg_img_input_ids,
+                        txt=neg_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                        txt_ids=neg_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                        y=neg_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                        timesteps=t_vec,
+                        guidance=guidance_vec,
+                        timestep_index=user_step,
+                        total_num_timesteps=total_steps,
+                        controlnet_double_block_residuals=None,
+                        controlnet_single_block_residuals=None,
+                        ip_adapter_extensions=neg_ip_adapter_extensions,
+                        regional_prompting_extension=neg_regional_prompting_extension,
+                    )
+
+                    if img_cond_seq is not None:
+                        neg_pred = neg_pred[:, :original_seq_len]
+                    pred = neg_pred + step_cfg_scale * (pred - neg_pred)
+
+                # Use scheduler.step() for the update
+                step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img)
+                img = step_output.prev_sample
+
+                # Get t_prev for inpainting (next sigma value)
+                if step_index + 1 < len(scheduler.sigmas):
+                    t_prev = scheduler.sigmas[step_index + 1].item()
+                else:
+                    t_prev = 0.0
+
+                if inpaint_extension is not None:
+                    img = inpaint_extension.merge_intermediate_latents_with_init_latents(img, t_prev)
+
+                # For Heun, only increment user step after second-order step completes
+                if is_heun:
+                    if not in_first_order:
+                        # Second order step completed
+                        user_step += 1
+                        # Only call step_callback if we haven't exceeded total_steps
+                        if user_step <= total_steps:
+                            pbar.update(1)
+                            preview_img = img - t_curr * pred
+                            if inpaint_extension is not None:
+                                preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(
+                                    preview_img, 0.0
+                                )
+                            step_callback(
+                                PipelineIntermediateState(
+                                    step=user_step,
+                                    order=2,
+                                    total_steps=total_steps,
+                                    timestep=int(t_curr * 1000),
+                                    latents=preview_img,
+                                ),
+                            )
+                else:
+                    # For LCM and other first-order schedulers
+                    user_step += 1
+                    # Only call step_callback if we haven't exceeded total_steps
+                    # (LCM scheduler may have more internal steps than user-facing steps)
+                    if user_step <= total_steps:
+                        pbar.update(1)
+                        preview_img = img - t_curr * pred
+                        if inpaint_extension is not None:
+                            preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(
+                                preview_img, 0.0
+                            )
+                        step_callback(
+                            PipelineIntermediateState(
+                                step=user_step,
+                                order=1,
+                                total_steps=total_steps,
+                                timestep=int(t_curr * 1000),
+                                latents=preview_img,
+                            ),
+                        )
+
+            pbar.close()
+            return img
+
+        # Original Euler implementation (when scheduler is None)
+        for step_index, (t_curr, t_prev) in tqdm(list(enumerate(zip(timesteps[:-1], timesteps[1:], strict=True)))):
+            # DyPE: Update step state for timestep-dependent scaling
+            if dype_extension is not None and dype_embedder is not None:
+                dype_extension.update_step_state(
+                    embedder=dype_embedder,
+                    sigma=t_curr,
+                )
+
+            t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+            # Run ControlNet models.
+            controlnet_residuals: list[ControlNetFluxOutput] = []
+            for controlnet_extension in controlnet_extensions:
+                controlnet_residuals.append(
+                    controlnet_extension.run_controlnet(
+                        timestep_index=step_index,
+                        total_num_timesteps=total_steps,
+                        img=img,
+                        img_ids=img_ids,
+                        txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                        txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                        y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                        timesteps=t_vec,
+                        guidance=guidance_vec,
+                    )
+                )
+
+            # Merge the ControlNet residuals from multiple ControlNets.
+            # TODO(ryand): We may want to calculate the sum just-in-time to keep peak memory low. Keep in mind, that the
+            # controlnet_residuals datastructure is efficient in that it likely contains multiple references to the same
+            # tensors. Calculating the sum materializes each tensor into its own instance.
+            merged_controlnet_residuals = sum_controlnet_flux_outputs(controlnet_residuals)
+
+            # Prepare input for model - concatenate fresh each step
+            img_input = img
+            img_input_ids = img_ids
+
+            # Add channel-wise conditioning (for ControlNet, FLUX Fill, etc.)
+            if img_cond is not None:
+                img_input = torch.cat((img_input, img_cond), dim=-1)
+
+            # Add sequence-wise conditioning (for Kontext)
+            if img_cond_seq is not None:
+                assert img_cond_seq_ids is not None, (
+                    "You need to provide either both or neither of the sequence conditioning"
+                )
+                img_input = torch.cat((img_input, img_cond_seq), dim=1)
+                img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids), dim=1)
+
+            pred = model(
+                img=img_input,
+                img_ids=img_input_ids,
+                txt=pos_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                txt_ids=pos_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                y=pos_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                timesteps=t_vec,
+                guidance=guidance_vec,
+                timestep_index=step_index,
+                total_num_timesteps=total_steps,
+                controlnet_double_block_residuals=merged_controlnet_residuals.double_block_residuals,
+                controlnet_single_block_residuals=merged_controlnet_residuals.single_block_residuals,
+                ip_adapter_extensions=pos_ip_adapter_extensions,
+                regional_prompting_extension=pos_regional_prompting_extension,
+            )
+
+            # Slice prediction to only include the main image tokens
+            if img_cond_seq is not None:
+                pred = pred[:, :original_seq_len]
+
+            step_cfg_scale = cfg_scale[step_index]
+
+            # If step_cfg_scale, is 1.0, then we don't need to run the negative prediction.
+            if not math.isclose(step_cfg_scale, 1.0):
+                # TODO(ryand): Add option to run positive and negative predictions in a single batch for better performance
+                # on systems with sufficient VRAM.
+
+                if neg_regional_prompting_extension is None:
+                    raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
+
+                # For negative prediction with Kontext, we need to include the reference images
+                # to maintain consistency between positive and negative passes. Without this,
+                # CFG would create artifacts as the attention mechanism would see different
+                # spatial structures in each pass
+                neg_img_input = img
+                neg_img_input_ids = img_ids
+
+                # Add channel-wise conditioning for negative pass if present
+                if img_cond is not None:
+                    neg_img_input = torch.cat((neg_img_input, img_cond), dim=-1)
+
+                # Add sequence-wise conditioning (Kontext) for negative pass
+                # This ensures reference images are processed consistently
+                if img_cond_seq is not None:
+                    neg_img_input = torch.cat((neg_img_input, img_cond_seq), dim=1)
+                    neg_img_input_ids = torch.cat((neg_img_input_ids, img_cond_seq_ids), dim=1)
+
+                neg_pred = model(
+                    img=neg_img_input,
+                    img_ids=neg_img_input_ids,
+                    txt=neg_regional_prompting_extension.regional_text_conditioning.t5_embeddings,
+                    txt_ids=neg_regional_prompting_extension.regional_text_conditioning.t5_txt_ids,
+                    y=neg_regional_prompting_extension.regional_text_conditioning.clip_embeddings,
+                    timesteps=t_vec,
+                    guidance=guidance_vec,
+                    timestep_index=step_index,
+                    total_num_timesteps=total_steps,
+                    controlnet_double_block_residuals=None,
+                    controlnet_single_block_residuals=None,
+                    ip_adapter_extensions=neg_ip_adapter_extensions,
+                    regional_prompting_extension=neg_regional_prompting_extension,
+                )
+
+                # Slice negative prediction to match main image tokens
+                if img_cond_seq is not None:
+                    neg_pred = neg_pred[:, :original_seq_len]
+                pred = neg_pred + step_cfg_scale * (pred - neg_pred)
+
+            preview_img = img - t_curr * pred
+            img = img + (t_prev - t_curr) * pred
+
+            if inpaint_extension is not None:
+                img = inpaint_extension.merge_intermediate_latents_with_init_latents(img, t_prev)
+                preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_img, 0.0)
+
+            step_callback(
+                PipelineIntermediateState(
+                    step=step_index + 1,
+                    order=1,
+                    total_steps=total_steps,
+                    timestep=int(t_curr),
+                    latents=preview_img,
+                ),
+            )
+
+        return img
+
+    finally:
+        # DyPE: Restore original position embedder
+        if original_pe_embedder is not None:
+            DyPEExtension.restore_model(model, original_pe_embedder)

invokeai/backend/flux/dype/__init__.py

@@ -0,0 +1,35 @@
+"""Dynamic Position Extrapolation (DyPE) for FLUX models.
+
+DyPE enables high-resolution image generation with pretrained FLUX models by
+dynamically modulating RoPE extrapolation during denoising.
+
+Based on the official DyPE project: https://github.com/guyyariv/DyPE
+"""
+
+from invokeai.backend.flux.dype.base import DyPEConfig
+from invokeai.backend.flux.dype.embed import DyPEEmbedND
+from invokeai.backend.flux.dype.presets import (
+    DYPE_PRESET_4K,
+    DYPE_PRESET_AREA,
+    DYPE_PRESET_AUTO,
+    DYPE_PRESET_LABELS,
+    DYPE_PRESET_MANUAL,
+    DYPE_PRESET_OFF,
+    DyPEPreset,
+    get_dype_config_for_area,
+    get_dype_config_for_resolution,
+)
+
+__all__ = [
+    "DyPEConfig",
+    "DyPEEmbedND",
+    "DyPEPreset",
+    "DYPE_PRESET_OFF",
+    "DYPE_PRESET_MANUAL",
+    "DYPE_PRESET_AUTO",
+    "DYPE_PRESET_AREA",
+    "DYPE_PRESET_4K",
+    "DYPE_PRESET_LABELS",
+    "get_dype_config_for_area",
+    "get_dype_config_for_resolution",
+]

invokeai/backend/flux/dype/base.py

@@ -0,0 +1,115 @@
+"""DyPE base configuration and utilities for FLUX vision_yarn RoPE."""
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class DyPEConfig:
+    """Configuration for Dynamic Position Extrapolation."""
+
+    enable_dype: bool = True
+    base_resolution: int = 1024  # Native training resolution
+    dype_scale: float = 2.0  # Magnitude λs (0.0-8.0)
+    dype_exponent: float = 2.0  # Decay speed λt (0.0-1000.0)
+    dype_start_sigma: float = 1.0  # When DyPE decay starts
+
+
+def get_timestep_kappa(
+    current_sigma: float,
+    dype_scale: float,
+    dype_exponent: float,
+    dype_start_sigma: float,
+) -> float:
+    """Calculate the paper-style DyPE scheduler value κ(t).
+
+    The key insight of DyPE: early steps focus on low frequencies (global structure),
+    late steps on high frequencies (details). DyPE expresses this as a direct
+    timestep scheduler over the positional extrapolation strength:
+
+        κ(t) = λs * t^λt
+
+    Args:
+        current_sigma: Current noise level (1.0 = full noise, 0.0 = clean)
+        dype_scale: DyPE magnitude (λs)
+        dype_exponent: DyPE decay speed (λt)
+        dype_start_sigma: Sigma threshold to start decay
+
+    Returns:
+        Timestep scheduler value κ(t)
+    """
+    if dype_scale <= 0.0 or dype_start_sigma <= 0.0:
+        return 0.0
+
+    t_normalized = max(0.0, min(current_sigma / dype_start_sigma, 1.0))
+    return dype_scale * (t_normalized**dype_exponent)
+
+
+def compute_vision_yarn_freqs(
+    pos: Tensor,
+    dim: int,
+    theta: int,
+    scale_h: float,
+    scale_w: float,
+    current_sigma: float,
+    dype_config: DyPEConfig,
+) -> tuple[Tensor, Tensor]:
+    """Compute RoPE frequencies using NTK-aware scaling for high-resolution.
+
+    This method extends FLUX's position encoding to handle resolutions beyond
+    the 1024px training resolution by scaling the base frequency (theta).
+
+    The NTK-aware approach smoothly interpolates frequencies to cover larger
+    position ranges without breaking the attention patterns.
+
+    DyPE (Dynamic Position Extrapolation) modulates the NTK scaling based on
+    the current timestep - stronger extrapolation in early steps (global structure),
+    weaker in late steps (fine details).
+
+    Args:
+        pos: Position tensor
+        dim: Embedding dimension
+        theta: RoPE base frequency
+        scale_h: Height scaling factor
+        scale_w: Width scaling factor
+        current_sigma: Current noise level (1.0 = full noise, 0.0 = clean)
+        dype_config: DyPE configuration
+
+    Returns:
+        Tuple of (cos, sin) frequency tensors
+    """
+    assert dim % 2 == 0
+
+    scale = max(scale_h, scale_w)
+
+    device = pos.device
+    dtype = torch.float64 if device.type != "mps" else torch.float32
+
+    # DyPE applies a direct timestep scheduler to the NTK extrapolation exponent.
+    # Early steps keep strong extrapolation; late steps relax smoothly back
+    # toward the training-time RoPE.
+    if scale > 1.0:
+        ntk_exponent = dim / (dim - 2)
+        kappa = get_timestep_kappa(
+            current_sigma=current_sigma,
+            dype_scale=dype_config.dype_scale,
+            dype_exponent=dype_config.dype_exponent,
+            dype_start_sigma=dype_config.dype_start_sigma,
+        )
+        scaled_theta = theta * (scale ** (ntk_exponent * kappa))
+    else:
+        scaled_theta = theta
+
+    # Standard RoPE frequency computation
+    freq_seq = torch.arange(0, dim, 2, dtype=dtype, device=device) / dim
+    freqs = 1.0 / (scaled_theta**freq_seq)
+
+    # Compute angles = position * frequency
+    angles = torch.einsum("...n,d->...nd", pos.to(dtype), freqs)
+
+    cos = torch.cos(angles)
+    sin = torch.sin(angles)
+
+    return cos.to(pos.dtype), sin.to(pos.dtype)

invokeai/backend/flux/dype/embed.py

@@ -0,0 +1,116 @@
+"""DyPE-enhanced position embedding module."""
+
+import torch
+from torch import Tensor, nn
+
+from invokeai.backend.flux.dype.base import DyPEConfig
+from invokeai.backend.flux.dype.rope import rope_dype
+
+
+class DyPEEmbedND(nn.Module):
+    """N-dimensional position embedding with DyPE support.
+
+    This class replaces the standard EmbedND from FLUX with a DyPE-aware version
+    that dynamically scales position embeddings based on resolution and timestep.
+
+    The key difference from EmbedND:
+    - Maintains step state (current_sigma, target dimensions)
+    - Uses rope_dype() instead of rope() for frequency computation
+    - Applies timestep-dependent scaling for better high-resolution generation
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        theta: int,
+        axes_dim: list[int],
+        dype_config: DyPEConfig,
+    ):
+        """Initialize DyPE position embedder.
+
+        Args:
+            dim: Total embedding dimension (sum of axes_dim)
+            theta: RoPE base frequency
+            axes_dim: Dimension allocation per axis (e.g., [16, 56, 56] for FLUX)
+            dype_config: DyPE configuration
+        """
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+        self.dype_config = dype_config
+
+        # Step state - updated before each denoising step
+        self._current_sigma: float = 1.0
+        self._target_height: int = 1024
+        self._target_width: int = 1024
+
+    def set_step_state(self, sigma: float, height: int, width: int) -> None:
+        """Update the step state before each denoising step.
+
+        This method should be called by the DyPE extension before each step
+        to update the current noise level and target dimensions.
+
+        Args:
+            sigma: Current noise level (timestep value, 1.0 = full noise)
+            height: Target image height in pixels
+            width: Target image width in pixels
+        """
+        self._current_sigma = sigma
+        self._target_height = height
+        self._target_width = width
+
+    def forward(self, ids: Tensor) -> Tensor:
+        """Compute position embeddings with DyPE scaling.
+
+        Args:
+            ids: Position indices tensor with shape (batch, seq_len, n_axes)
+                 For FLUX: n_axes=3 (time/channel, height, width)
+
+        Returns:
+            Position embedding tensor with shape (batch, 1, seq_len, dim)
+        """
+        n_axes = ids.shape[-1]
+
+        # Compute RoPE for each axis with DyPE scaling
+        embeddings = []
+        for i in range(n_axes):
+            axis_emb = rope_dype(
+                pos=ids[..., i],
+                dim=self.axes_dim[i],
+                theta=self.theta,
+                current_sigma=self._current_sigma,
+                target_height=self._target_height,
+                target_width=self._target_width,
+                dype_config=self.dype_config,
+            )
+            embeddings.append(axis_emb)
+
+        # Concatenate embeddings from all axes
+        emb = torch.cat(embeddings, dim=-3)
+
+        return emb.unsqueeze(1)
+
+    @classmethod
+    def from_embednd(
+        cls,
+        embed_nd: nn.Module,
+        dype_config: DyPEConfig,
+    ) -> "DyPEEmbedND":
+        """Create a DyPEEmbedND from an existing EmbedND.
+
+        This is a convenience method for patching an existing FLUX model.
+
+        Args:
+            embed_nd: Original EmbedND module from FLUX
+            dype_config: DyPE configuration
+
+        Returns:
+            New DyPEEmbedND with same parameters
+        """
+        return cls(
+            dim=embed_nd.dim,
+            theta=embed_nd.theta,
+            axes_dim=embed_nd.axes_dim,
+            dype_config=dype_config,
+        )

invokeai/backend/flux/dype/presets.py

@@ -0,0 +1,198 @@
+"""DyPE presets and automatic configuration."""
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+from invokeai.backend.flux.dype.base import DyPEConfig
+
+# DyPE preset type - using Literal for proper frontend dropdown support
+DyPEPreset = Literal["off", "manual", "auto", "area", "4k"]
+
+# Constants for preset values
+DYPE_PRESET_OFF: DyPEPreset = "off"
+DYPE_PRESET_MANUAL: DyPEPreset = "manual"
+DYPE_PRESET_AUTO: DyPEPreset = "auto"
+DYPE_PRESET_AREA: DyPEPreset = "area"
+DYPE_PRESET_4K: DyPEPreset = "4k"
+
+# Human-readable labels for the UI
+DYPE_PRESET_LABELS: dict[str, str] = {
+    "off": "Off",
+    "manual": "Manual",
+    "auto": "Auto (>1536px)",
+    "area": "Area (auto)",
+    "4k": "4K Optimized",
+}
+
+
+@dataclass
+class DyPEPresetConfig:
+    """Preset configuration values."""
+
+    base_resolution: int
+    dype_scale: float
+    dype_exponent: float
+    dype_start_sigma: float
+
+
+# Predefined preset configurations
+DYPE_PRESETS: dict[DyPEPreset, DyPEPresetConfig] = {
+    DYPE_PRESET_4K: DyPEPresetConfig(
+        base_resolution=1024,
+        dype_scale=2.0,
+        dype_exponent=2.0,
+        dype_start_sigma=1.0,
+    ),
+}
+
+
+def get_dype_config_for_resolution(
+    width: int,
+    height: int,
+    base_resolution: int = 1024,
+    activation_threshold: int = 1536,
+) -> DyPEConfig | None:
+    """Automatically determine DyPE config based on target resolution.
+
+    FLUX can handle resolutions up to ~1.5x natively without significant artifacts.
+    DyPE is only activated when the resolution exceeds the activation threshold.
+
+    Args:
+        width: Target image width in pixels
+        height: Target image height in pixels
+        base_resolution: Native training resolution of the model (for scale calculation)
+        activation_threshold: Resolution threshold above which DyPE is activated
+
+    Returns:
+        DyPEConfig if DyPE should be enabled, None otherwise
+    """
+    max_dim = max(width, height)
+
+    if max_dim <= activation_threshold:
+        return None  # FLUX can handle this natively
+
+    # Calculate scaling factor based on base_resolution
+    scale = max_dim / base_resolution
+
+    # Dynamic parameters based on scaling
+    # Higher resolution = higher dype_scale, capped at 8.0
+    dynamic_dype_scale = min(2.0 * scale, 8.0)
+
+    return DyPEConfig(
+        enable_dype=True,
+        base_resolution=base_resolution,
+        dype_scale=dynamic_dype_scale,
+        dype_exponent=2.0,
+        dype_start_sigma=1.0,
+    )
+
+
+def get_dype_config_for_area(
+    width: int,
+    height: int,
+    base_resolution: int = 1024,
+) -> DyPEConfig | None:
+    """Automatically determine DyPE config based on target area.
+
+    Uses sqrt(area/base_area) as an effective side-length ratio.
+    DyPE is enabled only when target area exceeds base area.
+
+    Returns:
+        DyPEConfig if DyPE should be enabled, None otherwise
+    """
+    area = width * height
+    base_area = base_resolution**2
+
+    if area <= base_area:
+        return None
+
+    area_ratio = area / base_area
+    effective_side_ratio = math.sqrt(area_ratio)
+    aspect_ratio = max(width, height) / min(width, height)
+    aspect_attenuation = 1.0 if aspect_ratio <= 2.0 else 2.0 / aspect_ratio
+
+    # Retune area mode to be "auto, but area-aware" instead of dramatically
+    # stronger than auto. This keeps it closer to the paper-style core DyPE.
+    dynamic_dype_scale = 2.4 * effective_side_ratio
+    dynamic_dype_scale *= aspect_attenuation
+    dynamic_dype_scale = max(0.0, min(dynamic_dype_scale, 8.0))
+
+    # Use a narrower, higher exponent range than the old area heuristic so the
+    # paper-style scheduler decays more conservatively and artifacts are reduced.
+    exponent_progress = max(0.0, min(effective_side_ratio - 1.0, 1.0))
+    dype_exponent = 1.25 + 0.75 * exponent_progress
+
+    return DyPEConfig(
+        enable_dype=True,
+        base_resolution=base_resolution,
+        dype_scale=dynamic_dype_scale,
+        dype_exponent=dype_exponent,
+        dype_start_sigma=1.0,
+    )
+
+
+def get_dype_config_from_preset(
+    preset: DyPEPreset,
+    width: int,
+    height: int,
+    custom_scale: float | None = None,
+    custom_exponent: float | None = None,
+) -> DyPEConfig | None:
+    """Get DyPE configuration from a preset or custom values.
+
+    Args:
+        preset: The DyPE preset to use
+        width: Target image width
+        height: Target image height
+        custom_scale: Optional custom dype_scale (only used with 'manual' preset)
+        custom_exponent: Optional custom dype_exponent (only used with 'manual' preset)
+
+    Returns:
+        DyPEConfig if DyPE should be enabled, None otherwise
+    """
+    if preset == DYPE_PRESET_OFF:
+        return None
+
+    if preset == DYPE_PRESET_MANUAL:
+        # Manual mode - custom values can override defaults
+        max_dim = max(width, height)
+        scale = max_dim / 1024
+        dynamic_dype_scale = min(2.0 * scale, 8.0)
+        return DyPEConfig(
+            enable_dype=True,
+            base_resolution=1024,
+            dype_scale=custom_scale if custom_scale is not None else dynamic_dype_scale,
+            dype_exponent=custom_exponent if custom_exponent is not None else 2.0,
+            dype_start_sigma=1.0,
+        )
+
+    if preset == DYPE_PRESET_AUTO:
+        # Auto preset - custom values are ignored
+        return get_dype_config_for_resolution(
+            width=width,
+            height=height,
+            base_resolution=1024,
+            activation_threshold=1536,
+        )
+
+    if preset == DYPE_PRESET_AREA:
+        # Area-based preset - custom values are ignored
+        return get_dype_config_for_area(
+            width=width,
+            height=height,
+            base_resolution=1024,
+        )
+
+    # Use preset configuration (4K etc.) - custom values are ignored
+    preset_config = DYPE_PRESETS.get(preset)
+    if preset_config is None:
+        return None
+
+    return DyPEConfig(
+        enable_dype=True,
+        base_resolution=preset_config.base_resolution,
+        dype_scale=preset_config.dype_scale,
+        dype_exponent=preset_config.dype_exponent,
+        dype_start_sigma=preset_config.dype_start_sigma,
+    )

invokeai/backend/flux/dype/rope.py

@@ -0,0 +1,86 @@
+"""DyPE-enhanced RoPE (Rotary Position Embedding) functions."""
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+from invokeai.backend.flux.dype.base import (
+    DyPEConfig,
+    compute_vision_yarn_freqs,
+)
+
+
+def rope_dype(
+    pos: Tensor,
+    dim: int,
+    theta: int,
+    current_sigma: float,
+    target_height: int,
+    target_width: int,
+    dype_config: DyPEConfig,
+) -> Tensor:
+    """Compute RoPE with Dynamic Position Extrapolation.
+
+    This is the core DyPE function that replaces the standard rope() function.
+    It applies resolution-aware and timestep-aware scaling to position embeddings.
+
+    Args:
+        pos: Position indices tensor
+        dim: Embedding dimension per axis
+        theta: RoPE base frequency (typically 10000)
+        current_sigma: Current noise level (1.0 = full noise, 0.0 = clean)
+        target_height: Target image height in pixels
+        target_width: Target image width in pixels
+        dype_config: DyPE configuration
+
+    Returns:
+        Rotary position embedding tensor with shape suitable for FLUX attention
+    """
+    assert dim % 2 == 0
+
+    # Calculate scaling factors
+    base_res = dype_config.base_resolution
+    scale_h = target_height / base_res
+    scale_w = target_width / base_res
+    scale = max(scale_h, scale_w)
+
+    # If no scaling needed and DyPE disabled, use base method
+    if not dype_config.enable_dype or scale <= 1.0:
+        return _rope_base(pos, dim, theta)
+
+    cos, sin = compute_vision_yarn_freqs(
+        pos=pos,
+        dim=dim,
+        theta=theta,
+        scale_h=scale_h,
+        scale_w=scale_w,
+        current_sigma=current_sigma,
+        dype_config=dype_config,
+    )
+
+    # Construct rotation matrix from cos/sin
+    # Output shape: (batch, seq_len, dim/2, 2, 2)
+    out = torch.stack([cos, -sin, sin, cos], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+
+    return out.to(dtype=pos.dtype, device=pos.device)
+
+
+def _rope_base(pos: Tensor, dim: int, theta: int) -> Tensor:
+    """Standard RoPE without DyPE scaling.
+
+    This matches the original rope() function from invokeai.backend.flux.math.
+    """
+    assert dim % 2 == 0
+
+    device = pos.device
+    dtype = torch.float64 if device.type != "mps" else torch.float32
+
+    scale = torch.arange(0, dim, 2, dtype=dtype, device=device) / dim
+    omega = 1.0 / (theta**scale)
+
+    out = torch.einsum("...n,d->...nd", pos.to(dtype), omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+
+    return out.to(dtype=pos.dtype, device=pos.device)

invokeai/backend/flux/extensions/base_controlnet_extension.py

@@ -0,0 +1,45 @@
+import math
+from abc import ABC, abstractmethod
+from typing import List, Union
+
+import torch
+
+from invokeai.backend.flux.controlnet.controlnet_flux_output import ControlNetFluxOutput
+
+
+class BaseControlNetExtension(ABC):
+    def __init__(
+        self,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+    ):
+        self._weight = weight
+        self._begin_step_percent = begin_step_percent
+        self._end_step_percent = end_step_percent
+
+    def _get_weight(self, timestep_index: int, total_num_timesteps: int) -> float:
+        first_step = math.floor(self._begin_step_percent * total_num_timesteps)
+        last_step = math.ceil(self._end_step_percent * total_num_timesteps)
+
+        if timestep_index < first_step or timestep_index > last_step:
+            return 0.0
+
+        if isinstance(self._weight, list):
+            return self._weight[timestep_index]
+
+        return self._weight
+
+    @abstractmethod
+    def run_controlnet(
+        self,
+        timestep_index: int,
+        total_num_timesteps: int,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt: torch.Tensor,
+        txt_ids: torch.Tensor,
+        y: torch.Tensor,
+        timesteps: torch.Tensor,
+        guidance: torch.Tensor | None,
+    ) -> ControlNetFluxOutput: ...

invokeai/backend/flux/extensions/dype_extension.py

@@ -0,0 +1,113 @@
+"""DyPE extension for FLUX denoising pipeline."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Sequence
+
+import torch
+
+from invokeai.backend.flux.dype.base import DyPEConfig
+from invokeai.backend.flux.dype.embed import DyPEEmbedND
+
+if TYPE_CHECKING:
+    from invokeai.backend.flux.model import Flux
+
+
+@dataclass
+class DyPEExtension:
+    """Extension for Dynamic Position Extrapolation in FLUX models.
+
+    This extension manages the patching of the FLUX model's position embedder
+    and updates the step state during denoising.
+
+    Usage:
+        1. Create extension with config and target dimensions
+        2. Call patch_model() to replace pe_embedder with DyPE version
+        3. Call update_step_state() before each denoising step
+        4. Call restore_model() after denoising to restore original embedder
+    """
+
+    config: DyPEConfig
+    target_height: int
+    target_width: int
+
+    def patch_model(self, model: "Flux") -> tuple[DyPEEmbedND, object]:
+        """Patch the model's position embedder with DyPE version.
+
+        Args:
+            model: The FLUX model to patch
+
+        Returns:
+            Tuple of (new DyPE embedder, original embedder for restoration)
+        """
+        original_embedder = model.pe_embedder
+
+        dype_embedder = DyPEEmbedND.from_embednd(
+            embed_nd=original_embedder,
+            dype_config=self.config,
+        )
+
+        # Set initial state
+        dype_embedder.set_step_state(
+            sigma=1.0,
+            height=self.target_height,
+            width=self.target_width,
+        )
+
+        # Replace the embedder
+        model.pe_embedder = dype_embedder
+
+        return dype_embedder, original_embedder
+
+    def update_step_state(
+        self,
+        embedder: DyPEEmbedND,
+        sigma: float,
+    ) -> None:
+        """Update the step state in the DyPE embedder.
+
+        This should be called before each denoising step to update the
+        current noise level for timestep-dependent scaling.
+
+        Args:
+            embedder: The DyPE embedder to update
+            sigma: Current noise level for the active denoising step
+        """
+        embedder.set_step_state(
+            sigma=sigma,
+            height=self.target_height,
+            width=self.target_width,
+        )
+
+    @staticmethod
+    def resolve_step_sigma(
+        fallback_sigma: float,
+        step_index: int,
+        scheduler_sigmas: Sequence[float] | torch.Tensor | None,
+    ) -> float:
+        """Resolve the actual sigma for the current denoising step.
+
+        Diffusers schedulers may expose both normalized timesteps and the underlying
+        sigma sequence. DyPE should follow the noise schedule, so prefer
+        ``scheduler.sigmas`` when available and fall back to the provided value
+        otherwise.
+        """
+        if scheduler_sigmas is None:
+            return fallback_sigma
+
+        if step_index >= len(scheduler_sigmas):
+            return fallback_sigma
+
+        sigma = scheduler_sigmas[step_index]
+        if isinstance(sigma, torch.Tensor):
+            return float(sigma.item())
+        return float(sigma)
+
+    @staticmethod
+    def restore_model(model: "Flux", original_embedder: object) -> None:
+        """Restore the original position embedder.
+
+        Args:
+            model: The FLUX model to restore
+            original_embedder: The original embedder saved from patch_model()
+        """
+        model.pe_embedder = original_embedder

invokeai/backend/flux/extensions/instantx_controlnet_extension.py

@@ -0,0 +1,194 @@
+import math
+from typing import List, Union
+
+import torch
+from PIL.Image import Image
+
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
+from invokeai.app.invocations.flux_vae_encode import FluxVaeEncodeInvocation
+from invokeai.app.util.controlnet_utils import CONTROLNET_RESIZE_VALUES, prepare_control_image
+from invokeai.backend.flux.controlnet.controlnet_flux_output import ControlNetFluxOutput
+from invokeai.backend.flux.controlnet.instantx_controlnet_flux import (
+    InstantXControlNetFlux,
+    InstantXControlNetFluxOutput,
+)
+from invokeai.backend.flux.extensions.base_controlnet_extension import BaseControlNetExtension
+from invokeai.backend.flux.sampling_utils import pack
+from invokeai.backend.model_manager.load.load_base import LoadedModel
+
+
+class InstantXControlNetExtension(BaseControlNetExtension):
+    def __init__(
+        self,
+        model: InstantXControlNetFlux,
+        controlnet_cond: torch.Tensor,
+        instantx_control_mode: torch.Tensor | None,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+    ):
+        super().__init__(
+            weight=weight,
+            begin_step_percent=begin_step_percent,
+            end_step_percent=end_step_percent,
+        )
+        self._model = model
+        # The VAE-encoded and 'packed' control image to pass to the ControlNet model.
+        self._controlnet_cond = controlnet_cond
+        # TODO(ryand): Should we define an enum for the instantx_control_mode? Is it likely to change for future models?
+        # The control mode for InstantX ControlNet union models.
+        # See the values defined here: https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union#control-mode
+        # Expected shape: (batch_size, 1), Expected dtype: torch.long
+        # If None, a zero-embedding will be used.
+        self._instantx_control_mode = instantx_control_mode
+
+        # TODO(ryand): Pass in these params if a new base transformer / InstantX ControlNet pair get released.
+        self._flux_transformer_num_double_blocks = 19
+        self._flux_transformer_num_single_blocks = 38
+
+    @classmethod
+    def prepare_controlnet_cond(
+        cls,
+        controlnet_image: Image,
+        vae_info: LoadedModel,
+        latent_height: int,
+        latent_width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        resize_mode: CONTROLNET_RESIZE_VALUES,
+    ):
+        image_height = latent_height * LATENT_SCALE_FACTOR
+        image_width = latent_width * LATENT_SCALE_FACTOR
+
+        resized_controlnet_image = prepare_control_image(
+            image=controlnet_image,
+            do_classifier_free_guidance=False,
+            width=image_width,
+            height=image_height,
+            device=device,
+            dtype=dtype,
+            control_mode="balanced",
+            resize_mode=resize_mode,
+        )
+
+        # Shift the image from [0, 1] to [-1, 1].
+        resized_controlnet_image = resized_controlnet_image * 2 - 1
+
+        # Run VAE encoder.
+        controlnet_cond = FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=resized_controlnet_image)
+        controlnet_cond = pack(controlnet_cond)
+
+        return controlnet_cond
+
+    @classmethod
+    def from_controlnet_image(
+        cls,
+        model: InstantXControlNetFlux,
+        controlnet_image: Image,
+        instantx_control_mode: torch.Tensor | None,
+        vae_info: LoadedModel,
+        latent_height: int,
+        latent_width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        resize_mode: CONTROLNET_RESIZE_VALUES,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+    ):
+        image_height = latent_height * LATENT_SCALE_FACTOR
+        image_width = latent_width * LATENT_SCALE_FACTOR
+
+        resized_controlnet_image = prepare_control_image(
+            image=controlnet_image,
+            do_classifier_free_guidance=False,
+            width=image_width,
+            height=image_height,
+            device=device,
+            dtype=dtype,
+            control_mode="balanced",
+            resize_mode=resize_mode,
+        )
+
+        # Shift the image from [0, 1] to [-1, 1].
+        resized_controlnet_image = resized_controlnet_image * 2 - 1
+
+        # Run VAE encoder.
+        controlnet_cond = FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=resized_controlnet_image)
+        controlnet_cond = pack(controlnet_cond)
+
+        return cls(
+            model=model,
+            controlnet_cond=controlnet_cond,
+            instantx_control_mode=instantx_control_mode,
+            weight=weight,
+            begin_step_percent=begin_step_percent,
+            end_step_percent=end_step_percent,
+        )
+
+    def _instantx_output_to_controlnet_output(
+        self, instantx_output: InstantXControlNetFluxOutput
+    ) -> ControlNetFluxOutput:
+        # The `interval_control` logic here is based on
+        # https://github.com/huggingface/diffusers/blob/31058cdaef63ca660a1a045281d156239fba8192/src/diffusers/models/transformers/transformer_flux.py#L507-L511
+
+        # Handle double block residuals.
+        double_block_residuals: list[torch.Tensor] = []
+        double_block_samples = instantx_output.controlnet_block_samples
+        if double_block_samples:
+            interval_control = self._flux_transformer_num_double_blocks / len(double_block_samples)
+            interval_control = int(math.ceil(interval_control))
+            for i in range(self._flux_transformer_num_double_blocks):
+                double_block_residuals.append(double_block_samples[i // interval_control])
+
+        # Handle single block residuals.
+        single_block_residuals: list[torch.Tensor] = []
+        single_block_samples = instantx_output.controlnet_single_block_samples
+        if single_block_samples:
+            interval_control = self._flux_transformer_num_single_blocks / len(single_block_samples)
+            interval_control = int(math.ceil(interval_control))
+            for i in range(self._flux_transformer_num_single_blocks):
+                single_block_residuals.append(single_block_samples[i // interval_control])
+
+        return ControlNetFluxOutput(
+            double_block_residuals=double_block_residuals or None,
+            single_block_residuals=single_block_residuals or None,
+        )
+
+    def run_controlnet(
+        self,
+        timestep_index: int,
+        total_num_timesteps: int,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt: torch.Tensor,
+        txt_ids: torch.Tensor,
+        y: torch.Tensor,
+        timesteps: torch.Tensor,
+        guidance: torch.Tensor | None,
+    ) -> ControlNetFluxOutput:
+        weight = self._get_weight(timestep_index=timestep_index, total_num_timesteps=total_num_timesteps)
+        if weight < 1e-6:
+            return ControlNetFluxOutput(single_block_residuals=None, double_block_residuals=None)
+
+        # Make sure inputs have correct device and dtype.
+        self._controlnet_cond = self._controlnet_cond.to(device=img.device, dtype=img.dtype)
+        self._instantx_control_mode = (
+            self._instantx_control_mode.to(device=img.device) if self._instantx_control_mode is not None else None
+        )
+
+        instantx_output: InstantXControlNetFluxOutput = self._model(
+            controlnet_cond=self._controlnet_cond,
+            controlnet_mode=self._instantx_control_mode,
+            img=img,
+            img_ids=img_ids,
+            txt=txt,
+            txt_ids=txt_ids,
+            timesteps=timesteps,
+            y=y,
+            guidance=guidance,
+        )
+
+        controlnet_output = self._instantx_output_to_controlnet_output(instantx_output)
+        controlnet_output.apply_weight(weight)
+        return controlnet_output

invokeai/backend/flux/extensions/kontext_extension.py

@@ -0,0 +1,218 @@
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+from einops import repeat
+
+from invokeai.app.invocations.fields import FluxKontextConditioningField
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+from invokeai.backend.flux.sampling_utils import pack
+from invokeai.backend.util.devices import TorchDevice
+
+
+def generate_img_ids_with_offset(
+    latent_height: int,
+    latent_width: int,
+    batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    idx_offset: int = 0,
+    h_offset: int = 0,
+    w_offset: int = 0,
+) -> torch.Tensor:
+    """Generate tensor of image position ids with optional index and spatial offsets.
+
+    Args:
+        latent_height (int): Height of image in latent space (after packing, this becomes h//2).
+        latent_width (int): Width of image in latent space (after packing, this becomes w//2).
+        batch_size (int): Number of images in the batch.
+        device (torch.device): Device to create tensors on.
+        dtype (torch.dtype): Data type for the tensors.
+        idx_offset (int): Offset to add to the first dimension of the image ids (default: 0).
+        h_offset (int): Spatial offset for height/y-coordinates in latent space (default: 0).
+        w_offset (int): Spatial offset for width/x-coordinates in latent space (default: 0).
+
+    Returns:
+        torch.Tensor: Image position ids with shape [batch_size, (latent_height//2 * latent_width//2), 3].
+    """
+
+    if device.type == "mps":
+        orig_dtype = dtype
+        dtype = torch.float16
+
+    # After packing, the spatial dimensions are halved due to the 2x2 patch structure
+    packed_height = latent_height // 2
+    packed_width = latent_width // 2
+
+    # Convert spatial offsets from latent space to packed space
+    packed_h_offset = h_offset // 2
+    packed_w_offset = w_offset // 2
+
+    # Create base tensor for position IDs with shape [packed_height, packed_width, 3]
+    # The 3 channels represent: [batch_offset, y_position, x_position]
+    img_ids = torch.zeros(packed_height, packed_width, 3, device=device, dtype=dtype)
+
+    # Set the batch offset for all positions
+    img_ids[..., 0] = idx_offset
+
+    # Create y-coordinate indices (vertical positions) with spatial offset
+    y_indices = torch.arange(packed_height, device=device, dtype=dtype) + packed_h_offset
+    # Broadcast y_indices to match the spatial dimensions [packed_height, 1]
+    img_ids[..., 1] = y_indices[:, None]
+
+    # Create x-coordinate indices (horizontal positions) with spatial offset
+    x_indices = torch.arange(packed_width, device=device, dtype=dtype) + packed_w_offset
+    # Broadcast x_indices to match the spatial dimensions [1, packed_width]
+    img_ids[..., 2] = x_indices[None, :]
+
+    # Expand to include batch dimension: [batch_size, (packed_height * packed_width), 3]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=batch_size)
+
+    if device.type == "mps":
+        img_ids = img_ids.to(orig_dtype)
+
+    return img_ids
+
+
+class KontextExtension:
+    """Applies FLUX Kontext (reference image) conditioning."""
+
+    def __init__(
+        self,
+        kontext_conditioning: list[FluxKontextConditioningField],
+        context: InvocationContext,
+        vae_field: VAEField,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        """
+        Initializes the KontextExtension, pre-processing the reference images
+        into latents and positional IDs.
+        """
+        self._context = context
+        self._device = device
+        self._dtype = dtype
+        self._vae_field = vae_field
+        self.kontext_conditioning = kontext_conditioning
+
+        # Pre-process and cache the kontext latents and ids upon initialization.
+        self.kontext_latents, self.kontext_ids = self._prepare_kontext()
+
+    def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encodes the reference images and prepares their concatenated latents and IDs with spatial tiling."""
+        all_latents = []
+        all_ids = []
+
+        # Track cumulative dimensions for spatial tiling
+        # These track the running extent of the virtual canvas in latent space
+        canvas_h = 0  # Running canvas height
+        canvas_w = 0  # Running canvas width
+
+        vae_info = self._context.models.load(self._vae_field.vae)
+
+        for idx, kontext_field in enumerate(self.kontext_conditioning):
+            image = self._context.images.get_pil(kontext_field.image.image_name)
+
+            # Convert to RGB
+            image = image.convert("RGB")
+
+            # Convert to tensor using torchvision transforms for consistency
+            transformation = T.Compose(
+                [
+                    T.ToTensor(),  # Converts PIL image to tensor and scales to [0, 1]
+                ]
+            )
+            image_tensor = transformation(image)
+            # Convert from [0, 1] to [-1, 1] range expected by VAE
+            image_tensor = image_tensor * 2.0 - 1.0
+            image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension
+            image_tensor = image_tensor.to(self._device)
+
+            # Continue with VAE encoding
+            # Don't sample from the distribution for reference images - use the mean (matching ComfyUI)
+            # Estimate working memory for encode operation (50% of decode memory requirements)
+            img_h = image_tensor.shape[-2]
+            img_w = image_tensor.shape[-1]
+            element_size = next(vae_info.model.parameters()).element_size()
+            scaling_constant = 1100  # 50% of decode scaling constant (2200)
+            estimated_working_memory = int(img_h * img_w * element_size * scaling_constant)
+
+            with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+                assert isinstance(vae, AutoEncoder)
+                vae_dtype = next(iter(vae.parameters())).dtype
+                image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
+                # Use sample=False to get the distribution mean without noise
+                kontext_latents_unpacked = vae.encode(image_tensor, sample=False)
+                TorchDevice.empty_cache()
+
+            # Extract tensor dimensions
+            batch_size, _, latent_height, latent_width = kontext_latents_unpacked.shape
+
+            # Pad latents to be compatible with patch_size=2
+            # This ensures dimensions are even for the pack() function
+            pad_h = (2 - latent_height % 2) % 2
+            pad_w = (2 - latent_width % 2) % 2
+            if pad_h > 0 or pad_w > 0:
+                kontext_latents_unpacked = F.pad(kontext_latents_unpacked, (0, pad_w, 0, pad_h), mode="circular")
+                # Update dimensions after padding
+                _, _, latent_height, latent_width = kontext_latents_unpacked.shape
+
+            # Pack the latents
+            kontext_latents_packed = pack(kontext_latents_unpacked).to(self._device, self._dtype)
+
+            # Determine spatial offsets for this reference image
+            h_offset = 0
+            w_offset = 0
+
+            if idx > 0:  # First image starts at (0, 0)
+                # Calculate potential canvas dimensions for each tiling option
+                # Option 1: Tile vertically (below existing content)
+                potential_h_vertical = canvas_h + latent_height
+
+                # Option 2: Tile horizontally (to the right of existing content)
+                potential_w_horizontal = canvas_w + latent_width
+
+                # Choose arrangement that minimizes the maximum dimension
+                # This keeps the canvas closer to square, optimizing attention computation
+                if potential_h_vertical > potential_w_horizontal:
+                    # Tile horizontally (to the right of existing images)
+                    w_offset = canvas_w
+                    canvas_w = canvas_w + latent_width
+                    canvas_h = max(canvas_h, latent_height)
+                else:
+                    # Tile vertically (below existing images)
+                    h_offset = canvas_h
+                    canvas_h = canvas_h + latent_height
+                    canvas_w = max(canvas_w, latent_width)
+            else:
+                # First image - just set canvas dimensions
+                canvas_h = latent_height
+                canvas_w = latent_width
+
+            # Generate IDs with both index offset and spatial offsets
+            kontext_ids = generate_img_ids_with_offset(
+                latent_height=latent_height,
+                latent_width=latent_width,
+                batch_size=batch_size,
+                device=self._device,
+                dtype=self._dtype,
+                idx_offset=1,  # All reference images use index=1 (matching ComfyUI implementation)
+                h_offset=h_offset,
+                w_offset=w_offset,
+            )
+
+            all_latents.append(kontext_latents_packed)
+            all_ids.append(kontext_ids)
+
+        # Concatenate all latents and IDs along the sequence dimension
+        concatenated_latents = torch.cat(all_latents, dim=1)  # Concatenate along sequence dimension
+        concatenated_ids = torch.cat(all_ids, dim=1)  # Concatenate along sequence dimension
+
+        return concatenated_latents, concatenated_ids
+
+    def ensure_batch_size(self, target_batch_size: int) -> None:
+        """Ensures the kontext latents and IDs match the target batch size by repeating if necessary."""
+        if self.kontext_latents.shape[0] != target_batch_size:
+            self.kontext_latents = self.kontext_latents.repeat(target_batch_size, 1, 1)
+            self.kontext_ids = self.kontext_ids.repeat(target_batch_size, 1, 1)

invokeai/backend/flux/extensions/regional_prompting_extension.py

@@ -0,0 +1,295 @@
+from typing import Optional
+
+import torch
+import torchvision
+
+from invokeai.backend.flux.text_conditioning import (
+    FluxReduxConditioning,
+    FluxRegionalTextConditioning,
+    FluxTextConditioning,
+)
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.mask import to_standard_float_mask
+
+
+class RegionalPromptingExtension:
+    """A class for managing regional prompting with FLUX.
+
+    This implementation is inspired by https://arxiv.org/pdf/2411.02395 (though there are significant differences).
+    """
+
+    def __init__(
+        self,
+        regional_text_conditioning: FluxRegionalTextConditioning,
+        restricted_attn_mask: torch.Tensor | None = None,
+    ):
+        self.regional_text_conditioning = regional_text_conditioning
+        self.restricted_attn_mask = restricted_attn_mask
+
+    def get_double_stream_attn_mask(self, block_index: int) -> torch.Tensor | None:
+        order = [self.restricted_attn_mask, None]
+        return order[block_index % len(order)]
+
+    def get_single_stream_attn_mask(self, block_index: int) -> torch.Tensor | None:
+        order = [self.restricted_attn_mask, None]
+        return order[block_index % len(order)]
+
+    @classmethod
+    def from_text_conditioning(
+        cls,
+        text_conditioning: list[FluxTextConditioning],
+        redux_conditioning: list[FluxReduxConditioning],
+        img_seq_len: int,
+    ):
+        """Create a RegionalPromptingExtension from a list of text conditionings.
+
+        Args:
+            text_conditioning (list[FluxTextConditioning]): The text conditionings to use for regional prompting.
+            img_seq_len (int): The image sequence length (i.e. packed_height * packed_width).
+        """
+        regional_text_conditioning = cls._concat_regional_text_conditioning(text_conditioning, redux_conditioning)
+        attn_mask_with_restricted_img_self_attn = cls._prepare_restricted_attn_mask(
+            regional_text_conditioning, img_seq_len
+        )
+        return cls(
+            regional_text_conditioning=regional_text_conditioning,
+            restricted_attn_mask=attn_mask_with_restricted_img_self_attn,
+        )
+
+    # Keeping _prepare_unrestricted_attn_mask for reference as an alternative masking strategy:
+    #
+    # @classmethod
+    # def _prepare_unrestricted_attn_mask(
+    #     cls,
+    #     regional_text_conditioning: FluxRegionalTextConditioning,
+    #     img_seq_len: int,
+    # ) -> torch.Tensor:
+    #     """Prepare an 'unrestricted' attention mask. In this context, 'unrestricted' means that:
+    #     - img self-attention is not masked.
+    #     - img regions attend to both txt within their own region and to global prompts.
+    #     """
+    #     device = TorchDevice.choose_torch_device()
+
+    #     # Infer txt_seq_len from the t5_embeddings tensor.
+    #     txt_seq_len = regional_text_conditioning.t5_embeddings.shape[1]
+
+    #     # In the attention blocks, the txt seq and img seq are concatenated and then attention is applied.
+    #     # Concatenation happens in the following order: [txt_seq, img_seq].
+    #     # There are 4 portions of the attention mask to consider as we prepare it:
+    #     # 1. txt attends to itself
+    #     # 2. txt attends to corresponding regional img
+    #     # 3. regional img attends to corresponding txt
+    #     # 4. regional img attends to itself
+
+    #     # Initialize empty attention mask.
+    #     regional_attention_mask = torch.zeros(
+    #         (txt_seq_len + img_seq_len, txt_seq_len + img_seq_len), device=device, dtype=torch.float16
+    #     )
+
+    #     for image_mask, t5_embedding_range in zip(
+    #         regional_text_conditioning.image_masks, regional_text_conditioning.t5_embedding_ranges, strict=True
+    #     ):
+    #         # 1. txt attends to itself
+    #         regional_attention_mask[
+    #             t5_embedding_range.start : t5_embedding_range.end, t5_embedding_range.start : t5_embedding_range.end
+    #         ] = 1.0
+
+    #         # 2. txt attends to corresponding regional img
+    #         # Note that we reshape to (1, img_seq_len) to ensure broadcasting works as desired.
+    #         fill_value = image_mask.view(1, img_seq_len) if image_mask is not None else 1.0
+    #         regional_attention_mask[t5_embedding_range.start : t5_embedding_range.end, txt_seq_len:] = fill_value
+
+    #         # 3. regional img attends to corresponding txt
+    #         # Note that we reshape to (img_seq_len, 1) to ensure broadcasting works as desired.
+    #         fill_value = image_mask.view(img_seq_len, 1) if image_mask is not None else 1.0
+    #         regional_attention_mask[txt_seq_len:, t5_embedding_range.start : t5_embedding_range.end] = fill_value
+
+    #     # 4. regional img attends to itself
+    #     # Allow unrestricted img self attention.
+    #     regional_attention_mask[txt_seq_len:, txt_seq_len:] = 1.0
+
+    #     # Convert attention mask to boolean.
+    #     regional_attention_mask = regional_attention_mask > 0.5
+
+    #     return regional_attention_mask
+
+    @classmethod
+    def _prepare_restricted_attn_mask(
+        cls,
+        regional_text_conditioning: FluxRegionalTextConditioning,
+        img_seq_len: int,
+    ) -> torch.Tensor | None:
+        """Prepare a 'restricted' attention mask. In this context, 'restricted' means that:
+        - img self-attention is only allowed within regions.
+        - img regions only attend to txt within their own region, not to global prompts.
+        """
+        # Identify background region. I.e. the region that is not covered by any region masks.
+        background_region_mask: None | torch.Tensor = None
+        for image_mask in regional_text_conditioning.image_masks:
+            if image_mask is not None:
+                if background_region_mask is None:
+                    background_region_mask = torch.ones_like(image_mask)
+                background_region_mask *= 1 - image_mask
+
+        if background_region_mask is None:
+            # There are no region masks, short-circuit and return None.
+            # TODO(ryand): We could restrict txt-txt attention across multiple global prompts, but this would
+            # is a rare use case and would make the logic here significantly more complicated.
+            return None
+
+        device = TorchDevice.choose_torch_device()
+
+        # Infer txt_seq_len from the t5_embeddings tensor.
+        txt_seq_len = regional_text_conditioning.t5_embeddings.shape[1]
+
+        # In the attention blocks, the txt seq and img seq are concatenated and then attention is applied.
+        # Concatenation happens in the following order: [txt_seq, img_seq].
+        # There are 4 portions of the attention mask to consider as we prepare it:
+        # 1. txt attends to itself
+        # 2. txt attends to corresponding regional img
+        # 3. regional img attends to corresponding txt
+        # 4. regional img attends to itself
+
+        # Initialize empty attention mask.
+        regional_attention_mask = torch.zeros(
+            (txt_seq_len + img_seq_len, txt_seq_len + img_seq_len), device=device, dtype=torch.float16
+        )
+
+        for image_mask, t5_embedding_range in zip(
+            regional_text_conditioning.image_masks, regional_text_conditioning.t5_embedding_ranges, strict=True
+        ):
+            # 1. txt attends to itself
+            regional_attention_mask[
+                t5_embedding_range.start : t5_embedding_range.end, t5_embedding_range.start : t5_embedding_range.end
+            ] = 1.0
+
+            if image_mask is not None:
+                # 2. txt attends to corresponding regional img
+                # Note that we reshape to (1, img_seq_len) to ensure broadcasting works as desired.
+                regional_attention_mask[t5_embedding_range.start : t5_embedding_range.end, txt_seq_len:] = (
+                    image_mask.view(1, img_seq_len)
+                )
+
+                # 3. regional img attends to corresponding txt
+                # Note that we reshape to (img_seq_len, 1) to ensure broadcasting works as desired.
+                regional_attention_mask[txt_seq_len:, t5_embedding_range.start : t5_embedding_range.end] = (
+                    image_mask.view(img_seq_len, 1)
+                )
+
+                # 4. regional img attends to itself
+                image_mask = image_mask.view(img_seq_len, 1)
+                regional_attention_mask[txt_seq_len:, txt_seq_len:] += image_mask @ image_mask.T
+            else:
+                # We don't allow attention between non-background image regions and global prompts. This helps to ensure
+                # that regions focus on their local prompts. We do, however, allow attention between background regions
+                # and global prompts. If we didn't do this, then the background regions would not attend to any txt
+                # embeddings, which we found experimentally to cause artifacts.
+
+                # 2. global txt attends to background region
+                # Note that we reshape to (1, img_seq_len) to ensure broadcasting works as desired.
+                regional_attention_mask[t5_embedding_range.start : t5_embedding_range.end, txt_seq_len:] = (
+                    background_region_mask.view(1, img_seq_len)
+                )
+
+                # 3. background region attends to global txt
+                # Note that we reshape to (img_seq_len, 1) to ensure broadcasting works as desired.
+                regional_attention_mask[txt_seq_len:, t5_embedding_range.start : t5_embedding_range.end] = (
+                    background_region_mask.view(img_seq_len, 1)
+                )
+
+        # Allow background regions to attend to themselves.
+        regional_attention_mask[txt_seq_len:, txt_seq_len:] += background_region_mask.view(img_seq_len, 1)
+        regional_attention_mask[txt_seq_len:, txt_seq_len:] += background_region_mask.view(1, img_seq_len)
+
+        # Convert attention mask to boolean.
+        regional_attention_mask = regional_attention_mask > 0.5
+
+        return regional_attention_mask
+
+    @classmethod
+    def _concat_regional_text_conditioning(
+        cls,
+        text_conditionings: list[FluxTextConditioning],
+        redux_conditionings: list[FluxReduxConditioning],
+    ) -> FluxRegionalTextConditioning:
+        """Concatenate regional text conditioning data into a single conditioning tensor (with associated masks)."""
+        concat_t5_embeddings: list[torch.Tensor] = []
+        concat_t5_embedding_ranges: list[Range] = []
+        image_masks: list[torch.Tensor | None] = []
+
+        # Choose global CLIP embedding.
+        # Use the first global prompt's CLIP embedding as the global CLIP embedding. If there is no global prompt, use
+        # the first prompt's CLIP embedding.
+        global_clip_embedding: torch.Tensor = text_conditionings[0].clip_embeddings
+        for text_conditioning in text_conditionings:
+            if text_conditioning.mask is None:
+                global_clip_embedding = text_conditioning.clip_embeddings
+                break
+
+        # Handle T5 text embeddings.
+        cur_t5_embedding_len = 0
+        for text_conditioning in text_conditionings:
+            concat_t5_embeddings.append(text_conditioning.t5_embeddings)
+            concat_t5_embedding_ranges.append(
+                Range(start=cur_t5_embedding_len, end=cur_t5_embedding_len + text_conditioning.t5_embeddings.shape[1])
+            )
+            image_masks.append(text_conditioning.mask)
+            cur_t5_embedding_len += text_conditioning.t5_embeddings.shape[1]
+
+        # Handle Redux embeddings.
+        for redux_conditioning in redux_conditionings:
+            concat_t5_embeddings.append(redux_conditioning.redux_embeddings)
+            concat_t5_embedding_ranges.append(
+                Range(
+                    start=cur_t5_embedding_len, end=cur_t5_embedding_len + redux_conditioning.redux_embeddings.shape[1]
+                )
+            )
+            image_masks.append(redux_conditioning.mask)
+            cur_t5_embedding_len += redux_conditioning.redux_embeddings.shape[1]
+
+        t5_embeddings = torch.cat(concat_t5_embeddings, dim=1)
+
+        # Initialize the txt_ids tensor.
+        pos_bs, pos_t5_seq_len, _ = t5_embeddings.shape
+        t5_txt_ids = torch.zeros(
+            pos_bs, pos_t5_seq_len, 3, dtype=t5_embeddings.dtype, device=TorchDevice.choose_torch_device()
+        )
+
+        return FluxRegionalTextConditioning(
+            t5_embeddings=t5_embeddings,
+            clip_embeddings=global_clip_embedding,
+            t5_txt_ids=t5_txt_ids,
+            image_masks=image_masks,
+            t5_embedding_ranges=concat_t5_embedding_ranges,
+        )
+
+    @staticmethod
+    def preprocess_regional_prompt_mask(
+        mask: Optional[torch.Tensor], packed_height: int, packed_width: int, dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        """Preprocess a regional prompt mask to match the target height and width.
+        If mask is None, returns a mask of all ones with the target height and width.
+        If mask is not None, resizes the mask to the target height and width using 'nearest' interpolation.
+
+        packed_height and packed_width are the target height and width of the mask in the 'packed' latent space.
+
+        Returns:
+            torch.Tensor: The processed mask. shape: (1, 1, packed_height * packed_width).
+        """
+
+        if mask is None:
+            return torch.ones((1, 1, packed_height * packed_width), dtype=dtype, device=device)
+
+        mask = to_standard_float_mask(mask, out_dtype=dtype)
+
+        tf = torchvision.transforms.Resize(
+            (packed_height, packed_width), interpolation=torchvision.transforms.InterpolationMode.NEAREST
+        )
+
+        # Add a batch dimension to the mask, because torchvision expects shape (batch, channels, h, w).
+        mask = mask.unsqueeze(0)  # Shape: (1, h, w) -> (1, 1, h, w)
+        resized_mask = tf(mask)
+
+        # Flatten the height and width dimensions into a single image_seq_len dimension.
+        return resized_mask.flatten(start_dim=2)

invokeai/backend/flux/extensions/xlabs_controlnet_extension.py

@@ -0,0 +1,150 @@
+from typing import List, Union
+
+import torch
+from PIL.Image import Image
+
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
+from invokeai.app.util.controlnet_utils import CONTROLNET_RESIZE_VALUES, prepare_control_image
+from invokeai.backend.flux.controlnet.controlnet_flux_output import ControlNetFluxOutput
+from invokeai.backend.flux.controlnet.xlabs_controlnet_flux import XLabsControlNetFlux, XLabsControlNetFluxOutput
+from invokeai.backend.flux.extensions.base_controlnet_extension import BaseControlNetExtension
+
+
+class XLabsControlNetExtension(BaseControlNetExtension):
+    def __init__(
+        self,
+        model: XLabsControlNetFlux,
+        controlnet_cond: torch.Tensor,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+    ):
+        super().__init__(
+            weight=weight,
+            begin_step_percent=begin_step_percent,
+            end_step_percent=end_step_percent,
+        )
+
+        self._model = model
+        # _controlnet_cond is the control image passed to the ControlNet model.
+        # Pixel values are in the range [-1, 1]. Shape: (batch_size, 3, height, width).
+        self._controlnet_cond = controlnet_cond
+
+        # TODO(ryand): Pass in these params if a new base transformer / XLabs ControlNet pair get released.
+        self._flux_transformer_num_double_blocks = 19
+        self._flux_transformer_num_single_blocks = 38
+
+    @classmethod
+    def prepare_controlnet_cond(
+        cls,
+        controlnet_image: Image,
+        latent_height: int,
+        latent_width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        resize_mode: CONTROLNET_RESIZE_VALUES,
+    ):
+        image_height = latent_height * LATENT_SCALE_FACTOR
+        image_width = latent_width * LATENT_SCALE_FACTOR
+
+        controlnet_cond = prepare_control_image(
+            image=controlnet_image,
+            do_classifier_free_guidance=False,
+            width=image_width,
+            height=image_height,
+            device=device,
+            dtype=dtype,
+            control_mode="balanced",
+            resize_mode=resize_mode,
+        )
+
+        # Map pixel values from [0, 1] to [-1, 1].
+        controlnet_cond = controlnet_cond * 2 - 1
+
+        return controlnet_cond
+
+    @classmethod
+    def from_controlnet_image(
+        cls,
+        model: XLabsControlNetFlux,
+        controlnet_image: Image,
+        latent_height: int,
+        latent_width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        resize_mode: CONTROLNET_RESIZE_VALUES,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+    ):
+        image_height = latent_height * LATENT_SCALE_FACTOR
+        image_width = latent_width * LATENT_SCALE_FACTOR
+
+        controlnet_cond = prepare_control_image(
+            image=controlnet_image,
+            do_classifier_free_guidance=False,
+            width=image_width,
+            height=image_height,
+            device=device,
+            dtype=dtype,
+            control_mode="balanced",
+            resize_mode=resize_mode,
+        )
+
+        # Map pixel values from [0, 1] to [-1, 1].
+        controlnet_cond = controlnet_cond * 2 - 1
+
+        return cls(
+            model=model,
+            controlnet_cond=controlnet_cond,
+            weight=weight,
+            begin_step_percent=begin_step_percent,
+            end_step_percent=end_step_percent,
+        )
+
+    def _xlabs_output_to_controlnet_output(self, xlabs_output: XLabsControlNetFluxOutput) -> ControlNetFluxOutput:
+        # The modulo index logic used here is based on:
+        # https://github.com/XLabs-AI/x-flux/blob/47495425dbed499be1e8e5a6e52628b07349cba2/src/flux/model.py#L198-L200
+
+        # Handle double block residuals.
+        double_block_residuals: list[torch.Tensor] = []
+        xlabs_double_block_residuals = xlabs_output.controlnet_double_block_residuals
+        if xlabs_double_block_residuals is not None:
+            for i in range(self._flux_transformer_num_double_blocks):
+                double_block_residuals.append(xlabs_double_block_residuals[i % len(xlabs_double_block_residuals)])
+
+        return ControlNetFluxOutput(
+            double_block_residuals=double_block_residuals,
+            single_block_residuals=None,
+        )
+
+    def run_controlnet(
+        self,
+        timestep_index: int,
+        total_num_timesteps: int,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt: torch.Tensor,
+        txt_ids: torch.Tensor,
+        y: torch.Tensor,
+        timesteps: torch.Tensor,
+        guidance: torch.Tensor | None,
+    ) -> ControlNetFluxOutput:
+        weight = self._get_weight(timestep_index=timestep_index, total_num_timesteps=total_num_timesteps)
+        if weight < 1e-6:
+            return ControlNetFluxOutput(single_block_residuals=None, double_block_residuals=None)
+
+        xlabs_output: XLabsControlNetFluxOutput = self._model(
+            img=img,
+            img_ids=img_ids,
+            controlnet_cond=self._controlnet_cond,
+            txt=txt,
+            txt_ids=txt_ids,
+            timesteps=timesteps,
+            y=y,
+            guidance=guidance,
+        )
+
+        controlnet_output = self._xlabs_output_to_controlnet_output(xlabs_output)
+        controlnet_output.apply_weight(weight)
+        return controlnet_output

invokeai/backend/flux/extensions/xlabs_ip_adapter_extension.py

@@ -0,0 +1,90 @@
+import math
+from typing import List, Union
+
+import einops
+import torch
+from PIL import Image
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+
+from invokeai.backend.flux.ip_adapter.xlabs_ip_adapter_flux import XlabsIpAdapterFlux
+from invokeai.backend.flux.modules.layers import DoubleStreamBlock
+from invokeai.backend.util.devices import TorchDevice
+
+
+class XLabsIPAdapterExtension:
+    def __init__(
+        self,
+        model: XlabsIpAdapterFlux,
+        image_prompt_clip_embed: torch.Tensor,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+    ):
+        self._model = model
+        self._image_prompt_clip_embed = image_prompt_clip_embed
+        self._weight = weight
+        self._begin_step_percent = begin_step_percent
+        self._end_step_percent = end_step_percent
+
+        self._image_proj: torch.Tensor | None = None
+
+    def _get_weight(self, timestep_index: int, total_num_timesteps: int) -> float:
+        first_step = math.floor(self._begin_step_percent * total_num_timesteps)
+        last_step = math.ceil(self._end_step_percent * total_num_timesteps)
+
+        if timestep_index < first_step or timestep_index > last_step:
+            return 0.0
+
+        if isinstance(self._weight, list):
+            return self._weight[timestep_index]
+
+        return self._weight
+
+    @staticmethod
+    def run_clip_image_encoder(
+        pil_image: List[Image.Image], image_encoder: CLIPVisionModelWithProjection
+    ) -> torch.Tensor:
+        clip_image_processor = CLIPImageProcessor()
+        clip_image: torch.Tensor = clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
+        clip_image = clip_image.to(device=TorchDevice.choose_torch_device(), dtype=image_encoder.dtype)
+        clip_image_embeds = image_encoder(clip_image).image_embeds
+        return clip_image_embeds
+
+    def run_image_proj(self, dtype: torch.dtype):
+        image_prompt_clip_embed = self._image_prompt_clip_embed.to(dtype=dtype)
+        self._image_proj = self._model.image_proj(image_prompt_clip_embed)
+
+    def run_ip_adapter(
+        self,
+        timestep_index: int,
+        total_num_timesteps: int,
+        block_index: int,
+        block: DoubleStreamBlock,
+        img_q: torch.Tensor,
+        img: torch.Tensor,
+    ) -> torch.Tensor:
+        """The logic in this function is based on:
+        https://github.com/XLabs-AI/x-flux/blob/47495425dbed499be1e8e5a6e52628b07349cba2/src/flux/modules/layers.py#L245-L301
+        """
+        weight = self._get_weight(timestep_index=timestep_index, total_num_timesteps=total_num_timesteps)
+        if weight < 1e-6:
+            return img
+
+        ip_adapter_block = self._model.ip_adapter_double_blocks.double_blocks[block_index]
+
+        ip_key = ip_adapter_block.ip_adapter_double_stream_k_proj(self._image_proj)
+        ip_value = ip_adapter_block.ip_adapter_double_stream_v_proj(self._image_proj)
+
+        # Reshape projections for multi-head attention.
+        ip_key = einops.rearrange(ip_key, "B L (H D) -> B H L D", H=block.num_heads)
+        ip_value = einops.rearrange(ip_value, "B L (H D) -> B H L D", H=block.num_heads)
+
+        # Compute attention between IP projections and the latent query.
+        ip_attn = torch.nn.functional.scaled_dot_product_attention(
+            img_q, ip_key, ip_value, dropout_p=0.0, is_causal=False
+        )
+        ip_attn = einops.rearrange(ip_attn, "B H L D -> B L (H D)", H=block.num_heads)
+
+        img = img + weight * ip_attn
+
+        return img

invokeai/backend/flux/flux_state_dict_utils.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+
+def get_flux_in_channels_from_state_dict(state_dict: dict[str | int, Any]) -> int | None:
+    """Gets the in channels from the state dict."""
+
+    # "Standard" FLUX models use "img_in.weight", but some community fine tunes use
+    # "model.diffusion_model.img_in.weight". Known models that use the latter key:
+    # - https://civitai.com/models/885098?modelVersionId=990775
+    # - https://civitai.com/models/1018060?modelVersionId=1596255
+    # - https://civitai.com/models/978314/ultrareal-fine-tune?modelVersionId=1413133
+
+    keys = {"img_in.weight", "model.diffusion_model.img_in.weight"}
+
+    for key in keys:
+        val = state_dict.get(key)
+        if val is not None:
+            return val.shape[1]
+
+    return None

invokeai/backend/flux/ip_adapter/ip_double_stream_block_processor.py

@@ -0,0 +1,93 @@
+# This file is based on:
+# https://github.com/XLabs-AI/x-flux/blob/47495425dbed499be1e8e5a6e52628b07349cba2/src/flux/modules/layers.py#L221
+import einops
+import torch
+
+from invokeai.backend.flux.math import attention
+from invokeai.backend.flux.modules.layers import DoubleStreamBlock
+
+
+class IPDoubleStreamBlockProcessor(torch.nn.Module):
+    """Attention processor for handling IP-adapter with double stream block."""
+
+    def __init__(self, context_dim: int, hidden_dim: int):
+        super().__init__()
+
+        # Ensure context_dim matches the dimension of image_proj
+        self.context_dim = context_dim
+        self.hidden_dim = hidden_dim
+
+        # Initialize projections for IP-adapter
+        self.ip_adapter_double_stream_k_proj = torch.nn.Linear(context_dim, hidden_dim, bias=True)
+        self.ip_adapter_double_stream_v_proj = torch.nn.Linear(context_dim, hidden_dim, bias=True)
+
+        torch.nn.init.zeros_(self.ip_adapter_double_stream_k_proj.weight)
+        torch.nn.init.zeros_(self.ip_adapter_double_stream_k_proj.bias)
+
+        torch.nn.init.zeros_(self.ip_adapter_double_stream_v_proj.weight)
+        torch.nn.init.zeros_(self.ip_adapter_double_stream_v_proj.bias)
+
+    def __call__(
+        self,
+        attn: DoubleStreamBlock,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        pe: torch.Tensor,
+        image_proj: torch.Tensor,
+        ip_scale: float = 1.0,
+    ):
+        # Prepare image for attention
+        img_mod1, img_mod2 = attn.img_mod(vec)
+        txt_mod1, txt_mod2 = attn.txt_mod(vec)
+
+        img_modulated = attn.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = attn.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = einops.rearrange(
+            img_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads, D=attn.head_dim
+        )
+        img_q, img_k = attn.img_attn.norm(img_q, img_k, img_v)
+
+        txt_modulated = attn.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = attn.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = einops.rearrange(
+            txt_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads, D=attn.head_dim
+        )
+        txt_q, txt_k = attn.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+
+        attn1 = attention(q, k, v, pe=pe)
+        txt_attn, img_attn = attn1[:, : txt.shape[1]], attn1[:, txt.shape[1] :]
+
+        # print(f"txt_attn shape: {txt_attn.size()}")
+        # print(f"img_attn shape: {img_attn.size()}")
+
+        img = img + img_mod1.gate * attn.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * attn.img_mlp((1 + img_mod2.scale) * attn.img_norm2(img) + img_mod2.shift)
+
+        txt = txt + txt_mod1.gate * attn.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * attn.txt_mlp((1 + txt_mod2.scale) * attn.txt_norm2(txt) + txt_mod2.shift)
+
+        # IP-adapter processing
+        ip_query = img_q  # latent sample query
+        ip_key = self.ip_adapter_double_stream_k_proj(image_proj)
+        ip_value = self.ip_adapter_double_stream_v_proj(image_proj)
+
+        # Reshape projections for multi-head attention
+        ip_key = einops.rearrange(ip_key, "B L (H D) -> B H L D", H=attn.num_heads, D=attn.head_dim)
+        ip_value = einops.rearrange(ip_value, "B L (H D) -> B H L D", H=attn.num_heads, D=attn.head_dim)
+
+        # Compute attention between IP projections and the latent query
+        ip_attention = torch.nn.functional.scaled_dot_product_attention(
+            ip_query, ip_key, ip_value, dropout_p=0.0, is_causal=False
+        )
+        ip_attention = einops.rearrange(ip_attention, "B H L D -> B L (H D)", H=attn.num_heads, D=attn.head_dim)
+
+        img = img + ip_scale * ip_attention
+
+        return img, txt

invokeai/backend/flux/ip_adapter/state_dict_utils.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+import torch
+
+from invokeai.backend.flux.ip_adapter.xlabs_ip_adapter_flux import XlabsIpAdapterParams
+
+
+def is_state_dict_xlabs_ip_adapter(sd: dict[str | int, Any]) -> bool:
+    """Is the state dict for an XLabs FLUX IP-Adapter model?
+
+    This is intended to be a reasonably high-precision detector, but it is not guaranteed to have perfect precision.
+    """
+    # If all of the expected keys are present, then this is very likely an XLabs IP-Adapter model.
+    expected_keys = {
+        "double_blocks.0.processor.ip_adapter_double_stream_k_proj.bias",
+        "double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight",
+        "double_blocks.0.processor.ip_adapter_double_stream_v_proj.bias",
+        "double_blocks.0.processor.ip_adapter_double_stream_v_proj.weight",
+        "ip_adapter_proj_model.norm.bias",
+        "ip_adapter_proj_model.norm.weight",
+        "ip_adapter_proj_model.proj.bias",
+        "ip_adapter_proj_model.proj.weight",
+    }
+
+    if expected_keys.issubset(sd.keys()):
+        return True
+    return False
+
+
+def infer_xlabs_ip_adapter_params_from_state_dict(state_dict: dict[str | int, torch.Tensor]) -> XlabsIpAdapterParams:
+    num_double_blocks = 0
+    context_dim = 0
+    hidden_dim = 0
+
+    # Count the number of double blocks.
+    double_block_index = 0
+    while f"double_blocks.{double_block_index}.processor.ip_adapter_double_stream_k_proj.weight" in state_dict:
+        double_block_index += 1
+    num_double_blocks = double_block_index
+
+    hidden_dim = state_dict["double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight"].shape[0]
+    context_dim = state_dict["double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight"].shape[1]
+    clip_embeddings_dim = state_dict["ip_adapter_proj_model.proj.weight"].shape[1]
+    clip_extra_context_tokens = state_dict["ip_adapter_proj_model.proj.weight"].shape[0] // context_dim
+
+    return XlabsIpAdapterParams(
+        num_double_blocks=num_double_blocks,
+        context_dim=context_dim,
+        hidden_dim=hidden_dim,
+        clip_embeddings_dim=clip_embeddings_dim,
+        clip_extra_context_tokens=clip_extra_context_tokens,
+    )

invokeai/backend/flux/ip_adapter/xlabs_ip_adapter_flux.py

@@ -0,0 +1,70 @@
+from dataclasses import dataclass
+
+import torch
+
+from invokeai.backend.ip_adapter.ip_adapter import ImageProjModel
+
+
+class IPDoubleStreamBlock(torch.nn.Module):
+    def __init__(self, context_dim: int, hidden_dim: int):
+        super().__init__()
+
+        self.context_dim = context_dim
+        self.hidden_dim = hidden_dim
+
+        self.ip_adapter_double_stream_k_proj = torch.nn.Linear(context_dim, hidden_dim, bias=True)
+        self.ip_adapter_double_stream_v_proj = torch.nn.Linear(context_dim, hidden_dim, bias=True)
+
+
+class IPAdapterDoubleBlocks(torch.nn.Module):
+    def __init__(self, num_double_blocks: int, context_dim: int, hidden_dim: int):
+        super().__init__()
+        self.double_blocks = torch.nn.ModuleList(
+            [IPDoubleStreamBlock(context_dim, hidden_dim) for _ in range(num_double_blocks)]
+        )
+
+
+@dataclass
+class XlabsIpAdapterParams:
+    num_double_blocks: int
+    context_dim: int
+    hidden_dim: int
+
+    clip_embeddings_dim: int
+    clip_extra_context_tokens: int
+
+
+class XlabsIpAdapterFlux(torch.nn.Module):
+    def __init__(self, params: XlabsIpAdapterParams):
+        super().__init__()
+        self.image_proj = ImageProjModel(
+            cross_attention_dim=params.context_dim,
+            clip_embeddings_dim=params.clip_embeddings_dim,
+            clip_extra_context_tokens=params.clip_extra_context_tokens,
+        )
+        self.ip_adapter_double_blocks = IPAdapterDoubleBlocks(
+            num_double_blocks=params.num_double_blocks, context_dim=params.context_dim, hidden_dim=params.hidden_dim
+        )
+
+    def load_xlabs_state_dict(self, state_dict: dict[str, torch.Tensor], assign: bool = False):
+        """We need this custom function to load state dicts rather than using .load_state_dict(...) because the model
+        structure does not match the state_dict structure.
+        """
+        # Split the state_dict into the image projection model and the double blocks.
+        image_proj_sd: dict[str, torch.Tensor] = {}
+        double_blocks_sd: dict[str, torch.Tensor] = {}
+        for k, v in state_dict.items():
+            if k.startswith("ip_adapter_proj_model."):
+                image_proj_sd[k] = v
+            elif k.startswith("double_blocks."):
+                double_blocks_sd[k] = v
+            else:
+                raise ValueError(f"Unexpected key: {k}")
+
+        # Initialize the image projection model.
+        image_proj_sd = {k.replace("ip_adapter_proj_model.", ""): v for k, v in image_proj_sd.items()}
+        self.image_proj.load_state_dict(image_proj_sd, assign=assign)
+
+        # Initialize the double blocks.
+        double_blocks_sd = {k.replace("processor.", ""): v for k, v in double_blocks_sd.items()}
+        self.ip_adapter_double_blocks.load_state_dict(double_blocks_sd, assign=assign)

invokeai/backend/flux/math.py

@@ -0,0 +1,35 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, attn_mask: Tensor | None = None) -> Tensor:
+    q, k = apply_rope(q, k, pe)
+
+    x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+    x = rearrange(x, "B H L D -> B L (H D)")
+
+    return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = (
+        torch.arange(0, dim, 2, dtype=torch.float32 if pos.device.type == "mps" else torch.float64, device=pos.device)
+        / dim
+    )
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.to(dtype=pos.dtype, device=pos.device)
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.view(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.view(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.view(*xq.shape).type_as(xq), xk_out.view(*xk.shape).type_as(xk)

invokeai/backend/flux/model.py

@@ -0,0 +1,168 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+
+from invokeai.backend.flux.custom_block_processor import (
+    CustomDoubleStreamBlockProcessor,
+    CustomSingleStreamBlockProcessor,
+)
+from invokeai.backend.flux.extensions.regional_prompting_extension import RegionalPromptingExtension
+from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
+from invokeai.backend.flux.modules.layers import (
+    DoubleStreamBlock,
+    EmbedND,
+    LastLayer,
+    MLPEmbedder,
+    SingleStreamBlock,
+    timestep_embedding,
+)
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list[int]
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+    out_channels: Optional[int] = None
+
+
+class Flux(nn.Module):
+    """
+    Transformer model for flow matching on sequences.
+    """
+
+    def __init__(self, params: FluxParams):
+        super().__init__()
+
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = params.out_channels or self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}")
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
+
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+    def forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor | None,
+        timestep_index: int,
+        total_num_timesteps: int,
+        controlnet_double_block_residuals: list[Tensor] | None,
+        controlnet_single_block_residuals: list[Tensor] | None,
+        ip_adapter_extensions: list[XLabsIPAdapterExtension],
+        regional_prompting_extension: RegionalPromptingExtension,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        # Validate double_block_residuals shape.
+        if controlnet_double_block_residuals is not None:
+            assert len(controlnet_double_block_residuals) == len(self.double_blocks)
+        for block_index, block in enumerate(self.double_blocks):
+            assert isinstance(block, DoubleStreamBlock)
+            img, txt = CustomDoubleStreamBlockProcessor.custom_double_block_forward(
+                timestep_index=timestep_index,
+                total_num_timesteps=total_num_timesteps,
+                block_index=block_index,
+                block=block,
+                img=img,
+                txt=txt,
+                vec=vec,
+                pe=pe,
+                ip_adapter_extensions=ip_adapter_extensions,
+                regional_prompting_extension=regional_prompting_extension,
+            )
+
+            if controlnet_double_block_residuals is not None:
+                img += controlnet_double_block_residuals[block_index]
+
+        img = torch.cat((txt, img), 1)
+
+        # Validate single_block_residuals shape.
+        if controlnet_single_block_residuals is not None:
+            assert len(controlnet_single_block_residuals) == len(self.single_blocks)
+
+        for block_index, block in enumerate(self.single_blocks):
+            assert isinstance(block, SingleStreamBlock)
+            img = CustomSingleStreamBlockProcessor.custom_single_block_forward(
+                timestep_index=timestep_index,
+                total_num_timesteps=total_num_timesteps,
+                block_index=block_index,
+                block=block,
+                img=img,
+                vec=vec,
+                pe=pe,
+                regional_prompting_extension=regional_prompting_extension,
+            )
+
+            if controlnet_single_block_residuals is not None:
+                img[:, txt.shape[1] :, ...] += controlnet_single_block_residuals[block_index]
+
+        img = img[:, txt.shape[1] :, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        return img

invokeai/backend/flux/modules/autoencoder.py

@@ -0,0 +1,324 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int
+    in_channels: int
+    ch: int
+    out_ch: int
+    ch_mult: list[int]
+    num_res_blocks: int
+    z_channels: int
+    scale_factor: float
+    shift_factor: float
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        # downsampling
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        # end
+        h = self.norm_out(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv_out(h)
+        return h
+
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, chunk_dim: int = 1):
+        super().__init__()
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor, sample: bool = True, generator: torch.Generator | None = None) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        if sample:
+            std = torch.exp(0.5 * logvar)
+            # Unfortunately, torch.randn_like(...) does not accept a generator argument at the time of writing, so we
+            # have to use torch.randn(...) instead.
+            return mean + std * torch.randn(size=mean.size(), generator=generator, dtype=mean.dtype, device=mean.device)
+        else:
+            return mean
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__()
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian()
+
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor, sample: bool = True, generator: torch.Generator | None = None) -> Tensor:
+        """Run VAE encoding on input tensor x.
+
+        Args:
+            x (Tensor): Input image tensor. Shape: (batch_size, in_channels, height, width).
+            sample (bool, optional): If True, sample from the encoded distribution, else, return the distribution mean.
+                Defaults to True.
+            generator (torch.Generator | None, optional): Optional random number generator for reproducibility.
+                Defaults to None.
+
+        Returns:
+            Tensor: Encoded latent tensor. Shape: (batch_size, z_channels, latent_height, latent_width).
+        """
+
+        z = self.reg(self.encoder(x), sample=sample, generator=generator)
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))

invokeai/backend/flux/modules/conditioner.py

@@ -0,0 +1,44 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from torch import Tensor, nn
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class HFEncoder(nn.Module):
+    def __init__(
+        self,
+        encoder: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
+        is_clip: bool,
+        max_length: int,
+    ):
+        super().__init__()
+        self.max_length = max_length
+        self.is_clip = is_clip
+        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
+        self.tokenizer = tokenizer
+        self.hf_module = encoder
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+    def forward(self, text: list[str]) -> Tensor:
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=False,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        # Move inputs to the same device as the model to support cpu_only models
+        model_device = get_effective_device(self.hf_module)
+
+        outputs = self.hf_module(
+            input_ids=batch_encoding["input_ids"].to(model_device),
+            attention_mask=None,
+            output_hidden_states=False,
+        )
+        return outputs[self.output_key]

invokeai/backend/flux/modules/layers.py

@@ -0,0 +1,250 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+import math
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+from invokeai.backend.flux.math import attention, rope
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    """
+    Create sinusoidal timestep embeddings.
+    :param t: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an (N, D) Tensor of positional embeddings.
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(t.device)
+
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor):
+        return torch.nn.functional.rms_norm(x, self.scale.shape, self.scale, eps=1e-6)
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        x = attention(q, k, v, pe=pe)
+        x = self.proj(x)
+        return x
+
+
+@dataclass
+class ModulationOut:
+    shift: Tensor
+    scale: Tensor
+    gate: Tensor
+
+
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
+
+        return (
+            ModulationOut(*out[:3]),
+            ModulationOut(*out[3:]) if self.is_double else None,
+        )
+
+
+class DoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.img_mod = Modulation(hidden_size, double=True)
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+        self.txt_mod = Modulation(hidden_size, double=True)
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = self.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        # prepare txt for attention
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+
+        attn = attention(q, k, v, pe=pe)
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+
+        # calculate the txt bloks
+        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+        return img, txt
+
+
+class SingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_scale: float | None = None,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        # qkv and mlp_in
+        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+        # proj and mlp_out
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+
+        self.norm = QKNorm(head_dim)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False)
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + mod.gate * output
+
+
+class LastLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.linear(x)
+        return x

invokeai/backend/flux/redux/flux_redux_model.py

@@ -0,0 +1,17 @@
+import torch
+
+# This model definition is based on:
+# https://github.com/black-forest-labs/flux/blob/716724eb276d94397be99710a0a54d352664e23b/src/flux/modules/image_embedders.py#L66
+
+
+class FluxReduxModel(torch.nn.Module):
+    def __init__(self, redux_dim: int = 1152, txt_in_features: int = 4096) -> None:
+        super().__init__()
+
+        self.redux_dim = redux_dim
+
+        self.redux_up = torch.nn.Linear(redux_dim, txt_in_features * 3)
+        self.redux_down = torch.nn.Linear(txt_in_features * 3, txt_in_features)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.redux_down(torch.nn.functional.silu(self.redux_up(x)))

invokeai/backend/flux/redux/flux_redux_state_dict_utils.py

@@ -0,0 +1,11 @@
+from typing import Any
+
+
+def is_state_dict_likely_flux_redux(state_dict: dict[str | int, Any]) -> bool:
+    """Checks if the provided state dict is likely a FLUX Redux model."""
+
+    expected_keys = {"redux_down.bias", "redux_down.weight", "redux_up.bias", "redux_up.weight"}
+    if set(state_dict.keys()) == expected_keys:
+        return True
+
+    return False

invokeai/backend/flux/sampling_utils.py

@@ -0,0 +1,186 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+import math
+from typing import Callable
+
+import torch
+from einops import rearrange, repeat
+
+
+def get_noise(
+    num_samples: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+):
+    # We always generate noise on the same device and dtype then cast to ensure consistency across devices/dtypes.
+    rand_device = "cpu"
+    rand_dtype = torch.float16
+    return torch.randn(
+        num_samples,
+        16,
+        # allow for packing
+        2 * math.ceil(height / 16),
+        2 * math.ceil(width / 16),
+        device=rand_device,
+        dtype=rand_dtype,
+        generator=torch.Generator(device=rand_device).manual_seed(seed),
+    ).to(device=device, dtype=dtype)
+
+
+def time_shift(mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+
+def get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15) -> Callable[[float], float]:
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def get_schedule(
+    num_steps: int,
+    image_seq_len: int,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+    shift: bool = True,
+) -> list[float]:
+    # extra step for zero
+    timesteps = torch.linspace(1, 0, num_steps + 1)
+
+    # shifting the schedule to favor high timesteps for higher signal images
+    if shift:
+        # estimate mu based on linear estimation between two points
+        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
+        timesteps = time_shift(mu, 1.0, timesteps)
+
+    return timesteps.tolist()
+
+
+def _find_last_index_ge_val(timesteps: list[float], val: float, eps: float = 1e-6) -> int:
+    """Find the last index in timesteps that is >= val.
+
+    We use epsilon-close equality to avoid potential floating point errors.
+    """
+    idx = len(list(filter(lambda t: t >= (val - eps), timesteps))) - 1
+    assert idx >= 0
+    return idx
+
+
+def clip_timestep_schedule(timesteps: list[float], denoising_start: float, denoising_end: float) -> list[float]:
+    """Clip the timestep schedule to the denoising range.
+
+    Args:
+        timesteps (list[float]): The original timestep schedule: [1.0, ..., 0.0].
+        denoising_start (float): A value in [0, 1] specifying the start of the denoising process. E.g. a value of 0.2
+            would mean that the denoising process start at the last timestep in the schedule >= 0.8.
+        denoising_end (float): A value in [0, 1] specifying the end of the denoising process. E.g. a value of 0.8 would
+            mean that the denoising process end at the last timestep in the schedule >= 0.2.
+
+    Returns:
+        list[float]: The clipped timestep schedule.
+    """
+    assert 0.0 <= denoising_start <= 1.0
+    assert 0.0 <= denoising_end <= 1.0
+    assert denoising_start <= denoising_end
+
+    t_start_val = 1.0 - denoising_start
+    t_end_val = 1.0 - denoising_end
+
+    t_start_idx = _find_last_index_ge_val(timesteps, t_start_val)
+    t_end_idx = _find_last_index_ge_val(timesteps, t_end_val)
+
+    clipped_timesteps = timesteps[t_start_idx : t_end_idx + 1]
+
+    return clipped_timesteps
+
+
+def clip_timestep_schedule_fractional(
+    timesteps: list[float], denoising_start: float, denoising_end: float
+) -> list[float]:
+    """Clip the timestep schedule to the denoising range. Insert new timesteps to exactly match the desired denoising
+    range. (A fractional version of clip_timestep_schedule().)
+
+    Args:
+        timesteps (list[float]): The original timestep schedule: [1.0, ..., 0.0].
+        denoising_start (float): A value in [0, 1] specifying the start of the denoising process. E.g. a value of 0.2
+            would mean that the denoising process start at t=0.8.
+        denoising_end (float): A value in [0, 1] specifying the end of the denoising process. E.g. a value of 0.8 would
+            mean that the denoising process ends at t=0.2.
+
+    Returns:
+        list[float]: The clipped timestep schedule.
+    """
+    assert 0.0 <= denoising_start <= 1.0
+    assert 0.0 <= denoising_end <= 1.0
+    assert denoising_start <= denoising_end
+
+    t_start_val = 1.0 - denoising_start
+    t_end_val = 1.0 - denoising_end
+
+    t_start_idx = _find_last_index_ge_val(timesteps, t_start_val)
+    t_end_idx = _find_last_index_ge_val(timesteps, t_end_val)
+
+    clipped_timesteps = timesteps[t_start_idx : t_end_idx + 1]
+
+    # We know that clipped_timesteps[0] >= t_start_val. Replace clipped_timesteps[0] with t_start_val.
+    clipped_timesteps[0] = t_start_val
+
+    # We know that clipped_timesteps[-1] >= t_end_val. If clipped_timesteps[-1] > t_end_val, add another step to
+    # t_end_val.
+    eps = 1e-6
+    if clipped_timesteps[-1] > t_end_val + eps:
+        clipped_timesteps.append(t_end_val)
+
+    return clipped_timesteps
+
+
+def unpack(x: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """Unpack flat array of patch embeddings to latent image."""
+    return rearrange(
+        x,
+        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+        h=math.ceil(height / 16),
+        w=math.ceil(width / 16),
+        ph=2,
+        pw=2,
+    )
+
+
+def pack(x: torch.Tensor) -> torch.Tensor:
+    """Pack latent image to flattented array of patch embeddings."""
+    # Pixel unshuffle with a scale of 2, and flatten the height/width dimensions to get an array of patches.
+    return rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+
+
+def generate_img_ids(h: int, w: int, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """Generate tensor of image position ids.
+
+    Args:
+        h (int): Height of image in latent space.
+        w (int): Width of image in latent space.
+        batch_size (int): Batch size.
+        device (torch.device): Device.
+        dtype (torch.dtype): dtype.
+
+    Returns:
+        torch.Tensor: Image position ids.
+    """
+
+    if device.type == "mps":
+        orig_dtype = dtype
+        dtype = torch.float16
+
+    img_ids = torch.zeros(h // 2, w // 2, 3, device=device, dtype=dtype)
+    # Set batch offset to 0 for main image tokens
+    img_ids[..., 0] = 0
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2, device=device, dtype=dtype)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2, device=device, dtype=dtype)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=batch_size)
+
+    if device.type == "mps":
+        img_ids = img_ids.to(orig_dtype)
+
+    return img_ids

invokeai/backend/flux/schedulers.py

@@ -0,0 +1,131 @@
+"""Flow Matching scheduler definitions and mapping.
+
+This module provides the scheduler types and mapping for Flow Matching models
+(Flux and Z-Image), supporting multiple schedulers from the diffusers library.
+"""
+
+from typing import Any, Literal, Type
+
+from diffusers import (
+    DPMSolverMultistepScheduler,
+    FlowMatchEulerDiscreteScheduler,
+    FlowMatchHeunDiscreteScheduler,
+)
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from invokeai.backend.rectified_flow.er_sde_scheduler import ERSDEScheduler
+
+# Note: FlowMatchLCMScheduler may not be available in all diffusers versions
+try:
+    from diffusers import FlowMatchLCMScheduler
+
+    _HAS_LCM = True
+except ImportError:
+    _HAS_LCM = False
+
+# Scheduler name literal type for type checking
+FLUX_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
+
+# Human-readable labels for the UI
+FLUX_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "lcm": "LCM",
+}
+
+# Mapping from scheduler names to scheduler classes
+FLUX_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
+    "euler": FlowMatchEulerDiscreteScheduler,
+    "heun": FlowMatchHeunDiscreteScheduler,
+}
+
+if _HAS_LCM:
+    FLUX_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler
+
+
+# Z-Image scheduler types (Flow Matching schedulers)
+# Note: Z-Image-Turbo is optimized for ~8 steps with Euler, LCM can also work.
+# Z-Image Base (undistilled) should only use Euler or Heun (LCM not supported for undistilled models).
+ZIMAGE_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
+
+# Human-readable labels for the UI
+ZIMAGE_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "lcm": "LCM",
+}
+
+# Mapping from scheduler names to scheduler classes
+ZIMAGE_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
+    "euler": FlowMatchEulerDiscreteScheduler,
+    "heun": FlowMatchHeunDiscreteScheduler,
+}
+
+if _HAS_LCM:
+    ZIMAGE_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler
+
+
+# Anima scheduler types.
+# Anima uses rectified flow with shift=3.0. The driver passes pre-shifted sigmas via
+# set_timesteps(sigmas=...) when the scheduler accepts that signature. For those, the
+# entry carries shift=1.0 to avoid double-shifting (the scheduler uses our sigmas verbatim).
+# Schedulers that don't accept sigmas= (Heun, DPM++ on diffusers 0.35.1) build their own
+# internal schedule, so they need shift=ANIMA_SHIFT/flow_shift=ANIMA_SHIFT in kwargs to match
+# Anima's reference loglinear schedule.
+
+# Fixed shift factor for the Anima rectified-flow noise schedule.
+ANIMA_SHIFT = 3.0
+
+ANIMA_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "dpmpp_2m", "dpmpp_2m_sde", "er_sde", "lcm"]
+
+ANIMA_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "dpmpp_2m": "DPM++ 2M",
+    "dpmpp_2m_sde": "DPM++ 2M SDE",
+    "er_sde": "ER-SDE",
+    "lcm": "LCM",
+}
+
+# When adding a new Anima scheduler: add to all three of NAME_VALUES, LABELS,
+# and this MAP. The MAP entry is `(SchedulerClass, scheduler_kwargs)`. For
+# rectified-flow schedulers, set `use_flow_sigmas=True` and use
+# `prediction_type="flow_prediction"`. If the scheduler accepts set_timesteps(sigmas=...),
+# use shift=1.0 (driver passes pre-shifted sigmas); otherwise use shift=ANIMA_SHIFT
+# so the scheduler builds the correct internal schedule.
+ANIMA_SCHEDULER_MAP: dict[str, tuple[Type[SchedulerMixin], dict[str, Any]]] = {
+    "euler": (FlowMatchEulerDiscreteScheduler, {"shift": 1.0}),
+    "heun": (FlowMatchHeunDiscreteScheduler, {"shift": ANIMA_SHIFT}),
+    "dpmpp_2m": (
+        DPMSolverMultistepScheduler,
+        {
+            "prediction_type": "flow_prediction",
+            "use_flow_sigmas": True,
+            "flow_shift": ANIMA_SHIFT,
+            "solver_order": 2,
+        },
+    ),
+    "dpmpp_2m_sde": (
+        DPMSolverMultistepScheduler,
+        {
+            "prediction_type": "flow_prediction",
+            "use_flow_sigmas": True,
+            "flow_shift": ANIMA_SHIFT,
+            "algorithm_type": "sde-dpmsolver++",
+            "solver_order": 2,
+        },
+    ),
+    "er_sde": (
+        ERSDEScheduler,
+        {
+            "prediction_type": "flow_prediction",
+            "use_flow_sigmas": True,
+            "flow_shift": ANIMA_SHIFT,
+            "solver_order": 3,
+            "stochastic": True,
+        },
+    ),
+}
+
+if _HAS_LCM:
+    ANIMA_SCHEDULER_MAP["lcm"] = (FlowMatchLCMScheduler, {"shift": 1.0})

invokeai/backend/flux/text_conditioning.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+import torch
+
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range
+
+
+@dataclass
+class FluxTextConditioning:
+    t5_embeddings: torch.Tensor
+    clip_embeddings: torch.Tensor
+    # If mask is None, the prompt is a global prompt.
+    mask: torch.Tensor | None
+
+
+@dataclass
+class FluxReduxConditioning:
+    redux_embeddings: torch.Tensor
+    # If mask is None, the prompt is a global prompt.
+    mask: torch.Tensor | None
+
+
+@dataclass
+class FluxRegionalTextConditioning:
+    # Concatenated text embeddings.
+    # Shape: (1, concatenated_txt_seq_len, 4096)
+    t5_embeddings: torch.Tensor
+    # Shape: (1, concatenated_txt_seq_len, 3)
+    t5_txt_ids: torch.Tensor
+
+    # Global CLIP embeddings.
+    # Shape: (1, 768)
+    clip_embeddings: torch.Tensor
+
+    # A binary mask indicating the regions of the image that the prompt should be applied to. If None, the prompt is a
+    # global prompt.
+    # image_masks[i] is the mask for the ith prompt.
+    # image_masks[i] has shape (1, image_seq_len) and dtype torch.bool.
+    image_masks: list[torch.Tensor | None]
+
+    # List of ranges that represent the embedding ranges for each mask.
+    # t5_embedding_ranges[i] contains the range of the t5 embeddings that correspond to image_masks[i].
+    t5_embedding_ranges: list[Range]

invokeai/backend/flux/util.py

@@ -0,0 +1,195 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from dataclasses import dataclass
+from typing import Literal
+
+from invokeai.backend.flux.model import FluxParams
+from invokeai.backend.flux.modules.autoencoder import AutoEncoderParams
+from invokeai.backend.model_manager.taxonomy import AnyVariant, Flux2VariantType, FluxVariantType
+
+
+@dataclass
+class ModelSpec:
+    params: FluxParams
+    ae_params: AutoEncoderParams
+    ckpt_path: str | None
+    ae_path: str | None
+    repo_id: str | None
+    repo_flow: str | None
+    repo_ae: str | None
+
+
+# Preferred resolutions for Kontext models to avoid tiling artifacts
+# These are the specific resolutions the model was trained on
+PREFERED_KONTEXT_RESOLUTIONS = [
+    (672, 1568),
+    (688, 1504),
+    (720, 1456),
+    (752, 1392),
+    (800, 1328),
+    (832, 1248),
+    (880, 1184),
+    (944, 1104),
+    (1024, 1024),
+    (1104, 944),
+    (1184, 880),
+    (1248, 832),
+    (1328, 800),
+    (1392, 752),
+    (1456, 720),
+    (1504, 688),
+    (1568, 672),
+]
+
+
+_flux_max_seq_lengths: dict[AnyVariant, Literal[256, 512]] = {
+    FluxVariantType.Dev: 512,
+    FluxVariantType.DevFill: 512,
+    FluxVariantType.Schnell: 256,
+    Flux2VariantType.Klein4B: 512,
+    Flux2VariantType.Klein9B: 512,
+}
+
+
+def get_flux_max_seq_length(variant: AnyVariant):
+    try:
+        return _flux_max_seq_lengths[variant]
+    except KeyError:
+        raise ValueError(f"Unknown variant for FLUX max seq len: {variant}")
+
+
+_flux_ae_params = AutoEncoderParams(
+    resolution=256,
+    in_channels=3,
+    ch=128,
+    out_ch=3,
+    ch_mult=[1, 2, 4, 4],
+    num_res_blocks=2,
+    z_channels=16,
+    scale_factor=0.3611,
+    shift_factor=0.1159,
+)
+
+
+def get_flux_ae_params() -> AutoEncoderParams:
+    return _flux_ae_params
+
+
+_flux_transformer_params: dict[AnyVariant, FluxParams] = {
+    FluxVariantType.Dev: FluxParams(
+        in_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=True,
+    ),
+    FluxVariantType.Schnell: FluxParams(
+        in_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
+    ),
+    FluxVariantType.DevFill: FluxParams(
+        in_channels=384,
+        out_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=True,
+    ),
+    # Flux2 Klein 4B uses Qwen3 4B text encoder with stacked embeddings from layers [9, 18, 27]
+    # The context_in_dim is 3 * hidden_size of Qwen3 (3 * 2560 = 7680)
+    Flux2VariantType.Klein4B: FluxParams(
+        in_channels=64,
+        vec_in_dim=2560,  # Qwen3-4B hidden size (used for pooled output)
+        context_in_dim=7680,  # 3 layers * 2560 = 7680 for Qwen3-4B
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
+    ),
+    # Flux2 Klein 4B Base is the undistilled foundation model. It shares the same
+    # architecture as Klein 4B (distilled) and reports guidance_embeds=False in its
+    # HF transformer config - classical CFG (external negative pass) is the guidance mechanism.
+    Flux2VariantType.Klein4BBase: FluxParams(
+        in_channels=64,
+        vec_in_dim=2560,  # Qwen3-4B hidden size (used for pooled output)
+        context_in_dim=7680,  # 3 layers * 2560 = 7680 for Qwen3-4B
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
+    ),
+    # Flux2 Klein 9B uses Qwen3 8B text encoder with stacked embeddings from layers [9, 18, 27]
+    # The context_in_dim is 3 * hidden_size of Qwen3 (3 * 4096 = 12288)
+    Flux2VariantType.Klein9B: FluxParams(
+        in_channels=64,
+        vec_in_dim=4096,  # Qwen3-8B hidden size (used for pooled output)
+        context_in_dim=12288,  # 3 layers * 4096 = 12288 for Qwen3-8B
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
+    ),
+    # Flux2 Klein 9B Base is the undistilled foundation model. It shares the same
+    # architecture as Klein 9B (distilled) and reports guidance_embeds=False in its
+    # HF transformer config - the guidance scalar is inert for all Klein variants.
+    Flux2VariantType.Klein9BBase: FluxParams(
+        in_channels=64,
+        vec_in_dim=4096,  # Qwen3-8B hidden size (used for pooled output)
+        context_in_dim=12288,  # 3 layers * 4096 = 12288 for Qwen3-8B
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
+    ),
+}
+
+
+def get_flux_transformers_params(variant: AnyVariant):
+    try:
+        return _flux_transformer_params[variant]
+    except KeyError:
+        raise ValueError(f"Unknown variant for FLUX transformer params: {variant}")

invokeai/backend/flux2/__init__.py

@@ -0,0 +1,4 @@
+"""FLUX.2 backend modules.
+
+This package contains modules specific to FLUX.2 models (e.g., Klein).
+"""

invokeai/backend/flux2/denoise.py

@@ -0,0 +1,306 @@
+"""Flux2 Klein Denoising Function.
+
+This module provides the denoising function for FLUX.2 Klein models,
+which use Qwen3 as the text encoder instead of CLIP+T5.
+"""
+
+import inspect
+import math
+from typing import Any, Callable
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
+from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
+
+
+def denoise(
+    model: torch.nn.Module,
+    # model input
+    img: torch.Tensor,
+    img_ids: torch.Tensor,
+    txt: torch.Tensor,
+    txt_ids: torch.Tensor,
+    # sampling parameters
+    timesteps: list[float],
+    step_callback: Callable[[PipelineIntermediateState], None],
+    guidance: float,
+    cfg_scale: list[float],
+    # Negative conditioning for CFG
+    neg_txt: torch.Tensor | None = None,
+    neg_txt_ids: torch.Tensor | None = None,
+    # Scheduler for stepping (e.g., FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler)
+    scheduler: Any = None,
+    # Dynamic shifting parameter for FLUX.2 Klein (computed from image resolution)
+    mu: float | None = None,
+    # Inpainting extension for merging latents during denoising
+    inpaint_extension: RectifiedFlowInpaintExtension | None = None,
+    # Reference image conditioning (multi-reference image editing)
+    img_cond_seq: torch.Tensor | None = None,
+    img_cond_seq_ids: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Denoise latents using a FLUX.2 Klein transformer model.
+
+    This is a simplified denoise function for FLUX.2 Klein models that uses
+    the diffusers Flux2Transformer2DModel interface.
+
+    All current FLUX.2 Klein variants (4B, 4B Base, 9B, 9B Base) have guidance_embeds=False
+    in their HF transformer config (or absent/zeroed projection weights), so the guidance
+    value is passed but effectively ignored by the model. The argument is retained for
+    node-graph compatibility and future variants that may ship trained guidance projections.
+    CFG is applied externally using negative conditioning when cfg_scale != 1.0.
+
+    Args:
+        model: The Flux2Transformer2DModel from diffusers.
+        img: Packed latent image tensor of shape (B, seq_len, channels).
+        img_ids: Image position IDs tensor.
+        txt: Text encoder hidden states (Qwen3 embeddings).
+        txt_ids: Text position IDs tensor.
+        timesteps: List of timesteps for denoising schedule (linear sigmas from 1.0 to 1/n).
+        step_callback: Callback function for progress updates.
+        guidance: Guidance strength. Inert for all current FLUX.2 Klein variants
+            (their guidance_embeds projection weights are absent/zero).
+        cfg_scale: List of CFG scale values per step.
+        neg_txt: Negative text embeddings for CFG (optional).
+        neg_txt_ids: Negative text position IDs (optional).
+        scheduler: Optional diffusers scheduler (Euler, Heun, LCM). If None, uses manual Euler.
+        mu: Dynamic shifting parameter computed from image resolution. Required when scheduler
+            has use_dynamic_shifting=True.
+
+    Returns:
+        Denoised latent tensor.
+    """
+    total_steps = len(timesteps) - 1
+
+    # Store original sequence length for extracting output later (before concatenating reference images)
+    original_seq_len = img.shape[1]
+
+    # Concatenate reference image conditioning if provided (multi-reference image editing)
+    if img_cond_seq is not None and img_cond_seq_ids is not None:
+        img = torch.cat([img, img_cond_seq], dim=1)
+        img_ids = torch.cat([img_ids, img_cond_seq_ids], dim=1)
+
+    # The transformer forward() requires a guidance tensor even when guidance_embeds=False,
+    # because the Flux2TimestepGuidanceEmbeddings forward signature takes it unconditionally.
+    # All current Klein variants have guidance_embeds=False, so the value is ignored internally.
+    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+
+    # Use scheduler if provided
+    use_scheduler = scheduler is not None
+    if use_scheduler:
+        # Set up scheduler with sigmas and mu for dynamic shifting
+        # Convert timesteps (0-1 range) to sigmas for the scheduler
+        # The scheduler will apply dynamic shifting internally using mu (if enabled in scheduler config)
+        sigmas = np.array(timesteps[:-1], dtype=np.float32)  # Exclude final 0.0
+
+        # Check if scheduler supports sigmas parameter using inspect.signature
+        # FlowMatchHeunDiscreteScheduler and FlowMatchLCMScheduler don't support sigmas
+        set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
+        supports_sigmas = "sigmas" in set_timesteps_sig.parameters
+        if supports_sigmas and mu is not None:
+            # Pass mu if provided - it will only be used if scheduler has use_dynamic_shifting=True
+            scheduler.set_timesteps(sigmas=sigmas.tolist(), mu=mu, device=img.device)
+        elif supports_sigmas:
+            scheduler.set_timesteps(sigmas=sigmas.tolist(), device=img.device)
+        else:
+            # Scheduler doesn't support sigmas (e.g., Heun, LCM) - use num_inference_steps
+            #
+            # Important for img2img callers: if the initial latent/noise blend was
+            # computed from a separate pre-scheduler schedule, that preblend may not
+            # match this scheduler's true first step exactly.
+            scheduler_kwargs: dict[str, Any] = {"num_inference_steps": len(sigmas), "device": img.device}
+            if mu is not None and "mu" in set_timesteps_sig.parameters:
+                scheduler_kwargs["mu"] = mu
+            scheduler.set_timesteps(**scheduler_kwargs)
+        num_scheduler_steps = len(scheduler.timesteps)
+        is_heun = hasattr(scheduler, "state_in_first_order")
+        user_step = 0
+
+        pbar = tqdm(total=total_steps, desc="Denoising")
+        for step_index in range(num_scheduler_steps):
+            timestep = scheduler.timesteps[step_index]
+            # Convert scheduler timestep (0-1000) to normalized (0-1) for the model
+            t_curr = timestep.item() / scheduler.config.num_train_timesteps
+            t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+            # Track if we're in first or second order step (for Heun)
+            in_first_order = scheduler.state_in_first_order if is_heun else True
+
+            # Run the transformer model (matching diffusers: guidance=guidance, return_dict=False)
+            output = model(
+                hidden_states=img,
+                encoder_hidden_states=txt,
+                timestep=t_vec,
+                img_ids=img_ids,
+                txt_ids=txt_ids,
+                guidance=guidance_vec,
+                return_dict=False,
+            )
+
+            # Extract the sample from the output (return_dict=False returns tuple)
+            pred = output[0] if isinstance(output, tuple) else output
+
+            step_cfg_scale = cfg_scale[min(user_step, len(cfg_scale) - 1)]
+
+            # Apply CFG if scale is not 1.0
+            if not math.isclose(step_cfg_scale, 1.0):
+                if neg_txt is None:
+                    raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
+
+                neg_output = model(
+                    hidden_states=img,
+                    encoder_hidden_states=neg_txt,
+                    timestep=t_vec,
+                    img_ids=img_ids,
+                    txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids,
+                    guidance=guidance_vec,
+                    return_dict=False,
+                )
+
+                neg_pred = neg_output[0] if isinstance(neg_output, tuple) else neg_output
+                pred = neg_pred + step_cfg_scale * (pred - neg_pred)
+
+            # Use scheduler.step() for the update
+            step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img)
+            img = step_output.prev_sample
+
+            # Get t_prev for inpainting (next sigma value)
+            if step_index + 1 < len(scheduler.sigmas):
+                t_prev = scheduler.sigmas[step_index + 1].item()
+            else:
+                t_prev = 0.0
+
+            # Apply inpainting merge at each step
+            if inpaint_extension is not None:
+                # Separate the generated latents from the reference conditioning
+                gen_img = img[:, :original_seq_len, :]
+                ref_img = img[:, original_seq_len:, :]
+
+                # Merge only the generated part
+                gen_img = inpaint_extension.merge_intermediate_latents_with_init_latents(gen_img, t_prev)
+
+                # Concatenate back together
+                img = torch.cat([gen_img, ref_img], dim=1)
+
+            # For Heun, only increment user step after second-order step completes
+            if is_heun:
+                if not in_first_order:
+                    user_step += 1
+                    if user_step <= total_steps:
+                        pbar.update(1)
+                        preview_img = img - t_curr * pred
+                        if inpaint_extension is not None:
+                            preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(
+                                preview_img, 0.0
+                            )
+                        # Extract only the generated image portion for preview (exclude reference images)
+                        callback_latents = (
+                            preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
+                        )
+                        step_callback(
+                            PipelineIntermediateState(
+                                step=user_step,
+                                order=2,
+                                total_steps=total_steps,
+                                timestep=int(t_curr * 1000),
+                                latents=callback_latents,
+                            ),
+                        )
+            else:
+                user_step += 1
+                if user_step <= total_steps:
+                    pbar.update(1)
+                    preview_img = img - t_curr * pred
+                    if inpaint_extension is not None:
+                        preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_img, 0.0)
+                    # Extract only the generated image portion for preview (exclude reference images)
+                    callback_latents = preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
+                    step_callback(
+                        PipelineIntermediateState(
+                            step=user_step,
+                            order=1,
+                            total_steps=total_steps,
+                            timestep=int(t_curr * 1000),
+                            latents=callback_latents,
+                        ),
+                    )
+
+        pbar.close()
+    else:
+        # Manual Euler stepping (original behavior)
+        for step_index, (t_curr, t_prev) in tqdm(list(enumerate(zip(timesteps[:-1], timesteps[1:], strict=True)))):
+            t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+
+            # Run the transformer model (matching diffusers: guidance=guidance, return_dict=False)
+            output = model(
+                hidden_states=img,
+                encoder_hidden_states=txt,
+                timestep=t_vec,
+                img_ids=img_ids,
+                txt_ids=txt_ids,
+                guidance=guidance_vec,
+                return_dict=False,
+            )
+
+            # Extract the sample from the output (return_dict=False returns tuple)
+            pred = output[0] if isinstance(output, tuple) else output
+
+            step_cfg_scale = cfg_scale[step_index]
+
+            # Apply CFG if scale is not 1.0
+            if not math.isclose(step_cfg_scale, 1.0):
+                if neg_txt is None:
+                    raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
+
+                neg_output = model(
+                    hidden_states=img,
+                    encoder_hidden_states=neg_txt,
+                    timestep=t_vec,
+                    img_ids=img_ids,
+                    txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids,
+                    guidance=guidance_vec,
+                    return_dict=False,
+                )
+
+                neg_pred = neg_output[0] if isinstance(neg_output, tuple) else neg_output
+                pred = neg_pred + step_cfg_scale * (pred - neg_pred)
+
+            # Euler step
+            preview_img = img - t_curr * pred
+            img = img + (t_prev - t_curr) * pred
+
+            # Apply inpainting merge at each step
+            if inpaint_extension is not None:
+                # Separate the generated latents from the reference conditioning
+                gen_img = img[:, :original_seq_len, :]
+                ref_img = img[:, original_seq_len:, :]
+
+                # Merge only the generated part
+                gen_img = inpaint_extension.merge_intermediate_latents_with_init_latents(gen_img, t_prev)
+
+                # Concatenate back together
+                img = torch.cat([gen_img, ref_img], dim=1)
+
+                # Handling preview images
+                preview_gen = preview_img[:, :original_seq_len, :]
+                preview_gen = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_gen, 0.0)
+
+            # Extract only the generated image portion for preview (exclude reference images)
+            callback_latents = preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
+            step_callback(
+                PipelineIntermediateState(
+                    step=step_index + 1,
+                    order=1,
+                    total_steps=total_steps,
+                    timestep=int(t_curr),
+                    latents=callback_latents,
+                ),
+            )
+
+    # Extract only the generated image portion (exclude concatenated reference images)
+    if img_cond_seq is not None:
+        img = img[:, :original_seq_len, :]
+
+    return img

invokeai/backend/flux2/ref_image_extension.py

@@ -0,0 +1,294 @@
+"""FLUX.2 Klein Reference Image Extension for multi-reference image editing.
+
+This module provides the Flux2RefImageExtension for FLUX.2 Klein models,
+which handles encoding reference images using the FLUX.2 VAE and
+generating the appropriate position IDs for multi-reference image editing.
+
+FLUX.2 Klein has built-in support for reference image editing (unlike FLUX.1
+which requires a separate Kontext model).
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+from einops import repeat
+from PIL import Image
+
+from invokeai.app.invocations.fields import FluxKontextConditioningField
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux2.sampling_utils import pack_flux2
+from invokeai.backend.util.devices import TorchDevice
+
+# Maximum pixel counts for reference images (matches BFL FLUX.2 sampling.py)
+# Single reference image: 2024² pixels, Multiple: 1024² pixels
+MAX_PIXELS_SINGLE_REF = 2024**2  # ~4.1M pixels
+MAX_PIXELS_MULTI_REF = 1024**2  # ~1M pixels
+
+
+def resize_image_to_max_pixels(image: Image.Image, max_pixels: int) -> Image.Image:
+    """Resize image to fit within max_pixels while preserving aspect ratio.
+
+    This matches the BFL FLUX.2 sampling.py cap_pixels() behavior.
+
+    Args:
+        image: PIL Image to resize.
+        max_pixels: Maximum total pixel count (width * height).
+
+    Returns:
+        Resized PIL Image (or original if already within bounds).
+    """
+    width, height = image.size
+    pixel_count = width * height
+
+    if pixel_count <= max_pixels:
+        return image
+
+    # Calculate scale factor to fit within max_pixels (BFL approach)
+    scale = math.sqrt(max_pixels / pixel_count)
+    new_width = int(width * scale)
+    new_height = int(height * scale)
+
+    # Ensure dimensions are at least 1
+    new_width = max(1, new_width)
+    new_height = max(1, new_height)
+
+    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+
+def generate_img_ids_flux2_with_offset(
+    latent_height: int,
+    latent_width: int,
+    batch_size: int,
+    device: torch.device,
+    idx_offset: int = 0,
+    h_offset: int = 0,
+    w_offset: int = 0,
+) -> torch.Tensor:
+    """Generate tensor of image position ids with optional offsets for FLUX.2.
+
+    FLUX.2 uses 4D position coordinates (T, H, W, L) for its rotary position embeddings.
+    Position IDs use int64 (long) dtype.
+
+    Args:
+        latent_height: Height of image in latent space (before packing).
+        latent_width: Width of image in latent space (before packing).
+        batch_size: Number of images in the batch.
+        device: Device to create tensors on.
+        idx_offset: Offset for T (time/index) coordinate - use 1 for reference images.
+        h_offset: Spatial offset for H coordinate in latent space.
+        w_offset: Spatial offset for W coordinate in latent space.
+
+    Returns:
+        Image position ids with shape [batch_size, (latent_height//2 * latent_width//2), 4].
+    """
+    # After packing, the spatial dimensions are halved due to the 2x2 patch structure
+    packed_height = latent_height // 2
+    packed_width = latent_width // 2
+
+    # Convert spatial offsets from latent space to packed space
+    packed_h_offset = h_offset // 2
+    packed_w_offset = w_offset // 2
+
+    # Create base tensor for position IDs with shape [packed_height, packed_width, 4]
+    # The 4 channels represent: [T, H, W, L]
+    img_ids = torch.zeros(packed_height, packed_width, 4, device=device, dtype=torch.long)
+
+    # Set T (time/index offset) for all positions - use 1 for reference images
+    img_ids[..., 0] = idx_offset
+
+    # Set H (height/y) coordinates with offset
+    h_coords = torch.arange(packed_height, device=device, dtype=torch.long) + packed_h_offset
+    img_ids[..., 1] = h_coords[:, None]
+
+    # Set W (width/x) coordinates with offset
+    w_coords = torch.arange(packed_width, device=device, dtype=torch.long) + packed_w_offset
+    img_ids[..., 2] = w_coords[None, :]
+
+    # L (layer) coordinate stays 0
+
+    # Expand to include batch dimension: [batch_size, (packed_height * packed_width), 4]
+    img_ids = img_ids.reshape(1, packed_height * packed_width, 4)
+    img_ids = repeat(img_ids, "1 s c -> b s c", b=batch_size)
+
+    return img_ids
+
+
+class Flux2RefImageExtension:
+    """Applies FLUX.2 Klein reference image conditioning.
+
+    This extension handles encoding reference images using the FLUX.2 VAE
+    and generating the appropriate 4D position IDs for multi-reference image editing.
+
+    FLUX.2 Klein has built-in support for reference image editing, unlike FLUX.1
+    which requires a separate Kontext model.
+    """
+
+    def __init__(
+        self,
+        ref_image_conditioning: list[FluxKontextConditioningField],
+        context: InvocationContext,
+        vae_field: VAEField,
+        device: torch.device,
+        dtype: torch.dtype,
+        bn_mean: torch.Tensor | None = None,
+        bn_std: torch.Tensor | None = None,
+    ):
+        """Initialize the Flux2RefImageExtension.
+
+        Args:
+            ref_image_conditioning: List of reference image conditioning fields.
+            context: The invocation context for loading models and images.
+            vae_field: The FLUX.2 VAE field for encoding images.
+            device: Target device for tensors.
+            dtype: Target dtype for tensors.
+            bn_mean: BN running mean for normalizing latents (shape: 128).
+            bn_std: BN running std for normalizing latents (shape: 128).
+        """
+        self._context = context
+        self._device = device
+        self._dtype = dtype
+        self._vae_field = vae_field
+        self._bn_mean = bn_mean
+        self._bn_std = bn_std
+        self.ref_image_conditioning = ref_image_conditioning
+
+        # Pre-process and cache the reference image latents and ids upon initialization
+        self.ref_image_latents, self.ref_image_ids = self._prepare_ref_images()
+
+    def _bn_normalize(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply BN normalization to packed latents.
+
+        BN formula (affine=False): y = (x - mean) / std
+
+        Args:
+            x: Packed latents of shape (B, seq, 128).
+
+        Returns:
+            Normalized latents of same shape.
+        """
+        assert self._bn_mean is not None and self._bn_std is not None
+        bn_mean = self._bn_mean.to(x.device, x.dtype)
+        bn_std = self._bn_std.to(x.device, x.dtype)
+        return (x - bn_mean) / bn_std
+
+    def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encode reference images and prepare their concatenated latents and IDs with spatial tiling."""
+        all_latents = []
+        all_ids = []
+
+        # Track cumulative dimensions for spatial tiling
+        canvas_h = 0
+        canvas_w = 0
+
+        vae_info = self._context.models.load(self._vae_field.vae)
+
+        # Determine max pixels based on number of reference images (BFL FLUX.2 approach)
+        num_refs = len(self.ref_image_conditioning)
+        max_pixels = MAX_PIXELS_SINGLE_REF if num_refs == 1 else MAX_PIXELS_MULTI_REF
+
+        for idx, ref_image_field in enumerate(self.ref_image_conditioning):
+            image = self._context.images.get_pil(ref_image_field.image.image_name)
+            image = image.convert("RGB")
+
+            # Resize large images to max pixel count (matches BFL FLUX.2 sampling.py)
+            image = resize_image_to_max_pixels(image, max_pixels)
+
+            # Convert to tensor using torchvision transforms
+            transformation = T.Compose([T.ToTensor()])
+            image_tensor = transformation(image)
+            # Convert from [0, 1] to [-1, 1] range expected by VAE
+            image_tensor = image_tensor * 2.0 - 1.0
+            image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension
+
+            # Encode using FLUX.2 VAE
+            with vae_info.model_on_device() as (_, vae):
+                vae_dtype = next(iter(vae.parameters())).dtype
+                image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
+
+                # FLUX.2 VAE uses diffusers API
+                latent_dist = vae.encode(image_tensor, return_dict=False)[0]
+
+                # Use mode() for deterministic encoding (no sampling)
+                if hasattr(latent_dist, "mode"):
+                    ref_image_latents_unpacked = latent_dist.mode()
+                elif hasattr(latent_dist, "sample"):
+                    ref_image_latents_unpacked = latent_dist.sample()
+                else:
+                    ref_image_latents_unpacked = latent_dist
+
+                TorchDevice.empty_cache()
+
+            # Extract tensor dimensions (B, 32, H, W for FLUX.2)
+            batch_size, _, latent_height, latent_width = ref_image_latents_unpacked.shape
+
+            # Pad latents to be compatible with patch_size=2
+            pad_h = (2 - latent_height % 2) % 2
+            pad_w = (2 - latent_width % 2) % 2
+            if pad_h > 0 or pad_w > 0:
+                ref_image_latents_unpacked = F.pad(ref_image_latents_unpacked, (0, pad_w, 0, pad_h), mode="circular")
+                _, _, latent_height, latent_width = ref_image_latents_unpacked.shape
+
+            # Pack the latents using FLUX.2 pack function (32 channels -> 128)
+            ref_image_latents_packed = pack_flux2(ref_image_latents_unpacked).to(self._device, self._dtype)
+
+            # Apply BN normalization to match the input latents scale
+            # This is critical - the transformer expects normalized latents
+            if self._bn_mean is not None and self._bn_std is not None:
+                ref_image_latents_packed = self._bn_normalize(ref_image_latents_packed)
+
+            # Determine spatial offsets for this reference image
+            h_offset = 0
+            w_offset = 0
+
+            if idx > 0:  # First image starts at (0, 0)
+                # Calculate potential canvas dimensions for each tiling option
+                potential_h_vertical = canvas_h + latent_height
+                potential_w_horizontal = canvas_w + latent_width
+
+                # Choose arrangement that minimizes the maximum dimension
+                if potential_h_vertical > potential_w_horizontal:
+                    # Tile horizontally (to the right)
+                    w_offset = canvas_w
+                    canvas_w = canvas_w + latent_width
+                    canvas_h = max(canvas_h, latent_height)
+                else:
+                    # Tile vertically (below)
+                    h_offset = canvas_h
+                    canvas_h = canvas_h + latent_height
+                    canvas_w = max(canvas_w, latent_width)
+            else:
+                canvas_h = latent_height
+                canvas_w = latent_width
+
+            # Generate position IDs with 4D format (T, H, W, L)
+            # Use T-coordinate offset with scale=10 like diffusers Flux2Pipeline:
+            # T = scale + scale * idx (so first ref image is T=10, second is T=20, etc.)
+            # The generated image uses T=0, so this clearly separates reference images
+            t_offset = 10 + 10 * idx  # scale=10 matches diffusers
+            ref_image_ids = generate_img_ids_flux2_with_offset(
+                latent_height=latent_height,
+                latent_width=latent_width,
+                batch_size=batch_size,
+                device=self._device,
+                idx_offset=t_offset,  # Reference images use T=10, 20, 30...
+                h_offset=h_offset,
+                w_offset=w_offset,
+            )
+
+            all_latents.append(ref_image_latents_packed)
+            all_ids.append(ref_image_ids)
+
+        # Concatenate all latents and IDs along the sequence dimension
+        concatenated_latents = torch.cat(all_latents, dim=1)
+        concatenated_ids = torch.cat(all_ids, dim=1)
+
+        return concatenated_latents, concatenated_ids
+
+    def ensure_batch_size(self, target_batch_size: int) -> None:
+        """Ensure the reference image latents and IDs match the target batch size."""
+        if self.ref_image_latents.shape[0] != target_batch_size:
+            self.ref_image_latents = self.ref_image_latents.repeat(target_batch_size, 1, 1)
+            self.ref_image_ids = self.ref_image_ids.repeat(target_batch_size, 1, 1)

invokeai/backend/flux2/sampling_utils.py

@@ -0,0 +1,206 @@
+"""FLUX.2 Klein Sampling Utilities.
+
+FLUX.2 Klein uses a 32-channel VAE (AutoencoderKLFlux2) instead of the 16-channel VAE
+used by FLUX.1. This module provides sampling utilities adapted for FLUX.2.
+"""
+
+import math
+
+import torch
+from einops import rearrange
+
+
+def get_noise_flux2(
+    num_samples: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+) -> torch.Tensor:
+    """Generate noise for FLUX.2 Klein (32 channels).
+
+    FLUX.2 uses a 32-channel VAE, so noise must have 32 channels.
+    The spatial dimensions are calculated to allow for packing.
+
+    Args:
+        num_samples: Batch size.
+        height: Target image height in pixels.
+        width: Target image width in pixels.
+        device: Target device.
+        dtype: Target dtype.
+        seed: Random seed.
+
+    Returns:
+        Noise tensor of shape (num_samples, 32, latent_h, latent_w).
+    """
+    # We always generate noise on the same device and dtype then cast to ensure consistency.
+    rand_device = "cpu"
+    rand_dtype = torch.float16
+
+    # FLUX.2 uses 32 latent channels
+    # Latent dimensions: height/8, width/8 (from VAE downsampling)
+    # Must be divisible by 2 for packing (patchify step)
+    latent_h = 2 * math.ceil(height / 16)
+    latent_w = 2 * math.ceil(width / 16)
+
+    return torch.randn(
+        num_samples,
+        32,  # FLUX.2 uses 32 latent channels (vs 16 for FLUX.1)
+        latent_h,
+        latent_w,
+        device=rand_device,
+        dtype=rand_dtype,
+        generator=torch.Generator(device=rand_device).manual_seed(seed),
+    ).to(device=device, dtype=dtype)
+
+
+def pack_flux2(x: torch.Tensor) -> torch.Tensor:
+    """Pack latent image to flattened array of patch embeddings for FLUX.2.
+
+    This performs the patchify + pack operation in one step:
+    1. Patchify: Group 2x2 spatial patches into channels (C*4)
+    2. Pack: Flatten spatial dimensions to sequence
+
+    For 32-channel input: (B, 32, H, W) -> (B, H/2*W/2, 128)
+
+    Args:
+        x: Latent tensor of shape (B, 32, H, W).
+
+    Returns:
+        Packed tensor of shape (B, H/2*W/2, 128).
+    """
+    # Same operation as FLUX.1 pack, but input has 32 channels -> output has 128
+    return rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+
+
+def unpack_flux2(x: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """Unpack flat array of patch embeddings back to latent image for FLUX.2.
+
+    This reverses the pack_flux2 operation:
+    1. Unpack: Restore spatial dimensions from sequence
+    2. Unpatchify: Restore 32 channels from 128
+
+    Args:
+        x: Packed tensor of shape (B, H/2*W/2, 128).
+        height: Target image height in pixels.
+        width: Target image width in pixels.
+
+    Returns:
+        Latent tensor of shape (B, 32, H, W).
+    """
+    # Calculate latent dimensions
+    latent_h = 2 * math.ceil(height / 16)
+    latent_w = 2 * math.ceil(width / 16)
+
+    # Packed dimensions (after patchify)
+    packed_h = latent_h // 2
+    packed_w = latent_w // 2
+
+    return rearrange(
+        x,
+        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+        h=packed_h,
+        w=packed_w,
+        ph=2,
+        pw=2,
+    )
+
+
+def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
+    """Compute mu for FLUX.2 schedule shifting.
+
+    Uses a fixed mu value of 2.02, matching ComfyUI's proven FLUX.2 configuration.
+
+    The previous implementation (from diffusers' FLUX.1 pipeline) computed mu as a
+    linear function of image_seq_len, which produced excessively high values at
+    high resolutions (e.g., mu=3.23 at 2048x2048). This over-shifted the sigma
+    schedule, compressing almost all values above 0.9 and forcing the model to
+    denoise everything in the final 1-2 steps, causing severe grid/diamond artifacts.
+
+    ComfyUI uses a fixed shift=2.02 for FLUX.2 Klein at all resolutions and produces
+    artifact-free images even at 2048x2048.
+
+    Args:
+        image_seq_len: Number of image tokens (packed_h * packed_w). Currently unused.
+        num_steps: Number of denoising steps. Currently unused.
+
+    Returns:
+        The mu value (fixed at 2.02).
+    """
+    return 2.02
+
+
+def get_schedule_flux2(
+    num_steps: int,
+    image_seq_len: int,
+) -> list[float]:
+    """Get linear timestep schedule for FLUX.2.
+
+    Returns a linear sigma schedule from 1.0 to 1/num_steps.
+    The actual schedule shifting is handled by the FlowMatchEulerDiscreteScheduler
+    using the mu parameter and use_dynamic_shifting=True.
+
+    Args:
+        num_steps: Number of denoising steps.
+        image_seq_len: Number of image tokens (packed_h * packed_w). Currently unused,
+            but kept for API compatibility. The scheduler computes shifting internally.
+
+    Returns:
+        List of linear sigmas from 1.0 to 1/num_steps, plus final 0.0.
+    """
+    import numpy as np
+
+    # Create linear sigmas from 1.0 to 1/num_steps
+    # The scheduler will apply dynamic shifting using mu parameter
+    sigmas = np.linspace(1.0, 1 / num_steps, num_steps)
+    sigmas_list = [float(s) for s in sigmas]
+
+    # Add final 0.0 for the last step (scheduler needs n+1 timesteps for n steps)
+    sigmas_list.append(0.0)
+
+    return sigmas_list
+
+
+def generate_img_ids_flux2(h: int, w: int, batch_size: int, device: torch.device) -> torch.Tensor:
+    """Generate tensor of image position ids for FLUX.2 with RoPE scaling.
+
+    FLUX.2 uses 4D position coordinates (T, H, W, L) for its rotary position embeddings.
+    This is different from FLUX.1 which uses 3D coordinates.
+
+    RoPE Scaling: For resolutions >1536x1536, position IDs are scaled down using
+    Position Interpolation to prevent RoPE degradation and diamond/grid artifacts.
+
+    IMPORTANT: Position IDs must use int64 (long) dtype like diffusers, not bfloat16.
+    Using floating point dtype for position IDs can cause NaN in rotary embeddings.
+
+    Args:
+        h: Height of image in latent space.
+        w: Width of image in latent space.
+        batch_size: Batch size.
+        device: Device.
+
+    Returns:
+        Image position ids tensor of shape (batch_size, h/2*w/2, 4) with int64 dtype.
+    """
+    # After packing, spatial dims are h/2 x w/2
+    packed_h = h // 2
+    packed_w = w // 2
+
+    # Create coordinate grids - 4D: (T, H, W, L)
+    # T = time/batch index, H = height, W = width, L = layer/channel
+    # Use int64 (long) dtype like diffusers
+    img_ids = torch.zeros(packed_h, packed_w, 4, device=device, dtype=torch.long)
+
+    # T (time/batch) coordinate - set to 0 (already initialized)
+    # H coordinates
+    img_ids[..., 1] = torch.arange(packed_h, device=device, dtype=torch.long)[:, None]
+    # W coordinates
+    img_ids[..., 2] = torch.arange(packed_w, device=device, dtype=torch.long)[None, :]
+    # L (layer) coordinate - set to 0 (already initialized)
+
+    # Flatten and expand for batch
+    img_ids = img_ids.reshape(1, packed_h * packed_w, 4)
+    img_ids = img_ids.expand(batch_size, -1, -1)
+
+    return img_ids

invokeai/backend/image_util/__init__.py

@@ -0,0 +1,12 @@
+"""
+Initialization file for invokeai.backend.image_util methods.
+"""
+
+from invokeai.backend.image_util.infill_methods.patchmatch import PatchMatch  # noqa: F401
+from invokeai.backend.image_util.pngwriter import (  # noqa: F401
+    PngWriter,
+    PromptFormatter,
+    retrieve_metadata,
+    write_metadata,
+)
+from invokeai.backend.image_util.util import InitImageResizer, make_grid  # noqa: F401

invokeai/backend/image_util/basicsr/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2018-2022 BasicSR Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

invokeai/backend/image_util/basicsr/__init__.py

@@ -0,0 +1,18 @@
+"""
+Adapted from https://github.com/XPixelGroup/BasicSR
+License: Apache-2.0
+
+As of Feb 2024, `basicsr` appears to be unmaintained. It imports a function from `torchvision` that is removed in
+`torchvision` 0.17. Here is the deprecation warning:
+
+    UserWarning: The torchvision.transforms.functional_tensor module is deprecated in 0.15 and will be **removed in
+    0.17**. Please don't rely on it. You probably just need to use APIs in torchvision.transforms.functional or in
+    torchvision.transforms.v2.functional.
+
+As a result, a dependency on `basicsr` means we cannot keep our `torchvision` dependency up to date.
+
+Because we only rely on a single class `RRDBNet` from `basicsr`, we've copied the relevant code here and removed the
+dependency on `basicsr`.
+
+The code is almost unchanged, only a few type annotations have been added. The license is also copied.
+"""

invokeai/backend/image_util/basicsr/arch_util.py

@@ -0,0 +1,75 @@
+from typing import Type
+
+import torch
+from torch import nn as nn
+from torch.nn import init as init
+from torch.nn.modules.batchnorm import _BatchNorm
+
+
+@torch.no_grad()
+def default_init_weights(
+    module_list: list[nn.Module] | nn.Module, scale: float = 1, bias_fill: float = 0, **kwargs
+) -> None:
+    """Initialize network weights.
+
+    Args:
+        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
+        scale (float): Scale initialized weights, especially for residual
+            blocks. Default: 1.
+        bias_fill (float): The value to fill bias. Default: 0
+        kwargs (dict): Other arguments for initialization function.
+    """
+    if not isinstance(module_list, list):
+        module_list = [module_list]
+    for module in module_list:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, _BatchNorm):
+                init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+
+
+def make_layer(basic_block: Type[nn.Module], num_basic_block: int, **kwarg) -> nn.Sequential:
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (Type[nn.Module]): nn.Module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+# TODO: may write a cpp file
+def pixel_unshuffle(x: torch.Tensor, scale: int) -> torch.Tensor:
+    """Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)

invokeai/backend/image_util/basicsr/rrdbnet_arch.py

@@ -0,0 +1,125 @@
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from invokeai.backend.image_util.basicsr.arch_util import default_init_weights, make_layer, pixel_unshuffle
+
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block.
+
+    Used in RRDB block in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat: int = 64, num_grow_ch: int = 32) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block.
+
+    Used in RRDB-Net in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat: int, num_grow_ch: int = 32) -> None:
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """Networks consisting of Residual in Residual Dense Block, which is used
+    in ESRGAN.
+
+    ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
+
+    We extend ESRGAN for scale x2 and scale x1.
+    Note: This is one option for scale 1, scale 2 in RRDBNet.
+    We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size
+    and enlarge the channel size before feeding inputs into the main ESRGAN architecture.
+
+    Args:
+        num_in_ch (int): Channel number of inputs.
+        num_out_ch (int): Channel number of outputs.
+        num_feat (int): Channel number of intermediate features.
+            Default: 64
+        num_block (int): Block number in the trunk network. Defaults: 23
+        num_grow_ch (int): Channels for each growth. Default: 32.
+    """
+
+    def __init__(
+        self,
+        num_in_ch: int,
+        num_out_ch: int,
+        scale: int = 4,
+        num_feat: int = 64,
+        num_block: int = 23,
+        num_grow_ch: int = 32,
+    ) -> None:
+        super(RRDBNet, self).__init__()
+        self.scale = scale
+        if scale == 2:
+            num_in_ch = num_in_ch * 4
+        elif scale == 1:
+            num_in_ch = num_in_ch * 16
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # upsample
+        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.scale == 2:
+            feat = pixel_unshuffle(x, scale=2)
+        elif self.scale == 1:
+            feat = pixel_unshuffle(x, scale=4)
+        else:
+            feat = x
+        feat = self.conv_first(feat)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        # upsample
+        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode="nearest")))
+        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode="nearest")))
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out

invokeai/backend/image_util/canny.py

@@ -0,0 +1,41 @@
+import cv2
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    cv2_to_pil,
+    normalize_image_channel_count,
+    pil_to_cv2,
+    resize_image_to_resolution,
+)
+
+
+def get_canny_edges(
+    image: Image.Image, low_threshold: int, high_threshold: int, detect_resolution: int, image_resolution: int
+) -> Image.Image:
+    """Returns the edges of an image using the Canny edge detection algorithm.
+
+    Adapted from https://github.com/huggingface/controlnet_aux (Apache-2.0 license).
+
+    Args:
+        image: The input image.
+        low_threshold: The lower threshold for the hysteresis procedure.
+        high_threshold: The upper threshold for the hysteresis procedure.
+        input_resolution: The resolution of the input image. The image will be resized to this resolution before edge detection.
+        output_resolution: The resolution of the output image. The edges will be resized to this resolution before returning.
+
+    Returns:
+        The Canny edges of the input image.
+    """
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    np_image = pil_to_cv2(image)
+    np_image = normalize_image_channel_count(np_image)
+    np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+    edge_map = cv2.Canny(np_image, low_threshold, high_threshold)
+    edge_map = normalize_image_channel_count(edge_map)
+    edge_map = resize_image_to_resolution(edge_map, image_resolution)
+
+    return cv2_to_pil(edge_map)

invokeai/backend/image_util/composition.py

@@ -0,0 +1,122 @@
+# TODO: Improve blend modes
+# TODO: Add nodes like Hue Adjust for Saturation/Contrast/etc... ?
+# TODO: Continue implementing more blend modes/color spaces(?)
+# TODO: Custom ICC profiles with PIL.ImageCms?
+# TODO: Blend multiple layers all crammed into a tensor(?) or list
+
+# Copyright (c) 2023 Darren Ringer <dwringer@gmail.com>
+# Parts based on Oklab: Copyright (c) 2021 Bj�rn Ottosson <https://bottosson.github.io/>
+# HSL code based on CPython: Copyright (c) 2001-2023 Python Software Foundation; All Rights Reserved
+from math import pi as PI
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from invokeai.backend.image_util.color_conversion import (
+    gamut_clip_tensor,
+)
+from invokeai.backend.image_util.color_conversion import (
+    srgb_from_linear_srgb as shared_srgb_from_linear_srgb,
+)
+from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
+
+MAX_FLOAT = torch.finfo(torch.tensor(1.0).dtype).max
+
+# CIE Lab to Uniform Perceptual Lab profile is copyright © 2003 Bruce Justin Lindbloom. All rights reserved. <http://www.brucelindbloom.com>
+CIELAB_TO_UPLAB_ICC_PATH = Path(__file__).parent / "assets" / "CIELab_to_UPLab.icc"
+
+
+def equivalent_achromatic_lightness(lch_tensor: torch.Tensor):
+    """Calculate Equivalent Achromatic Lightness accounting for Helmholtz-Kohlrausch effect"""
+    # As described by High, Green, and Nussbaum (2023): https://doi.org/10.1002/col.22839
+
+    k = [0.1644, 0.0603, 0.1307, 0.0060]
+
+    h_minus_90 = torch.sub(lch_tensor[2, :, :], PI / 2.0)
+    h_minus_90 = torch.sub(torch.remainder(torch.add(h_minus_90, 3 * PI), 2 * PI), PI)
+
+    f_by = torch.add(k[0] * torch.abs(torch.sin(torch.div(h_minus_90, 2.0))), k[1])
+    f_r_0 = torch.add(k[2] * torch.abs(torch.cos(lch_tensor[2, :, :])), k[3])
+
+    f_r = torch.zeros(lch_tensor[0, :, :].shape)
+    mask_hi = torch.ge(lch_tensor[2, :, :], -1 * (PI / 2.0))
+    mask_lo = torch.le(lch_tensor[2, :, :], PI / 2.0)
+    mask = torch.logical_and(mask_hi, mask_lo)
+    f_r[mask] = f_r_0[mask]
+
+    l_max = torch.ones(lch_tensor[0, :, :].shape)
+    l_min = torch.zeros(lch_tensor[0, :, :].shape)
+    l_adjustment = torch.tensordot(torch.add(f_by, f_r), lch_tensor[1, :, :], dims=([0, 1], [0, 1]))
+    l_max = torch.add(l_max, l_adjustment)
+    l_min = torch.add(l_min, l_adjustment)
+    l_eal_tensor = torch.add(lch_tensor[0, :, :], l_adjustment)
+
+    l_eal_tensor = torch.add(
+        lch_tensor[0, :, :], torch.tensordot(torch.add(f_by, f_r), lch_tensor[1, :, :], dims=([0, 1], [0, 1]))
+    )
+    l_eal_tensor = torch.div(torch.sub(l_eal_tensor, l_min.min()), l_max.max() - l_min.min())
+
+    return l_eal_tensor
+
+
+def srgb_from_linear_srgb(linear_srgb_tensor: torch.Tensor, alpha: float = 0.0, steps: int = 1):
+    """Get gamma-corrected sRGB from a linear-light sRGB image tensor"""
+
+    if 0.0 < alpha:
+        linear_srgb_tensor = gamut_clip_tensor(linear_srgb_tensor, alpha=alpha, steps=steps)
+    return shared_srgb_from_linear_srgb(linear_srgb_tensor)
+
+
+def remove_nans(tensor: torch.Tensor, replace_with: float = MAX_FLOAT):
+    return torch.where(torch.isnan(tensor), replace_with, tensor)
+
+
+def tensor_from_pil_image(img: Image.Image, normalize: bool = False):
+    return image_resized_to_grid_as_tensor(img, normalize=normalize, multiple_of=1)
+
+
+# PSF LICENSE AGREEMENT FOR PYTHON 3.11.5
+
+# 1. This LICENSE AGREEMENT is between the Python Software Foundation ("PSF"), and
+#    the Individual or Organization ("Licensee") accessing and otherwise using Python
+#    3.11.5 software in source or binary form and its associated documentation.
+
+# 2. Subject to the terms and conditions of this License Agreement, PSF hereby
+#    grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce,
+#    analyze, test, perform and/or display publicly, prepare derivative works,
+#    distribute, and otherwise use Python 3.11.5 alone or in any derivative
+#    version, provided, however, that PSF's License Agreement and PSF's notice of
+#    copyright, i.e., "Copyright (c) 2001-2023 Python Software Foundation; All Rights
+#    Reserved" are retained in Python 3.11.5 alone or in any derivative version
+#    prepared by Licensee.
+
+# 3. In the event Licensee prepares a derivative work that is based on or
+#    incorporates Python 3.11.5 or any part thereof, and wants to make the
+#    derivative work available to others as provided herein, then Licensee hereby
+#    agrees to include in any such work a brief summary of the changes made to Python
+#    3.11.5.
+
+# 4. PSF is making Python 3.11.5 available to Licensee on an "AS IS" basis.
+#    PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR IMPLIED.  BY WAY OF
+#    EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND DISCLAIMS ANY REPRESENTATION OR
+#    WARRANTY OF MERCHANTABILITY OR FITNESS FOR ANY PARTICULAR PURPOSE OR THAT THE
+#    USE OF PYTHON 3.11.5 WILL NOT INFRINGE ANY THIRD PARTY RIGHTS.
+
+# 5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON 3.11.5
+#    FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS A RESULT OF
+#    MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON 3.11.5, OR ANY DERIVATIVE
+#    THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF.
+
+# 6. This License Agreement will automatically terminate upon a material breach of
+#    its terms and conditions.
+
+# 7. Nothing in this License Agreement shall be deemed to create any relationship
+#    of agency, partnership, or joint venture between PSF and Licensee.  This License
+#    Agreement does not grant permission to use PSF trademarks or trade name in a
+#    trademark sense to endorse or promote products or services of Licensee, or any
+#    third party.
+
+# 8. By copying, installing or otherwise using Python 3.11.5, Licensee agrees
+#    to be bound by the terms and conditions of this License Agreement.
+######################################################################################/

invokeai/backend/image_util/content_shuffle.py

@@ -0,0 +1,40 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np
+
+
+def make_noise_disk(H, W, C, F):
+    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
+    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
+    noise = noise[F : F + H, F : F + W]
+    noise -= np.min(noise)
+    noise /= np.max(noise)
+    if C == 1:
+        noise = noise[:, :, None]
+    return noise
+
+
+def content_shuffle(input_image: Image.Image, scale_factor: int | None = None) -> Image.Image:
+    """Shuffles the content of an image using a disk noise pattern, similar to a 'liquify' effect."""
+
+    np_img = pil_to_np(input_image)
+
+    height, width, _channels = np_img.shape
+
+    if scale_factor is None:
+        scale_factor = 256
+
+    x = make_noise_disk(height, width, 1, scale_factor) * float(width - 1)
+    y = make_noise_disk(height, width, 1, scale_factor) * float(height - 1)
+
+    flow = np.concatenate([x, y], axis=2).astype(np.float32)
+
+    shuffled_img = cv2.remap(np_img, flow, None, cv2.INTER_LINEAR)
+
+    output_img = np_to_pil(shuffled_img)
+
+    return output_img

invokeai/backend/image_util/controlnet_processor.py

@@ -0,0 +1,188 @@
+"""Utilities for processing images with ControlNet processors."""
+
+from datetime import datetime
+from typing import Any, Optional
+
+from invokeai.app.invocations.fields import ImageField
+from invokeai.app.services.invoker import InvocationServices
+from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
+from invokeai.app.services.shared.graph import Graph, GraphExecutionState
+from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context
+
+
+def _get_processor_invocation_class(processor_type: str):
+    """Get the invocation class for a processor type."""
+    # Import processor invocation classes on demand
+    processor_class_map = {
+        "canny_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.canny", fromlist=["CannyEdgeDetectionInvocation"]
+            ).CannyEdgeDetectionInvocation
+        ),
+        "hed_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.hed", fromlist=["HEDEdgeDetectionInvocation"]
+            ).HEDEdgeDetectionInvocation
+        ),
+        "mlsd_image_processor": lambda: (
+            __import__("invokeai.app.invocations.mlsd", fromlist=["MLSDDetectionInvocation"]).MLSDDetectionInvocation
+        ),
+        "depth_anything_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.depth_anything", fromlist=["DepthAnythingDepthEstimationInvocation"]
+            ).DepthAnythingDepthEstimationInvocation
+        ),
+        "normalbae_image_processor": lambda: (
+            __import__("invokeai.app.invocations.normal_bae", fromlist=["NormalMapInvocation"]).NormalMapInvocation
+        ),
+        "pidi_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.pidi", fromlist=["PiDiNetEdgeDetectionInvocation"]
+            ).PiDiNetEdgeDetectionInvocation
+        ),
+        "lineart_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.lineart", fromlist=["LineartEdgeDetectionInvocation"]
+            ).LineartEdgeDetectionInvocation
+        ),
+        "lineart_anime_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.lineart_anime", fromlist=["LineartAnimeEdgeDetectionInvocation"]
+            ).LineartAnimeEdgeDetectionInvocation
+        ),
+        "content_shuffle_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.content_shuffle", fromlist=["ContentShuffleInvocation"]
+            ).ContentShuffleInvocation
+        ),
+        "dw_openpose_image_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.dw_openpose", fromlist=["DWOpenposeDetectionInvocation"]
+            ).DWOpenposeDetectionInvocation
+        ),
+        "mediapipe_face_processor": lambda: (
+            __import__(
+                "invokeai.app.invocations.mediapipe_face", fromlist=["MediaPipeFaceDetectionInvocation"]
+            ).MediaPipeFaceDetectionInvocation
+        ),
+        # Note: zoe_depth_image_processor doesn't have a processor invocation implementation
+        "color_map_image_processor": lambda: (
+            __import__("invokeai.app.invocations.color_map", fromlist=["ColorMapInvocation"]).ColorMapInvocation
+        ),
+    }
+
+    if processor_type in processor_class_map:
+        return processor_class_map[processor_type]()
+    return None
+
+
+# Map processor type names to their default parameters
+PROCESSOR_DEFAULT_PARAMS = {
+    "canny_image_processor": {"low_threshold": 100, "high_threshold": 200},
+    "hed_image_processor": {"scribble": False},
+    "mlsd_image_processor": {"detect_resolution": 512, "thr_v": 0.1, "thr_d": 0.1},
+    "depth_anything_image_processor": {"model_size": "small"},
+    "normalbae_image_processor": {"detect_resolution": 512},
+    "pidi_image_processor": {"detect_resolution": 512, "safe": False},
+    "lineart_image_processor": {"detect_resolution": 512, "coarse": False},
+    "lineart_anime_image_processor": {"detect_resolution": 512},
+    "content_shuffle": {},
+    "dw_openpose_image_processor": {"draw_body": True, "draw_face": True, "draw_hands": True},
+    "mediapipe_face_processor": {"max_faces": 1, "min_confidence": 0.5},
+    "zoe_depth_image_processor": {},
+    "color_map_image_processor": {"color_map_tile_size": 64},
+}
+
+
+def process_controlnet_image(image_name: str, model_key: str, services: InvocationServices) -> Optional[dict[str, Any]]:
+    """
+    Process a controlnet image using the appropriate processor based on the model's default settings.
+
+    Args:
+        image_name: The filename of the image to process
+        model_key: The model key to look up default processor settings
+        services: The invocation services providing access to models and images
+
+    Returns:
+        A dictionary with the processed image data (image_name, width, height) or None if processing fails
+    """
+    logger = services.logger
+
+    try:
+        # Get model config to find default processor
+        model_record = services.model_manager.store.get_model(model_key)
+        if not model_record or not model_record.default_settings:
+            logger.info(f"No default processor settings found for model {model_key}")
+            return None
+
+        preprocessor = model_record.default_settings.preprocessor
+        if not preprocessor:
+            logger.info(f"No preprocessor configured for model {model_key}")
+            return None
+
+        # Get the invocation class for this processor
+        invocation_class = _get_processor_invocation_class(preprocessor)
+        if not invocation_class:
+            logger.info(f"No processor mapping found for preprocessor '{preprocessor}'")
+            return None
+
+        # Get default parameters for this processor
+        default_params = PROCESSOR_DEFAULT_PARAMS.get(preprocessor, {})
+        logger.info(f"Processing image {image_name} with processor {preprocessor}")
+
+        # Create a minimal context to run the invocation
+        # We need a fake queue item and session for the context
+        fake_session = GraphExecutionState(graph=Graph())
+        now = datetime.now()
+
+        # Create invocation instance first so we have its ID
+        invocation_params = {"image": ImageField(image_name=image_name), **default_params}
+        invocation = invocation_class(**invocation_params)
+
+        # Add the invocation ID to the session's prepared_source_mapping
+        # This is required for the invocation context to emit progress events
+        fake_session.prepared_source_mapping[invocation.id] = invocation.id
+
+        fake_queue_item = SessionQueueItem(
+            item_id=0,
+            session_id=fake_session.id,
+            queue_id="default",
+            batch_id="recall_processor",
+            field_values=None,
+            session=fake_session,
+            status="in_progress",
+            created_at=now,
+            updated_at=now,
+            started_at=now,
+            completed_at=None,
+        )
+
+        context_data = InvocationContextData(
+            invocation=invocation,
+            source_invocation_id=invocation.id,
+            queue_item=fake_queue_item,
+        )
+
+        context = build_invocation_context(
+            data=context_data,
+            services=services,
+            is_canceled=lambda: False,
+        )
+
+        # Invoke the processor
+        output = invocation.invoke(context)
+
+        # Get the processed image DTO
+        processed_image_dto = services.images.get_dto(output.image.image_name)
+
+        logger.info(f"Successfully processed image {image_name} -> {processed_image_dto.image_name}")
+
+        return {
+            "image_name": processed_image_dto.image_name,
+            "width": processed_image_dto.width,
+            "height": processed_image_dto.height,
+        }
+
+    except Exception as e:
+        logger.error(f"Error processing controlnet image {image_name}: {e}", exc_info=True)
+        return None

invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py

@@ -0,0 +1,41 @@
+import pathlib
+from typing import Optional
+
+import torch
+from PIL import Image
+from transformers import pipeline
+from transformers.pipelines import DepthEstimationPipeline
+
+from invokeai.backend.raw_model import RawModel
+
+
+class DepthAnythingPipeline(RawModel):
+    """Custom wrapper for the Depth Estimation pipeline from transformers adding compatibility
+    for Invoke's Model Management System"""
+
+    def __init__(self, pipeline: DepthEstimationPipeline) -> None:
+        self._pipeline = pipeline
+
+    def generate_depth(self, image: Image.Image) -> Image.Image:
+        depth_map = self._pipeline(image)["depth"]
+        assert isinstance(depth_map, Image.Image)
+        return depth_map
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._pipeline.model.to(device=device, dtype=dtype)
+        self._pipeline.device = self._pipeline.model.device
+
+    def calc_size(self) -> int:
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._pipeline.model)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path):
+        """Load the model from the given path and return a DepthAnythingPipeline instance."""
+
+        depth_anything_pipeline = pipeline(model=str(model_path), task="depth-estimation", local_files_only=True)
+        assert isinstance(depth_anything_pipeline, DepthEstimationPipeline)
+        return cls(depth_anything_pipeline)

invokeai/backend/image_util/dw_openpose/__init__.py

@@ -0,0 +1,152 @@
+from pathlib import Path
+from typing import Dict
+
+import huggingface_hub
+import numpy as np
+import onnxruntime as ort
+import torch
+from PIL import Image
+
+from invokeai.backend.image_util.dw_openpose.onnxdet import inference_detector
+from invokeai.backend.image_util.dw_openpose.onnxpose import inference_pose
+from invokeai.backend.image_util.dw_openpose.utils import NDArrayInt, draw_bodypose, draw_facepose, draw_handpose
+from invokeai.backend.image_util.util import np_to_pil
+from invokeai.backend.util.devices import TorchDevice
+
+
+class DWOpenposeDetector:
+    """
+    Code from the original implementation of the DW Openpose Detector.
+    Credits: https://github.com/IDEA-Research/DWPose
+    """
+
+    hf_repo_id = "yzd-v/DWPose"
+    hf_filename_onnx_det = "yolox_l.onnx"
+    hf_filename_onnx_pose = "dw-ll_ucoco_384.onnx"
+
+    @classmethod
+    def get_model_url_det(cls) -> str:
+        """Returns the URL for the detection model."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_onnx_det)
+
+    @classmethod
+    def get_model_url_pose(cls) -> str:
+        """Returns the URL for the pose model."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_onnx_pose)
+
+    @staticmethod
+    def create_onnx_inference_session(model_path: Path) -> ort.InferenceSession:
+        """Creates an ONNX Inference Session for the given model path, using the appropriate execution provider based on
+        the device type."""
+
+        device = TorchDevice.choose_torch_device()
+        providers = ["CUDAExecutionProvider"] if device.type == "cuda" else ["CPUExecutionProvider"]
+        return ort.InferenceSession(path_or_bytes=model_path, providers=providers)
+
+    def __init__(self, session_det: ort.InferenceSession, session_pose: ort.InferenceSession):
+        self.session_det = session_det
+        self.session_pose = session_pose
+
+    def pose_estimation(self, np_image: np.ndarray):
+        """Does the pose estimation on the given image and returns the keypoints and scores."""
+
+        det_result = inference_detector(self.session_det, np_image)
+        keypoints, scores = inference_pose(self.session_pose, det_result, np_image)
+
+        keypoints_info = np.concatenate((keypoints, scores[..., None]), axis=-1)
+        # compute neck joint
+        neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
+        # neck score when visualizing pred
+        neck[:, 2:4] = np.logical_and(keypoints_info[:, 5, 2:4] > 0.3, keypoints_info[:, 6, 2:4] > 0.3).astype(int)
+        new_keypoints_info = np.insert(keypoints_info, 17, neck, axis=1)
+        mmpose_idx = [17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3]
+        openpose_idx = [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17]
+        new_keypoints_info[:, openpose_idx] = new_keypoints_info[:, mmpose_idx]
+        keypoints_info = new_keypoints_info
+
+        keypoints, scores = keypoints_info[..., :2], keypoints_info[..., 2]
+
+        return keypoints, scores
+
+    def run(
+        self,
+        image: Image.Image,
+        draw_face: bool = False,
+        draw_body: bool = True,
+        draw_hands: bool = False,
+    ) -> Image.Image:
+        """Detects the pose in the given image and returns an solid black image with pose drawn on top, suitable for
+        use with a ControlNet."""
+
+        np_image = np.array(image)
+        H, W, C = np_image.shape
+
+        with torch.no_grad():
+            candidate, subset = self.pose_estimation(np_image)
+            nums, keys, locs = candidate.shape
+            candidate[..., 0] /= float(W)
+            candidate[..., 1] /= float(H)
+            body = candidate[:, :18].copy()
+            body = body.reshape(nums * 18, locs)
+            score = subset[:, :18]
+            for i in range(len(score)):
+                for j in range(len(score[i])):
+                    if score[i][j] > 0.3:
+                        score[i][j] = int(18 * i + j)
+                    else:
+                        score[i][j] = -1
+
+            un_visible = subset < 0.3
+            candidate[un_visible] = -1
+
+            # foot = candidate[:, 18:24]
+
+            faces = candidate[:, 24:92]
+
+            hands = candidate[:, 92:113]
+            hands = np.vstack([hands, candidate[:, 113:]])
+
+            bodies = {"candidate": body, "subset": score}
+            pose = {"bodies": bodies, "hands": hands, "faces": faces}
+
+            return DWOpenposeDetector.draw_pose(
+                pose, H, W, draw_face=draw_face, draw_hands=draw_hands, draw_body=draw_body
+            )
+
+    @staticmethod
+    def draw_pose(
+        pose: Dict[str, NDArrayInt | Dict[str, NDArrayInt]],
+        H: int,
+        W: int,
+        draw_face: bool = True,
+        draw_body: bool = True,
+        draw_hands: bool = True,
+    ) -> Image.Image:
+        """Draws the pose on a black image and returns it as a PIL Image."""
+
+        bodies = pose["bodies"]
+        faces = pose["faces"]
+        hands = pose["hands"]
+
+        assert isinstance(bodies, dict)
+        candidate = bodies["candidate"]
+
+        assert isinstance(bodies, dict)
+        subset = bodies["subset"]
+
+        canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+
+        if draw_body:
+            canvas = draw_bodypose(canvas, candidate, subset)
+
+        if draw_hands:
+            assert isinstance(hands, np.ndarray)
+            canvas = draw_handpose(canvas, hands)
+
+        if draw_face:
+            assert isinstance(hands, np.ndarray)
+            canvas = draw_facepose(canvas, faces)  # type: ignore
+
+        dwpose_image = np_to_pil(canvas)
+
+        return dwpose_image

invokeai/backend/image_util/dw_openpose/onnxdet.py

@@ -0,0 +1,128 @@
+# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
+
+import cv2
+import numpy as np
+
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= nms_thr)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def multiclass_nms(boxes, scores, nms_thr, score_thr):
+    """Multiclass NMS implemented in Numpy. Class-aware version."""
+    final_dets = []
+    num_classes = scores.shape[1]
+    for cls_ind in range(num_classes):
+        cls_scores = scores[:, cls_ind]
+        valid_score_mask = cls_scores > score_thr
+        if valid_score_mask.sum() == 0:
+            continue
+        else:
+            valid_scores = cls_scores[valid_score_mask]
+            valid_boxes = boxes[valid_score_mask]
+            keep = nms(valid_boxes, valid_scores, nms_thr)
+            if len(keep) > 0:
+                cls_inds = np.ones((len(keep), 1)) * cls_ind
+                dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
+                final_dets.append(dets)
+    if len(final_dets) == 0:
+        return None
+    return np.concatenate(final_dets, 0)
+
+
+def demo_postprocess(outputs, img_size, p6=False):
+    grids = []
+    expanded_strides = []
+    strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
+
+    hsizes = [img_size[0] // stride for stride in strides]
+    wsizes = [img_size[1] // stride for stride in strides]
+
+    for hsize, wsize, stride in zip(hsizes, wsizes, strides, strict=False):
+        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
+        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+        grids.append(grid)
+        shape = grid.shape[:2]
+        expanded_strides.append(np.full((*shape, 1), stride))
+
+    grids = np.concatenate(grids, 1)
+    expanded_strides = np.concatenate(expanded_strides, 1)
+    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
+    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
+
+    return outputs
+
+
+def preprocess(img, input_size, swap=(2, 0, 1)):
+    if len(img.shape) == 3:
+        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+    else:
+        padded_img = np.ones(input_size, dtype=np.uint8) * 114
+
+    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+    resized_img = cv2.resize(
+        img,
+        (int(img.shape[1] * r), int(img.shape[0] * r)),
+        interpolation=cv2.INTER_LINEAR,
+    ).astype(np.uint8)
+    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+    padded_img = padded_img.transpose(swap)
+    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+    return padded_img, r
+
+
+def inference_detector(session, oriImg):
+    input_shape = (640, 640)
+    img, ratio = preprocess(oriImg, input_shape)
+
+    ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
+    output = session.run(None, ort_inputs)
+    predictions = demo_postprocess(output[0], input_shape)[0]
+
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4:5] * predictions[:, 5:]
+
+    boxes_xyxy = np.ones_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
+    boxes_xyxy /= ratio
+    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
+    if dets is not None:
+        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+        isscore = final_scores > 0.3
+        iscat = final_cls_inds == 0
+        isbbox = [i and j for (i, j) in zip(isscore, iscat, strict=False)]
+        final_boxes = final_boxes[isbbox]
+    else:
+        final_boxes = np.array([])
+
+    return final_boxes

invokeai/backend/image_util/dw_openpose/onnxpose.py

@@ -0,0 +1,361 @@
+# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
+
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+
+def preprocess(
+    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Do preprocessing for RTMPose model inference.
+
+    Args:
+        img (np.ndarray): Input image in shape.
+        input_size (tuple): Input image size in shape (w, h).
+
+    Returns:
+        tuple:
+        - resized_img (np.ndarray): Preprocessed image.
+        - center (np.ndarray): Center of image.
+        - scale (np.ndarray): Scale of image.
+    """
+    # get shape of image
+    img_shape = img.shape[:2]
+    out_img, out_center, out_scale = [], [], []
+    if len(out_bbox) == 0:
+        out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
+    for i in range(len(out_bbox)):
+        x0 = out_bbox[i][0]
+        y0 = out_bbox[i][1]
+        x1 = out_bbox[i][2]
+        y1 = out_bbox[i][3]
+        bbox = np.array([x0, y0, x1, y1])
+
+        # get center and scale
+        center, scale = bbox_xyxy2cs(bbox, padding=1.25)
+
+        # do affine transformation
+        resized_img, scale = top_down_affine(input_size, scale, center, img)
+
+        # normalize image
+        mean = np.array([123.675, 116.28, 103.53])
+        std = np.array([58.395, 57.12, 57.375])
+        resized_img = (resized_img - mean) / std
+
+        out_img.append(resized_img)
+        out_center.append(center)
+        out_scale.append(scale)
+
+    return out_img, out_center, out_scale
+
+
+def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray:
+    """Inference RTMPose model.
+
+    Args:
+        sess (ort.InferenceSession): ONNXRuntime session.
+        img (np.ndarray): Input image in shape.
+
+    Returns:
+        outputs (np.ndarray): Output of RTMPose model.
+    """
+    all_out = []
+    # build input
+    for i in range(len(img)):
+        input = [img[i].transpose(2, 0, 1)]
+
+        # build output
+        sess_input = {sess.get_inputs()[0].name: input}
+        sess_output = []
+        for out in sess.get_outputs():
+            sess_output.append(out.name)
+
+        # run model
+        outputs = sess.run(sess_output, sess_input)
+        all_out.append(outputs)
+
+    return all_out
+
+
+def postprocess(
+    outputs: List[np.ndarray],
+    model_input_size: Tuple[int, int],
+    center: Tuple[int, int],
+    scale: Tuple[int, int],
+    simcc_split_ratio: float = 2.0,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Postprocess for RTMPose model output.
+
+    Args:
+        outputs (np.ndarray): Output of RTMPose model.
+        model_input_size (tuple): RTMPose model Input image size.
+        center (tuple): Center of bbox in shape (x, y).
+        scale (tuple): Scale of bbox in shape (w, h).
+        simcc_split_ratio (float): Split ratio of simcc.
+
+    Returns:
+        tuple:
+        - keypoints (np.ndarray): Rescaled keypoints.
+        - scores (np.ndarray): Model predict scores.
+    """
+    all_key = []
+    all_score = []
+    for i in range(len(outputs)):
+        # use simcc to decode
+        simcc_x, simcc_y = outputs[i]
+        keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
+
+        # rescale keypoints
+        keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
+        all_key.append(keypoints[0])
+        all_score.append(scores[0])
+
+    return np.array(all_key), np.array(all_score)
+
+
+def bbox_xyxy2cs(bbox: np.ndarray, padding: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
+    """Transform the bbox format from (x,y,w,h) into (center, scale)
+
+    Args:
+        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
+            as (left, top, right, bottom)
+        padding (float): BBox padding factor that will be multilied to scale.
+            Default: 1.0
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
+            (n, 2)
+        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
+            (n, 2)
+    """
+    # convert single bbox from (4, ) to (1, 4)
+    dim = bbox.ndim
+    if dim == 1:
+        bbox = bbox[None, :]
+
+    # get bbox center and scale
+    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
+    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
+    scale = np.hstack([x2 - x1, y2 - y1]) * padding
+
+    if dim == 1:
+        center = center[0]
+        scale = scale[0]
+
+    return center, scale
+
+
+def _fix_aspect_ratio(bbox_scale: np.ndarray, aspect_ratio: float) -> np.ndarray:
+    """Extend the scale to match the given aspect ratio.
+
+    Args:
+        scale (np.ndarray): The image scale (w, h) in shape (2, )
+        aspect_ratio (float): The ratio of ``w/h``
+
+    Returns:
+        np.ndarray: The reshaped image scale in (2, )
+    """
+    w, h = np.hsplit(bbox_scale, [1])
+    bbox_scale = np.where(w > h * aspect_ratio, np.hstack([w, w / aspect_ratio]), np.hstack([h * aspect_ratio, h]))
+    return bbox_scale
+
+
+def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
+    """Rotate a point by an angle.
+
+    Args:
+        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
+        angle_rad (float): rotation angle in radian
+
+    Returns:
+        np.ndarray: Rotated point in shape (2, )
+    """
+    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
+    rot_mat = np.array([[cs, -sn], [sn, cs]])
+    return rot_mat @ pt
+
+
+def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """To calculate the affine matrix, three pairs of points are required. This
+    function is used to get the 3rd point, given 2D points a & b.
+
+    The 3rd point is defined by rotating vector `a - b` by 90 degrees
+    anticlockwise, using b as the rotation center.
+
+    Args:
+        a (np.ndarray): The 1st point (x,y) in shape (2, )
+        b (np.ndarray): The 2nd point (x,y) in shape (2, )
+
+    Returns:
+        np.ndarray: The 3rd point.
+    """
+    direction = a - b
+    c = b + np.r_[-direction[1], direction[0]]
+    return c
+
+
+def get_warp_matrix(
+    center: np.ndarray,
+    scale: np.ndarray,
+    rot: float,
+    output_size: Tuple[int, int],
+    shift: Tuple[float, float] = (0.0, 0.0),
+    inv: bool = False,
+) -> np.ndarray:
+    """Calculate the affine transformation matrix that can warp the bbox area
+    in the input image to the output size.
+
+    Args:
+        center (np.ndarray[2, ]): Center of the bounding box (x, y).
+        scale (np.ndarray[2, ]): Scale of the bounding box
+            wrt [width, height].
+        rot (float): Rotation angle (degree).
+        output_size (np.ndarray[2, ] | list(2,)): Size of the
+            destination heatmaps.
+        shift (0-100%): Shift translation ratio wrt the width/height.
+            Default (0., 0.).
+        inv (bool): Option to inverse the affine transform direction.
+            (inv=False: src->dst or inv=True: dst->src)
+
+    Returns:
+        np.ndarray: A 2x3 transformation matrix
+    """
+    shift = np.array(shift)
+    src_w = scale[0]
+    dst_w = output_size[0]
+    dst_h = output_size[1]
+
+    # compute transformation matrix
+    rot_rad = np.deg2rad(rot)
+    src_dir = _rotate_point(np.array([0.0, src_w * -0.5]), rot_rad)
+    dst_dir = np.array([0.0, dst_w * -0.5])
+
+    # get four corners of the src rectangle in the original image
+    src = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale * shift
+    src[1, :] = center + src_dir + scale * shift
+    src[2, :] = _get_3rd_point(src[0, :], src[1, :])
+
+    # get four corners of the dst rectangle in the input image
+    dst = np.zeros((3, 2), dtype=np.float32)
+    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return warp_mat
+
+
+def top_down_affine(
+    input_size: dict, bbox_scale: dict, bbox_center: dict, img: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Get the bbox image as the model input by affine transform.
+
+    Args:
+        input_size (dict): The input size of the model.
+        bbox_scale (dict): The bbox scale of the img.
+        bbox_center (dict): The bbox center of the img.
+        img (np.ndarray): The original image.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: img after affine transform.
+        - np.ndarray[float32]: bbox scale after affine transform.
+    """
+    w, h = input_size
+    warp_size = (int(w), int(h))
+
+    # reshape bbox to fixed aspect ratio
+    bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
+
+    # get the affine matrix
+    center = bbox_center
+    scale = bbox_scale
+    rot = 0
+    warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
+
+    # do affine transform
+    img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
+
+    return img, bbox_scale
+
+
+def get_simcc_maximum(simcc_x: np.ndarray, simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Get maximum response location and value from simcc representations.
+
+    Note:
+        instance number: N
+        num_keypoints: K
+        heatmap height: H
+        heatmap width: W
+
+    Args:
+        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
+        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
+
+    Returns:
+        tuple:
+        - locs (np.ndarray): locations of maximum heatmap responses in shape
+            (K, 2) or (N, K, 2)
+        - vals (np.ndarray): values of maximum heatmap responses in shape
+            (K,) or (N, K)
+    """
+    N, K, Wx = simcc_x.shape
+    simcc_x = simcc_x.reshape(N * K, -1)
+    simcc_y = simcc_y.reshape(N * K, -1)
+
+    # get maximum value locations
+    x_locs = np.argmax(simcc_x, axis=1)
+    y_locs = np.argmax(simcc_y, axis=1)
+    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
+    max_val_x = np.amax(simcc_x, axis=1)
+    max_val_y = np.amax(simcc_y, axis=1)
+
+    # get maximum value across x and y axis
+    mask = max_val_x > max_val_y
+    max_val_x[mask] = max_val_y[mask]
+    vals = max_val_x
+    locs[vals <= 0.0] = -1
+
+    # reshape
+    locs = locs.reshape(N, K, 2)
+    vals = vals.reshape(N, K)
+
+    return locs, vals
+
+
+def decode(simcc_x: np.ndarray, simcc_y: np.ndarray, simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]:
+    """Modulate simcc distribution with Gaussian.
+
+    Args:
+        simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
+        simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
+        simcc_split_ratio (int): The split ratio of simcc.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
+        - np.ndarray[float32]: scores in shape (K,) or (n, K)
+    """
+    keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
+    keypoints /= simcc_split_ratio
+
+    return keypoints, scores
+
+
+def inference_pose(session, out_bbox, oriImg):
+    h, w = session.get_inputs()[0].shape[2:]
+    model_input_size = (w, h)
+    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
+    outputs = inference(session, resized_img)
+    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
+
+    return keypoints, scores

invokeai/backend/image_util/dw_openpose/utils.py

@@ -0,0 +1,158 @@
+# Code from the original DWPose Implementation: https://github.com/IDEA-Research/DWPose
+
+import math
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+
+eps = 0.01
+NDArrayInt = npt.NDArray[np.uint8]
+
+
+def draw_bodypose(canvas: NDArrayInt, candidate: NDArrayInt, subset: NDArrayInt) -> NDArrayInt:
+    H, W, C = canvas.shape
+    candidate = np.array(candidate)
+    subset = np.array(subset)
+
+    stickwidth = 4
+
+    limbSeq = [
+        [2, 3],
+        [2, 6],
+        [3, 4],
+        [4, 5],
+        [6, 7],
+        [7, 8],
+        [2, 9],
+        [9, 10],
+        [10, 11],
+        [2, 12],
+        [12, 13],
+        [13, 14],
+        [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],
+        [3, 17],
+        [6, 18],
+    ]
+
+    colors = [
+        [255, 0, 0],
+        [255, 85, 0],
+        [255, 170, 0],
+        [255, 255, 0],
+        [170, 255, 0],
+        [85, 255, 0],
+        [0, 255, 0],
+        [0, 255, 85],
+        [0, 255, 170],
+        [0, 255, 255],
+        [0, 170, 255],
+        [0, 85, 255],
+        [0, 0, 255],
+        [85, 0, 255],
+        [170, 0, 255],
+        [255, 0, 255],
+        [255, 0, 170],
+        [255, 0, 85],
+    ]
+
+    for i in range(17):
+        for n in range(len(subset)):
+            index = subset[n][np.array(limbSeq[i]) - 1]
+            if -1 in index:
+                continue
+            Y = candidate[index.astype(int), 0] * float(W)
+            X = candidate[index.astype(int), 1] * float(H)
+            mX = np.mean(X)
+            mY = np.mean(Y)
+            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+            cv2.fillConvexPoly(canvas, polygon, colors[i])
+
+    canvas = (canvas * 0.6).astype(np.uint8)
+
+    for i in range(18):
+        for n in range(len(subset)):
+            index = int(subset[n][i])
+            if index == -1:
+                continue
+            x, y = candidate[index][0:2]
+            x = int(x * W)
+            y = int(y * H)
+            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
+
+    return canvas
+
+
+def draw_handpose(canvas: NDArrayInt, all_hand_peaks: NDArrayInt) -> NDArrayInt:
+    H, W, C = canvas.shape
+
+    edges = [
+        [0, 1],
+        [1, 2],
+        [2, 3],
+        [3, 4],
+        [0, 5],
+        [5, 6],
+        [6, 7],
+        [7, 8],
+        [0, 9],
+        [9, 10],
+        [10, 11],
+        [11, 12],
+        [0, 13],
+        [13, 14],
+        [14, 15],
+        [15, 16],
+        [0, 17],
+        [17, 18],
+        [18, 19],
+        [19, 20],
+    ]
+
+    for peaks in all_hand_peaks:
+        peaks = np.array(peaks)
+
+        for ie, e in enumerate(edges):
+            x1, y1 = peaks[e[0]]
+            x2, y2 = peaks[e[1]]
+            x1 = int(x1 * W)
+            y1 = int(y1 * H)
+            x2 = int(x2 * W)
+            y2 = int(y2 * H)
+            if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
+                hsv_color = np.array([[[ie / float(len(edges)) * 180, 255, 255]]], dtype=np.uint8)
+                rgb_color = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2RGB)[0, 0]
+                cv2.line(
+                    canvas,
+                    (x1, y1),
+                    (x2, y2),
+                    rgb_color.tolist(),
+                    thickness=2,
+                )
+
+        for _, keyponit in enumerate(peaks):
+            x, y = keyponit
+            x = int(x * W)
+            y = int(y * H)
+            if x > eps and y > eps:
+                cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
+    return canvas
+
+
+def draw_facepose(canvas: NDArrayInt, all_lmks: NDArrayInt) -> NDArrayInt:
+    H, W, C = canvas.shape
+    for lmks in all_lmks:
+        lmks = np.array(lmks)
+        for lmk in lmks:
+            x, y = lmk
+            x = int(x * W)
+            y = int(y * H)
+            if x > eps and y > eps:
+                cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
+    return canvas

invokeai/backend/image_util/grounding_dino/detection_result.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class BoundingBox(BaseModel):
+    """Bounding box helper class."""
+
+    xmin: int
+    ymin: int
+    xmax: int
+    ymax: int
+
+
+class DetectionResult(BaseModel):
+    """Detection result from Grounding DINO."""
+
+    score: float
+    label: str
+    box: BoundingBox
+    model_config = ConfigDict(
+        # Allow arbitrary types for mask, since it will be a numpy array.
+        arbitrary_types_allowed=True
+    )

invokeai/backend/image_util/grounding_dino/grounding_dino_pipeline.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+from PIL import Image
+from transformers.pipelines import ZeroShotObjectDetectionPipeline
+
+from invokeai.backend.image_util.grounding_dino.detection_result import DetectionResult
+from invokeai.backend.raw_model import RawModel
+
+
+class GroundingDinoPipeline(RawModel):
+    """A wrapper class for a ZeroShotObjectDetectionPipeline that makes it compatible with the model manager's memory
+    management system.
+    """
+
+    def __init__(self, pipeline: ZeroShotObjectDetectionPipeline):
+        self._pipeline = pipeline
+
+    def detect(self, image: Image.Image, candidate_labels: list[str], threshold: float = 0.1) -> list[DetectionResult]:
+        results = self._pipeline(image=image, candidate_labels=candidate_labels, threshold=threshold)
+        assert results is not None
+        results = [DetectionResult.model_validate(result) for result in results]
+        return results
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        # HACK(ryand): The GroundingDinoPipeline does not work on MPS devices. We only allow it to be moved to CPU or
+        # CUDA.
+        if device is not None and device.type not in {"cpu", "cuda"}:
+            device = None
+        self._pipeline.model.to(device=device, dtype=dtype)
+        self._pipeline.device = self._pipeline.model.device
+
+    def calc_size(self) -> int:
+        # HACK(ryand): Fix the circular import issue.
+        from invokeai.backend.model_manager.load.model_util import calc_module_size
+
+        return calc_module_size(self._pipeline.model)

invokeai/backend/image_util/hed.py

@@ -0,0 +1,217 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import pathlib
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    nms,
+    normalize_image_channel_count,
+    np_to_pil,
+    pil_to_np,
+    resize_image_to_resolution,
+    safe_step,
+)
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class DoubleConvBlock(torch.nn.Module):
+    def __init__(self, input_channel, output_channel, layer_number):
+        super().__init__()
+        self.convs = torch.nn.Sequential()
+        self.convs.append(
+            torch.nn.Conv2d(
+                in_channels=input_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1
+            )
+        )
+        for _i in range(1, layer_number):
+            self.convs.append(
+                torch.nn.Conv2d(
+                    in_channels=output_channel,
+                    out_channels=output_channel,
+                    kernel_size=(3, 3),
+                    stride=(1, 1),
+                    padding=1,
+                )
+            )
+        self.projection = torch.nn.Conv2d(
+            in_channels=output_channel, out_channels=1, kernel_size=(1, 1), stride=(1, 1), padding=0
+        )
+
+    def __call__(self, x, down_sampling=False):
+        h = x
+        if down_sampling:
+            h = torch.nn.functional.max_pool2d(h, kernel_size=(2, 2), stride=(2, 2))
+        for conv in self.convs:
+            h = conv(h)
+            h = torch.nn.functional.relu(h)
+        return h, self.projection(h)
+
+
+class ControlNetHED_Apache2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
+        self.block1 = DoubleConvBlock(input_channel=3, output_channel=64, layer_number=2)
+        self.block2 = DoubleConvBlock(input_channel=64, output_channel=128, layer_number=2)
+        self.block3 = DoubleConvBlock(input_channel=128, output_channel=256, layer_number=3)
+        self.block4 = DoubleConvBlock(input_channel=256, output_channel=512, layer_number=3)
+        self.block5 = DoubleConvBlock(input_channel=512, output_channel=512, layer_number=3)
+
+    def __call__(self, x):
+        h = x - self.norm
+        h, projection1 = self.block1(h)
+        h, projection2 = self.block2(h, down_sampling=True)
+        h, projection3 = self.block3(h, down_sampling=True)
+        h, projection4 = self.block4(h, down_sampling=True)
+        h, projection5 = self.block5(h, down_sampling=True)
+        return projection1, projection2, projection3, projection4, projection5
+
+
+class HEDProcessor:
+    """Holistically-Nested Edge Detection.
+
+    On instantiation, loads the HED model from the HuggingFace Hub.
+    """
+
+    def __init__(self):
+        model_path = hf_hub_download("lllyasviel/Annotators", "ControlNetHED.pth")
+        self.network = ControlNetHED_Apache2()
+        self.network.load_state_dict(torch.load(model_path, map_location="cpu"))
+        self.network.float().eval()
+
+    def to(self, device: torch.device):
+        self.network.to(device)
+        return self
+
+    def run(
+        self,
+        input_image: Image.Image,
+        detect_resolution: int = 512,
+        image_resolution: int = 512,
+        safe: bool = False,
+        scribble: bool = False,
+    ) -> Image.Image:
+        """Processes an image and returns the detected edges.
+
+        Args:
+            input_image: The input image.
+            detect_resolution: The resolution to fit the image to before edge detection.
+            image_resolution: The resolution to fit the edges to before returning.
+            safe: Whether to apply safe step to the detected edges.
+            scribble: Whether to apply non-maximum suppression and Gaussian blur to the detected edges.
+
+        Returns:
+            The detected edges.
+        """
+        device = get_effective_device(self.network)
+        np_image = pil_to_np(input_image)
+        np_image = normalize_image_channel_count(np_image)
+        np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+        assert np_image.ndim == 3
+        height, width, _channels = np_image.shape
+        with torch.no_grad():
+            image_hed = torch.from_numpy(np_image.copy()).float().to(device)
+            image_hed = rearrange(image_hed, "h w c -> 1 c h w")
+            edges = self.network(image_hed)
+            edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
+            edges = [cv2.resize(e, (width, height), interpolation=cv2.INTER_LINEAR) for e in edges]
+            edges = np.stack(edges, axis=2)
+            edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
+            if safe:
+                edge = safe_step(edge)
+            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = edge
+        detected_map = normalize_image_channel_count(detected_map)
+
+        img = resize_image_to_resolution(np_image, image_resolution)
+        height, width, _channels = img.shape
+
+        detected_map = cv2.resize(detected_map, (width, height), interpolation=cv2.INTER_LINEAR)
+
+        if scribble:
+            detected_map = nms(detected_map, 127, 3.0)
+            detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
+            detected_map[detected_map > 4] = 255
+            detected_map[detected_map < 255] = 0
+
+        return np_to_pil(detected_map)
+
+
+class HEDEdgeDetector:
+    """Simple wrapper around the HED model for detecting edges in an image."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename = "ControlNetHED.pth"
+
+    def __init__(self, model: ControlNetHED_Apache2):
+        self.model = model
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> ControlNetHED_Apache2:
+        """Load the model from a file."""
+        model = ControlNetHED_Apache2()
+        model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        model.float().eval()
+        return model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image, safe: bool = False, scribble: bool = False) -> Image.Image:
+        """Processes an image and returns the detected edges.
+
+        Args:
+            image: The input image.
+            safe: Whether to apply safe step to the detected edges.
+            scribble: Whether to apply non-maximum suppression and Gaussian blur to the detected edges.
+
+        Returns:
+            The detected edges.
+        """
+
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(image)
+
+        height, width, _channels = np_image.shape
+
+        with torch.no_grad():
+            image_hed = torch.from_numpy(np_image.copy()).float().to(device)
+            image_hed = rearrange(image_hed, "h w c -> 1 c h w")
+            edges = self.model(image_hed)
+            edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
+            edges = [cv2.resize(e, (width, height), interpolation=cv2.INTER_LINEAR) for e in edges]
+            edges = np.stack(edges, axis=2)
+            edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
+            if safe:
+                edge = safe_step(edge)
+            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = edge
+
+        detected_map = cv2.resize(detected_map, (width, height), interpolation=cv2.INTER_LINEAR)
+
+        if scribble:
+            detected_map = nms(detected_map, 127, 3.0)
+            detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
+            detected_map[detected_map > 4] = 255
+            detected_map[detected_map < 255] = 0
+
+        output = np_to_pil(detected_map)
+
+        return output

invokeai/backend/image_util/imwatermark/vendor.py

@@ -0,0 +1,310 @@
+# This file is vendored from https://github.com/ShieldMnt/invisible-watermark
+#
+# `invisible-watermark` is MIT licensed as of August 23, 2025, when the code was copied into this repo.
+#
+# Why we vendored it in:
+# `invisible-watermark` has a dependency on `opencv-python`, which conflicts with Invoke's dependency on
+# `opencv-contrib-python`. It's easier to copy the code over than complicate the installation process by
+# requiring an extra post-install step of removing `opencv-python` and installing `opencv-contrib-python`.
+
+import struct
+import uuid
+import base64
+import cv2
+import numpy as np
+import pywt
+
+
+class WatermarkEncoder(object):
+    def __init__(self, content=b""):
+        seq = np.array([n for n in content], dtype=np.uint8)
+        self._watermarks = list(np.unpackbits(seq))
+        self._wmLen = len(self._watermarks)
+        self._wmType = "bytes"
+
+    def set_by_ipv4(self, addr):
+        bits = []
+        ips = addr.split(".")
+        for ip in ips:
+            bits += list(np.unpackbits(np.array([ip % 255], dtype=np.uint8)))
+        self._watermarks = bits
+        self._wmLen = len(self._watermarks)
+        self._wmType = "ipv4"
+        assert self._wmLen == 32
+
+    def set_by_uuid(self, uid):
+        u = uuid.UUID(uid)
+        self._wmType = "uuid"
+        seq = np.array([n for n in u.bytes], dtype=np.uint8)
+        self._watermarks = list(np.unpackbits(seq))
+        self._wmLen = len(self._watermarks)
+
+    def set_by_bytes(self, content):
+        self._wmType = "bytes"
+        seq = np.array([n for n in content], dtype=np.uint8)
+        self._watermarks = list(np.unpackbits(seq))
+        self._wmLen = len(self._watermarks)
+
+    def set_by_b16(self, b16):
+        content = base64.b16decode(b16)
+        self.set_by_bytes(content)
+        self._wmType = "b16"
+
+    def set_by_bits(self, bits=None):
+        if bits is None:
+            bits = []
+        self._watermarks = [int(bit) % 2 for bit in bits]
+        self._wmLen = len(self._watermarks)
+        self._wmType = "bits"
+
+    def set_watermark(self, wmType="bytes", content=""):
+        if wmType == "ipv4":
+            self.set_by_ipv4(content)
+        elif wmType == "uuid":
+            self.set_by_uuid(content)
+        elif wmType == "bits":
+            self.set_by_bits(content)
+        elif wmType == "bytes":
+            self.set_by_bytes(content)
+        elif wmType == "b16":
+            self.set_by_b16(content)
+        else:
+            raise NameError("%s is not supported" % wmType)
+
+    def get_length(self):
+        return self._wmLen
+
+    # @classmethod
+    # def loadModel(cls):
+    #     RivaWatermark.loadModel()
+
+    def encode(self, cv2Image, method="dwtDct", **configs):
+        (r, c, channels) = cv2Image.shape
+        if r * c < 256 * 256:
+            raise RuntimeError("image too small, should be larger than 256x256")
+
+        if method == "dwtDct":
+            embed = EmbedMaxDct(self._watermarks, wmLen=self._wmLen, **configs)
+            return embed.encode(cv2Image)
+        # elif method == 'dwtDctSvd':
+        #     embed = EmbedDwtDctSvd(self._watermarks, wmLen=self._wmLen, **configs)
+        #     return embed.encode(cv2Image)
+        # elif method == 'rivaGan':
+        #     embed = RivaWatermark(self._watermarks, self._wmLen)
+        #     return embed.encode(cv2Image)
+        else:
+            raise NameError("%s is not supported" % method)
+
+
+class WatermarkDecoder(object):
+    def __init__(self, wm_type="bytes", length=0):
+        self._wmType = wm_type
+        if wm_type == "ipv4":
+            self._wmLen = 32
+        elif wm_type == "uuid":
+            self._wmLen = 128
+        elif wm_type == "bytes":
+            self._wmLen = length
+        elif wm_type == "bits":
+            self._wmLen = length
+        elif wm_type == "b16":
+            self._wmLen = length
+        else:
+            raise NameError("%s is unsupported" % wm_type)
+
+    def reconstruct_ipv4(self, bits):
+        ips = [str(ip) for ip in list(np.packbits(bits))]
+        return ".".join(ips)
+
+    def reconstruct_uuid(self, bits):
+        nums = np.packbits(bits)
+        bstr = b""
+        for i in range(16):
+            bstr += struct.pack(">B", nums[i])
+
+        return str(uuid.UUID(bytes=bstr))
+
+    def reconstruct_bits(self, bits):
+        # return ''.join([str(b) for b in bits])
+        return bits
+
+    def reconstruct_b16(self, bits):
+        bstr = self.reconstruct_bytes(bits)
+        return base64.b16encode(bstr)
+
+    def reconstruct_bytes(self, bits):
+        nums = np.packbits(bits)
+        bstr = b""
+        for i in range(self._wmLen // 8):
+            bstr += struct.pack(">B", nums[i])
+        return bstr
+
+    def reconstruct(self, bits):
+        if len(bits) != self._wmLen:
+            raise RuntimeError("bits are not matched with watermark length")
+
+        if self._wmType == "ipv4":
+            return self.reconstruct_ipv4(bits)
+        elif self._wmType == "uuid":
+            return self.reconstruct_uuid(bits)
+        elif self._wmType == "bits":
+            return self.reconstruct_bits(bits)
+        elif self._wmType == "b16":
+            return self.reconstruct_b16(bits)
+        else:
+            return self.reconstruct_bytes(bits)
+
+    def decode(self, cv2Image, method="dwtDct", **configs):
+        (r, c, channels) = cv2Image.shape
+        if r * c < 256 * 256:
+            raise RuntimeError("image too small, should be larger than 256x256")
+
+        bits = []
+        if method == "dwtDct":
+            embed = EmbedMaxDct(watermarks=[], wmLen=self._wmLen, **configs)
+            bits = embed.decode(cv2Image)
+        # elif method == 'dwtDctSvd':
+        #     embed = EmbedDwtDctSvd(watermarks=[], wmLen=self._wmLen, **configs)
+        #     bits = embed.decode(cv2Image)
+        # elif method == 'rivaGan':
+        #     embed = RivaWatermark(watermarks=[], wmLen=self._wmLen, **configs)
+        #     bits = embed.decode(cv2Image)
+        else:
+            raise NameError("%s is not supported" % method)
+        return self.reconstruct(bits)
+
+    # @classmethod
+    # def loadModel(cls):
+    #     RivaWatermark.loadModel()
+
+
+class EmbedMaxDct(object):
+    def __init__(self, watermarks=None, wmLen=8, scales=None, block=4):
+        if watermarks is None:
+            watermarks = []
+        if scales is None:
+            scales = [0, 36, 36]
+        self._watermarks = watermarks
+        self._wmLen = wmLen
+        self._scales = scales
+        self._block = block
+
+    def encode(self, bgr):
+        (row, col, channels) = bgr.shape
+
+        yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
+
+        for channel in range(2):
+            if self._scales[channel] <= 0:
+                continue
+
+            ca1, (h1, v1, d1) = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar")
+            self.encode_frame(ca1, self._scales[channel])
+
+            yuv[: row // 4 * 4, : col // 4 * 4, channel] = pywt.idwt2((ca1, (v1, h1, d1)), "haar")
+
+        bgr_encoded = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
+        return bgr_encoded
+
+    def decode(self, bgr):
+        (row, col, channels) = bgr.shape
+
+        yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
+
+        scores = [[] for i in range(self._wmLen)]
+        for channel in range(2):
+            if self._scales[channel] <= 0:
+                continue
+
+            ca1, (h1, v1, d1) = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar")
+
+            scores = self.decode_frame(ca1, self._scales[channel], scores)
+
+        avgScores = list(map(lambda l: np.array(l).mean(), scores))
+
+        bits = np.array(avgScores) * 255 > 127
+        return bits
+
+    def decode_frame(self, frame, scale, scores):
+        (row, col) = frame.shape
+        num = 0
+
+        for i in range(row // self._block):
+            for j in range(col // self._block):
+                block = frame[
+                    i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
+                ]
+
+                score = self.infer_dct_matrix(block, scale)
+                # score = self.infer_dct_svd(block, scale)
+                wmBit = num % self._wmLen
+                scores[wmBit].append(score)
+                num = num + 1
+
+        return scores
+
+    def diffuse_dct_svd(self, block, wmBit, scale):
+        u, s, v = np.linalg.svd(cv2.dct(block))
+
+        s[0] = (s[0] // scale + 0.25 + 0.5 * wmBit) * scale
+        return cv2.idct(np.dot(u, np.dot(np.diag(s), v)))
+
+    def infer_dct_svd(self, block, scale):
+        u, s, v = np.linalg.svd(cv2.dct(block))
+
+        score = 0
+        score = int((s[0] % scale) > scale * 0.5)
+        return score
+        if score >= 0.5:
+            return 1.0
+        else:
+            return 0.0
+
+    def diffuse_dct_matrix(self, block, wmBit, scale):
+        pos = np.argmax(abs(block.flatten()[1:])) + 1
+        i, j = pos // self._block, pos % self._block
+        val = block[i][j]
+        if val >= 0.0:
+            block[i][j] = (val // scale + 0.25 + 0.5 * wmBit) * scale
+        else:
+            val = abs(val)
+            block[i][j] = -1.0 * (val // scale + 0.25 + 0.5 * wmBit) * scale
+        return block
+
+    def infer_dct_matrix(self, block, scale):
+        pos = np.argmax(abs(block.flatten()[1:])) + 1
+        i, j = pos // self._block, pos % self._block
+
+        val = block[i][j]
+        if val < 0:
+            val = abs(val)
+
+        if (val % scale) > 0.5 * scale:
+            return 1
+        else:
+            return 0
+
+    def encode_frame(self, frame, scale):
+        """
+        frame is a matrix (M, N)
+
+        we get K (watermark bits size) blocks (self._block x self._block)
+
+        For i-th block, we encode watermark[i] bit into it
+        """
+        (row, col) = frame.shape
+        num = 0
+        for i in range(row // self._block):
+            for j in range(col // self._block):
+                block = frame[
+                    i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
+                ]
+                wmBit = self._watermarks[(num % self._wmLen)]
+
+                diffusedBlock = self.diffuse_dct_matrix(block, wmBit, scale)
+                # diffusedBlock = self.diffuse_dct_svd(block, wmBit, scale)
+                frame[
+                    i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
+                ] = diffusedBlock
+
+                num = num + 1

invokeai/backend/image_util/infill_methods/cv2_inpaint.py

@@ -0,0 +1,20 @@
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def cv2_inpaint(image: Image.Image) -> Image.Image:
+    # Prepare Image
+    image_array = np.array(image.convert("RGB"))
+    image_cv = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
+
+    # Prepare Mask From Alpha Channel
+    mask = image.split()[3].convert("RGB")
+    mask_array = np.array(mask)
+    mask_cv = cv2.cvtColor(mask_array, cv2.COLOR_BGR2GRAY)
+    mask_inv = cv2.bitwise_not(mask_cv)
+
+    # Inpaint Image
+    inpainted_result = cv2.inpaint(image_cv, mask_inv, 3, cv2.INPAINT_TELEA)
+    inpainted_image = Image.fromarray(cv2.cvtColor(inpainted_result, cv2.COLOR_BGR2RGB))
+    return inpainted_image

invokeai/backend/image_util/infill_methods/lama.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+import invokeai.backend.util.logging as logger
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+from invokeai.backend.model_manager.taxonomy import AnyModel
+
+
+def norm_img(np_img):
+    if len(np_img.shape) == 2:
+        np_img = np_img[:, :, np.newaxis]
+    np_img = np.transpose(np_img, (2, 0, 1))
+    np_img = np_img.astype("float32") / 255
+    return np_img
+
+
+class LaMA:
+    def __init__(self, model: AnyModel):
+        self._model = model
+
+    def __call__(self, input_image: Image.Image, *args: Any, **kwds: Any) -> Any:
+        image = np.asarray(input_image.convert("RGB"))
+        image = norm_img(image)
+
+        mask = input_image.split()[-1]
+        mask = np.asarray(mask)
+        mask = np.invert(mask)
+        mask = norm_img(mask)
+        mask = (mask > 0) * 1
+
+        device = get_effective_device(self._model)
+        image = torch.from_numpy(image).unsqueeze(0).to(device)
+        mask = torch.from_numpy(mask).unsqueeze(0).to(device)
+
+        with torch.inference_mode():
+            infilled_image = self._model(image, mask)
+
+        infilled_image = infilled_image[0].permute(1, 2, 0).detach().cpu().numpy()
+        infilled_image = np.clip(infilled_image * 255, 0, 255).astype("uint8")
+        infilled_image = Image.fromarray(infilled_image)
+
+        return infilled_image
+
+    @staticmethod
+    def load_jit_model(url_or_path: str | Path, device: torch.device | str = "cpu") -> torch.nn.Module:
+        model_path = url_or_path
+        logger.info(f"Loading model from: {model_path}")
+        model: torch.nn.Module = torch.jit.load(model_path, map_location="cpu").to(device)  # type: ignore
+        model.eval()
+        return model

invokeai/backend/image_util/infill_methods/mosaic.py

@@ -0,0 +1,60 @@
+from typing import Tuple
+
+import numpy as np
+from PIL import Image
+
+
+def infill_mosaic(
+    image: Image.Image,
+    tile_shape: Tuple[int, int] = (64, 64),
+    min_color: Tuple[int, int, int, int] = (0, 0, 0, 0),
+    max_color: Tuple[int, int, int, int] = (255, 255, 255, 0),
+) -> Image.Image:
+    """
+    image:PIL - A PIL Image
+    tile_shape: Tuple[int,int] - Tile width & Tile Height
+    min_color: Tuple[int,int,int] - RGB values for the lowest color to clip to (0-255)
+    max_color: Tuple[int,int,int] - RGB values for the highest color to clip to (0-255)
+    """
+
+    np_image = np.array(image)  # Convert image to np array
+    alpha = np_image[:, :, 3]  # Get the mask from the alpha channel of the image
+    non_transparent_pixels = np_image[alpha != 0, :3]  # List of non-transparent pixels
+
+    # Create color tiles to paste in the empty areas of the image
+    tile_width, tile_height = tile_shape
+
+    # Clip the range of colors in the image to a particular spectrum only
+    r_min, g_min, b_min, _ = min_color
+    r_max, g_max, b_max, _ = max_color
+    non_transparent_pixels[:, 0] = np.clip(non_transparent_pixels[:, 0], r_min, r_max)
+    non_transparent_pixels[:, 1] = np.clip(non_transparent_pixels[:, 1], g_min, g_max)
+    non_transparent_pixels[:, 2] = np.clip(non_transparent_pixels[:, 2], b_min, b_max)
+
+    tiles = []
+    for _ in range(256):
+        color = non_transparent_pixels[np.random.randint(len(non_transparent_pixels))]
+        tile = np.zeros((tile_height, tile_width, 3), dtype=np.uint8)
+        tile[:, :] = color
+        tiles.append(tile)
+
+    # Fill the transparent area with tiles
+    filled_image = np.zeros((image.height, image.width, 3), dtype=np.uint8)
+
+    for x in range(image.width):
+        for y in range(image.height):
+            tile = tiles[np.random.randint(len(tiles))]
+            try:
+                filled_image[
+                    y - (y % tile_height) : y - (y % tile_height) + tile_height,
+                    x - (x % tile_width) : x - (x % tile_width) + tile_width,
+                ] = tile
+            except ValueError:
+                # Need to handle edge cases - literally
+                pass
+
+    filled_image = Image.fromarray(filled_image)  # Convert the filled tiles image to PIL
+    image = Image.composite(
+        image, filled_image, image.split()[-1]
+    )  # Composite the original image on top of the filled tiles
+    return image

invokeai/backend/image_util/infill_methods/patchmatch.py

@@ -0,0 +1,67 @@
+"""
+This module defines a singleton object, "patchmatch" that
+wraps the actual patchmatch object. It respects the global
+"try_patchmatch" attribute, so that patchmatch loading can
+be suppressed or deferred
+"""
+
+import numpy as np
+from PIL import Image
+
+import invokeai.backend.util.logging as logger
+from invokeai.app.services.config.config_default import get_config
+
+
+class PatchMatch:
+    """
+    Thin class wrapper around the patchmatch function.
+    """
+
+    patch_match = None
+    tried_load: bool = False
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def _load_patch_match(cls):
+        if cls.tried_load:
+            return
+        if get_config().patchmatch:
+            from patchmatch import patch_match as pm
+
+            if pm.patchmatch_available:
+                logger.info("Patchmatch initialized")
+                cls.patch_match = pm
+            else:
+                logger.info("Patchmatch not loaded (nonfatal)")
+        else:
+            logger.info("Patchmatch loading disabled")
+        cls.tried_load = True
+
+    @classmethod
+    def patchmatch_available(cls) -> bool:
+        cls._load_patch_match()
+        if not cls.patch_match:
+            return False
+        return cls.patch_match.patchmatch_available
+
+    @classmethod
+    def inpaint(cls, image: Image.Image) -> Image.Image:
+        if cls.patch_match is None or not cls.patchmatch_available():
+            return image
+
+        np_image = np.array(image)
+        mask = 255 - np_image[:, :, 3]
+        infilled = cls.patch_match.inpaint(np_image[:, :, :3], mask, patch_size=3)
+        return Image.fromarray(infilled, mode="RGB")
+
+
+def infill_patchmatch(image: Image.Image) -> Image.Image:
+    IS_PATCHMATCH_AVAILABLE = PatchMatch.patchmatch_available()
+
+    if not IS_PATCHMATCH_AVAILABLE:
+        logger.warning("PatchMatch is not available on this system")
+        return image
+
+    return PatchMatch.inpaint(image)

invokeai/backend/image_util/infill_methods/tile.ipynb

@@ -0,0 +1,95 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Smoke test for the tile infill\"\"\"\n",
+    "\n",
+    "from pathlib import Path\n",
+    "from typing import Optional\n",
+    "from PIL import Image\n",
+    "from invokeai.backend.image_util.infill_methods.tile import infill_tile\n",
+    "\n",
+    "images: list[tuple[str, Image.Image]] = []\n",
+    "\n",
+    "for i in sorted(Path(\"./test_images/\").glob(\"*.webp\")):\n",
+    "    images.append((i.name, Image.open(i)))\n",
+    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_LEFT_RIGHT)))\n",
+    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_TOP_BOTTOM)))\n",
+    "    images.append((i.name, Image.open(i).resize((512, 512))))\n",
+    "    images.append((i.name, Image.open(i).resize((1234, 461))))\n",
+    "\n",
+    "outputs: list[tuple[str, Image.Image, Image.Image, Optional[Image.Image]]] = []\n",
+    "\n",
+    "for name, image in images:\n",
+    "    try:\n",
+    "        output = infill_tile(image, seed=0, tile_size=32)\n",
+    "        outputs.append((name, image, output.infilled, output.tile_image))\n",
+    "    except ValueError as e:\n",
+    "        print(f\"Skipping image {name}: {e}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Display the images in jupyter notebook\n",
+    "import matplotlib.pyplot as plt\n",
+    "from PIL import ImageOps\n",
+    "\n",
+    "fig, axes = plt.subplots(len(outputs), 3, figsize=(10, 3 * len(outputs)))\n",
+    "plt.subplots_adjust(hspace=0)\n",
+    "\n",
+    "for i, (name, original, infilled, tile_image) in enumerate(outputs):\n",
+    "    # Add a border to each image, helps to see the edges\n",
+    "    size = original.size\n",
+    "    original = ImageOps.expand(original, border=5, fill=\"red\")\n",
+    "    filled = ImageOps.expand(infilled, border=5, fill=\"red\")\n",
+    "    if tile_image:\n",
+    "        tile_image = ImageOps.expand(tile_image, border=5, fill=\"red\")\n",
+    "\n",
+    "    axes[i, 0].imshow(original)\n",
+    "    axes[i, 0].axis(\"off\")\n",
+    "    axes[i, 0].set_title(f\"Original ({name} - {size})\")\n",
+    "\n",
+    "    if tile_image:\n",
+    "        axes[i, 1].imshow(tile_image)\n",
+    "        axes[i, 1].axis(\"off\")\n",
+    "        axes[i, 1].set_title(\"Tile Image\")\n",
+    "    else:\n",
+    "        axes[i, 1].axis(\"off\")\n",
+    "        axes[i, 1].set_title(\"NO TILES GENERATED (NO TRANSPARENCY)\")\n",
+    "\n",
+    "    axes[i, 2].imshow(filled)\n",
+    "    axes[i, 2].axis(\"off\")\n",
+    "    axes[i, 2].set_title(\"Filled\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".invokeai",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

invokeai/backend/image_util/infill_methods/tile.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from PIL import Image
+
+
+def create_tile_pool(img_array: np.ndarray, tile_size: tuple[int, int]) -> list[np.ndarray]:
+    """
+    Create a pool of tiles from non-transparent areas of the image by systematically walking through the image.
+
+    Args:
+        img_array: numpy array of the image.
+        tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
+
+    Returns:
+        A list of numpy arrays, each representing a tile.
+    """
+    tiles: list[np.ndarray] = []
+    rows, cols = img_array.shape[:2]
+    tile_width, tile_height = tile_size
+
+    for y in range(0, rows - tile_height + 1, tile_height):
+        for x in range(0, cols - tile_width + 1, tile_width):
+            tile = img_array[y : y + tile_height, x : x + tile_width]
+            # Check if the image has an alpha channel and the tile is completely opaque
+            if img_array.shape[2] == 4 and np.all(tile[:, :, 3] == 255):
+                tiles.append(tile)
+            elif img_array.shape[2] == 3:  # If no alpha channel, append the tile
+                tiles.append(tile)
+
+    if not tiles:
+        raise ValueError(
+            "Not enough opaque pixels to generate any tiles. Use a smaller tile size or a different image."
+        )
+
+    return tiles
+
+
+def create_filled_image(
+    img_array: np.ndarray, tile_pool: list[np.ndarray], tile_size: tuple[int, int], seed: int
+) -> np.ndarray:
+    """
+    Create an image of the same dimensions as the original, filled entirely with tiles from the pool.
+
+    Args:
+        img_array: numpy array of the original image.
+        tile_pool: A list of numpy arrays, each representing a tile.
+        tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
+
+    Returns:
+        A numpy array representing the filled image.
+    """
+
+    rows, cols, _ = img_array.shape
+    tile_width, tile_height = tile_size
+
+    # Prep an empty RGB image
+    filled_img_array = np.zeros((rows, cols, 3), dtype=img_array.dtype)
+
+    # Make the random tile selection reproducible
+    rng = np.random.default_rng(seed)
+
+    for y in range(0, rows, tile_height):
+        for x in range(0, cols, tile_width):
+            # Pick a random tile from the pool
+            tile = tile_pool[rng.integers(len(tile_pool))]
+
+            # Calculate the space available (may be less than tile size near the edges)
+            space_y = min(tile_height, rows - y)
+            space_x = min(tile_width, cols - x)
+
+            # Crop the tile if necessary to fit into the available space
+            cropped_tile = tile[:space_y, :space_x, :3]
+
+            # Fill the available space with the (possibly cropped) tile
+            filled_img_array[y : y + space_y, x : x + space_x, :3] = cropped_tile
+
+    return filled_img_array
+
+
+@dataclass
+class InfillTileOutput:
+    infilled: Image.Image
+    tile_image: Optional[Image.Image] = None
+
+
+def infill_tile(image_to_infill: Image.Image, seed: int, tile_size: int) -> InfillTileOutput:
+    """Infills an image with random tiles from the image itself.
+
+    If the image is not an RGBA image, it is returned untouched.
+
+    Args:
+        image: The image to infill.
+        tile_size: The size of the tiles to use for infilling.
+
+    Raises:
+        ValueError: If there are not enough opaque pixels to generate any tiles.
+    """
+
+    if image_to_infill.mode != "RGBA":
+        return InfillTileOutput(infilled=image_to_infill)
+
+    # Internally, we want a tuple of (tile_width, tile_height). In the future, the tile size can be any rectangle.
+    _tile_size = (tile_size, tile_size)
+    np_image = np.array(image_to_infill, dtype=np.uint8)
+
+    # Create the pool of tiles that we will use to infill
+    tile_pool = create_tile_pool(np_image, _tile_size)
+
+    # Create an image from the tiles, same size as the original
+    tile_np_image = create_filled_image(np_image, tile_pool, _tile_size, seed)
+
+    # Paste the OG image over the tile image, effectively infilling the area
+    tile_image = Image.fromarray(tile_np_image, "RGB")
+    infilled = tile_image.copy()
+    infilled.paste(image_to_infill, (0, 0), image_to_infill.split()[-1])
+
+    # I think we want this to be "RGBA"?
+    infilled.convert("RGBA")
+
+    return InfillTileOutput(infilled=infilled, tile_image=tile_image)

invokeai/backend/image_util/invisible_watermark.py

@@ -0,0 +1,49 @@
+"""
+This module defines a singleton object, "invisible_watermark" that
+wraps the invisible watermark model. It respects the global "invisible_watermark"
+configuration variable, that allows the watermarking to be supressed.
+"""
+
+import cv2
+import numpy as np
+from PIL import Image
+
+import invokeai.backend.util.logging as logger
+from invokeai.backend.image_util.imwatermark.vendor import WatermarkDecoder, WatermarkEncoder
+
+
+class InvisibleWatermark:
+    """
+    Wrapper around InvisibleWatermark module.
+    """
+
+    @classmethod
+    def add_watermark(cls, image: Image.Image, watermark_text: str) -> Image.Image:
+        logger.debug(f'Applying invisible watermark "{watermark_text}"')
+        bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
+        encoder = WatermarkEncoder()
+        encoder.set_watermark("bytes", watermark_text.encode("utf-8"))
+        bgr_encoded = encoder.encode(bgr, "dwtDct")
+        return Image.fromarray(cv2.cvtColor(bgr_encoded, cv2.COLOR_BGR2RGB)).convert("RGBA")
+
+    @classmethod
+    def decode_watermark(cls, image: Image.Image, length: int = 8) -> str:
+        """Attempt to decode an invisible watermark from an image.
+
+        Args:
+            image: The PIL Image to decode the watermark from.
+            length: The expected watermark length in bytes. Must match the length used when encoding.
+                The WatermarkDecoder requires the length in bits; this value is multiplied by 8 internally.
+
+        Returns:
+            The decoded watermark text, or an empty string if no watermark is detected or decoding fails.
+        """
+        logger.debug("Attempting to decode invisible watermark")
+        try:
+            bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
+            decoder = WatermarkDecoder("bytes", length * 8)
+            watermark_bytes = decoder.decode(bgr, "dwtDct")
+            return watermark_bytes.decode("utf-8", errors="ignore").rstrip("\x00")
+        except Exception:
+            logger.debug("Failed to decode invisible watermark")
+            return ""

invokeai/backend/image_util/lineart.py

@@ -0,0 +1,228 @@
+"""Adapted from https://github.com/huggingface/controlnet_aux (Apache-2.0 license)."""
+
+import pathlib
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    normalize_image_channel_count,
+    np_to_pil,
+    pil_to_np,
+    resize_image_to_resolution,
+)
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+
+        conv_block = [
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+        ]
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+
+class Generator(nn.Module):
+    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
+        super(Generator, self).__init__()
+
+        # Initial convolution block
+        model0 = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, 7), nn.InstanceNorm2d(64), nn.ReLU(inplace=True)]
+        self.model0 = nn.Sequential(*model0)
+
+        # Downsampling
+        model1 = []
+        in_features = 64
+        out_features = in_features * 2
+        for _ in range(2):
+            model1 += [
+                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+            out_features = in_features * 2
+        self.model1 = nn.Sequential(*model1)
+
+        model2 = []
+        # Residual blocks
+        for _ in range(n_residual_blocks):
+            model2 += [ResidualBlock(in_features)]
+        self.model2 = nn.Sequential(*model2)
+
+        # Upsampling
+        model3 = []
+        out_features = in_features // 2
+        for _ in range(2):
+            model3 += [
+                nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+            out_features = in_features // 2
+        self.model3 = nn.Sequential(*model3)
+
+        # Output layer
+        model4 = [nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, 7)]
+        if sigmoid:
+            model4 += [nn.Sigmoid()]
+
+        self.model4 = nn.Sequential(*model4)
+
+    def forward(self, x, cond=None):
+        out = self.model0(x)
+        out = self.model1(out)
+        out = self.model2(out)
+        out = self.model3(out)
+        out = self.model4(out)
+
+        return out
+
+
+class LineartProcessor:
+    """Processor for lineart detection."""
+
+    def __init__(self):
+        model_path = hf_hub_download("lllyasviel/Annotators", "sk_model.pth")
+        self.model = Generator(3, 1, 3)
+        self.model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+        self.model.eval()
+
+        coarse_model_path = hf_hub_download("lllyasviel/Annotators", "sk_model2.pth")
+        self.model_coarse = Generator(3, 1, 3)
+        self.model_coarse.load_state_dict(torch.load(coarse_model_path, map_location=torch.device("cpu")))
+        self.model_coarse.eval()
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        self.model_coarse.to(device)
+        return self
+
+    def run(
+        self, input_image: Image.Image, coarse: bool = False, detect_resolution: int = 512, image_resolution: int = 512
+    ) -> Image.Image:
+        """Processes an image to detect lineart.
+
+        Args:
+            input_image: The input image.
+            coarse: Whether to use the coarse model.
+            detect_resolution: The resolution to fit the image to before edge detection.
+            image_resolution: The resolution of the output image.
+
+        Returns:
+            The detected lineart.
+        """
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(input_image)
+        np_image = normalize_image_channel_count(np_image)
+        np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+        model = self.model_coarse if coarse else self.model
+        assert np_image.ndim == 3
+        image = np_image
+        with torch.no_grad():
+            image = torch.from_numpy(image).float().to(device)
+            image = image / 255.0
+            image = rearrange(image, "h w c -> 1 c h w")
+            line = model(image)[0][0]
+
+            line = line.cpu().numpy()
+            line = (line * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = line
+
+        detected_map = normalize_image_channel_count(detected_map)
+
+        img = resize_image_to_resolution(np_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+        detected_map = 255 - detected_map
+
+        return np_to_pil(detected_map)
+
+
+class LineartEdgeDetector:
+    """Simple wrapper around the fine and coarse lineart models for detecting edges in an image."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename_fine = "sk_model.pth"
+    hf_filename_coarse = "sk_model2.pth"
+
+    @classmethod
+    def get_model_url(cls, coarse: bool = False) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        if coarse:
+            return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_coarse)
+        else:
+            return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename_fine)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> Generator:
+        """Load the model from a file."""
+        model = Generator(3, 1, 3)
+        model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        model.float().eval()
+        return model
+
+    def __init__(self, model: Generator) -> None:
+        self.model = model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image) -> Image.Image:
+        """Detects edges in the input image with the selected lineart model.
+
+        Args:
+            input: The input image.
+            coarse: Whether to use the coarse model.
+
+        Returns:
+            The detected edges.
+        """
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(image)
+
+        with torch.no_grad():
+            np_image = torch.from_numpy(np_image).float().to(device)
+            np_image = np_image / 255.0
+            np_image = rearrange(np_image, "h w c -> 1 c h w")
+            line = self.model(np_image)[0][0]
+
+            line = line.cpu().numpy()
+            line = (line * 255.0).clip(0, 255).astype(np.uint8)
+
+        detected_map = 255 - line
+
+        # The lineart model often outputs a lot of almost-black noise. SD1.5 ControlNets seem to be OK with this, but
+        # SDXL ControlNets are not - they need a cleaner map. 12 was experimentally determined to be a good threshold,
+        # eliminating all the noise while keeping the actual edges. Other approaches to thresholding may be better,
+        # for example stretching the contrast or removing noise.
+        detected_map[detected_map < 12] = 0
+
+        output = np_to_pil(detected_map)
+
+        return output

invokeai/backend/image_util/lineart_anime.py

@@ -0,0 +1,274 @@
+"""Adapted from https://github.com/huggingface/controlnet_aux (Apache-2.0 license)."""
+
+import functools
+import pathlib
+from typing import Optional
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from invokeai.backend.image_util.util import (
+    normalize_image_channel_count,
+    np_to_pil,
+    pil_to_np,
+    resize_image_to_resolution,
+)
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+class UnetGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(
+        self,
+        input_nc: int,
+        output_nc: int,
+        num_downs: int,
+        ngf: int = 64,
+        norm_layer=nn.BatchNorm2d,
+        use_dropout: bool = False,
+    ):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(UnetGenerator, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True
+        )  # add the innermost layer
+        for _ in range(num_downs - 5):  # add intermediate layers with ngf * 8 filters
+            unet_block = UnetSkipConnectionBlock(
+                ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout
+            )
+        # gradually reduce the number of filters from ngf * 8 to ngf
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+        )
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+        )
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        self.model = UnetSkipConnectionBlock(
+            output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer
+        )  # add the outermost layer
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class UnetSkipConnectionBlock(nn.Module):
+    """Defines the Unet submodule with skip connection.
+    X -------------------identity----------------------
+    |-- downsampling -- |submodule| -- upsampling --|
+    """
+
+    def __init__(
+        self,
+        outer_nc: int,
+        inner_nc: int,
+        input_nc: Optional[int] = None,
+        submodule=None,
+        outermost: bool = False,
+        innermost: bool = False,
+        norm_layer=nn.BatchNorm2d,
+        use_dropout: bool = False,
+    ):
+        """Construct a Unet submodule with skip connections.
+        Parameters:
+            outer_nc (int) -- the number of filters in the outer conv layer
+            inner_nc (int) -- the number of filters in the inner conv layer
+            input_nc (int) -- the number of channels in input images/features
+            submodule (UnetSkipConnectionBlock) -- previously defined submodules
+            outermost (bool)    -- if this module is the outermost module
+            innermost (bool)    -- if this module is the innermost module
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+        """
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if isinstance(norm_layer, functools.partial):
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:  # add skip connections
+            return torch.cat([x, self.model(x)], 1)
+
+
+class LineartAnimeProcessor:
+    """Processes an image to detect lineart."""
+
+    def __init__(self):
+        model_path = hf_hub_download("lllyasviel/Annotators", "netG.pth")
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+        self.model = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)
+        ckpt = torch.load(model_path)
+        for key in list(ckpt.keys()):
+            if "module." in key:
+                ckpt[key.replace("module.", "")] = ckpt[key]
+                del ckpt[key]
+        self.model.load_state_dict(ckpt)
+        self.model.eval()
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, input_image: Image.Image, detect_resolution: int = 512, image_resolution: int = 512) -> Image.Image:
+        """Processes an image to detect lineart.
+
+        Args:
+            input_image: The input image.
+            detect_resolution: The resolution to use for detection.
+            image_resolution: The resolution to use for the output image.
+
+        Returns:
+            The detected lineart.
+        """
+        device = get_effective_device(self.model)
+        np_image = pil_to_np(input_image)
+
+        np_image = normalize_image_channel_count(np_image)
+        np_image = resize_image_to_resolution(np_image, detect_resolution)
+
+        H, W, C = np_image.shape
+        Hn = 256 * int(np.ceil(float(H) / 256.0))
+        Wn = 256 * int(np.ceil(float(W) / 256.0))
+        img = cv2.resize(np_image, (Wn, Hn), interpolation=cv2.INTER_CUBIC)
+        with torch.no_grad():
+            image_feed = torch.from_numpy(img).float().to(device)
+            image_feed = image_feed / 127.5 - 1.0
+            image_feed = rearrange(image_feed, "h w c -> 1 c h w")
+
+            line = self.model(image_feed)[0, 0] * 127.5 + 127.5
+            line = line.cpu().numpy()
+
+            line = cv2.resize(line, (W, H), interpolation=cv2.INTER_CUBIC)
+            line = line.clip(0, 255).astype(np.uint8)
+
+        detected_map = line
+
+        detected_map = normalize_image_channel_count(detected_map)
+
+        img = resize_image_to_resolution(np_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
+        detected_map = 255 - detected_map
+
+        return np_to_pil(detected_map)
+
+
+class LineartAnimeEdgeDetector:
+    """Simple wrapper around the Lineart Anime model for detecting edges in an image."""
+
+    hf_repo_id = "lllyasviel/Annotators"
+    hf_filename = "netG.pth"
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> UnetGenerator:
+        """Load the model from a file."""
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+        model = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)
+        ckpt = torch.load(model_path)
+        for key in list(ckpt.keys()):
+            if "module." in key:
+                ckpt[key.replace("module.", "")] = ckpt[key]
+                del ckpt[key]
+        model.load_state_dict(ckpt)
+        model.eval()
+        return model
+
+    def __init__(self, model: UnetGenerator) -> None:
+        self.model = model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image) -> Image.Image:
+        """Processes an image and returns the detected edges."""
+        device = get_effective_device(self.model)
+
+        np_image = pil_to_np(image)
+
+        height, width, _channels = np_image.shape
+        new_height = 256 * int(np.ceil(float(height) / 256.0))
+        new_width = 256 * int(np.ceil(float(width) / 256.0))
+
+        resized_img = cv2.resize(np_image, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
+
+        with torch.no_grad():
+            image_feed = torch.from_numpy(resized_img).float().to(device)
+            image_feed = image_feed / 127.5 - 1.0
+            image_feed = rearrange(image_feed, "h w c -> 1 c h w")
+
+            line = self.model(image_feed)[0, 0] * 127.5 + 127.5
+            line = line.cpu().numpy()
+
+            line = cv2.resize(line, (width, height), interpolation=cv2.INTER_CUBIC)
+            line = line.clip(0, 255).astype(np.uint8)
+
+        detected_map = 255 - line
+
+        # The lineart model often outputs a lot of almost-black noise. SD1.5 ControlNets seem to be OK with this, but
+        # SDXL ControlNets are not - they need a cleaner map. 12 was experimentally determined to be a good threshold,
+        # eliminating all the noise while keeping the actual edges. Other approaches to thresholding may be better,
+        # for example stretching the contrast or removing noise.
+        detected_map[detected_map < 12] = 0
+
+        output = np_to_pil(detected_map)
+
+        return output

invokeai/backend/image_util/mediapipe_face/__init__.py

@@ -0,0 +1,15 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+from PIL import Image
+
+from invokeai.backend.image_util.mediapipe_face.mediapipe_face_common import generate_annotation
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np
+
+
+def detect_faces(image: Image.Image, max_faces: int = 1, min_confidence: float = 0.5) -> Image.Image:
+    """Detects faces in an image using MediaPipe."""
+
+    np_img = pil_to_np(image)
+    detected_map = generate_annotation(np_img, max_faces, min_confidence)
+    detected_map_pil = np_to_pil(detected_map)
+    return detected_map_pil

invokeai/backend/image_util/mediapipe_face/mediapipe_face_common.py

@@ -0,0 +1,149 @@
+from typing import Mapping
+
+import mediapipe as mp
+import numpy
+
+mp_drawing = mp.solutions.drawing_utils
+mp_drawing_styles = mp.solutions.drawing_styles
+mp_face_detection = mp.solutions.face_detection  # Only for counting faces.
+mp_face_mesh = mp.solutions.face_mesh
+mp_face_connections = mp.solutions.face_mesh_connections.FACEMESH_TESSELATION
+mp_hand_connections = mp.solutions.hands_connections.HAND_CONNECTIONS
+mp_body_connections = mp.solutions.pose_connections.POSE_CONNECTIONS
+
+DrawingSpec = mp.solutions.drawing_styles.DrawingSpec
+PoseLandmark = mp.solutions.drawing_styles.PoseLandmark
+
+min_face_size_pixels: int = 64
+f_thick = 2
+f_rad = 1
+right_iris_draw = DrawingSpec(color=(10, 200, 250), thickness=f_thick, circle_radius=f_rad)
+right_eye_draw = DrawingSpec(color=(10, 200, 180), thickness=f_thick, circle_radius=f_rad)
+right_eyebrow_draw = DrawingSpec(color=(10, 220, 180), thickness=f_thick, circle_radius=f_rad)
+left_iris_draw = DrawingSpec(color=(250, 200, 10), thickness=f_thick, circle_radius=f_rad)
+left_eye_draw = DrawingSpec(color=(180, 200, 10), thickness=f_thick, circle_radius=f_rad)
+left_eyebrow_draw = DrawingSpec(color=(180, 220, 10), thickness=f_thick, circle_radius=f_rad)
+mouth_draw = DrawingSpec(color=(10, 180, 10), thickness=f_thick, circle_radius=f_rad)
+head_draw = DrawingSpec(color=(10, 200, 10), thickness=f_thick, circle_radius=f_rad)
+
+# mp_face_mesh.FACEMESH_CONTOURS has all the items we care about.
+face_connection_spec = {}
+for edge in mp_face_mesh.FACEMESH_FACE_OVAL:
+    face_connection_spec[edge] = head_draw
+for edge in mp_face_mesh.FACEMESH_LEFT_EYE:
+    face_connection_spec[edge] = left_eye_draw
+for edge in mp_face_mesh.FACEMESH_LEFT_EYEBROW:
+    face_connection_spec[edge] = left_eyebrow_draw
+# for edge in mp_face_mesh.FACEMESH_LEFT_IRIS:
+#    face_connection_spec[edge] = left_iris_draw
+for edge in mp_face_mesh.FACEMESH_RIGHT_EYE:
+    face_connection_spec[edge] = right_eye_draw
+for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
+    face_connection_spec[edge] = right_eyebrow_draw
+# for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
+#    face_connection_spec[edge] = right_iris_draw
+for edge in mp_face_mesh.FACEMESH_LIPS:
+    face_connection_spec[edge] = mouth_draw
+iris_landmark_spec = {468: right_iris_draw, 473: left_iris_draw}
+
+
+def draw_pupils(image, landmark_list, drawing_spec, halfwidth: int = 2):
+    """We have a custom function to draw the pupils because the mp.draw_landmarks method requires a parameter for all
+    landmarks.  Until our PR is merged into mediapipe, we need this separate method."""
+    if len(image.shape) != 3:
+        raise ValueError("Input image must be H,W,C.")
+    image_rows, image_cols, image_channels = image.shape
+    if image_channels != 3:  # BGR channels
+        raise ValueError("Input image must contain three channel bgr data.")
+    for idx, landmark in enumerate(landmark_list.landmark):
+        if (landmark.HasField("visibility") and landmark.visibility < 0.9) or (
+            landmark.HasField("presence") and landmark.presence < 0.5
+        ):
+            continue
+        if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
+            continue
+        image_x = int(image_cols * landmark.x)
+        image_y = int(image_rows * landmark.y)
+        draw_color = None
+        if isinstance(drawing_spec, Mapping):
+            if drawing_spec.get(idx) is None:
+                continue
+            else:
+                draw_color = drawing_spec[idx].color
+        elif isinstance(drawing_spec, DrawingSpec):
+            draw_color = drawing_spec.color
+        image[image_y - halfwidth : image_y + halfwidth, image_x - halfwidth : image_x + halfwidth, :] = draw_color
+
+
+def reverse_channels(image):
+    """Given a numpy array in RGB form, convert to BGR.  Will also convert from BGR to RGB."""
+    # im[:,:,::-1] is a neat hack to convert BGR to RGB by reversing the indexing order.
+    # im[:,:,::[2,1,0]] would also work but makes a copy of the data.
+    return image[:, :, ::-1]
+
+
+def generate_annotation(img_rgb, max_faces: int, min_confidence: float):
+    """
+    Find up to 'max_faces' inside the provided input image.
+    If min_face_size_pixels is provided and nonzero it will be used to filter faces that occupy less than this many
+    pixels in the image.
+    """
+    with mp_face_mesh.FaceMesh(
+        static_image_mode=True,
+        max_num_faces=max_faces,
+        refine_landmarks=True,
+        min_detection_confidence=min_confidence,
+    ) as facemesh:
+        img_height, img_width, img_channels = img_rgb.shape
+        assert img_channels == 3
+
+        results = facemesh.process(img_rgb).multi_face_landmarks
+
+        if results is None:
+            print("No faces detected in controlnet image for Mediapipe face annotator.")
+            return numpy.zeros_like(img_rgb)
+
+        # Filter faces that are too small
+        filtered_landmarks = []
+        for lm in results:
+            landmarks = lm.landmark
+            face_rect = [
+                landmarks[0].x,
+                landmarks[0].y,
+                landmarks[0].x,
+                landmarks[0].y,
+            ]  # Left, up, right, down.
+            for i in range(len(landmarks)):
+                face_rect[0] = min(face_rect[0], landmarks[i].x)
+                face_rect[1] = min(face_rect[1], landmarks[i].y)
+                face_rect[2] = max(face_rect[2], landmarks[i].x)
+                face_rect[3] = max(face_rect[3], landmarks[i].y)
+            if min_face_size_pixels > 0:
+                face_width = abs(face_rect[2] - face_rect[0])
+                face_height = abs(face_rect[3] - face_rect[1])
+                face_width_pixels = face_width * img_width
+                face_height_pixels = face_height * img_height
+                face_size = min(face_width_pixels, face_height_pixels)
+                if face_size >= min_face_size_pixels:
+                    filtered_landmarks.append(lm)
+            else:
+                filtered_landmarks.append(lm)
+
+        # Annotations are drawn in BGR for some reason, but we don't need to flip a zero-filled image at the start.
+        empty = numpy.zeros_like(img_rgb)
+
+        # Draw detected faces:
+        for face_landmarks in filtered_landmarks:
+            mp_drawing.draw_landmarks(
+                empty,
+                face_landmarks,
+                connections=face_connection_spec.keys(),
+                landmark_drawing_spec=None,
+                connection_drawing_spec=face_connection_spec,
+            )
+            draw_pupils(empty, face_landmarks, iris_landmark_spec, 2)
+
+        # Flip BGR back to RGB.
+        empty = reverse_channels(empty).copy()
+
+        return empty

invokeai/backend/image_util/mlsd/__init__.py

@@ -0,0 +1,66 @@
+# Adapted from https://github.com/huggingface/controlnet_aux
+
+import pathlib
+
+import cv2
+import huggingface_hub
+import numpy as np
+import torch
+from PIL import Image
+
+from invokeai.backend.image_util.mlsd.models.mbv2_mlsd_large import MobileV2_MLSD_Large
+from invokeai.backend.image_util.mlsd.utils import pred_lines
+from invokeai.backend.image_util.util import np_to_pil, pil_to_np, resize_to_multiple
+
+
+class MLSDDetector:
+    """Simple wrapper around a MLSD model for detecting edges as line segments in an image."""
+
+    hf_repo_id = "lllyasviel/ControlNet"
+    hf_filename = "annotator/ckpts/mlsd_large_512_fp32.pth"
+
+    @classmethod
+    def get_model_url(cls) -> str:
+        """Get the URL to download the model from the Hugging Face Hub."""
+
+        return huggingface_hub.hf_hub_url(cls.hf_repo_id, cls.hf_filename)
+
+    @classmethod
+    def load_model(cls, model_path: pathlib.Path) -> MobileV2_MLSD_Large:
+        """Load the model from a file."""
+
+        model = MobileV2_MLSD_Large()
+        model.load_state_dict(torch.load(model_path), strict=True)
+        model.eval()
+        return model
+
+    def __init__(self, model: MobileV2_MLSD_Large) -> None:
+        self.model = model
+
+    def to(self, device: torch.device):
+        self.model.to(device)
+        return self
+
+    def run(self, image: Image.Image, score_threshold: float = 0.1, distance_threshold: float = 20.0) -> Image.Image:
+        """Processes an image and returns the detected edges."""
+
+        np_img = pil_to_np(image)
+
+        height, width, _channels = np_img.shape
+
+        # This model requires the input image to have a resolution that is a multiple of 64
+        np_img = resize_to_multiple(np_img, 64)
+        img_output = np.zeros_like(np_img)
+
+        with torch.no_grad():
+            lines = pred_lines(np_img, self.model, [np_img.shape[0], np_img.shape[1]], score_threshold, distance_threshold)
+            for line in lines:
+                x_start, y_start, x_end, y_end = [int(val) for val in line]
+                cv2.line(img_output, (x_start, y_start), (x_end, y_end), [255, 255, 255], 1)
+
+        detected_map = img_output[:, :, 0]
+
+        # Back to the original size
+        output_image = cv2.resize(detected_map, (width, height), interpolation=cv2.INTER_LINEAR)
+
+        return np_to_pil(output_image)

invokeai/backend/image_util/mlsd/models/mbv2_mlsd_large.py

@@ -0,0 +1,290 @@
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from torch.nn import functional as F
+
+
+class BlockTypeA(nn.Module):
+    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
+        super(BlockTypeA, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c2, out_c2, kernel_size=1),
+            nn.BatchNorm2d(out_c2),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c1, out_c1, kernel_size=1),
+            nn.BatchNorm2d(out_c1),
+            nn.ReLU(inplace=True)
+        )
+        self.upscale = upscale
+
+    def forward(self, a, b):
+        b = self.conv1(b)
+        a = self.conv2(a)
+        if self.upscale:
+             b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
+        return torch.cat((a, b), dim=1)
+
+
+class BlockTypeB(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeB, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_c),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        x = self.conv1(x) + x
+        x = self.conv2(x)
+        return x
+
+class BlockTypeC(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeC, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+        self.channel_pad = out_planes - in_planes
+        self.stride = stride
+        #padding = (kernel_size - 1) // 2
+
+        # TFLite uses slightly different padding than PyTorch
+        if stride == 2:
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+
+
+    def forward(self, x):
+        # TFLite uses  different padding
+        if self.stride == 2:
+            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
+            #print(x.shape)
+
+        for module in self:
+            if not isinstance(module, nn.MaxPool2d):
+                x = module(x)
+        return x
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, pretrained=True):
+        """
+        MobileNet V2 main class
+        Args:
+            num_classes (int): Number of classes
+            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
+            inverted_residual_setting: Network structure
+            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
+            Set to 1 to turn off rounding
+            block: Module specifying inverted residual building block for mobilenet
+        """
+        super(MobileNetV2, self).__init__()
+
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        width_mult = 1.0
+        round_nearest = 8
+
+        inverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            #[6, 160, 3, 2],
+            #[6, 320, 1, 1],
+        ]
+
+        # only check the first element, assuming user knows t,c,n,s are required
+        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
+            raise ValueError("inverted_residual_setting should be non-empty "
+                             "or a 4-element list, got {}".format(inverted_residual_setting))
+
+        # building first layer
+        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
+        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
+        features = [ConvBNReLU(4, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s in inverted_residual_setting:
+            output_channel = _make_divisible(c * width_mult, round_nearest)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+                input_channel = output_channel
+
+        self.features = nn.Sequential(*features)
+        self.fpn_selected = [1, 3, 6, 10, 13]
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+        if pretrained:
+           self._load_pretrained_model()
+
+    def _forward_impl(self, x):
+        # This exists since TorchScript doesn't support inheritance, so the superclass method
+        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
+        fpn_features = []
+        for i, f in enumerate(self.features):
+            if i > self.fpn_selected[-1]:
+                break
+            x = f(x)
+            if i in self.fpn_selected:
+                fpn_features.append(x)
+
+        c1, c2, c3, c4, c5 = fpn_features
+        return c1, c2, c3, c4, c5
+
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+class MobileV2_MLSD_Large(nn.Module):
+    def __init__(self):
+        super(MobileV2_MLSD_Large, self).__init__()
+
+        self.backbone = MobileNetV2(pretrained=False)
+        ## A, B
+        self.block15 = BlockTypeA(in_c1= 64, in_c2= 96,
+                                  out_c1= 64, out_c2=64,
+                                  upscale=False)
+        self.block16 = BlockTypeB(128, 64)
+
+        ## A, B
+        self.block17 = BlockTypeA(in_c1 = 32,  in_c2 = 64,
+                                  out_c1= 64,  out_c2= 64)
+        self.block18 = BlockTypeB(128, 64)
+
+        ## A, B
+        self.block19 = BlockTypeA(in_c1=24, in_c2=64,
+                                  out_c1=64, out_c2=64)
+        self.block20 = BlockTypeB(128, 64)
+
+        ## A, B, C
+        self.block21 = BlockTypeA(in_c1=16, in_c2=64,
+                                  out_c1=64, out_c2=64)
+        self.block22 = BlockTypeB(128, 64)
+
+        self.block23 = BlockTypeC(64, 16)
+
+    def forward(self, x):
+        c1, c2, c3, c4, c5 = self.backbone(x)
+
+        x = self.block15(c4, c5)
+        x = self.block16(x)
+
+        x = self.block17(c3, x)
+        x = self.block18(x)
+
+        x = self.block19(c2, x)
+        x = self.block20(x)
+
+        x = self.block21(c1, x)
+        x = self.block22(x)
+        x = self.block23(x)
+        x = x[:, 7:, :, :]
+
+        return x

invokeai/backend/image_util/mlsd/models/mbv2_mlsd_tiny.py

@@ -0,0 +1,273 @@
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from torch.nn import functional as F
+
+
+class BlockTypeA(nn.Module):
+    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
+        super(BlockTypeA, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c2, out_c2, kernel_size=1),
+            nn.BatchNorm2d(out_c2),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c1, out_c1, kernel_size=1),
+            nn.BatchNorm2d(out_c1),
+            nn.ReLU(inplace=True)
+        )
+        self.upscale = upscale
+
+    def forward(self, a, b):
+        b = self.conv1(b)
+        a = self.conv2(a)
+        b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
+        return torch.cat((a, b), dim=1)
+
+
+class BlockTypeB(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeB, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_c),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        x = self.conv1(x) + x
+        x = self.conv2(x)
+        return x
+
+class BlockTypeC(nn.Module):
+    def __init__(self, in_c, out_c):
+        super(BlockTypeC, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
+            nn.BatchNorm2d(in_c),
+            nn.ReLU()
+        )
+        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+        self.channel_pad = out_planes - in_planes
+        self.stride = stride
+        #padding = (kernel_size - 1) // 2
+
+        # TFLite uses slightly different padding than PyTorch
+        if stride == 2:
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+
+
+    def forward(self, x):
+        # TFLite uses  different padding
+        if self.stride == 2:
+            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
+            #print(x.shape)
+
+        for module in self:
+            if not isinstance(module, nn.MaxPool2d):
+                x = module(x)
+        return x
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, pretrained=True):
+        """
+        MobileNet V2 main class
+        Args:
+            num_classes (int): Number of classes
+            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
+            inverted_residual_setting: Network structure
+            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
+            Set to 1 to turn off rounding
+            block: Module specifying inverted residual building block for mobilenet
+        """
+        super(MobileNetV2, self).__init__()
+
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        width_mult = 1.0
+        round_nearest = 8
+
+        inverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            #[6, 96, 3, 1],
+            #[6, 160, 3, 2],
+            #[6, 320, 1, 1],
+        ]
+
+        # only check the first element, assuming user knows t,c,n,s are required
+        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
+            raise ValueError("inverted_residual_setting should be non-empty "
+                             "or a 4-element list, got {}".format(inverted_residual_setting))
+
+        # building first layer
+        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
+        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
+        features = [ConvBNReLU(4, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s in inverted_residual_setting:
+            output_channel = _make_divisible(c * width_mult, round_nearest)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+                input_channel = output_channel
+        self.features = nn.Sequential(*features)
+
+        self.fpn_selected = [3, 6, 10]
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+
+        #if pretrained:
+        #    self._load_pretrained_model()
+
+    def _forward_impl(self, x):
+        # This exists since TorchScript doesn't support inheritance, so the superclass method
+        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
+        fpn_features = []
+        for i, f in enumerate(self.features):
+            if i > self.fpn_selected[-1]:
+                break
+            x = f(x)
+            if i in self.fpn_selected:
+                fpn_features.append(x)
+
+        c2, c3, c4 = fpn_features
+        return c2, c3, c4
+
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+class MobileV2_MLSD_Tiny(nn.Module):
+    def __init__(self):
+        super(MobileV2_MLSD_Tiny, self).__init__()
+
+        self.backbone = MobileNetV2(pretrained=True)
+
+        self.block12 = BlockTypeA(in_c1= 32, in_c2= 64,
+                                  out_c1= 64, out_c2=64)
+        self.block13 = BlockTypeB(128, 64)
+
+        self.block14 = BlockTypeA(in_c1 = 24,  in_c2 = 64,
+                                  out_c1= 32,  out_c2= 32)
+        self.block15 = BlockTypeB(64, 64)
+
+        self.block16 = BlockTypeC(64, 16)
+
+    def forward(self, x):
+        c2, c3, c4 = self.backbone(x)
+
+        x = self.block12(c3, c4)
+        x = self.block13(x)
+        x = self.block14(c2, x)
+        x = self.block15(x)
+        x = self.block16(x)
+        x = x[:, 7:, :, :]
+        #print(x.shape)
+        x = F.interpolate(x, scale_factor=2.0, mode='bilinear', align_corners=True)
+
+        return x

invokeai/backend/image_util/mlsd/utils.py

@@ -0,0 +1,589 @@
+'''
+modified by  lihaoweicv
+pytorch version
+'''
+
+'''
+M-LSD
+Copyright 2021-present NAVER Corp.
+Apache License v2.0
+'''
+
+import cv2
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
+
+def deccode_output_score_and_ptss(tpMap, topk_n = 200, ksize = 5):
+    '''
+    tpMap:
+    center: tpMap[1, 0, :, :]
+    displacement: tpMap[1, 1:5, :, :]
+    '''
+    b, c, h, w = tpMap.shape
+    assert  b==1, 'only support bsize==1'
+    displacement = tpMap[:, 1:5, :, :][0]
+    center = tpMap[:, 0, :, :]
+    heat = torch.sigmoid(center)
+    hmax = F.max_pool2d( heat, (ksize, ksize), stride=1, padding=(ksize-1)//2)
+    keep = (hmax == heat).float()
+    heat = heat * keep
+    heat = heat.reshape(-1, )
+
+    scores, indices = torch.topk(heat, topk_n, dim=-1, largest=True)
+    yy = torch.floor_divide(indices, w).unsqueeze(-1)
+    xx = torch.fmod(indices, w).unsqueeze(-1)
+    ptss = torch.cat((yy, xx),dim=-1)
+
+    ptss   = ptss.detach().cpu().numpy()
+    scores = scores.detach().cpu().numpy()
+    displacement = displacement.detach().cpu().numpy()
+    displacement = displacement.transpose((1,2,0))
+    return  ptss, scores, displacement
+
+
+def pred_lines(image, model,
+               input_shape=[512, 512],
+               score_thr=0.10,
+               dist_thr=20.0):
+    h, w, _ = image.shape
+
+    device = get_effective_device(model)
+    h_ratio, w_ratio = [h / input_shape[0], w / input_shape[1]]
+
+    resized_image = np.concatenate([cv2.resize(image, (input_shape[1], input_shape[0]), interpolation=cv2.INTER_AREA),
+                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
+
+    resized_image = resized_image.transpose((2,0,1))
+    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
+    batch_image = (batch_image / 127.5) - 1.0
+
+    batch_image = torch.from_numpy(batch_image).float()
+    batch_image = batch_image.to(device)
+    outputs = model(batch_image)
+    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
+    start = vmap[:, :, :2]
+    end = vmap[:, :, 2:]
+    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
+
+    segments_list = []
+    for center, score in zip(pts, pts_score, strict=False):
+        y, x = center
+        distance = dist_map[y, x]
+        if score > score_thr and distance > dist_thr:
+            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
+            x_start = x + disp_x_start
+            y_start = y + disp_y_start
+            x_end = x + disp_x_end
+            y_end = y + disp_y_end
+            segments_list.append([x_start, y_start, x_end, y_end])
+
+    if segments_list:
+        lines = 2 * np.array(segments_list)  # 256 > 512
+        lines[:, 0] = lines[:, 0] * w_ratio
+        lines[:, 1] = lines[:, 1] * h_ratio
+        lines[:, 2] = lines[:, 2] * w_ratio
+        lines[:, 3] = lines[:, 3] * h_ratio
+    else:
+        # No segments detected - return empty array
+        lines = np.array([])
+
+    return lines
+
+
+def pred_squares(image,
+                 model,
+                 input_shape=[512, 512],
+                 params={'score': 0.06,
+                         'outside_ratio': 0.28,
+                         'inside_ratio': 0.45,
+                         'w_overlap': 0.0,
+                         'w_degree': 1.95,
+                         'w_length': 0.0,
+                         'w_area': 1.86,
+                         'w_center': 0.14}):
+    '''
+    shape = [height, width]
+    '''
+    h, w, _ = image.shape
+    original_shape = [h, w]
+    device = get_effective_device(model)
+
+    resized_image = np.concatenate([cv2.resize(image, (input_shape[0], input_shape[1]), interpolation=cv2.INTER_AREA),
+                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
+    resized_image = resized_image.transpose((2, 0, 1))
+    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
+    batch_image = (batch_image / 127.5) - 1.0
+
+    batch_image = torch.from_numpy(batch_image).float().to(device)
+    outputs = model(batch_image)
+
+    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
+    start = vmap[:, :, :2]  # (x, y)
+    end = vmap[:, :, 2:]  # (x, y)
+    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
+
+    junc_list = []
+    segments_list = []
+    for junc, score in zip(pts, pts_score, strict=False):
+        y, x = junc
+        distance = dist_map[y, x]
+        if score > params['score'] and distance > 20.0:
+            junc_list.append([x, y])
+            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
+            d_arrow = 1.0
+            x_start = x + d_arrow * disp_x_start
+            y_start = y + d_arrow * disp_y_start
+            x_end = x + d_arrow * disp_x_end
+            y_end = y + d_arrow * disp_y_end
+            segments_list.append([x_start, y_start, x_end, y_end])
+
+    segments = np.array(segments_list)
+
+    ####### post processing for squares
+    # 1. get unique lines
+    point = np.array([[0, 0]])
+    point = point[0]
+    start = segments[:, :2]
+    end = segments[:, 2:]
+    diff = start - end
+    a = diff[:, 1]
+    b = -diff[:, 0]
+    c = a * start[:, 0] + b * start[:, 1]
+
+    d = np.abs(a * point[0] + b * point[1] - c) / np.sqrt(a ** 2 + b ** 2 + 1e-10)
+    theta = np.arctan2(diff[:, 0], diff[:, 1]) * 180 / np.pi
+    theta[theta < 0.0] += 180
+    hough = np.concatenate([d[:, None], theta[:, None]], axis=-1)
+
+    d_quant = 1
+    theta_quant = 2
+    hough[:, 0] //= d_quant
+    hough[:, 1] //= theta_quant
+    _, indices, counts = np.unique(hough, axis=0, return_index=True, return_counts=True)
+
+    acc_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='float32')
+    idx_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='int32') - 1
+    yx_indices = hough[indices, :].astype('int32')
+    acc_map[yx_indices[:, 0], yx_indices[:, 1]] = counts
+    idx_map[yx_indices[:, 0], yx_indices[:, 1]] = indices
+
+    acc_map_np = acc_map
+    # acc_map = acc_map[None, :, :, None]
+    #
+    # ### fast suppression using tensorflow op
+    # acc_map = tf.constant(acc_map, dtype=tf.float32)
+    # max_acc_map = tf.keras.layers.MaxPool2D(pool_size=(5, 5), strides=1, padding='same')(acc_map)
+    # acc_map = acc_map * tf.cast(tf.math.equal(acc_map, max_acc_map), tf.float32)
+    # flatten_acc_map = tf.reshape(acc_map, [1, -1])
+    # topk_values, topk_indices = tf.math.top_k(flatten_acc_map, k=len(pts))
+    # _, h, w, _ = acc_map.shape
+    # y = tf.expand_dims(topk_indices // w, axis=-1)
+    # x = tf.expand_dims(topk_indices % w, axis=-1)
+    # yx = tf.concat([y, x], axis=-1)
+
+    ### fast suppression using pytorch op
+    acc_map = torch.from_numpy(acc_map_np).unsqueeze(0).unsqueeze(0)
+    _,_, h, w = acc_map.shape
+    max_acc_map = F.max_pool2d(acc_map,kernel_size=5, stride=1, padding=2)
+    acc_map = acc_map * ( (acc_map == max_acc_map).float() )
+    flatten_acc_map = acc_map.reshape([-1, ])
+
+    scores, indices = torch.topk(flatten_acc_map, len(pts), dim=-1, largest=True)
+    yy = torch.div(indices, w, rounding_mode='floor').unsqueeze(-1)
+    xx = torch.fmod(indices, w).unsqueeze(-1)
+    yx = torch.cat((yy, xx), dim=-1)
+
+    yx = yx.detach().cpu().numpy()
+
+    topk_values = scores.detach().cpu().numpy()
+    indices = idx_map[yx[:, 0], yx[:, 1]]
+    basis = 5 // 2
+
+    merged_segments = []
+    for yx_pt, max_indice, value in zip(yx, indices, topk_values, strict=False):
+        y, x = yx_pt
+        if max_indice == -1 or value == 0:
+            continue
+        segment_list = []
+        for y_offset in range(-basis, basis + 1):
+            for x_offset in range(-basis, basis + 1):
+                indice = idx_map[y + y_offset, x + x_offset]
+                cnt = int(acc_map_np[y + y_offset, x + x_offset])
+                if indice != -1:
+                    segment_list.append(segments[indice])
+                if cnt > 1:
+                    check_cnt = 1
+                    current_hough = hough[indice]
+                    for new_indice, new_hough in enumerate(hough):
+                        if (current_hough == new_hough).all() and indice != new_indice:
+                            segment_list.append(segments[new_indice])
+                            check_cnt += 1
+                        if check_cnt == cnt:
+                            break
+        group_segments = np.array(segment_list).reshape([-1, 2])
+        sorted_group_segments = np.sort(group_segments, axis=0)
+        x_min, y_min = sorted_group_segments[0, :]
+        x_max, y_max = sorted_group_segments[-1, :]
+
+        deg = theta[max_indice]
+        if deg >= 90:
+            merged_segments.append([x_min, y_max, x_max, y_min])
+        else:
+            merged_segments.append([x_min, y_min, x_max, y_max])
+
+    # 2. get intersections
+    new_segments = np.array(merged_segments)  # (x1, y1, x2, y2)
+    start = new_segments[:, :2]  # (x1, y1)
+    end = new_segments[:, 2:]  # (x2, y2)
+    new_centers = (start + end) / 2.0
+    diff = start - end
+    dist_segments = np.sqrt(np.sum(diff ** 2, axis=-1))
+
+    # ax + by = c
+    a = diff[:, 1]
+    b = -diff[:, 0]
+    c = a * start[:, 0] + b * start[:, 1]
+    pre_det = a[:, None] * b[None, :]
+    det = pre_det - np.transpose(pre_det)
+
+    pre_inter_y = a[:, None] * c[None, :]
+    inter_y = (pre_inter_y - np.transpose(pre_inter_y)) / (det + 1e-10)
+    pre_inter_x = c[:, None] * b[None, :]
+    inter_x = (pre_inter_x - np.transpose(pre_inter_x)) / (det + 1e-10)
+    inter_pts = np.concatenate([inter_x[:, :, None], inter_y[:, :, None]], axis=-1).astype('int32')
+
+    # 3. get corner information
+    # 3.1 get distance
+    '''
+    dist_segments:
+        | dist(0), dist(1), dist(2), ...|
+    dist_inter_to_segment1:
+        | dist(inter,0), dist(inter,0), dist(inter,0), ... |
+        | dist(inter,1), dist(inter,1), dist(inter,1), ... |
+        ...
+    dist_inter_to_semgnet2:
+        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
+        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
+        ...
+    '''
+
+    dist_inter_to_segment1_start = np.sqrt(
+        np.sum(((inter_pts - start[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment1_end = np.sqrt(
+        np.sum(((inter_pts - end[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment2_start = np.sqrt(
+        np.sum(((inter_pts - start[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+    dist_inter_to_segment2_end = np.sqrt(
+        np.sum(((inter_pts - end[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
+
+    # sort ascending
+    dist_inter_to_segment1 = np.sort(
+        np.concatenate([dist_inter_to_segment1_start, dist_inter_to_segment1_end], axis=-1),
+        axis=-1)  # [n_batch, n_batch, 2]
+    dist_inter_to_segment2 = np.sort(
+        np.concatenate([dist_inter_to_segment2_start, dist_inter_to_segment2_end], axis=-1),
+        axis=-1)  # [n_batch, n_batch, 2]
+
+    # 3.2 get degree
+    inter_to_start = new_centers[:, None, :] - inter_pts
+    deg_inter_to_start = np.arctan2(inter_to_start[:, :, 1], inter_to_start[:, :, 0]) * 180 / np.pi
+    deg_inter_to_start[deg_inter_to_start < 0.0] += 360
+    inter_to_end = new_centers[None, :, :] - inter_pts
+    deg_inter_to_end = np.arctan2(inter_to_end[:, :, 1], inter_to_end[:, :, 0]) * 180 / np.pi
+    deg_inter_to_end[deg_inter_to_end < 0.0] += 360
+
+    '''
+    B -- G
+    |    |
+    C -- R
+    B : blue / G: green / C: cyan / R: red
+
+    0 -- 1
+    |    |
+    3 -- 2
+    '''
+    # rename variables
+    deg1_map, deg2_map = deg_inter_to_start, deg_inter_to_end
+    # sort deg ascending
+    deg_sort = np.sort(np.concatenate([deg1_map[:, :, None], deg2_map[:, :, None]], axis=-1), axis=-1)
+
+    deg_diff_map = np.abs(deg1_map - deg2_map)
+    # we only consider the smallest degree of intersect
+    deg_diff_map[deg_diff_map > 180] = 360 - deg_diff_map[deg_diff_map > 180]
+
+    # define available degree range
+    deg_range = [60, 120]
+
+    corner_dict = {corner_info: [] for corner_info in range(4)}
+    inter_points = []
+    for i in range(inter_pts.shape[0]):
+        for j in range(i + 1, inter_pts.shape[1]):
+            # i, j > line index, always i < j
+            x, y = inter_pts[i, j, :]
+            deg1, deg2 = deg_sort[i, j, :]
+            deg_diff = deg_diff_map[i, j]
+
+            check_degree = deg_diff > deg_range[0] and deg_diff < deg_range[1]
+
+            outside_ratio = params['outside_ratio']  # over ratio >>> drop it!
+            inside_ratio = params['inside_ratio']  # over ratio >>> drop it!
+            check_distance = ((dist_inter_to_segment1[i, j, 1] >= dist_segments[i] and \
+                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * outside_ratio) or \
+                              (dist_inter_to_segment1[i, j, 1] <= dist_segments[i] and \
+                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * inside_ratio)) and \
+                             ((dist_inter_to_segment2[i, j, 1] >= dist_segments[j] and \
+                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * outside_ratio) or \
+                              (dist_inter_to_segment2[i, j, 1] <= dist_segments[j] and \
+                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * inside_ratio))
+
+            if check_degree and check_distance:
+                corner_info = None
+
+                if (deg1 >= 0 and deg1 <= 45 and deg2 >= 45 and deg2 <= 120) or \
+                        (deg2 >= 315 and deg1 >= 45 and deg1 <= 120):
+                    corner_info, color_info = 0, 'blue'
+                elif (deg1 >= 45 and deg1 <= 125 and deg2 >= 125 and deg2 <= 225):
+                    corner_info, color_info = 1, 'green'
+                elif (deg1 >= 125 and deg1 <= 225 and deg2 >= 225 and deg2 <= 315):
+                    corner_info, color_info = 2, 'black'
+                elif (deg1 >= 0 and deg1 <= 45 and deg2 >= 225 and deg2 <= 315) or \
+                        (deg2 >= 315 and deg1 >= 225 and deg1 <= 315):
+                    corner_info, color_info = 3, 'cyan'
+                else:
+                    corner_info, color_info = 4, 'red'  # we don't use it
+                    continue
+
+                corner_dict[corner_info].append([x, y, i, j])
+                inter_points.append([x, y])
+
+    square_list = []
+    connect_list = []
+    segments_list = []
+    for corner0 in corner_dict[0]:
+        for corner1 in corner_dict[1]:
+            connect01 = False
+            for corner0_line in corner0[2:]:
+                if corner0_line in corner1[2:]:
+                    connect01 = True
+                    break
+            if connect01:
+                for corner2 in corner_dict[2]:
+                    connect12 = False
+                    for corner1_line in corner1[2:]:
+                        if corner1_line in corner2[2:]:
+                            connect12 = True
+                            break
+                    if connect12:
+                        for corner3 in corner_dict[3]:
+                            connect23 = False
+                            for corner2_line in corner2[2:]:
+                                if corner2_line in corner3[2:]:
+                                    connect23 = True
+                                    break
+                            if connect23:
+                                for corner3_line in corner3[2:]:
+                                    if corner3_line in corner0[2:]:
+                                        # SQUARE!!!
+                                        '''
+                                        0 -- 1
+                                        |    |
+                                        3 -- 2
+                                        square_list:
+                                            order: 0 > 1 > 2 > 3
+                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
+                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
+                                            ...
+                                        connect_list:
+                                            order: 01 > 12 > 23 > 30
+                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
+                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
+                                            ...
+                                        segments_list:
+                                            order: 0 > 1 > 2 > 3
+                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
+                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
+                                            ...
+                                        '''
+                                        square_list.append(corner0[:2] + corner1[:2] + corner2[:2] + corner3[:2])
+                                        connect_list.append([corner0_line, corner1_line, corner2_line, corner3_line])
+                                        segments_list.append(corner0[2:] + corner1[2:] + corner2[2:] + corner3[2:])
+
+    def check_outside_inside(segments_info, connect_idx):
+        # return 'outside or inside', min distance, cover_param, peri_param
+        if connect_idx == segments_info[0]:
+            check_dist_mat = dist_inter_to_segment1
+        else:
+            check_dist_mat = dist_inter_to_segment2
+
+        i, j = segments_info
+        min_dist, max_dist = check_dist_mat[i, j, :]
+        connect_dist = dist_segments[connect_idx]
+        if max_dist > connect_dist:
+            return 'outside', min_dist, 0, 1
+        else:
+            return 'inside', min_dist, -1, -1
+
+    top_square = None
+
+    try:
+        map_size = input_shape[0] / 2
+        squares = np.array(square_list).reshape([-1, 4, 2])
+        score_array = []
+        connect_array = np.array(connect_list)
+        segments_array = np.array(segments_list).reshape([-1, 4, 2])
+
+        # get degree of corners:
+        squares_rollup = np.roll(squares, 1, axis=1)
+        squares_rolldown = np.roll(squares, -1, axis=1)
+        vec1 = squares_rollup - squares
+        normalized_vec1 = vec1 / (np.linalg.norm(vec1, axis=-1, keepdims=True) + 1e-10)
+        vec2 = squares_rolldown - squares
+        normalized_vec2 = vec2 / (np.linalg.norm(vec2, axis=-1, keepdims=True) + 1e-10)
+        inner_products = np.sum(normalized_vec1 * normalized_vec2, axis=-1)  # [n_squares, 4]
+        squares_degree = np.arccos(inner_products) * 180 / np.pi  # [n_squares, 4]
+
+        # get square score
+        overlap_scores = []
+        degree_scores = []
+        length_scores = []
+
+        for connects, segments, square, degree in zip(connect_array, segments_array, squares, squares_degree, strict=False):
+            '''
+            0 -- 1
+            |    |
+            3 -- 2
+
+            # segments: [4, 2]
+            # connects: [4]
+            '''
+
+            ###################################### OVERLAP SCORES
+            cover = 0
+            perimeter = 0
+            # check 0 > 1 > 2 > 3
+            square_length = []
+
+            for start_idx in range(4):
+                end_idx = (start_idx + 1) % 4
+
+                connect_idx = connects[start_idx]  # segment idx of segment01
+                start_segments = segments[start_idx]
+                end_segments = segments[end_idx]
+
+                start_point = square[start_idx]
+                end_point = square[end_idx]
+
+                # check whether outside or inside
+                start_position, start_min, start_cover_param, start_peri_param = check_outside_inside(start_segments,
+                                                                                                      connect_idx)
+                end_position, end_min, end_cover_param, end_peri_param = check_outside_inside(end_segments, connect_idx)
+
+                cover += dist_segments[connect_idx] + start_cover_param * start_min + end_cover_param * end_min
+                perimeter += dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min
+
+                square_length.append(
+                    dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min)
+
+            overlap_scores.append(cover / perimeter)
+            ######################################
+            ###################################### DEGREE SCORES
+            '''
+            deg0 vs deg2
+            deg1 vs deg3
+            '''
+            deg0, deg1, deg2, deg3 = degree
+            deg_ratio1 = deg0 / deg2
+            if deg_ratio1 > 1.0:
+                deg_ratio1 = 1 / deg_ratio1
+            deg_ratio2 = deg1 / deg3
+            if deg_ratio2 > 1.0:
+                deg_ratio2 = 1 / deg_ratio2
+            degree_scores.append((deg_ratio1 + deg_ratio2) / 2)
+            ######################################
+            ###################################### LENGTH SCORES
+            '''
+            len0 vs len2
+            len1 vs len3
+            '''
+            len0, len1, len2, len3 = square_length
+            len_ratio1 = len0 / len2 if len2 > len0 else len2 / len0
+            len_ratio2 = len1 / len3 if len3 > len1 else len3 / len1
+            length_scores.append((len_ratio1 + len_ratio2) / 2)
+
+            ######################################
+
+        overlap_scores = np.array(overlap_scores)
+        overlap_scores /= np.max(overlap_scores)
+
+        degree_scores = np.array(degree_scores)
+        # degree_scores /= np.max(degree_scores)
+
+        length_scores = np.array(length_scores)
+
+        ###################################### AREA SCORES
+        area_scores = np.reshape(squares, [-1, 4, 2])
+        area_x = area_scores[:, :, 0]
+        area_y = area_scores[:, :, 1]
+        correction = area_x[:, -1] * area_y[:, 0] - area_y[:, -1] * area_x[:, 0]
+        area_scores = np.sum(area_x[:, :-1] * area_y[:, 1:], axis=-1) - np.sum(area_y[:, :-1] * area_x[:, 1:], axis=-1)
+        area_scores = 0.5 * np.abs(area_scores + correction)
+        area_scores /= (map_size * map_size)  # np.max(area_scores)
+        ######################################
+
+        ###################################### CENTER SCORES
+        centers = np.array([[256 // 2, 256 // 2]], dtype='float32')  # [1, 2]
+        # squares: [n, 4, 2]
+        square_centers = np.mean(squares, axis=1)  # [n, 2]
+        center2center = np.sqrt(np.sum((centers - square_centers) ** 2))
+        center_scores = center2center / (map_size / np.sqrt(2.0))
+
+        '''
+        score_w = [overlap, degree, area, center, length]
+        '''
+        score_w = [0.0, 1.0, 10.0, 0.5, 1.0]
+        score_array = params['w_overlap'] * overlap_scores \
+                      + params['w_degree'] * degree_scores \
+                      + params['w_area'] * area_scores \
+                      - params['w_center'] * center_scores \
+                      + params['w_length'] * length_scores
+
+        best_square = []
+
+        sorted_idx = np.argsort(score_array)[::-1]
+        score_array = score_array[sorted_idx]
+        squares = squares[sorted_idx]
+
+    except Exception:
+        pass
+
+    '''return list
+    merged_lines, squares, scores
+    '''
+
+    try:
+        new_segments[:, 0] = new_segments[:, 0] * 2 / input_shape[1] * original_shape[1]
+        new_segments[:, 1] = new_segments[:, 1] * 2 / input_shape[0] * original_shape[0]
+        new_segments[:, 2] = new_segments[:, 2] * 2 / input_shape[1] * original_shape[1]
+        new_segments[:, 3] = new_segments[:, 3] * 2 / input_shape[0] * original_shape[0]
+    except Exception:
+        new_segments = []
+
+    try:
+        squares[:, :, 0] = squares[:, :, 0] * 2 / input_shape[1] * original_shape[1]
+        squares[:, :, 1] = squares[:, :, 1] * 2 / input_shape[0] * original_shape[0]
+    except Exception:
+        squares = []
+        score_array = []
+
+    try:
+        inter_points = np.array(inter_points)
+        inter_points[:, 0] = inter_points[:, 0] * 2 / input_shape[1] * original_shape[1]
+        inter_points[:, 1] = inter_points[:, 1] * 2 / input_shape[0] * original_shape[0]
+    except Exception:
+        inter_points = []
+
+    return new_segments, squares, score_array, inter_points