DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-10 19:06:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add control-net z-image (#2187)

Browse Source
This commit is contained in:
SolicTous authored and GitHub committed 2026-10-25 03:11:57 +08:00
1 parent 4c870cc60c
commit b29cad540a
6 files changed
+470 -21
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/diffusion/model.hpp
+7 -1
View File
@@ -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;
src/model/diffusion/z_image.hpp
+352 -14
View File
@@ -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() {
src/pipeline/diffusion_engine.cpp
+35
View File
@@ -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{};
}
src/pipeline/image.cpp
+66 -5
View File
@@ -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,
src/pipeline/model_builders.cpp
+3
View File
@@ -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;
}
src/stable-diffusion.cpp
+7 -1
View File
@@ -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) {