From 9893126cdeaa492b3bfb710ec99cfcdf0520fff8 Mon Sep 17 00:00:00 2001 From: SolicTous Date: Wed, 21 Oct 2026 10:59:58 +0500 Subject: [PATCH 1/2] feat: add control-net z-image --- src/model/diffusion/model.hpp | 8 +- src/model/diffusion/z_image.hpp | 366 ++++++++++++++++++++++++++++-- src/pipeline/diffusion_engine.cpp | 35 +++ src/pipeline/image.cpp | 36 ++- src/pipeline/model_builders.cpp | 3 + 5 files changed, 429 insertions(+), 19 deletions(-) diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 955c5277a..cf3d23797 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -149,6 +149,11 @@ struct MingImageDiffusionExtra { const sd::Tensor* direct_context = nullptr; }; +struct ZImageDiffusionExtra { + const sd::Tensor* control_latent = nullptr; + float strength = 1.f; +}; + using DiffusionExtraParams = std::variant; + MingImageDiffusionExtra, + ZImageDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index ed2644047..b47950731 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(hidden_size, num_heads * head_dim, false); + blocks["attention.to_k"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["attention.to_v"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["attention.to_out.0"] = std::make_shared(num_heads * head_dim, hidden_size, false); + if (qk_norm) { + blocks["attention.norm_q"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + blocks["attention.norm_k"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + } + blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["attention_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["feed_forward"] = std::make_shared(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier); + blocks["adaLN_modulation.0"] = std::make_shared(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size); + if (with_before_proj) { + blocks["before_proj"] = std::make_shared(hidden_size, hidden_size); + } + blocks["after_proj"] = std::make_shared(hidden_size, hidden_size); + } + + ggml_tensor* forward_attention(GGMLRunnerContext* ctx, + ggml_tensor* c, + ggml_tensor* pe) { + auto to_q = std::dynamic_pointer_cast(blocks["attention.to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["attention.to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["attention.to_v"]); + auto out_proj = std::dynamic_pointer_cast(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(blocks["attention.norm_q"])->forward(ctx, q); + k = std::dynamic_pointer_cast(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 forward(GGMLRunnerContext* ctx, + ggml_tensor* c, + ggml_tensor* pe, + ggml_tensor* skip, + ggml_tensor* t_emb) { + auto attention_norm1 = std::dynamic_pointer_cast(blocks["attention_norm1"]); + auto ffn_norm1 = std::dynamic_pointer_cast(blocks["ffn_norm1"]); + auto attention_norm2 = std::dynamic_pointer_cast(blocks["attention_norm2"]); + auto ffn_norm2 = std::dynamic_pointer_cast(blocks["ffn_norm2"]); + auto feed_forward = std::dynamic_pointer_cast(blocks["feed_forward"]); + auto adaLN_modulation_0 = std::dynamic_pointer_cast(blocks["adaLN_modulation.0"]); + auto after_proj = std::dynamic_pointer_cast(blocks["after_proj"]); + + if (block_id == 0 && has_before_proj) { + GGML_ASSERT(skip != nullptr); + auto before_proj = std::dynamic_pointer_cast(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(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); blocks["t_embedder"] = std::make_shared(MIN(config.hidden_size, 1024), 256, 256); @@ -484,13 +623,154 @@ namespace ZImage { } blocks["final_layer"] = std::make_shared(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(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(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(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& 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& ref_hints) { + auto all_x_embedder = std::dynamic_pointer_cast(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(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(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& ref_hints) { + GGML_ASSERT(ref_hints.size() == static_cast(config.num_refiner_layers)); + + auto c = img; + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(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 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 main_hints; + main_hints.reserve(control_num_layers); + + for (int i = 0; i < control_num_layers; i++) { + auto block = std::dynamic_pointer_cast(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(blocks["x_embedder"]); auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); auto cap_embedder_0 = std::dynamic_pointer_cast(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 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(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 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(control_num_layers)); + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(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(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 ref_latents = {}) { + std::vector 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 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& tensors) { + z_image.get_control_param_tensors(tensors); + } + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, - const std::vector>& ref_latents_tensor = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + const std::vector>& ref_latents_tensor = {}, + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + const sd::Tensor& 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(x->ne[1]), static_cast(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& timesteps, const sd::Tensor& context, const std::vector>& ref_latents = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + const sd::Tensor& 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> empty_ref_latents; + static const sd::Tensor empty_control_latent; + + const auto* z_image_extra = std::get_if(&diffusion_params.extra); + const sd::Tensor& 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 22ef3ac45..ed7df036c 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -41,6 +41,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" @@ -296,6 +297,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; @@ -1178,6 +1188,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(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 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; } @@ -2710,6 +2741,10 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptrencode_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) { + sd::Tensor vae_output = sd->first_stage_model->encode(sd->n_threads, + control_image_tensor, + sd->vae_tiling_params, + sd->circular_x, + sd->circular_y); + if (vae_output.empty()) { + LOG_ERROR("failed to encode control image"); + return std::nullopt; + } + sd::Tensor ctrl_latent = 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::nullopt; + } + std::vector padded_shape = ctrl_latent.shape(); + padded_shape[2] = 33; + sd::Tensor padded(padded_shape); + sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent); + control_image_tensor = std::move(padded); + } 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; + } } } diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 4b435170f..9163d3fb1 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -694,6 +694,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; } From 8cedc64f17697caa750fa54c97934f206f859491 Mon Sep 17 00:00:00 2001 From: SolicTous Date: Fri, 23 Oct 2026 17:39:56 +0500 Subject: [PATCH 2/2] fix: control-net z-image integration (hires fix, TAE encode, has_control_net) --- src/pipeline/image.cpp | 71 +++++++++++++++++++++++++++++----------- src/stable-diffusion.cpp | 8 ++++- 2 files changed, 59 insertions(+), 20 deletions(-) diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index c28205d0e..492f30143 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -117,6 +117,32 @@ namespace sd::pipeline { channels); } + static sd::Tensor encode_z_image_control_latent(StableDiffusionGGML* sd, const sd::Tensor& control_rgb) { + sd::Tensor 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 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 padded_shape = ctrl_latent.shape(); + padded_shape[2] = 33; + sd::Tensor padded(padded_shape); + sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent); + return padded; + } + static std::optional prepare_image_generation_latents(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, GenerationRequest* request, @@ -253,26 +279,10 @@ namespace sd::pipeline { 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) { - sd::Tensor vae_output = sd->first_stage_model->encode(sd->n_threads, - control_image_tensor, - sd->vae_tiling_params, - sd->circular_x, - sd->circular_y); - if (vae_output.empty()) { - LOG_ERROR("failed to encode control image"); + control_image_tensor = encode_z_image_control_latent(sd, control_image_tensor); + if (control_image_tensor.empty()) { return std::nullopt; } - sd::Tensor ctrl_latent = 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::nullopt; - } - std::vector padded_shape = ctrl_latent.shape(); - padded_shape[2] = 33; - sd::Tensor padded(padded_shape); - sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent); - control_image_tensor = std::move(padded); } 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 = {}; @@ -989,6 +999,29 @@ namespace sd::pipeline { std::vector> hires_final_latents; int64_t hires_denoise_start = ggml_time_ms(); + sd::Tensor 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 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"); @@ -1026,7 +1059,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/stable-diffusion.cpp b/src/stable-diffusion.cpp index d476f92d6..e98482c40 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -729,7 +729,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) {