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

feat: add PuLID-Flux identity-injection support (#1595)

Browse Source
This commit is contained in:
RapidMark authored and GitHub committed 2026-06-15 23:33:50 +08:00
1 parent 6e66a1a4a4
commit 93527fda74
11 files changed
+878 -13
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/diffusion/flux.hpp
+133 -13
View File
@@ -4,6 +4,7 @@
#include <memory>
#include <vector>
#include "model/adapter/pulid.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/model.hpp"
@@ -49,6 +50,13 @@ 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;
static FluxConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
SDVersion version = VERSION_FLUX) {
@@ -138,6 +146,13 @@ 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;
}
}
if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != config.patch_size) {
GGML_ASSERT(config.patch_size == 2 * actual_radiance_patch_size);
@@ -957,6 +972,29 @@ namespace Flux {
blocks["double_stream_modulation_txt"] = std::make_shared<Modulation>(config.hidden_size, true, !config.disable_bias);
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_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_head=*/PuLIDPerceiverAttentionCA::DEFAULT_DIM_HEAD,
/*heads=*/ PuLIDPerceiverAttentionCA::DEFAULT_HEADS,
/*kv_dim=*/ PuLIDPerceiverAttentionCA::DEFAULT_KV_DIM));
}
}
}
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
@@ -967,7 +1005,9 @@ namespace Flux {
ggml_tensor* guidance,
ggml_tensor* pe,
ggml_tensor* mod_index_arange = nullptr,
std::vector<int> skip_layers = {}) {
std::vector<int> skip_layers = {},
ggml_tensor* pulid_id = nullptr,
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"]);
@@ -1044,6 +1084,23 @@ 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;
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()) {
continue;
@@ -1056,9 +1113,19 @@ namespace Flux {
txt = img_txt.second; // [N, n_txt_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(img, "flux.double_blocks." + std::to_string(i), "img");
sd::ggml_graph_cut::mark_graph_cut(txt, "flux.double_blocks." + std::to_string(i), "txt");
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));
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
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;
@@ -1067,6 +1134,31 @@ namespace Flux {
txt_img = block->forward(ctx, txt_img, vec, pe, txt_img_mask, ss_mods);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.single_blocks." + std::to_string(i), "txt_img");
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);
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);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.pulid_ca." + std::to_string(ca_idx), "txt_img");
ca_idx++;
}
}
img = ggml_view_3d(ctx->ggml_ctx,
@@ -1105,7 +1197,9 @@ namespace Flux {
ggml_tensor* mod_index_arange = nullptr,
ggml_tensor* dct = nullptr,
std::vector<ggml_tensor*> ref_latents = {},
std::vector<int> skip_layers = {}) {
std::vector<int> skip_layers = {},
ggml_tensor* pulid_id = nullptr,
float pulid_id_weight = 1.0f) {
GGML_ASSERT(x->ne[3] == 1);
int64_t W = x->ne[0];
@@ -1131,7 +1225,8 @@ namespace Flux {
img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], img->ne[3]); // [N, hidden_size, H/patch_size*W/patch_size]
img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3)); // [N, H/patch_size*W/patch_size, hidden_size]
auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers); // [N, n_img_token, hidden_size]
auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers,
pulid_id, pulid_id_weight); // [N, n_img_token, hidden_size]
// nerf decode
auto nerf_image_embedder = std::dynamic_pointer_cast<NerfEmbedder>(blocks["nerf_image_embedder"]);
@@ -1179,7 +1274,9 @@ namespace Flux {
ggml_tensor* mod_index_arange = nullptr,
ggml_tensor* dct = nullptr,
std::vector<ggml_tensor*> ref_latents = {},
std::vector<int> skip_layers = {}) {
std::vector<int> skip_layers = {},
ggml_tensor* pulid_id = nullptr,
float pulid_id_weight = 1.0f) {
GGML_ASSERT(x->ne[3] == 1);
int64_t W = x->ne[0];
@@ -1226,7 +1323,8 @@ namespace Flux {
}
}
auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers); // [N, num_tokens, C * patch_size * patch_size]
auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers,
pulid_id, pulid_id_weight); // [N, num_tokens, C * patch_size * patch_size]
if (out->ne[1] > img_tokens) {
out = ggml_view_3d(ctx->ggml_ctx, out, out->ne[0], img_tokens, out->ne[2], out->nb[1], out->nb[2], 0);
@@ -1248,7 +1346,9 @@ namespace Flux {
ggml_tensor* mod_index_arange = nullptr,
ggml_tensor* dct = nullptr,
std::vector<ggml_tensor*> ref_latents = {},
std::vector<int> skip_layers = {}) {
std::vector<int> skip_layers = {},
ggml_tensor* pulid_id = nullptr,
float pulid_id_weight = 1.0f) {
// Forward pass of DiT.
// x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
// timestep: (N,) tensor of diffusion timesteps
@@ -1271,7 +1371,9 @@ namespace Flux {
mod_index_arange,
dct,
ref_latents,
skip_layers);
skip_layers,
pulid_id,
pulid_id_weight);
} else {
return forward_flux_chroma(ctx,
x,
@@ -1284,7 +1386,9 @@ namespace Flux {
mod_index_arange,
dct,
ref_latents,
skip_layers);
skip_layers,
pulid_id,
pulid_id_weight);
}
}
};
@@ -1384,7 +1488,9 @@ namespace Flux {
const sd::Tensor<float>& guidance_tensor = {},
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
bool increase_ref_index = false,
std::vector<int> skip_layers = {}) {
std::vector<int> skip_layers = {},
const sd::Tensor<float>& pulid_id_tensor = {},
float pulid_id_weight = 1.0f) {
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_optional_input(context_tensor);
@@ -1461,6 +1567,13 @@ 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);
auto runner_ctx = get_context();
ggml_tensor* out = flux.forward(&runner_ctx,
@@ -1474,7 +1587,9 @@ namespace Flux {
mod_index_arange,
dct,
ref_latents,
skip_layers);
skip_layers,
pulid_id,
pulid_id_weight);
ggml_build_forward_expand(gf, out);
@@ -1490,14 +1605,17 @@ namespace Flux {
const sd::Tensor<float>& guidance = {},
const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false,
std::vector<int> skip_layers = std::vector<int>()) {
std::vector<int> skip_layers = std::vector<int>(),
const sd::Tensor<float>& pulid_id = {},
float pulid_id_weight = 1.0f) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
// y: [N, adm_in_channels] or [1, adm_in_channels]
// guidance: [N, ]
// pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers);
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers, pulid_id, pulid_id_weight);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -1520,7 +1638,9 @@ namespace Flux {
tensor_or_empty(extra->guidance),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index,
extra->skip_layers ? *extra->skip_layers : empty_skip_layers);
extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
tensor_or_empty(extra->pulid_id),
extra->pulid_id_weight);
}
void test() {