Skip to content

Commit c79d24b

Browse files
stduhpfstduhpf
authored
feat: add Krea2OstrisEdit support (#1775)
1 parent 1b04283 commit c79d24b

3 files changed

Lines changed: 162 additions & 28 deletions

File tree

‎src/conditioning/conditioner.hpp‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2141,6 +2141,41 @@ struct LLMEmbedder : public Conditioner {
21412141
out_layers = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35};
21422142

21432143
prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
2144+
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
2145+
std::string img_prompt = "";
2146+
const std::string placeholder = "<|image_pad|>";
2147+
2148+
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
2149+
const auto& image = (*conditioner_params.ref_images)[i];
2150+
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
2151+
int height = static_cast<int>(image.shape()[1]);
2152+
int width = static_cast<int>(image.shape()[0]);
2153+
double beta = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
2154+
int h_bar = std::max(static_cast<int>(factor),
2155+
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
2156+
int w_bar = std::max(static_cast<int>(factor),
2157+
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
2158+
2159+
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
2160+
2161+
auto resized_image = clip_preprocess(image, w_bar, h_bar);
2162+
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
2163+
GGML_ASSERT(!image_embed.empty());
2164+
2165+
std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
2166+
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
2167+
image_embeds.emplace_back(image_embed_idx, image_embed);
2168+
2169+
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
2170+
int64_t num_image_tokens = image_embed.shape()[1];
2171+
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
2172+
for (int j = 0; j < num_image_tokens; j++) {
2173+
img_prompt += placeholder;
2174+
}
2175+
img_prompt += "<|vision_end|>";
2176+
}
2177+
prompt += img_prompt;
2178+
}
21442179

21452180
prompt_attn_range.first = static_cast<int>(prompt.size());
21462181
prompt += conditioner_params.text;

‎src/model/diffusion/krea2.hpp‎

Lines changed: 126 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -421,29 +421,88 @@ namespace Krea2 {
421421
ggml_tensor* forward(GGMLRunnerContext* ctx,
422422
ggml_tensor* x,
423423
ggml_tensor* vec,
424-
ggml_tensor* pe) {
424+
ggml_tensor* pe,
425+
ggml_tensor* vec_refs = nullptr,
426+
int64_t ref_start = -1) {
425427
auto mod = std::dynamic_pointer_cast<KreaDoubleSharedModulation>(blocks["mod"]);
426428
auto prenorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["prenorm"]);
427429
auto postnorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["postnorm"]);
428430
auto attn = std::dynamic_pointer_cast<KreaAttention>(blocks["attn"]);
429431
auto mlp = std::dynamic_pointer_cast<KreaSwiGLU>(blocks["mlp"]);
430432

431-
auto mods = mod->forward(ctx, vec);
432-
auto attn_input = Flux::modulate(ctx->ggml_ctx,
433-
prenorm->forward(ctx, x),
434-
mods[1],
435-
mods[0],
436-
true);
437-
auto attn_out = attn->forward(ctx, attn_input, pe);
438-
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
439-
440-
auto mlp_input = Flux::modulate(ctx->ggml_ctx,
441-
postnorm->forward(ctx, x),
442-
mods[4],
443-
mods[3],
444-
true);
445-
auto mlp_out = mlp->forward(ctx, mlp_input);
446-
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
433+
if (ref_start >= 0 && vec_refs) {
434+
// same as normal, but since vec is different for refs and the rest, needs a lot of views and concats
435+
auto mods_main = mod->forward(ctx, vec);
436+
auto mods_refs = mod->forward(ctx, vec_refs);
437+
438+
int64_t D = x->ne[0];
439+
int64_t N = x->ne[1];
440+
int64_t B = x->ne[2];
441+
size_t nb1 = x->nb[1];
442+
size_t nb2 = x->nb[2];
443+
444+
int64_t len_main = ref_start;
445+
int64_t len_refs = N - ref_start;
446+
447+
auto pre_x = prenorm->forward(ctx, x);
448+
449+
auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0);
450+
auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1);
451+
452+
auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true);
453+
auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
454+
455+
auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1);
456+
457+
auto attn_out = attn->forward(ctx, attn_input, pe);
458+
459+
auto attn_out_main = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_main, B, attn_out->nb[1], attn_out->nb[2], 0);
460+
auto attn_out_refs = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_refs, B, attn_out->nb[1], attn_out->nb[2], len_main * attn_out->nb[1]);
461+
462+
auto res_main = ggml_mul(ctx->ggml_ctx, attn_out_main, mods_main[2]);
463+
auto res_refs = ggml_mul(ctx->ggml_ctx, attn_out_refs, mods_refs[2]);
464+
465+
auto attn_res = ggml_concat(ctx->ggml_ctx, res_main, res_refs, 1);
466+
467+
x = ggml_add(ctx->ggml_ctx, x, attn_res);
468+
469+
auto post_x = postnorm->forward(ctx, x);
470+
471+
auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0);
472+
auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]);
473+
474+
auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true);
475+
auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
476+
477+
auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1);
478+
auto mlp_out = mlp->forward(ctx, mlp_input);
479+
480+
auto mlp_out_main = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_main, B, mlp_out->nb[1], mlp_out->nb[2], 0);
481+
auto mlp_out_refs = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_refs, B, mlp_out->nb[1], mlp_out->nb[2], len_main * mlp_out->nb[1]);
482+
483+
auto mlp_res_main = ggml_mul(ctx->ggml_ctx, mlp_out_main, mods_main[5]);
484+
auto mlp_res_refs = ggml_mul(ctx->ggml_ctx, mlp_out_refs, mods_refs[5]);
485+
486+
auto mlp_res = ggml_concat(ctx->ggml_ctx, mlp_res_main, mlp_res_refs, 1);
487+
x = ggml_add(ctx->ggml_ctx, x, mlp_res);
488+
} else {
489+
auto mods = mod->forward(ctx, vec);
490+
auto attn_input = Flux::modulate(ctx->ggml_ctx,
491+
prenorm->forward(ctx, x),
492+
mods[1],
493+
mods[0],
494+
true);
495+
auto attn_out = attn->forward(ctx, attn_input, pe);
496+
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
497+
498+
auto mlp_input = Flux::modulate(ctx->ggml_ctx,
499+
postnorm->forward(ctx, x),
500+
mods[4],
501+
mods[3],
502+
true);
503+
auto mlp_out = mlp->forward(ctx, mlp_input);
504+
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
505+
}
447506
return x;
448507
}
449508

