upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-07-30 23:00:41 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix: simplify PuLID ID extraction setup (#1664)

Browse Source
This commit is contained in:
leejet
2026-06-15 23:55:38 +08:00
committed by GitHub
parent 93527fda74
commit 146b6cc49e
10 changed files with 85 additions and 257 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
src/model/diffusion/flux.hpp
View File

@@ -50,12 +50,9 @@ namespace Flux {
float ref_index_scale = 1.f;
ChromaRadianceConfig chroma_radiance_params;
// PuLID-Flux identity injection. Turned on by the runner when a
// --pulid-weights path is provided. The intervals are fixed by the
// PuLID v0.9.1 architecture (every 2nd double, every 4th single).
bool pulid_enabled = false;
int pulid_double_interval = 2;
int pulid_single_interval = 4;
bool pulid_enabled = false;
int pulid_double_interval = 2;
int pulid_single_interval = 4;
static FluxConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
@@ -146,10 +143,6 @@ namespace Flux {
if (ends_with(name, "double_blocks.0.txt_attn.norm.key_norm.scale")) {
head_dim = tensor_storage.ne[0];
}
// PuLID weights live alongside the diffusion model under the same
// prefix (pulid_ca.<i>.<sub>) when the pulid loader merges them in
// (see stable-diffusion.cpp). Spotting any pulid_ca.* key flips the
// flag so the Flux ctor builds the pulid_ca.<i> child blocks.
if (name.find("pulid_ca.") != std::string::npos) {
config.pulid_enabled = true;
}
@@ -973,26 +966,17 @@ namespace Flux {
blocks["single_stream_modulation"] = std::make_shared<Modulation>(config.hidden_size, false, !config.disable_bias);
}
// PuLID-Flux identity-injection cross-attention modules. Only constructed
// when config.pulid_enabled is set (turned on by the runner after seeing a
// --pulid-weights path during model load). Counts come straight from PuLID
// v0.9.1's pipeline_flux.py: every `pulid_double_interval` double block
// (=2) and every `pulid_single_interval` single block (=4). For a stock
// Flux Dev (depth=19, depth_single_blocks=38), this means 10 + 10 = 20
// hook points... but the reference uses ceil-rounding so the actual count
// is `ceil(depth/2) + ceil(depth_single_blocks/4)` = 10 + 10 = 20. PuLID
// v0.9.1 trained weights have 20 entries.
if (config.pulid_enabled) {
int num_double_ca = (config.depth + config.pulid_double_interval - 1) / config.pulid_double_interval;
int num_single_ca = (config.depth_single_blocks + config.pulid_single_interval - 1) / config.pulid_single_interval;
int num_double_ca = (config.depth + config.pulid_double_interval - 1) / config.pulid_double_interval;
int num_single_ca = (config.depth_single_blocks + config.pulid_single_interval - 1) / config.pulid_single_interval;
int num_ca = num_double_ca + num_single_ca;
for (int i = 0; i < num_ca; i++) {
blocks["pulid_ca." + std::to_string(i)] =
std::shared_ptr<GGMLBlock>(new PuLIDPerceiverAttentionCA(
/*dim=*/ config.hidden_size,
/*dim=*/config.hidden_size,
/*dim_head=*/PuLIDPerceiverAttentionCA::DEFAULT_DIM_HEAD,
/*heads=*/ PuLIDPerceiverAttentionCA::DEFAULT_HEADS,
/*kv_dim=*/ PuLIDPerceiverAttentionCA::DEFAULT_KV_DIM));
/*heads=*/PuLIDPerceiverAttentionCA::DEFAULT_HEADS,
/*kv_dim=*/PuLIDPerceiverAttentionCA::DEFAULT_KV_DIM));
}
}
}
@@ -1007,7 +991,7 @@ namespace Flux {
ggml_tensor* mod_index_arange = nullptr,
std::vector<int> skip_layers = {},
ggml_tensor* pulid_id = nullptr,
float pulid_id_weight = 1.0f) {
float pulid_id_weight = 1.0f) {
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
auto txt_in = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
auto final_layer = std::dynamic_pointer_cast<LastLayer>(blocks["final_layer"]);
@@ -1084,22 +1068,12 @@ namespace Flux {
sd::ggml_graph_cut::mark_graph_cut(txt, "flux.prelude", "txt");
sd::ggml_graph_cut::mark_graph_cut(vec, "flux.prelude", "vec");
// PuLID identity injection: mirrors ToTheBeginning/PuLID
// pulid/encoders_transformer.py + flux/model.py. The CA layers
// run *between* transformer blocks, with their output added to
// img (scaled by id_weight) at every `pulid_double_interval`-th
// double_block and every `pulid_single_interval`-th single_block.
//
// skip_layers + PuLID is NOT a supported combination -- skipping
// a block at a PuLID-aligned index would either misalign the
// ca_idx assignment (silent quality regression) or require us
// to invent a non-reference index policy. Refuse early instead.
const bool pulid_active = config.pulid_enabled && pulid_id != nullptr;
if (pulid_active && !skip_layers.empty()) {
LOG_WARN("PuLID + skip_layers is not supported; disabling PuLID for this generation.");
}
const bool pulid_run = pulid_active && skip_layers.empty();
int ca_idx = 0;
int ca_idx = 0;
for (int i = 0; i < config.depth; i++) {
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
@@ -1117,15 +1091,15 @@ namespace Flux {
if (pulid_run && (i % config.pulid_double_interval == 0)) {
auto pulid_ca = std::dynamic_pointer_cast<PuLIDPerceiverAttentionCA>(
blocks["pulid_ca." + std::to_string(ca_idx)]);
ggml_tensor* ca_out = pulid_ca->forward(ctx, pulid_id, img); // [N, n_img_token, hidden_size]
img = ggml_add(ctx->ggml_ctx, img, ggml_scale(ctx->ggml_ctx, ca_out, pulid_id_weight));
ggml_tensor* ca_out = pulid_ca->forward(ctx, pulid_id, img); // [N, n_img_token, hidden_size]
img = ggml_add(ctx->ggml_ctx, img, ggml_scale(ctx->ggml_ctx, ca_out, pulid_id_weight));
sd::ggml_graph_cut::mark_graph_cut(img, "flux.pulid_ca." + std::to_string(ca_idx), "img");
ca_idx++;
}
}
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
const int64_t n_txt_tok = txt->ne[1]; // for splitting back into img portion below
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
const int64_t n_txt_tok = txt->ne[1];
for (int i = 0; i < config.depth_single_blocks; i++) {
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + config.depth) != skip_layers.end()) {
continue;
@@ -1138,24 +1112,22 @@ namespace Flux {
if (pulid_run && (i % config.pulid_single_interval == 0)) {
auto pulid_ca = std::dynamic_pointer_cast<PuLIDPerceiverAttentionCA>(
blocks["pulid_ca." + std::to_string(ca_idx)]);
// Split txt_img into [txt | img], inject ID into the img portion
// only, then concatenate back. Matches the PyTorch reference.
ggml_tensor* txt_part = ggml_view_3d(ctx->ggml_ctx, txt_img,
txt_img->ne[0], n_txt_tok, txt_img->ne[2],
txt_img->nb[1], txt_img->nb[2],
0);
txt_img->ne[0], n_txt_tok, txt_img->ne[2],
txt_img->nb[1], txt_img->nb[2],
0);
ggml_tensor* img_part = ggml_view_3d(ctx->ggml_ctx, txt_img,
txt_img->ne[0],
txt_img->ne[1] - n_txt_tok,
txt_img->ne[2],
txt_img->nb[1],
txt_img->nb[2],
n_txt_tok * txt_img->nb[1]);
txt_part = ggml_cont(ctx->ggml_ctx, txt_part);
img_part = ggml_cont(ctx->ggml_ctx, img_part);
ggml_tensor* ca_out = pulid_ca->forward(ctx, pulid_id, img_part);
img_part = ggml_add(ctx->ggml_ctx, img_part, ggml_scale(ctx->ggml_ctx, ca_out, pulid_id_weight));
txt_img = ggml_concat(ctx->ggml_ctx, txt_part, img_part, 1);
txt_img->ne[0],
txt_img->ne[1] - n_txt_tok,
txt_img->ne[2],
txt_img->nb[1],
txt_img->nb[2],
n_txt_tok * txt_img->nb[1]);
txt_part = ggml_cont(ctx->ggml_ctx, txt_part);
img_part = ggml_cont(ctx->ggml_ctx, img_part);
ggml_tensor* ca_out = pulid_ca->forward(ctx, pulid_id, img_part);
img_part = ggml_add(ctx->ggml_ctx, img_part, ggml_scale(ctx->ggml_ctx, ca_out, pulid_id_weight));
txt_img = ggml_concat(ctx->ggml_ctx, txt_part, img_part, 1);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.pulid_ca." + std::to_string(ca_idx), "txt_img");
ca_idx++;
}
@@ -1567,12 +1539,9 @@ namespace Flux {
set_backend_tensor_data(dct, dct_vec.data());
}
// Materialize the PuLID id embedding into the compute graph when
// pulid_id_tensor is non-empty. forward() accepts nullptr for the
// no-injection case.
ggml_tensor* pulid_id = pulid_id_tensor.empty()
? nullptr
: make_input(pulid_id_tensor);
? nullptr
: make_input(pulid_id_tensor);
auto runner_ctx = get_context();

3
src/model/diffusion/model.hpp
View File

@@ -22,9 +22,6 @@ struct SkipLayerDiffusionExtra {
struct FluxDiffusionExtra {
const sd::Tensor<float>* guidance = nullptr;
const std::vector<int>* skip_layers = nullptr;
// PuLID-Flux: precomputed (N=1, num_tokens=32, kv_dim=2048) identity embedding
// produced by runtime-scripts/pulid_extract_id.py. nullptr when PuLID is
// disabled. id_weight is per-job (typical 0.7-1.2; default 1.0).
const sd::Tensor<float>* pulid_id = nullptr;
float pulid_id_weight = 1.0f;
};