feat: ControlNet Union support for SDXL (#2284)

examples/common/common.cpp

@@ -1187,6 +1187,13 @@ ArgOptions SDGenerationParams::get_options() {
          "under --hires-upscalers-dir (default: Latent)",
          0,
          &hires_upscaler},
+        {"",
+         "--control-type",
+         "kind of the control image, one of [openpose, depth, scribble, canny, normal, segment, tile, repaint]. "
+         "Required by an SDXL Control Net Union model (tile and repaint only exist in its ProMax variant), ignored by any other Control Net. "
+         "Also accepted: hed, softedge, pidi, ted (scribble); lineart, anime_lineart, mlsd (canny); seg (segment); inpaint, outpaint (repaint)",
+         0,
+         &control_type},
         {"",
          "--extra-sample-args",
          "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
@@ -2249,6 +2256,7 @@ bool SDGenerationParams::from_json_str(
 
     load_if_exists("strength", strength);
     load_if_exists("control_strength", control_strength);
+    load_if_exists("control_type", control_type);
     load_if_exists("ip_adapter_strength", ip_adapter_strength);
     load_if_exists("moe_boundary", moe_boundary);
     load_if_exists("vace_strength", vace_strength);
@@ -2638,6 +2646,15 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s
         sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100);
     }
 
+    resolved_control_type = SD_CONTROL_TYPE_NONE;
+    if (!control_type.empty()) {
+        resolved_control_type = str_to_sd_control_type(control_type.c_str());
+        if (resolved_control_type == SD_CONTROL_TYPE_COUNT) {
+            LOG_ERROR("invalid control type '%s' (openpose, depth, scribble, canny, normal, segment, tile, repaint)", control_type.c_str());
+            return false;
+        }
+    }
+
     hires_upscaler_model_path.clear();
     if (hires_enabled) {
         if (hires_upscaler.empty()) {
@@ -2877,6 +2894,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.qwen_image_layers   = qwen_image_layers;
     params.control_image       = control_image.get();
     params.control_strength    = control_strength;
+    params.control_type        = resolved_control_type;
     params.ip_adapter_image    = ip_adapter_image.get();
     params.ip_adapter_strength = ip_adapter_strength;
     params.pm_params           = pm_params;
@@ -3082,6 +3100,7 @@ std::string SDGenerationParams::to_string() const {
         << "  vace_strength: " << vace_strength << ",\n"
         << "  strength: " << strength << ",\n"
         << "  control_strength: " << control_strength << ",\n"
+        << "  control_type: \"" << control_type << "\",\n"
         << "  seed: " << seed << ",\n"
         << "  upscale_repeats: " << upscale_repeats << ",\n"
         << "  upscale_tile_size: " << upscale_tile_size << ",\n"

examples/common/common.h

@@ -227,6 +227,7 @@ struct SDGenerationParams {
     std::string end_image_path;
     std::string mask_image_path;
     std::string control_image_path;
+    std::string control_type;
     std::string ip_adapter_image_path;
     std::vector<std::string> ref_image_paths;
     std::vector<std::string> ref_video_paths;
@@ -291,6 +292,7 @@ struct SDGenerationParams {
     std::string prompt_with_lora;  // for metadata record only
     std::vector<sd_lora_t> lora_vec;
     sd_hires_upscaler_t resolved_hires_upscaler;
+    sd_control_type_t resolved_control_type = SD_CONTROL_TYPE_NONE;
 
     // Owned execution payload.
     SDImageOwner init_image;

examples/server/api.md

@@ -936,6 +936,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH
 | `ref_image_args` | `string` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `control_type` | `string` |
 | `ip_adapter_strength` | `number` |
 | `hires` | `object` |
 | `hires.enabled` | `boolean` |
@@ -1116,6 +1117,7 @@ Example:
   "ref_image_args": "",
   "increase_ref_index": false,
   "control_strength": 0.9,
+  "control_type": "",
   "ip_adapter_strength": 1.0,
   "embed_image_metadata": true,
 
@@ -1241,6 +1243,7 @@ Top-level scalar fields:
 | `image_preprocess` | `string \| array<string>` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `control_type` | `string` |
 | `ip_adapter_strength` | `number` |
 | `embed_image_metadata` | `boolean` |
 | `preview` | `string` |

examples/server/routes_sdcpp.cpp

@@ -154,6 +154,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"image_preprocess", defaults.image_preprocess},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
+        {"control_type", defaults.control_type},
         {"ip_adapter_strength", defaults.ip_adapter_strength},
         {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers, defaults.custom_sigmas)},
         {"hires", make_hires_json(defaults)},

include/stable-diffusion.h

@@ -400,6 +400,19 @@ enum sd_hires_upscaler_t {
     SD_HIRES_UPSCALER_COUNT,
 };
 
+enum sd_control_type_t {
+    SD_CONTROL_TYPE_NONE,
+    SD_CONTROL_TYPE_OPENPOSE,
+    SD_CONTROL_TYPE_DEPTH,
+    SD_CONTROL_TYPE_SCRIBBLE,  // thick lines: scribble, hed, softedge, ted
+    SD_CONTROL_TYPE_CANNY,     // thin lines: canny, lineart, anime lineart, mlsd
+    SD_CONTROL_TYPE_NORMAL,
+    SD_CONTROL_TYPE_SEGMENT,
+    SD_CONTROL_TYPE_TILE,
+    SD_CONTROL_TYPE_REPAINT,  // inpaint, outpaint
+    SD_CONTROL_TYPE_COUNT,
+};
+
 typedef struct {
     bool enabled;
     enum sd_hires_upscaler_t upscaler;
@@ -444,6 +457,7 @@ typedef struct {
     bool circular_x;
     bool circular_y;
     sd_image_preprocess_params_t image_preprocess;
+    enum sd_control_type_t control_type;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -531,6 +545,10 @@ SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode);
 SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str);
 SD_API const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler);
 SD_API enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str);
+SD_API const char* sd_control_type_name(enum sd_control_type_t control_type);
+// Also accepts the names of the kinds a type stands for (hed, lineart, mlsd, ...);
+// returns SD_CONTROL_TYPE_COUNT when the name is unknown.
+SD_API enum sd_control_type_t str_to_sd_control_type(const char* str);
 
 SD_API void sd_cache_params_init(sd_cache_params_t* cache_params);
 SD_API void sd_hires_params_init(sd_hires_params_t* hires_params);

src/generation_info.cpp

@@ -127,8 +127,11 @@ namespace sd {
         auto info = common_info(p, "image");
         if (p.init_image.data)
             info["strength"] = p.strength;
-        if (p.control_image.data)
+        if (p.control_image.data) {
             info["control_strength"] = p.control_strength;
+            if (p.control_type != SD_CONTROL_TYPE_NONE)
+                info["control_type"] = sd_control_type_name(p.control_type);
+        }
         if (p.ip_adapter_image.data)
             info["ip_adapter_strength"] = p.ip_adapter_strength;
         if (p.pm_params.id_images && p.pm_params.id_images_count > 0) {
@@ -188,7 +191,7 @@ namespace sd {
         auto inpaint        = image_info(inpaint_params);
         inpaint["strength"] = inpaint_params.strength;
         // ADetailer clears these inputs before generating each crop.
-        for (const auto* key : {"control_strength", "pm_params", "pulid_params", "hires", "image_preprocess"})
+        for (const auto* key : {"control_strength", "control_type", "pm_params", "pulid_params", "hires", "image_preprocess"})
             inpaint.erase(key);
         sd_img_gen_params_t defaults;
         sd_img_gen_params_init(&defaults);

src/model/common/block.hpp

@@ -576,6 +576,33 @@ public:
     }
 };
 
+class ResidualAttentionBlock : public GGMLBlock {
+public:
+    ResidualAttentionBlock(int64_t channels, int64_t n_head) {
+        blocks["ln_1"]       = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
+        blocks["attn"]       = std::shared_ptr<GGMLBlock>(new MultiheadAttention(channels, n_head, true, true, true));
+        blocks["ln_2"]       = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
+        blocks["mlp.c_fc"]   = std::shared_ptr<GGMLBlock>(new Linear(channels, channels * 4));
+        blocks["mlp.c_proj"] = std::shared_ptr<GGMLBlock>(new Linear(channels * 4, channels));
+    }
+
+    // x: [N, n_token, channels]
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        auto ln_1       = std::dynamic_pointer_cast<LayerNorm>(blocks["ln_1"]);
+        auto attn       = std::dynamic_pointer_cast<MultiheadAttention>(blocks["attn"]);
+        auto ln_2       = std::dynamic_pointer_cast<LayerNorm>(blocks["ln_2"]);
+        auto mlp_c_fc   = std::dynamic_pointer_cast<Linear>(blocks["mlp.c_fc"]);
+        auto mlp_c_proj = std::dynamic_pointer_cast<Linear>(blocks["mlp.c_proj"]);
+
+        x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, ln_1->forward(ctx, x)));
+
+        auto h = mlp_c_fc->forward(ctx, ln_2->forward(ctx, x));
+        h      = ggml_ext_gelu_quick(ctx->ggml_ctx, h);
+        h      = mlp_c_proj->forward(ctx, h);
+        return ggml_add(ctx->ggml_ctx, x, h);
+    }
+};
+
 class AlphaBlender : public GGMLBlock {
 protected:
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {

src/model/diffusion/control.hpp

@@ -2,10 +2,47 @@
 #define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
 #include "model/common/block.hpp"
+#include "stable-diffusion.h"
 
 // Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
 #define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
 
+/*
+    =============================== ControlNet Union ================================
+    Reference: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/controlnets/controlnet_union.py
+
+*/
+struct ControlNetUnionConfig {
+    int num_control_types = 0;  // 0: not a union model
+    int num_fusion_layers = 0;
+
+    static constexpr int CONTROL_TYPE_EMBED_DIM = 256;
+    static constexpr int FUSION_HEADS           = 8;
+
+    static ControlNetUnionConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                     const std::string& prefix,
+                                                     SDVersion version = VERSION_SD1) {
+        ControlNetUnionConfig config;
+        if (!sd_version_is_sdxl(version)) {
+            return config;
+        }
+        const std::string p = prefix.empty() ? "" : prefix + ".";
+        auto embedding      = tensor_storage_map.find(p + "control_add_embedding.linear_1.weight");
+        if (embedding == tensor_storage_map.end()) {
+            return config;
+        }
+        // control_add_embedding takes CONTROL_TYPE_EMBED_DIM inputs per control type
+        config.num_control_types = static_cast<int>(embedding->second.logical_ne0() / CONTROL_TYPE_EMBED_DIM);
+        while (tensor_storage_map.find(p + "transformer_layes." + std::to_string(config.num_fusion_layers) + ".ln_1.weight") != tensor_storage_map.end()) {
+            config.num_fusion_layers++;
+        }
+        LOG_VERBOSE("controlnet union: num_control_types = %d, num_fusion_layers = %d",
+                    config.num_control_types,
+                    config.num_fusion_layers);
+        return config;
+    }
+};
+
 /*
     =================================== ControlNet ===================================
     Reference: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/cldm/cldm.py
@@ -13,6 +50,8 @@
 */
 class ControlNetBlock : public GGMLBlock {
 protected:
+    ControlNetUnionConfig config;
+
     SDVersion version = VERSION_SD1;
     // network hparams
     int in_channels                        = 4;
@@ -28,12 +67,19 @@ protected:
     int context_dim                        = 768;  // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
     bool use_linear_projection             = false;
 
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        if (is_union()) {
+            // [num_control_types, model_channels] in the checkpoint
+            params["task_embedding"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, model_channels, config.num_control_types);
+        }
+    }
+
 public:
     int model_channels  = 320;
     int adm_in_channels = 2816;  // only for VERSION_SDXL
 
-    ControlNetBlock(SDVersion version = VERSION_SD1)
-        : version(version) {
+    ControlNetBlock(SDVersion version = VERSION_SD1, ControlNetUnionConfig config = {})
+        : config(config), version(version) {
         if (sd_version_is_sd2(version)) {
             context_dim       = 1024;
             num_head_channels = 64;
@@ -162,6 +208,76 @@ public:
 
         // middle_block_out
         blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
+
+        if (is_union()) {
+            blocks["control_add_embedding.linear_1"] = std::shared_ptr<GGMLBlock>(new Linear(ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, time_embed_dim));
+            // nn.SiLU()
+            blocks["control_add_embedding.linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
+            for (int i = 0; i < config.num_fusion_layers; i++) {
+                blocks["transformer_layes." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualAttentionBlock(model_channels, ControlNetUnionConfig::FUSION_HEADS));
+            }
+            blocks["spatial_ch_projs"] = std::shared_ptr<GGMLBlock>(new Linear(model_channels, model_channels));
+        }
+    }
+
+    bool is_union() const {
+        return config.num_control_types > 0;
+    }
+
+    // control_types: [N*num_control_types], 1 for the types in use and 0 for the others
+    ggml_tensor* control_type_embedding(GGMLRunnerContext* ctx, ggml_tensor* control_types, int64_t batch) {
+        auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["control_add_embedding.linear_1"]);
+        auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["control_add_embedding.linear_2"]);
+
+        // [N*num_control_types, CONTROL_TYPE_EMBED_DIM]
+        auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM);
+        e      = ggml_reshape_2d(ctx->ggml_ctx, e, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, batch);
+        e      = linear_1->forward(ctx, e);
+        e      = ggml_silu_inplace(ctx->ggml_ctx, e);
+        return linear_2->forward(ctx, e);  // [N, time_embed_dim]
+    }
+
+    // [N, C, H, W] => [N, C]
+    static ggml_tensor* spatial_mean(ggml_context* ctx, ggml_tensor* x) {
+        auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
+        return ggml_reshape_2d(ctx, ggml_mean(ctx, flat), x->ne[2], x->ne[3]);
+    }
+
+    // sample: [N, C, H, W], hint: [N, C, H, W] or [1, C, H, W]
+    ggml_tensor* fuse_condition(GGMLRunnerContext* ctx,
+                                ggml_tensor* sample,
+                                ggml_tensor* hint,
+                                int control_type_index) {
+        auto c           = ctx->ggml_ctx;
+        int64_t channels = sample->ne[2];
+        int64_t batch    = sample->ne[3];
+
+        auto task_embedding = params["task_embedding"];  // [num_control_types, C]
+        auto task           = ggml_view_2d(c, task_embedding, channels, 1, task_embedding->nb[1], control_type_index * task_embedding->nb[1]);
+
+        auto hint_feat = ggml_add(c, spatial_mean(c, hint), task);  // [N, C]
+        if (hint_feat->ne[1] != batch) {
+            hint_feat = ggml_repeat(c, hint_feat, ggml_new_tensor_2d(c, GGML_TYPE_F32, channels, batch));
+        }
+        auto sample_feat = spatial_mean(c, sample);  // [N, C]
+
+        // as in the reference (nn.MultiheadAttention without batch_first) the attention is over
+        // the batch, and the 2 features (hint, sample) are its batch: [2, N, C]
+        auto x = ggml_concat(c,
+                             ggml_reshape_3d(c, hint_feat, channels, batch, 1),
+                             ggml_reshape_3d(c, sample_feat, channels, batch, 1),
+                             2);
+        for (int i = 0; i < config.num_fusion_layers; i++) {
+            auto layer = std::dynamic_pointer_cast<ResidualAttentionBlock>(blocks["transformer_layes." + std::to_string(i)]);
+            x          = layer->forward(ctx, x);
+        }
+
+        auto projection = std::dynamic_pointer_cast<Linear>(blocks["spatial_ch_projs"]);
+        auto alpha      = ggml_reshape_2d(c, ggml_ext_slice(c, x, 2, 0, 1), channels, batch);
+        alpha           = projection->forward(ctx, alpha);
+        alpha           = ggml_reshape_4d(c, alpha, 1, 1, channels, batch);
+
+        return ggml_add(c, ggml_add(c, sample, hint), alpha);
     }
 
     ggml_tensor* resblock_forward(std::string name,
@@ -204,7 +320,9 @@ public:
                                       ggml_tensor* guided_hint,
                                       ggml_tensor* timesteps,
                                       ggml_tensor* context,
-                                      ggml_tensor* y = nullptr) {
+                                      ggml_tensor* y             = nullptr,
+                                      ggml_tensor* control_types = nullptr,
+                                      int control_type_index     = 0) {
         // x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
         // timesteps: [N,]
         // context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
@@ -246,6 +364,11 @@ public:
             emb = ggml_add(ctx->ggml_ctx, emb, label_emb);  // [N, time_embed_dim]
         }
 
+        if (is_union()) {
+            GGML_ASSERT(control_types != nullptr);
+            emb = ggml_add(ctx->ggml_ctx, emb, control_type_embedding(ctx, control_types, x->ne[3]));
+        }
+
         std::vector<ggml_tensor*> outs;
 
         if (guided_hint == nullptr) {
@@ -257,7 +380,11 @@ public:
 
         // input block 0
         auto h = input_blocks_0_0->forward(ctx, x);
-        h      = ggml_add(ctx->ggml_ctx, h, guided_hint);
+        if (is_union()) {
+            h = fuse_condition(ctx, h, guided_hint, control_type_index);
+        } else {
+            h = ggml_add(ctx->ggml_ctx, h, guided_hint);
+        }
         outs.push_back(zero_convs_0->forward(ctx, h));
 
         // input block 1-11
@@ -308,9 +435,15 @@ public:
 
 struct ControlNet : public GGMLRunner {
     SDVersion version = VERSION_SD1;
+    ControlNetUnionConfig config;
     ControlNetBlock control_net;
     std::string weight_prefix;
 
+    // ControlNet Union: index of the control type of the control image (-1: none yet)
+    // and the 0/1 vector of the types in use that the graph reads
+    int control_type_index = -1;
+    sd::Tensor<float> control_type_tensor;
+
     std::vector<ggml_tensor*> control_outputs_ggml;
     ggml_tensor* guided_hint_output_ggml = nullptr;
     std::vector<sd::Tensor<float>> controls;
@@ -320,12 +453,46 @@ struct ControlNet : public GGMLRunner {
         return "controlnet.guided_hint";
     }
 
+    bool is_union() const {
+        return config.num_control_types > 0;
+    }
+
+    // set the kind of the control image of the next generations, false (after saying why) when it
+    // can't be used: a union model needs one of the types it was trained on
+    bool set_control_type(sd_control_type_t type) {
+        if (!is_union()) {
+            if (type != SD_CONTROL_TYPE_NONE) {
+                LOG_WARN("control type '%s' ignored: the ControlNet is not a union model", sd_control_type_name(type));
+            }
+            control_type_index = -1;
+            return true;
+        }
+        if (type == SD_CONTROL_TYPE_NONE) {
+            LOG_ERROR("this ControlNet is a union model and needs a control type (openpose, depth, scribble, canny, normal, segment%s)",
+                      config.num_control_types > 6 ? ", tile, repaint" : "");
+            return false;
+        }
+        // the types of a union model, in training order: the enum without NONE
+        const int index = static_cast<int>(type) - 1;
+        if (index < 0 || index >= config.num_control_types) {
+            LOG_ERROR("control type '%s' is not supported by this ControlNet Union (it has %d control types)",
+                      sd_control_type_name(type), config.num_control_types);
+            return false;
+        }
+        control_type_index = index;
+        return true;
+    }
+
     ControlNet(ggml_backend_t backend,
                const String2TensorStorage& tensor_storage_map      = {},
                SDVersion version                                   = VERSION_SD1,
                const std::string& prefix                           = "",
                std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-        : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix) {
+        : GGMLRunner(backend, weight_manager),
+          version(version),
+          config(ControlNetUnionConfig::detect_from_weights(tensor_storage_map, prefix, version)),
+          control_net(version, config),
+          weight_prefix(prefix) {
         control_net.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -376,6 +543,16 @@ struct ControlNet : public GGMLRunner {
             hint = make_input(hint_tensor);
         }
 
+        ggml_tensor* control_types = nullptr;
+        if (is_union()) {
+            const int64_t batch = x_tensor.dim() >= 4 ? x_tensor.shape()[3] : 1;
+            control_type_tensor = sd::zeros<float>({config.num_control_types * batch});
+            for (int64_t i = 0; i < batch; i++) {
+                control_type_tensor[i * config.num_control_types + control_type_index] = 1.f;
+            }
+            control_types = make_input(control_type_tensor);
+        }
+
         auto runner_ctx = get_context(gf);
 
         auto outs = control_net.forward(&runner_ctx,
@@ -384,7 +561,9 @@ struct ControlNet : public GGMLRunner {
                                         guided_hint_input,
                                         timesteps,
                                         context,
-                                        y);
+                                        y,
+                                        control_types,
+                                        control_type_index);
 
         if (guided_hint_input == nullptr && !outs.empty()) {
             guided_hint_output_ggml = outs[0];
@@ -413,6 +592,10 @@ struct ControlNet : public GGMLRunner {
         // timesteps: [N, ]
         // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
+        if (is_union() && control_type_index < 0) {
+            LOG_ERROR("ControlNet Union: no control type set");
+            return std::nullopt;
+        }
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x, hint, timesteps, context, y);
         };

src/model/diffusion/unet.hpp

@@ -684,7 +684,8 @@ public:
         }
         sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
         if (controls.size() > 0) {
-            auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true);
+            // not in place, it segfaults the Vulkan backend with SDXL ControlNets (seen on RADV)
+            auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, false);
             h       = ggml_add(ctx->ggml_ctx, h, cs);  // middle control
         }
         int control_offset = static_cast<int>(controls.size() - 2);
@@ -697,7 +698,8 @@ public:
                 hs.pop_back();
 
                 if (controls.size() > 0) {
-                    auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, true);
+                    // not in place, see the middle control
+                    auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, false);
                     h_skip  = ggml_add(ctx->ggml_ctx, h_skip, cs);  // control net condition
                     control_offset--;
                 }

