Commit fdfc913 for stable-diffusion.cpp
commit fdfc913c515ad749cedd84a118e7c23a05ac1c65
Author: Daniele <57776841+daniandtheweb@users.noreply.github.com>
Date: Sat Nov 7 15:33:41 2026 +0000
feat: ControlNet Union support for SDXL (#2284)
diff --git a/examples/common/common.cpp b/examples/common/common.cpp
index d5496f6..02c3088 100644
--- a/examples/common/common.cpp
+++ b/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"
diff --git a/examples/common/common.h b/examples/common/common.h
index 39d6011..4ecbef4 100644
--- a/examples/common/common.h
+++ b/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;
diff --git a/examples/server/api.md b/examples/server/api.md
index 97cbd66..9c8d8c8 100644
--- a/examples/server/api.md
+++ b/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` |
diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp
index 0dda9a9..1d11b89 100644
--- a/examples/server/routes_sdcpp.cpp
+++ b/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)},
diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h
index f4ea0ea..1c61e4b 100644
--- a/include/stable-diffusion.h
+++ b/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);
diff --git a/src/generation_info.cpp b/src/generation_info.cpp
index 2f49054..cf40bc2 100644
--- a/src/generation_info.cpp
+++ b/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);
diff --git a/src/model/common/block.hpp b/src/model/common/block.hpp
index 66180fa..5be333d 100644
--- a/src/model/common/block.hpp
+++ b/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 {
diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp
index 3b08ed1..7848892 100644
--- a/src/model/diffusion/control.hpp
+++ b/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);
};
diff --git a/src/model/diffusion/unet.hpp b/src/model/diffusion/unet.hpp
index 675d942..837163f 100644
--- a/src/model/diffusion/unet.hpp
+++ b/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--;
}
diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp
index 2e14734..1f19b11 100644
--- a/src/name_conversion.cpp
+++ b/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 = {
diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp
index 4875858..a8eed8e 100644
--- a/src/pipeline/diffusion_engine.cpp
+++ b/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,
diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h
index e067781..97e3aca 100644
--- a/src/pipeline/diffusion_engine.h
+++ b/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,
diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp
index 0763e38..03a6707 100644
--- a/src/pipeline/image.cpp
+++ b/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()) {
diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp
index d333fbb..739b8e4 100644
--- a/src/pipeline/request.cpp
+++ b/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;
diff --git a/src/pipeline/request.h b/src/pipeline/request.h
index 1cf1cee..b1ab640 100644
--- a/src/pipeline/request.h
+++ b/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 = {};
diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp
index bca4b79..436328d 100644
--- a/src/stable-diffusion.cpp
+++ b/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),