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

feat: add configurable reference image processing for edit models (#1780)

Browse Source
This commit is contained in:
stduhpf
2026-07-14 16:58:22 +02:00
committed by GitHub
parent b5d812008e
commit 74bce049d0
15 changed files with 551 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/model/common/rope.hpp
View File

@@ -18,10 +18,6 @@ namespace Rope {
DECREASE,
};
__STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
}
template <class T>
__STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
std::vector<T> result(num);

61
src/model/diffusion/anima.hpp
View File

@@ -484,10 +484,11 @@ namespace Anima {
ggml_tensor* timestep,
ggml_tensor* encoder_hidden_states,
ggml_tensor* image_pe,
ggml_tensor* t5_ids = nullptr,
ggml_tensor* t5_weights = nullptr,
ggml_tensor* adapter_q_pe = nullptr,
ggml_tensor* adapter_k_pe = nullptr) {
ggml_tensor* t5_ids = nullptr,
ggml_tensor* t5_weights = nullptr,
ggml_tensor* adapter_q_pe = nullptr,
ggml_tensor* adapter_k_pe = nullptr,
std::vector<ggml_tensor*> ref_latents = {}) {
GGML_ASSERT(x->ne[3] == 1);
auto x_embedder = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
@@ -502,8 +503,16 @@ namespace Anima {
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W]
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
int64_t img_len = x->ne[1];
if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) {
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
ref = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2); // [N, C + 1, H, W]
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
x = ggml_concat(ctx->ggml_ctx, x, ref, 1);
}
}
x = x_embedder->forward(ctx, x);
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
@@ -543,6 +552,7 @@ namespace Anima {
x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
}
x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len);
x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C]
@@ -602,8 +612,8 @@ namespace Anima {
const std::vector<int>& axes_dim,
float h_extrapolation_ratio,
float w_extrapolation_ratio,
float t_extrapolation_ratio) {
static const std::vector<ggml_tensor*> empty_ref_latents;
float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h,
w,
patch_size,
@@ -611,7 +621,7 @@ namespace Anima {
static_cast<int>(axes_dim.size()),
0,
{},
empty_ref_latents,
ref_latents,
Rope::RefIndexMode::FIXED,
1.0f,
false);
@@ -626,14 +636,20 @@ namespace Anima {
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor = {},
const sd::Tensor<int32_t>& t5_ids_tensor = {},
const sd::Tensor<float>& t5_weights_tensor = {}) {
const sd::Tensor<float>& context_tensor = {},
const sd::Tensor<int32_t>& t5_ids_tensor = {},
const sd::Tensor<float>& t5_weights_tensor = {},
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* t5_ids = make_optional_input(t5_ids_tensor);
ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
std::vector<ggml_tensor*> ref_latents;
ref_latents.reserve(ref_latents_tensor.size());
for (const auto& ref_latent_tensor : ref_latents_tensor) {
ref_latents.push_back(make_input(ref_latent_tensor));
}
GGML_ASSERT(x->ne[3] == 1);
ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
@@ -650,7 +666,8 @@ namespace Anima {
config.axes_dim,
4.0f,
4.0f,
1.0f);
1.0f,
ref_latents);
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
set_backend_tensor_data(image_pe, image_pe_vec.data());
@@ -682,7 +699,8 @@ namespace Anima {
t5_ids,
t5_weights,
adapter_q_pe,
adapter_k_pe);
adapter_k_pe,
ref_latents);
ggml_build_forward_expand(gf, out);
return gf;
@@ -691,11 +709,13 @@ namespace Anima {
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context = {},
const sd::Tensor<int32_t>& t5_ids = {},
const sd::Tensor<float>& t5_weights = {}) {
const sd::Tensor<float>& context = {},
const sd::Tensor<int32_t>& t5_ids = {},
const sd::Tensor<float>& t5_weights = {},
const std::vector<sd::Tensor<float>>& ref_latents = {},
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("cosmos_reference")) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, t5_ids, t5_weights);
return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
@@ -705,12 +725,15 @@ namespace Anima {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
static const std::vector<sd::Tensor<float>> empty_ref_latents;
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
tensor_or_empty(extra->t5_ids),
tensor_or_empty(extra->t5_weights));
tensor_or_empty(extra->t5_weights),
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_image_params);
}
};
} // namespace Anima

2
src/model/diffusion/boogu.hpp
View File

@@ -827,7 +827,7 @@ namespace Boogu {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents);
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
}
};
} // namespace Boogu

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

@@ -1642,8 +1642,8 @@ namespace Flux {
tensor_or_empty(diffusion_params.c_concat),
tensor_or_empty(diffusion_params.y),
tensor_or_empty(extra->guidance),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode,
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,
extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
tensor_or_empty(extra->pulid_id),
extra->pulid_id_weight);

19
src/model/diffusion/krea2.hpp
View File

@@ -615,7 +615,8 @@ namespace Krea2 {
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {}) {
std::vector<ggml_tensor*> ref_latents = {},
bool zero_timestep_refs = false) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t N = x->ne[3];
@@ -645,7 +646,7 @@ namespace Krea2 {
auto tvec = tproj->forward(ctx, t);
ggml_tensor* tvec_0 = nullptr;
if (ref_latents.size() > 0) {
if (ref_latents.size() > 0 && zero_timestep_refs) {
// "index_timestep_zero" mode: use timestep = 0 for ref latents
auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
@@ -719,7 +720,7 @@ namespace Krea2 {
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 RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -741,13 +742,13 @@ namespace Krea2 {
config.theta,
config.axes_dim,
ref_latents,
ref_index_mode);
ref_image_params.ref_index_mode);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);
ggml_build_forward_expand(gf, out);
return gf;
}
@@ -757,9 +758,9 @@ namespace Krea2 {
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) {
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
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_image_params);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
@@ -773,8 +774,8 @@ namespace Krea2 {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode);
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_image_params);
}
};
} // namespace Krea2

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

@@ -10,6 +10,36 @@
#include "model/common/rope.hpp"
#include "model_manager.h"
enum class RefImageResizeMode {
NONE,
LONGEST_SIDE,
AREA,
};
struct RefImageParams {
bool pass_to_vlm = false;
bool pass_to_dit = true;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
bool force_ref_timestep_zero = false;
bool resize_before_vae = true;
int vae_input_max_pixels = -1;
RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
int vlm_min_size = -1;
int vlm_max_size = -1;
};
const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
{"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
{"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
{"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
{"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
};
struct UNetDiffusionExtra {
int num_video_frames = -1;
const std::vector<sd::Tensor<float>>* controls = nullptr;
@@ -74,7 +104,7 @@ struct DiffusionParams {
const sd::Tensor<float>* c_concat = nullptr;
const sd::Tensor<float>* y = nullptr;
const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
RefImageParams ref_image_params = {false, false, Rope::RefIndexMode::FIXED, false};
DiffusionExtraParams extra = std::monostate{};
};

4
src/model/diffusion/qwen_image.hpp
View File

@@ -715,8 +715,8 @@ namespace Qwen {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode);
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);
}
void test() {

4
src/model/diffusion/z_image.hpp
View File

@@ -648,8 +648,8 @@ namespace ZImage {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode);
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);
}
void test() {