@@ -555,7 +614,8 @@ namespace Krea2 {
555614
ggml_tensor* x,
556615
ggml_tensor* timestep,
557616
ggml_tensor* context,
558-
ggml_tensor* pe) {
617+
ggml_tensor* pe,
618+
std::vector<ggml_tensor*> ref_latents = {}) {
559619
int64_t W = x->ne[0];
560620
int64_t H = x->ne[1];
561621
int64_t N = x->ne[3];
@@ -570,26 +630,44 @@ namespace Krea2 {
570630

571631
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, true);
572632
int64_t img_len = img->ne[1];
633+
if (ref_latents.size() > 0) {
634+
for (ggml_tensor* ref : ref_latents) {
635+
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size, true);
636+
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
637+
}
638+
}
639+
int64_t ref_len = img->ne[1] - img_len;
573640
img = first->forward(ctx, img);
574641

575642
auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
576643
t = tmlp->forward(ctx, t);
577644
t = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]);
578645
auto tvec = tproj->forward(ctx, t);
579646

647+
ggml_tensor* tvec_0 = nullptr;
648+
if (ref_latents.size() > 0) {
649+
// "index_timestep_zero" mode: use timestep = 0 for ref latents
650+
auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
651+
auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
652+
t_0 = tmlp->forward(ctx, t_0);
653+
t_0 = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]);
654+
tvec_0 = tproj->forward(ctx, t_0);
655+
}
656+
580657
auto txt = txtfusion->forward(ctx, context);
581658
txt = txtmlp->forward(ctx, txt);
582659
int64_t txt_len = txt->ne[1];
583660

584661
auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, img, 1);
662+
int64_t ref_start = hidden_states->ne[1] - ref_len;
585663
for (int i = 0; i < config.layers; ++i) {
586664
auto block = std::dynamic_pointer_cast<KreaSingleStreamBlock>(blocks["blocks." + std::to_string(i)]);
587-
hidden_states = block->forward(ctx, hidden_states, tvec, pe);
665+
hidden_states = block->forward(ctx, hidden_states, tvec, pe, tvec_0, ref_start);
588666
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states");
589667
}
590668

