Commit b29cad5 for stable-diffusion.cpp
commit b29cad540a1aa957c24c99209ce54682912f06e4
Author: SolicTous <52211516+SolicTous@users.noreply.github.com>
Date: Sun Oct 25 00:11:57 2026 +0500
feat: add control-net z-image (#2187)
diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp
index ee0964b..036c7a4 100644
--- a/src/model/diffusion/model.hpp
+++ b/src/model/diffusion/model.hpp
@@ -171,6 +171,11 @@ struct MingImageDiffusionExtra {
const sd::Tensor<float>* direct_context = nullptr;
};
+struct ZImageDiffusionExtra {
+ const sd::Tensor<float>* control_latent = nullptr;
+ float strength = 1.f;
+};
+
using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra,
SkipLayerDiffusionExtra,
@@ -186,7 +191,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
SenseNovaU1DiffusionExtra,
HunyuanVideoDiffusionExtra,
LLaDAImageDiffusionExtra,
- MingImageDiffusionExtra>;
+ MingImageDiffusionExtra,
+ ZImageDiffusionExtra>;
struct DiffusionParams {
const sd::Tensor<float>* x = nullptr;
diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp
index ed26440..b479507 100644
--- a/src/model/diffusion/z_image.hpp
+++ b/src/model/diffusion/z_image.hpp
@@ -23,9 +23,10 @@
#endif
namespace ZImage {
- constexpr int Z_IMAGE_GRAPH_SIZE = 20480;
- constexpr int ADALN_EMBED_DIM = 256;
- constexpr int SEQ_MULTI_OF = 32;
+ constexpr int Z_IMAGE_GRAPH_SIZE = 24576;
+ constexpr int ADALN_EMBED_DIM = 256;
+ constexpr int SEQ_MULTI_OF = 32;
+ constexpr int CONTROL_IN_CHANNELS = 33;
struct ZImageConfig {
int patch_size = 2;
@@ -391,6 +392,135 @@ namespace ZImage {
}
};
+ struct ZImageControlTransformerBlock : public GGMLBlock {
+ protected:
+ int block_id;
+ bool has_before_proj;
+ int64_t hidden_size;
+ int64_t head_dim;
+ int64_t num_heads;
+ int64_t num_kv_heads;
+ bool qk_norm;
+
+ public:
+ ZImageControlTransformerBlock(int block_id,
+ int64_t hidden_size,
+ int64_t head_dim,
+ int64_t num_heads,
+ int64_t num_kv_heads,
+ int64_t multiple_of,
+ float ffn_dim_multiplier,
+ float norm_eps,
+ bool qk_norm,
+ bool with_before_proj = false,
+ bool norm_elementwise_affine = true,
+ float qk_norm_eps = 1e-6f)
+ : block_id(block_id),
+ has_before_proj(with_before_proj),
+ hidden_size(hidden_size),
+ head_dim(head_dim),
+ num_heads(num_heads),
+ num_kv_heads(num_kv_heads),
+ qk_norm(qk_norm) {
+ blocks["attention.to_q"] = std::make_shared<Linear>(hidden_size, num_heads * head_dim, false);
+ blocks["attention.to_k"] = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
+ blocks["attention.to_v"] = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
+ blocks["attention.to_out.0"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false);
+ if (qk_norm) {
+ blocks["attention.norm_q"] = std::make_shared<RMSNorm>(head_dim, qk_norm_eps, norm_elementwise_affine);
+ blocks["attention.norm_k"] = std::make_shared<RMSNorm>(head_dim, qk_norm_eps, norm_elementwise_affine);
+ }
+ blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+ blocks["ffn_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+ blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+ blocks["ffn_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+ blocks["feed_forward"] = std::make_shared<FeedForward>(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier);
+ blocks["adaLN_modulation.0"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size);
+ if (with_before_proj) {
+ blocks["before_proj"] = std::make_shared<Linear>(hidden_size, hidden_size);
+ }
+ blocks["after_proj"] = std::make_shared<Linear>(hidden_size, hidden_size);
+ }
+
+ ggml_tensor* forward_attention(GGMLRunnerContext* ctx,
+ ggml_tensor* c,
+ ggml_tensor* pe) {
+ auto to_q = std::dynamic_pointer_cast<Linear>(blocks["attention.to_q"]);
+ auto to_k = std::dynamic_pointer_cast<Linear>(blocks["attention.to_k"]);
+ auto to_v = std::dynamic_pointer_cast<Linear>(blocks["attention.to_v"]);
+ auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["attention.to_out.0"]);
+
+ if (out_proj->has_quantized_weight()) {
+ out_proj->set_scale(1.f / 16.f);
+ }
+
+ int64_t n_token = c->ne[1];
+ int64_t N = c->ne[2];
+
+ if (sd_backend_is(ctx->backend, "ROCm")) {
+ out_proj->set_force_prec_f32(true);
+ to_q->set_force_prec_f32(true);
+ to_k->set_force_prec_f32(true);
+ to_v->set_force_prec_f32(true);
+ }
+
+ auto q = ggml_reshape_4d(ctx->ggml_ctx, to_q->forward(ctx, c), head_dim, num_heads, n_token, N);
+ auto k = ggml_reshape_4d(ctx->ggml_ctx, to_k->forward(ctx, c), head_dim, num_kv_heads, n_token, N);
+ auto v = ggml_reshape_4d(ctx->ggml_ctx, to_v->forward(ctx, c), head_dim, num_kv_heads, n_token, N);
+
+ if (qk_norm) {
+ q = std::dynamic_pointer_cast<RMSNorm>(blocks["attention.norm_q"])->forward(ctx, q);
+ k = std::dynamic_pointer_cast<RMSNorm>(blocks["attention.norm_k"])->forward(ctx, k);
+ }
+
+ auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.f / 128.f);
+ return out_proj->forward(ctx, out);
+ }
+
+ std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+ ggml_tensor* c,
+ ggml_tensor* pe,
+ ggml_tensor* skip,
+ ggml_tensor* t_emb) {
+ auto attention_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
+ auto ffn_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm1"]);
+ auto attention_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm2"]);
+ auto ffn_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm2"]);
+ auto feed_forward = std::dynamic_pointer_cast<FeedForward>(blocks["feed_forward"]);
+ auto adaLN_modulation_0 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
+ auto after_proj = std::dynamic_pointer_cast<Linear>(blocks["after_proj"]);
+
+ if (block_id == 0 && has_before_proj) {
+ GGML_ASSERT(skip != nullptr);
+ auto before_proj = std::dynamic_pointer_cast<Linear>(blocks["before_proj"]);
+ c = ggml_add(ctx->ggml_ctx, before_proj->forward(ctx, c), skip);
+ }
+
+ auto m = adaLN_modulation_0->forward(ctx, t_emb);
+ auto mods = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0);
+ auto scale_msa = mods[0];
+ auto gate_msa = mods[1];
+ auto scale_mlp = mods[2];
+ auto gate_mlp = mods[3];
+
+ auto residual = c;
+ c = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, c), scale_msa);
+ c = forward_attention(ctx, c, pe);
+ c = attention_norm2->forward(ctx, c);
+ c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_msa));
+ c = ggml_add(ctx->ggml_ctx, c, residual);
+
+ residual = c;
+ c = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, c), scale_mlp);
+ c = feed_forward->forward(ctx, c);
+ c = ffn_norm2->forward(ctx, c);
+ c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_mlp));
+ c = ggml_add(ctx->ggml_ctx, c, residual);
+
+ return {after_proj->forward(ctx, c), c};
+ }
+ };
+
struct FinalLayer : public GGMLBlock {
public:
FinalLayer(int64_t hidden_size,
@@ -423,15 +553,24 @@ namespace ZImage {
class ZImageModel : public GGMLBlock {
protected:
ZImageConfig config;
+ GGMLBlockMap control_blocks;
+ int control_num_layers = 0;
+ int control_num_refiner_layers = 0;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
+
+ for (auto& [name, block] : control_blocks) {
+ block->init(ctx, tensor_storage_map, name);
+ }
}
public:
+ bool control_enabled = false;
+
ZImageModel() = default;
- ZImageModel(ZImageConfig config)
+ ZImageModel(ZImageConfig config, const String2TensorStorage& tensor_storage_map = {})
: config(config) {
blocks["x_embedder"] = std::make_shared<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, 256);
@@ -484,13 +623,154 @@ namespace ZImage {
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
+
+ if (tensor_storage_map.find("control_all_x_embedder.2-1.weight") != tensor_storage_map.end()) {
+ control_enabled = true;
+
+ for (const auto& pair : tensor_storage_map) {
+ if (starts_with(pair.first, "control_layers.")) {
+ auto items = split_string(pair.first, '.');
+ if (items.size() > 1) {
+ control_num_layers = std::max(control_num_layers, atoi(items[1].c_str()) + 1);
+ }
+ } else if (starts_with(pair.first, "control_noise_refiner.")) {
+ auto items = split_string(pair.first, '.');
+ if (items.size() > 1) {
+ control_num_refiner_layers = std::max(control_num_refiner_layers, atoi(items[1].c_str()) + 1);
+ }
+ }
+ }
+
+ GGML_ASSERT(control_num_layers > 0);
+ GGML_ASSERT(control_num_refiner_layers == config.num_refiner_layers);
+
+ control_blocks["control_all_x_embedder.2-1"] = std::make_shared<Linear>(config.patch_size * config.patch_size * CONTROL_IN_CHANNELS, config.hidden_size);
+
+ for (int i = 0; i < control_num_refiner_layers; i++) {
+ std::string name = "control_noise_refiner." + std::to_string(i);
+ auto block = std::make_shared<ZImageControlTransformerBlock>(i,
+ config.hidden_size,
+ config.head_dim,
+ config.num_heads,
+ config.num_kv_heads,
+ config.multiple_of,
+ config.ffn_dim_multiplier,
+ config.norm_eps,
+ config.qk_norm,
+ tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end());
+
+ control_blocks[name] = block;
+ }
+
+ for (int i = 0; i < control_num_layers; i++) {
+ std::string name = "control_layers." + std::to_string(i);
+ auto block = std::make_shared<ZImageControlTransformerBlock>(i,
+ config.hidden_size,
+ config.head_dim,
+ config.num_heads,
+ config.num_kv_heads,
+ config.multiple_of,
+ config.ffn_dim_multiplier,
+ config.norm_eps,
+ config.qk_norm,
+ tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end());
+
+ control_blocks[name] = block;
+ }
+ }
+ }
+
+ void get_control_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+ for (auto& [name, block] : control_blocks) {
+ block->get_param_tensors(tensors, name);
+ }
+ }
+
+ ggml_tensor* forward_control_refiner(GGMLRunnerContext* ctx,
+ ggml_tensor* control,
+ ggml_tensor* img,
+ ggml_tensor* img_pe,
+ ggml_tensor* t_emb,
+ float strength,
+ std::vector<ggml_tensor*>& ref_hints) {
+ auto all_x_embedder = std::dynamic_pointer_cast<Linear>(control_blocks["control_all_x_embedder.2-1"]);
+
+ auto c = all_x_embedder->forward(ctx, control);
+ sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.prelude", "control");
+
+ int64_t N = c->ne[2];
+ int64_t n_token = c->ne[1];
+
+ int64_t n_pad_token = Rope::bound_mod(static_cast<int>(n_token), SEQ_MULTI_OF);
+ if (n_pad_token > 0) {
+ auto img_pad_token = params["x_pad_token"];
+ auto ctrl_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_pad_token, N, 1);
+ c = ggml_concat(ctx->ggml_ctx, c, ctrl_pad_tokens, 1);
+ }
+
+ for (int i = 0; i < control_num_refiner_layers; i++) {
+ auto block = std::dynamic_pointer_cast<ZImageControlTransformerBlock>(control_blocks["control_noise_refiner." + std::to_string(i)]);
+
+ auto out = block->forward(ctx, c, img_pe, img, t_emb);
+ c = out.second;
+ ref_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength));
+ sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.noise_refiner." + std::to_string(i), "control");
+ }
+
+ return c;
+ }
+
+ ggml_tensor* forward_noise_refiner_hinted(GGMLRunnerContext* ctx,
+ ggml_tensor* img,
+ ggml_tensor* img_pe,
+ ggml_tensor* t_emb,
+ const std::vector<ggml_tensor*>& ref_hints) {
+ GGML_ASSERT(ref_hints.size() == static_cast<size_t>(config.num_refiner_layers));
+
+ auto c = img;
+ for (int i = 0; i < config.num_refiner_layers; i++) {
+ auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
+
+ c = block->forward(ctx, c, img_pe, nullptr, t_emb);
+ c = ggml_add(ctx->ggml_ctx, c, ref_hints[i]);
+ sd::ggml_graph_cut::mark_graph_cut(c, "z_image.noise_refiner_hinted." + std::to_string(i), "img");
+ }
+
+ return c;
+ }
+
+ std::vector<ggml_tensor*> forward_control_main(GGMLRunnerContext* ctx,
+ ggml_tensor* control_ctx,
+ ggml_tensor* img_c,
+ ggml_tensor* txt,
+ ggml_tensor* pe,
+ ggml_tensor* t_emb,
+ float strength) {
+ auto unified_c = ggml_concat(ctx->ggml_ctx, txt, img_c, 1);
+ auto ctrl_unified = ggml_concat(ctx->ggml_ctx, txt, control_ctx, 1);
+
+ std::vector<ggml_tensor*> main_hints;
+ main_hints.reserve(control_num_layers);
+
+ for (int i = 0; i < control_num_layers; i++) {
+ auto block = std::dynamic_pointer_cast<ZImageControlTransformerBlock>(control_blocks["control_layers." + std::to_string(i)]);
+
+ auto out = block->forward(ctx, ctrl_unified, pe, unified_c, t_emb);
+ ctrl_unified = out.second;
+ main_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength));
+ sd::ggml_graph_cut::mark_graph_cut(ctrl_unified, "z_image.control.layers." + std::to_string(i), "control");
+ }
+
+ return main_hints;
}
ggml_tensor* forward_core(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
- ggml_tensor* pe) {
+ ggml_tensor* pe,
+ ggml_tensor* control = nullptr,
+ float strength = 0.f) {
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
@@ -530,6 +810,14 @@ namespace ZImage {
auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
+ std::vector<ggml_tensor*> ref_hints;
+ ggml_tensor* control_ctx = nullptr;
+ if (control != nullptr) {
+ control_ctx = forward_control_refiner(ctx, control, img, img_pe, t_emb, strength, ref_hints);
+ }
+
+ ggml_tensor* img_pre_refiner = img;
+
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
@@ -544,13 +832,27 @@ namespace ZImage {
sd::ggml_graph_cut::mark_graph_cut(img, "z_image.noise_refiner." + std::to_string(i), "img");
}
+ ggml_tensor* img_c = nullptr;
+ if (control_ctx != nullptr) {
+ img_c = forward_noise_refiner_hinted(ctx, img_pre_refiner, img_pe, t_emb, ref_hints);
+ }
+
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img");
+ std::vector<ggml_tensor*> main_hints;
+ if (control_ctx != nullptr) {
+ main_hints = forward_control_main(ctx, control_ctx, img_c, txt, pe, t_emb, strength);
+ }
+ GGML_ASSERT(main_hints.empty() || main_hints.size() == static_cast<size_t>(control_num_layers));
+
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
+ if (i % 2 == 0 && (i / 2) < static_cast<int>(main_hints.size())) {
+ txt_img = ggml_add(ctx->ggml_ctx, txt_img, main_hints[i / 2]);
+ }
sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.layers." + std::to_string(i), "txt_img");
}
@@ -566,7 +868,9 @@ namespace ZImage {
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
- std::vector<ggml_tensor*> ref_latents = {}) {
+ std::vector<ggml_tensor*> ref_latents = {},
+ ggml_tensor* control_latent = nullptr,
+ float control_strength = 0.f) {
// Forward pass of DiT.
// x: [N, C, H, W]
// timestep: [N,]
@@ -591,7 +895,16 @@ namespace ZImage {
}
}
- auto out = forward_core(ctx, img, timestep, context, pe);
+ ggml_tensor* control = nullptr;
+ if (control_latent != nullptr) {
+ GGML_ASSERT(control_enabled);
+ GGML_ASSERT(control_latent->ne[2] == CONTROL_IN_CHANNELS);
+ GGML_ASSERT(control_latent->ne[0] == W && control_latent->ne[1] == H && control_latent->ne[3] == N);
+ GGML_ASSERT(ref_latents.empty());
+ control = DiT::pad_and_patchify(ctx, control_latent, patch_size, patch_size, false);
+ }
+
+ auto out = forward_core(ctx, img, timestep, context, pe, control, control_strength);
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token); // [N, n_img_token, ph*pw*C]
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false); // [N, C, H, W]
@@ -617,7 +930,7 @@ namespace ZImage {
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
- z_image = ZImageModel(config);
+ z_image = ZImageModel(config, tensor_storage_map);
z_image.init(params_ctx, tensor_storage_map, prefix);
}
@@ -629,11 +942,21 @@ namespace ZImage {
z_image.get_param_tensors(tensors, prefix);
}
+ bool control_enabled() const {
+ return z_image.control_enabled;
+ }
+
+ void get_control_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+ z_image.get_control_param_tensors(tensors);
+ }
+
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
- const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
- Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
+ const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
+ Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
+ const sd::Tensor<float>& control_latent_tensor = {},
+ float control_strength = 0.f) {
ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -646,6 +969,8 @@ namespace ZImage {
ref_latents.push_back(make_input(ref_latent_tensor));
}
+ ggml_tensor* control_latent = make_optional_input(control_latent_tensor);
+
pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
@@ -670,7 +995,9 @@ namespace ZImage {
timesteps,
context,
pe,
- ref_latents);
+ ref_latents,
+ control_latent,
+ control_strength);
ggml_build_forward_expand(gf, out);
@@ -682,12 +1009,14 @@ namespace ZImage {
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {},
- Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
+ Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
+ const sd::Tensor<float>& control_latent = {},
+ float control_strength = 0.f) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* {
- return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
+ return build_graph(x, timesteps, context, ref_latents, ref_index_mode, control_latent, control_strength);
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
@@ -698,12 +1027,21 @@ namespace ZImage {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
static const std::vector<sd::Tensor<float>> empty_ref_latents;
+ static const sd::Tensor<float> empty_control_latent;
+
+ const auto* z_image_extra = std::get_if<ZImageDiffusionExtra>(&diffusion_params.extra);
+ const sd::Tensor<float>& control_latent = (z_image_extra != nullptr && z_image_extra->control_latent != nullptr) ? *z_image_extra->control_latent
+ : empty_control_latent;
+ float control_strength = z_image_extra != nullptr ? z_image_extra->strength : 0.f;
+
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
- diffusion_params.ref_image_params.ref_index_mode);
+ diffusion_params.ref_image_params.ref_index_mode,
+ control_latent,
+ control_strength);
}
void test() {
diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp
index b0dc84b..6a1ff35 100644
--- a/src/pipeline/diffusion_engine.cpp
+++ b/src/pipeline/diffusion_engine.cpp
@@ -42,6 +42,7 @@
#include "model/diffusion/animatediff.hpp"
#include "model/diffusion/control.hpp"
#include "model/diffusion/model.hpp"
+#include "model/diffusion/z_image.hpp"
#include "model/vae/audio_vae.hpp"
#include "model/vae/ltx_vae.hpp"
#include "model/vae/vae.hpp"
@@ -234,6 +235,15 @@ bool StableDiffusionGGML::apply_model_update(ModelLoader candidate,
}
}
}
+ if (sd_version_is_z_image(version) && sd_version_is_z_image(next_version)) {
+ if (groups.count(RunnerGroup::ControlNet) > 0) {
+ groups.insert(RunnerGroup::Core);
+ }
+ if (groups.count(RunnerGroup::Core) > 0 &&
+ model_manager->registered_params_size({ModelComponent::ControlNet}) > 0) {
+ groups.insert(RunnerGroup::ControlNet);
+ }
+ }
runner_state_.ready = false;
if (!reset_runners(groups)) {
return false;
@@ -1116,6 +1126,27 @@ bool StableDiffusionGGML::build_control_net_runner() {
if (config_->control_net_file == 0) {
return true;
}
+ if (sd_version_is_z_image(version)) {
+ auto z_image = std::dynamic_pointer_cast<ZImage::ZImageRunner>(diffusion_model);
+ if (z_image == nullptr || !z_image->control_enabled()) {
+ LOG_WARN("z-image control net file present but the diffusion model has no control weights; ignoring");
+ return true;
+ }
+ std::map<std::string, ggml_tensor*> control_tensors;
+ z_image->get_control_param_tensors(control_tensors);
+ if (control_tensors.empty()) {
+ LOG_ERROR("z-image control net has no param tensors to register");
+ return false;
+ }
+ ModelManager::ResidencyMode residency_mode = backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION)
+ ? ModelManager::ResidencyMode::Disk
+ : ModelManager::ResidencyMode::ParamBackend;
+ return model_manager->register_param_tensors(ModelComponent::ControlNet,
+ std::move(control_tensors),
+ residency_mode,
+ backend_for(SDBackendModule::DIFFUSION),
+ params_backend_for(SDBackendModule::DIFFUSION));
+ }
if (!sd::model_builders::build_control_net_runner(model_build_context(), control_net)) {
return false;
}
@@ -2568,6 +2599,10 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
} else if (sd_version_is_sensenova_u1(version)) {
diffusion_params.extra = SenseNovaU1DiffusionExtra{
condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids};
+ } else if (sd_version_is_z_image(version)) {
+ if (!control_image.empty()) {
+ diffusion_params.extra = ZImageDiffusionExtra{&control_image, control_strength};
+ }
} else {
diffusion_params.extra = std::monostate{};
}
diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp
index 38ca95c..046e9cf 100644
--- a/src/pipeline/image.cpp
+++ b/src/pipeline/image.cpp
@@ -117,6 +117,32 @@ namespace sd::pipeline {
channels);
}
+ static sd::Tensor<float> encode_z_image_control_latent(StableDiffusionGGML* sd, const sd::Tensor<float>& control_rgb) {
+ sd::Tensor<float> vae_output = sd->first_stage_model->encode(sd->n_threads,
+ control_rgb,
+ sd->vae_tiling_params,
+ sd->circular_x,
+ sd->circular_y);
+ if (vae_output.empty()) {
+ LOG_ERROR("failed to encode control image");
+ return {};
+ }
+ // TAE encode already returns the latent; KL/Flux VAE returns a mean/logvar pair.
+ sd::Tensor<float> ctrl_latent = sd->first_stage_model->get_desc() == "taesd"
+ ? vae_output
+ : sd::ops::chunk(vae_output, 2, 2)[0];
+ ctrl_latent = sd->first_stage_model->vae_to_diffusion_latents(ctrl_latent);
+ if (ctrl_latent.empty() || ctrl_latent.dim() != 4 || ctrl_latent.shape()[2] > 33) {
+ LOG_ERROR("failed to encode control image");
+ return {};
+ }
+ std::vector<int64_t> padded_shape = ctrl_latent.shape();
+ padded_shape[2] = 33;
+ sd::Tensor<float> padded(padded_shape);
+ sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent);
+ return padded;
+ }
+
static std::optional<ImageGenerationLatents> prepare_image_generation_latents(StableDiffusionGGML* sd,
const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request,
@@ -250,10 +276,22 @@ namespace sd::pipeline {
}
if (!control_image_tensor.empty()) {
- control_latent = sd->encode_first_stage(control_image_tensor);
- if (control_latent.empty()) {
- LOG_ERROR("failed to encode control image");
- return std::nullopt;
+ const bool z_image_control = sd_version_is_z_image(sd->version) &&
+ sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0;
+ if (z_image_control) {
+ control_image_tensor = encode_z_image_control_latent(sd, control_image_tensor);
+ if (control_image_tensor.empty()) {
+ return std::nullopt;
+ }
+ } else if (sd_version_is_z_image(sd->version)) {
+ LOG_WARN("control image provided but the z-image model has no control net; ignoring");
+ control_image_tensor = {};
+ } else {
+ control_latent = sd->encode_first_stage(control_image_tensor);
+ if (control_latent.empty()) {
+ LOG_ERROR("failed to encode control image");
+ return std::nullopt;
+ }
}
}
@@ -969,6 +1007,29 @@ namespace sd::pipeline {
std::vector<sd::Tensor<float>> hires_final_latents;
int64_t hires_denoise_start = ggml_time_ms();
+ sd::Tensor<float> hires_control_image;
+ if (!latents.control_image.empty() &&
+ sd_version_is_z_image(sd->version) &&
+ sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0) {
+ const int64_t hires_latent_w = request.hires.target_width / request.vae_scale_factor;
+ const int64_t hires_latent_h = request.hires.target_height / request.vae_scale_factor;
+ if (latents.control_image.shape()[0] != hires_latent_w ||
+ latents.control_image.shape()[1] != hires_latent_h) {
+ LOG_INFO("hires fix: re-encoding control latent for %dx%d", request.hires.target_width, request.hires.target_height);
+ hires_control_image = encode_z_image_control_latent(sd,
+ sd_image_to_tensor(sd_img_gen_params->control_image,
+ request.hires.target_width,
+ request.hires.target_height));
+ if (hires_control_image.empty()) {
+ std::vector<int64_t> ctrl_shape = latents.control_image.shape();
+ ctrl_shape[0] = hires_latent_w;
+ ctrl_shape[1] = hires_latent_h;
+ hires_control_image = sd::ops::interpolate(latents.control_image, ctrl_shape, sd::ops::InterpolateMode::Bicubic);
+ }
+ } else {
+ hires_control_image = latents.control_image;
+ }
+ }
for (int b = 0; b < (int)final_latents.size(); b++) {
if (sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling generation during hires fix");
@@ -1006,7 +1067,7 @@ namespace sd::pipeline {
embeds.cond,
embeds.uncond,
embeds.img_uncond,
- latents.control_image,
+ !hires_control_image.empty() ? hires_control_image : latents.control_image,
request.control_strength,
request.guidance,
plan.eta,
diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp
index b630a72..8b85f9b 100644
--- a/src/pipeline/model_builders.cpp
+++ b/src/pipeline/model_builders.cpp
@@ -247,6 +247,9 @@ namespace sd::model_builders {
const auto& tensor_storage_map = ctx.tensor_storage_map;
const auto version = ctx.version;
const auto& weight_manager = ctx.weight_manager;
+ if (sd_version_is_z_image(version)) {
+ return true;
+ }
if (!ensure_backend_pair(ctx.backends, SDBackendModule::CONTROL_NET)) {
return false;
}
diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp
index 3d59181..00c5e8d 100644
--- a/src/stable-diffusion.cpp
+++ b/src/stable-diffusion.cpp
@@ -726,7 +726,13 @@ SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) {
if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
return false;
}
- return sd_ctx->sd->control_net != nullptr;
+ if (sd_ctx->sd->control_net != nullptr) {
+ return true;
+ }
+ // z-image keeps control net weights inside the diffusion runner, registered
+ // under ModelComponent::ControlNet instead of a separate control_net runner.
+ return sd_ctx->sd->model_manager != nullptr &&
+ sd_ctx->sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0;
}
const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx) {