upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-07-29 06:10:52 -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

219
src/conditioning/conditioner.hpp
View File

@@ -7,6 +7,7 @@
#include "core/tensor_ggml.hpp"
#include "core/util.h"
#include "model/diffusion/model.hpp"
#include "model/te/clip.hpp"
#include "model/te/llm.hpp"
#include "model/te/t5.hpp"
@@ -106,6 +107,7 @@ struct ConditionerParams {
int height = -1;
bool zero_out_masked = false;
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
RefImageParams ref_image_params;
};
struct Conditioner {
@@ -1993,6 +1995,54 @@ struct LLMEmbedder : public Conditioner {
return new_hidden_states;
}
void resize_image_dims(int height, int width, int& h_bar, int& w_bar, int factor, int min_size, int max_size, RefImageResizeMode mode) {
if (min_size > 0 && min_size == max_size) {
if (mode == RefImageResizeMode::AREA) {
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
int current_max_side = std::max(height, width);
double beta = static_cast<double>(min_size) / current_max_side;
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
}
return;
}
if (mode == RefImageResizeMode::AREA) {
double current_area = static_cast<double>(h_bar) * w_bar;
if (max_size > 0 && current_area > max_size) {
double beta = std::sqrt((static_cast<double>(height) * width) / static_cast<double>(max_size));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (min_size > 0 && current_area < min_size) {
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
int current_max_side = std::max(height, width);
if (max_size > 0 && current_max_side > max_size) {
double beta = static_cast<double>(max_size) / current_max_side;
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width * beta / factor)) * static_cast<int>(factor));
} else if (min_size > 0 && current_max_side < min_size) {
double beta = static_cast<double>(min_size) / current_max_side;
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
}
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
std::string prompt;
@@ -2007,7 +2057,8 @@ struct LLMEmbedder : public Conditioner {
bool spell_quotes = false;
std::set<int> out_layers;
int64_t t0 = ggml_time_ms();
int64_t t0 = ggml_time_ms();
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
if (sd_version_is_lingbot_video(version)) {
const int pad_token = 151643;
@@ -2040,28 +2091,35 @@ struct LLMEmbedder : public Conditioner {
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int min_pixels = static_cast<int>(4 * factor * factor);
int max_pixels = static_cast<int>(16384 * factor * factor);
int h_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(height / factor) * factor));
int w_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(width / factor) * factor));
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels = static_cast<int>(4 * factor * factor);
} else {
min_pixels = static_cast<int>(2 * factor);
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels = static_cast<int>(16384 * factor * factor);
} else {
max_pixels = static_cast<int>(128 * factor);
}
}
int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor) * factor));
int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor) * factor));
if (std::max(height, width) > 200 * std::min(height, width)) {
LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
}
if (h_bar * w_bar > max_pixels) {
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (h_bar * w_bar < min_pixels) {
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
@@ -2092,30 +2150,33 @@ struct LLMEmbedder : public Conditioner {
prompt_template_encode_start_idx = 64;
int image_embed_idx = 64 + 6;
int min_pixels = 384 * 384;
int max_pixels = 560 * 560;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 560;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor);
int w_bar = static_cast<int>(std::round(width / factor) * factor);
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
if (static_cast<double>(h_bar) * w_bar > max_pixels) {
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@@ -2170,16 +2231,33 @@ struct LLMEmbedder : public Conditioner {
std::string img_prompt;
const std::string placeholder = "<|image_pad|>";
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
double beta = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
int h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
int w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
int h_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@@ -2221,17 +2299,33 @@ struct LLMEmbedder : public Conditioner {
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
std::string img_prompt = "";
const std::string placeholder = "<|image_pad|>";
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 1024;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
double beta = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
int h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
int w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
int h_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@@ -2268,30 +2362,33 @@ struct LLMEmbedder : public Conditioner {
min_length = 512 + prompt_template_encode_start_idx;
int image_embed_idx = 36 + 6;
int min_pixels = 384 * 384;
int max_pixels = 560 * 560;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 560;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor);
int w_bar = static_cast<int>(std::round(width / factor) * factor);
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
if (static_cast<double>(h_bar) * w_bar > max_pixels) {
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);

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() {

193
src/stable-diffusion.cpp
View File

@@ -2297,7 +2297,7 @@ public:
const char* extra_sample_args,
const std::vector<float>& sigmas,
const std::vector<sd::Tensor<float>>& ref_latents,
bool increase_ref_index,
const RefImageParams& ref_image_params,
const sd::Tensor<float>& denoise_mask,
const sd::Tensor<float>& vace_context,
float vace_strength,
@@ -2458,9 +2458,9 @@ public:
sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
std::vector<sd::Tensor<float>> controls;
DiffusionParams diffusion_params;
diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.ref_image_params = ref_image_params;
sd::guidance::GuidanceInput step_guidance_input;
step_guidance_input.step = step;
step_guidance_input.schedule_size = sigmas.size();
@@ -2840,6 +2840,126 @@ public:
auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
return !!flow_denoiser;
}
std::string get_default_ref_image_preset(SDVersion version) const {
if (sd_version_is_longcat(version)) {
return "longcat";
} else if (sd_version_is_flux(version)) {
return "flux_kontext";
} else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
return "flux2";
} else if (version == VERSION_QWEN_IMAGE_LAYERED) {
return "qwen_layered";
} else if (sd_version_is_qwen_image(version)) {
return "qwen";
} else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) {
return "z_image_omni";
} else if (sd_version_is_krea2(version)) {
// have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit)
// and "krea2_ostris_edit" (for krea2 ostris edit)
// since krea2 ostris edit support predates, it should probably be default
return "krea2_ostris_edit";
} else if (sd_version_is_anima(version)) {
return "cosmos_reference";
}
return "default";
}
RefImageParams resolve_ref_image_params(const char* ref_image_args) const {
RefImageParams params;
std::string preset_name = get_default_ref_image_preset(version);
for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
if (key == "preset") {
std::string requested_preset_name = value;
if (REF_IMAGE_PRESETS.count(requested_preset_name)) {
preset_name = requested_preset_name;
} else if (value != "default") {
std::string valid_list;
for (auto const& [name, _] : REF_IMAGE_PRESETS) {
valid_list += (valid_list.empty() ? "" : ", ") + name;
}
LOG_WARN("ignoring invalid reference image preset '%s'. Valid options: [%s]", value.c_str(), valid_list.c_str());
}
break;
}
}
if (preset_name != "default") {
LOG_INFO("Using '%s' preset for reference images", preset_name.c_str());
params = REF_IMAGE_PRESETS.at(preset_name);
}
for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
if (key == "pass_to_vlm") {
if (!parse_strict_bool(value, params.pass_to_vlm)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "pass_to_dit") {
if (!parse_strict_bool(value, params.pass_to_dit)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "ref_index_mode") {
if (value == "fixed") {
params.ref_index_mode = Rope::RefIndexMode::FIXED;
} else if (value == "increase") {
params.ref_index_mode = Rope::RefIndexMode::INCREASE;
} else if (value == "decrease") {
params.ref_index_mode = Rope::RefIndexMode::DECREASE;
} else {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "force_ref_timestep_zero") {
if (!parse_strict_bool(value, params.force_ref_timestep_zero)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "resize_before_vae") {
if (!parse_strict_bool(value, params.resize_before_vae)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "vae_input_max_pixels") {
if (!parse_strict_int(value, params.vae_input_max_pixels)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "vlm_resize_mode") {
if (value == "longest_side") {
params.vlm_resize_mode = RefImageResizeMode::LONGEST_SIDE;
} else if (value == "area") {
params.vlm_resize_mode = RefImageResizeMode::AREA;
} else if (value == "none") {
params.vlm_resize_mode = RefImageResizeMode::NONE;
} else {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "vlm_max_size") {
if (!parse_strict_int(value, params.vlm_max_size)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "vlm_min_size") {
if (!parse_strict_int(value, params.vlm_min_size)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key != "preset" && key != "vlm_size") {
LOG_WARN("ignoring unknown reference image arg '%s'", key.c_str());
}
}
for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
if (key == "vlm_size") {
int vlm_size;
if (!parse_strict_int(value, vlm_size)) {
LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
} else {
LOG_INFO("vlm_size override: setting both min and max size to %ld", (long)vlm_size);
params.vlm_min_size = vlm_size;
params.vlm_max_size = vlm_size;
}
break;
}
}
if (params.force_ref_timestep_zero && !sd_version_is_krea2(version)) {
LOG_WARN("force_ref_timestep_zero is only supported by Krea2 architecture for now");
}
return params;
}
};
/*================================================= SD API ==================================================*/
@@ -3301,6 +3421,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
sd_sample_params_init(&sd_img_gen_params->sample_params);
sd_img_gen_params->clip_skip = -1;
sd_img_gen_params->ref_images_count = 0;
sd_img_gen_params->ref_image_args = "";
sd_img_gen_params->width = 512;
sd_img_gen_params->height = 512;
sd_img_gen_params->strength = 0.75f;
@@ -3338,8 +3459,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
"batch_count: %d\n"
"qwen_image_layers: %d\n"
"ref_images_count: %d\n"
"auto_resize_ref_image: %s\n"
"increase_ref_index: %s\n"
"ref_image_args: %s\n"
"control_strength: %.2f\n"
"photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
"VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n"
@@ -3357,8 +3477,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
sd_img_gen_params->batch_count,
sd_img_gen_params->qwen_image_layers,
sd_img_gen_params->ref_images_count,
BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
BOOL_STR(sd_img_gen_params->increase_ref_index),
SAFE_STR(sd_img_gen_params->ref_image_args),
sd_img_gen_params->control_strength,
sd_img_gen_params->pm_params.style_strength,
sd_img_gen_params->pm_params.id_images_count,
@@ -3655,8 +3774,6 @@ struct GenerationRequest {
float strength = 1.f;
float control_strength = 0.f;
float eta = 0.f;
bool increase_ref_index = false;
bool auto_resize_ref_image = false;
sd_guidance_params_t guidance = {};
sd_guidance_params_t high_noise_guidance = {};
sd_pm_params_t pm_params = {};
@@ -3682,8 +3799,6 @@ struct GenerationRequest {
strength = sd_img_gen_params->strength;
control_strength = sd_img_gen_params->control_strength;
eta = sd_img_gen_params->sample_params.eta;
increase_ref_index = sd_img_gen_params->increase_ref_index;
auto_resize_ref_image = sd_img_gen_params->auto_resize_ref_image;
has_ref_images = sd_img_gen_params->ref_images_count > 0;
guidance = sd_img_gen_params->sample_params.guidance;
pm_params = sd_img_gen_params->pm_params;
@@ -4431,7 +4546,8 @@ static sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, i
static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request,
SamplePlan* plan) {
SamplePlan* plan,
const RefImageParams& ref_image_params) {
int64_t prepare_start_ms = ggml_time_ms();
sd::Tensor<float> init_image_tensor;
@@ -4574,9 +4690,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
continue;
}
sd::Tensor<float> ref_latent;
if (request->auto_resize_ref_image && !sd_version_is_pid(sd_ctx->sd->version)) {
if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) {
LOG_DEBUG("auto resize ref images");
int vae_image_size = std::min(1024 * 1024, request->width * request->height);
int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024;
int vae_image_size = std::min(target_pixels, request->width * request->height);
double vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
double vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
@@ -4703,15 +4820,20 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request,
SamplePlan* plan,
ImageGenerationLatents* latents) {
ImageGenerationLatents* latents,
const RefImageParams& ref_image_params) {
ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()};
ConditionerParams condition_params;
condition_params.text = request->prompt;
condition_params.clip_skip = request->clip_skip;
condition_params.width = request->width;
condition_params.height = request->height;
condition_params.ref_images = &latents->ref_images;
condition_params.text = request->prompt;
condition_params.clip_skip = request->clip_skip;
condition_params.width = request->width;
condition_params.height = request->height;
if (ref_image_params.pass_to_vlm) {
condition_params.ref_images = &latents->ref_images;
}
condition_params.ref_image_params = ref_image_params;
sd_ctx->sd->prepare_generation_extensions(request->pm_params,
request->pulid_params,
@@ -4721,7 +4843,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
condition_params.zero_out_masked = false;
auto cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params);
if (cond.c_concat.empty()) {
if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
cond.c_concat = latents->concat_latent; // TODO: optimize
}
@@ -4750,7 +4872,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params);
}
if (uncond.c_concat.empty()) {
if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
uncond.c_concat = latents->concat_latent; // TODO: optimize
}
}
@@ -4774,7 +4896,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
}
img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params);
if (img_uncond.c_concat.empty()) {
if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize
}
}
@@ -5091,13 +5213,16 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count);
apply_circular_axes_to_diffusion(sd_ctx, sd_img_gen_params->circular_x, sd_img_gen_params->circular_y);
const RefImageParams ref_image_params = sd_ctx->sd->resolve_ref_image_params(sd_img_gen_params->ref_image_args);
ImageVaeAxesGuard axes_guard(sd_ctx, sd_img_gen_params, request);
SamplePlan plan(sd_ctx, sd_img_gen_params, request);
auto latents_opt = prepare_image_generation_latents(sd_ctx,
sd_img_gen_params,
&request,
&plan);
&plan,
ref_image_params);
if (!latents_opt.has_value()) {
return false;
}
@@ -5107,7 +5232,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
sd_img_gen_params,
&request,
&plan,
&latents);
&latents,
ref_image_params);
if (!embeds_opt.has_value()) {
return false;
}
@@ -5153,7 +5279,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
plan.extra_sample_args,
plan.sigmas,
latents.ref_latents,
request.increase_ref_index,
ref_image_params,
latents.denoise_mask,
sd::Tensor<float>(),
1.f,
@@ -5274,7 +5400,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
plan.extra_sample_args,
hires_sigma_sched,
latents.ref_latents,
request.increase_ref_index,
ref_image_params,
hires_denoise_mask,
sd::Tensor<float>(),
1.f,
@@ -5665,6 +5791,9 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
condition_params.text = request.prompt;
condition_params.zero_out_masked = true;
condition_params.ref_images = &latents.ref_images;
if (sd_version_is_lingbot_video(sd_ctx->sd->version)) {
condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA;
}
int64_t prepare_start_ms = ggml_time_ms();
embeds.cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
@@ -5945,6 +6074,8 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
sd_ctx->sd->reset_cancel_flag();
const RefImageParams ref_image_params;
if (num_frames_out != nullptr) {
*num_frames_out = 0;
}
@@ -6034,7 +6165,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
plan.high_noise_extra_sample_args,
high_noise_sigmas,
std::vector<sd::Tensor<float>>{},
false,
ref_image_params,
latents.denoise_mask,
latents.vace_context,
request.vace_strength,
@@ -6076,7 +6207,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
plan.extra_sample_args,
plan.sigmas,
std::vector<sd::Tensor<float>>{},
false,
ref_image_params,
latents.denoise_mask,
latents.vace_context,
request.vace_strength,
@@ -6214,7 +6345,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
plan.extra_sample_args,
hires_sigma_sched,
std::vector<sd::Tensor<float>>{},
false,
ref_image_params,
hires_denoise_mask,
sd::Tensor<float>(),
hires_request.vace_strength,