591-
hidden_states = last->forward(ctx, hidden_states, t);
592669
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len);
670+
hidden_states = last->forward(ctx, hidden_states, t);
593671
hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
594672
return hidden_states;
595673
}
@@ -601,10 +679,16 @@ namespace Krea2 {
601679
int bs,
602680
int context_len,
603681
float theta,
604-
const std::vector<int>& axes_dim) {
682+
const std::vector<int>& axes_dim,
683+
const std::vector<ggml_tensor*>& ref_latents,
684+
Rope::RefIndexMode ref_index_mode) {
605685
auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
606686
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
607687
auto ids = Rope::concat_ids(txt_ids, img_ids, bs);
688+
if (ref_latents.size() > 0) {
689+
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0);
690+
ids = Rope::concat_ids(ids, refs_ids, bs);
691+
}
608692
return Rope::embed_nd(ids, bs, theta, axes_dim);
609693
}
610694

@@ -633,37 +717,49 @@ namespace Krea2 {
633717

634718
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
635719
const sd::Tensor<float>& timesteps_tensor,
636-
const sd::Tensor<float>& context_tensor) {
720+
const sd::Tensor<float>& context_tensor,
721+
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
722+
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
637723
ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE);
638724
ggml_tensor* x = make_input(x_tensor);
639725
ggml_tensor* timesteps = make_input(timesteps_tensor);
640726
GGML_ASSERT(x->ne[3] == 1);
641727
GGML_ASSERT(!context_tensor.empty());
642728
ggml_tensor* context = make_input(context_tensor);
643729

730+
std::vector<ggml_tensor*> ref_latents;
731+
ref_latents.reserve(ref_latents_tensor.size());
732+
for (const auto& ref_latent_tensor : ref_latents_tensor) {
733+
ref_latents.push_back(make_input(ref_latent_tensor));
734+
}
735+
644736
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]),
645737
static_cast<int>(x->ne[0]),
646738
config.patch_size,
647739
static_cast<int>(x->ne[3]),
648740
static_cast<int>(context->ne[1]),
649741
config.theta,
650-
config.axes_dim);
742+
config.axes_dim,
743+
ref_latents,
744+
ref_index_mode);
651745
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
652746
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
653747
set_backend_tensor_data(pe, pe_vec.data());
654748

655749
auto runner_ctx = get_context();
656-
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe);
750+
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
657751
ggml_build_forward_expand(gf, out);
658752
return gf;
659753
}
660754

661755
sd::Tensor<float> compute(int n_threads,
662756
const sd::Tensor<float>& x,
663757
const sd::Tensor<float>& timesteps,
664-
const sd::Tensor<float>& context) {
758+
const sd::Tensor<float>& context,
759+
const std::vector<sd::Tensor<float>>& ref_latents = {},
760+
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
665761
auto get_graph = [&]() -> ggml_cgraph* {
666-
return build_graph(x, timesteps, context);
762+
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
667763
};
668764
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
669765
}
@@ -672,10 +768,13 @@ namespace Krea2 {
672768
const DiffusionParams& diffusion_params) override {
673769
GGML_ASSERT(diffusion_params.x != nullptr);
674770
GGML_ASSERT(diffusion_params.timesteps != nullptr);
771+
static const std::vector<sd::Tensor<float>> empty_ref_latents;
675772
return compute(n_threads,
676773
*diffusion_params.x,
677774
*diffusion_params.timesteps,
678-
tensor_or_empty(diffusion_params.context));
775+
tensor_or_empty(diffusion_params.context),
776+
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
777+
diffusion_params.ref_index_mode);
679778
}
680779
};
681780
} // namespace Krea2

‎src/stable-diffusion.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -981,7 +981,7 @@ class StableDiffusionGGML {
981981
tensor_storage_map,
982982
version,
983983
"",
984-
false,
984+
true,
985985
model_manager);
986986
diffusion_model = std::make_shared<Krea2::Krea2Runner>(backend_for(SDBackendModule::DIFFUSION),
987987
tensor_storage_map,

0 commit comments

Comments
 (0)