src/name_conversion.cpp

@@ -1840,6 +1840,15 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
         if (sd_version_is_sdxl(version) && is_diffusers_controlnet_name(name)) {
             name = convert_diffusers_controlnet_to_original_sdxl(name);
         }
+        // controlnet union fusion: nn.MultiheadAttention has in_proj_weight/in_proj_bias,
+        // MultiheadAttention here has in_proj.weight/in_proj.bias
+        if (starts_with(name, "transformer_layes.")) {
+            const std::string from = "attn.in_proj_";
+            size_t pos             = name.find(from);
+            if (pos != std::string::npos) {
+                name.replace(pos, from.size(), "attn.in_proj.");
+            }
+        }
     }
 
     static const std::vector<std::pair<std::string, std::string>> generic_name_map = {

src/pipeline/diffusion_engine.cpp

@@ -2270,6 +2270,10 @@ static sd::Tensor<float> batch_two_condition_tensors(const sd::Tensor<float>& a,
     return sd::ops::concat(a, b, static_cast<size_t>(batch_dim));
 }
 
+bool StableDiffusionGGML::set_control_type(sd_control_type_t type) {
+    return control_net == nullptr || control_net->set_control_type(type);
+}
+
 void StableDiffusionGGML::compute_sample_controls(const sd::Tensor<float>& control_image,
                                                   const sd::Tensor<float>& noised_input,
                                                   const sd::Tensor<float>& timesteps_tensor,

src/pipeline/diffusion_engine.h

@@ -440,6 +440,10 @@ public:
                                 bool terminal_sigma_is_zero,
                                 int64_t* last_progress_us);
 
+    // set the kind of the control image of the next generations (ControlNet Union),
+    // false (after saying why) when the loaded ControlNet can't use it
+    bool set_control_type(sd_control_type_t type);
+
     void compute_sample_controls(const sd::Tensor<float>& control_image,
                                  const sd::Tensor<float>& noised_input,
                                  const sd::Tensor<float>& timesteps_tensor,

src/pipeline/image.cpp

@@ -253,6 +253,9 @@ namespace sd::pipeline {
 
         if (sd_img_gen_params->control_image.data != nullptr) {
             control_image_tensor = sd_image_to_tensor(sd_img_gen_params->control_image, request->width, request->height);
+            if (!sd->set_control_type(request->control_type)) {
+                return std::nullopt;
+            }
         }
 
         if (init_image_tensor.empty() || mask_image_tensor.empty()) {

src/pipeline/request.cpp

@@ -135,6 +135,7 @@ namespace sd::pipeline {
         shifted_timestep            = sd_img_gen_params->sample_params.shifted_timestep;
         strength                    = sd_img_gen_params->strength;
         control_strength            = sd_img_gen_params->control_strength;
+        control_type                = sd_img_gen_params->control_type;
         eta                         = sd_img_gen_params->sample_params.eta;
         has_ref_images              = sd_img_gen_params->ref_images_count > 0;
         guidance                    = sd_img_gen_params->sample_params.guidance;

src/pipeline/request.h

@@ -47,6 +47,7 @@ namespace sd::pipeline {
         int shifted_timestep                        = 0;
         float strength                              = 1.f;
         float control_strength                      = 0.f;
+        sd_control_type_t control_type              = SD_CONTROL_TYPE_NONE;
         float eta                                   = 0.f;
         sd_sample_params_t sample_params            = {};
         sd_sample_params_t high_noise_sample_params = {};

src/stable-diffusion.cpp

@@ -267,6 +267,57 @@ enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) {
     return SD_HIRES_UPSCALER_COUNT;
 }
 
+const char* control_type_to_str[] = {
+    "none",
+    "openpose",
+    "depth",
+    "scribble",
+    "canny",
+    "normal",
+    "segment",
+    "tile",
+    "repaint",
+};
+
+static_assert(SD_CONTROL_TYPE_COUNT == sizeof(control_type_to_str) / sizeof(control_type_to_str[0]),
+              "\nnumber of elements in control_type_to_str[] != SD_CONTROL_TYPE_COUNT");
+
+const char* sd_control_type_name(enum sd_control_type_t control_type) {
+    if (control_type >= SD_CONTROL_TYPE_NONE && control_type < SD_CONTROL_TYPE_COUNT) {
+        return control_type_to_str[control_type];
+    }
+    return NONE_STR;
+}
+
+enum sd_control_type_t str_to_sd_control_type(const char* str) {
+    for (int i = 0; i < SD_CONTROL_TYPE_COUNT; i++) {
+        if (!strcmp(str, control_type_to_str[i])) {
+            return (enum sd_control_type_t)i;
+        }
+    }
+    static const struct {
+        const char* alias;
+        enum sd_control_type_t type;
+    } aliases[] = {
+        {"hed", SD_CONTROL_TYPE_SCRIBBLE},
+        {"softedge", SD_CONTROL_TYPE_SCRIBBLE},
+        {"pidi", SD_CONTROL_TYPE_SCRIBBLE},
+        {"ted", SD_CONTROL_TYPE_SCRIBBLE},
+        {"lineart", SD_CONTROL_TYPE_CANNY},
+        {"anime_lineart", SD_CONTROL_TYPE_CANNY},
+        {"mlsd", SD_CONTROL_TYPE_CANNY},
+        {"seg", SD_CONTROL_TYPE_SEGMENT},
+        {"inpaint", SD_CONTROL_TYPE_REPAINT},
+        {"outpaint", SD_CONTROL_TYPE_REPAINT},
+    };
+    for (const auto& entry : aliases) {
+        if (!strcmp(str, entry.alias)) {
+            return entry.type;
+        }
+    }
+    return SD_CONTROL_TYPE_COUNT;
+}
+
 const char* sd_vae_format_name(enum sd_vae_format_t format) {
     switch (format) {
         case SD_VAE_FORMAT_AUTO:
@@ -561,6 +612,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "ref_images_count: %d\n"
              "ref_image_args: %s\n"
              "control_strength: %.2f\n"
+             "control_type: %s\n"
              "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
              "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n"
              "circular_x: %s\n"
@@ -579,6 +631,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->ref_images_count,
              SAFE_STR(sd_img_gen_params->ref_image_args),
              sd_img_gen_params->control_strength,
+             sd_control_type_name(sd_img_gen_params->control_type),
              sd_img_gen_params->pm_params.style_strength,
              sd_img_gen_params->pm_params.id_images_count,
              SAFE_STR(sd_img_gen_params->pm_params.id_embed_path),