mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-24 20:20:37 +00:00
526 lines
28 KiB
C++
526 lines
28 KiB
C++
#ifndef __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
|
|
#define __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
|
|
|
|
#include <algorithm>
|
|
#include <cinttypes>
|
|
|
|
#include "core/ggml_extend.h"
|
|
#include "core/ggml_runner.h"
|
|
#include "core/util.h"
|
|
#include "model/common/ggml_block.hpp"
|
|
#include "model/diffusion/model.hpp"
|
|
#include "model/diffusion/z_image.hpp"
|
|
#include "model_loader.h"
|
|
|
|
// Ref: https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py
|
|
//
|
|
// The denoiser is Lumina2/z_image's NextDiT with identical hyperparameters, so the blocks are
|
|
// reused from ZImage. Two things differ: every norm here is non-parametric (the checkpoint
|
|
// carries no norm weights at all), and latents arrive already patchified from the Flux2 VAE,
|
|
// so patch_size is 1 over 128 channels.
|
|
|
|
namespace LLaDAImage {
|
|
constexpr int LLADA_IMAGE_GRAPH_SIZE = 20480;
|
|
|
|
struct LLaDAImageConfig {
|
|
int patch_size = 1;
|
|
int64_t hidden_size = 3840;
|
|
int64_t in_channels = 128;
|
|
int64_t out_channels = 128;
|
|
int64_t num_layers = 30;
|
|
int64_t num_refiner_layers = 2;
|
|
int64_t head_dim = 128;
|
|
int64_t num_heads = 30;
|
|
int64_t num_kv_heads = 30;
|
|
int64_t multiple_of = 256;
|
|
float ffn_dim_multiplier = 8.0f / 3.0f;
|
|
float norm_eps = 1e-5f;
|
|
bool qk_norm = true;
|
|
int64_t cap_feat_dim = 2560;
|
|
int64_t semantic_feat_dim = 4096;
|
|
int theta = 256;
|
|
std::vector<int> axes_dim = {32, 48, 48};
|
|
int64_t axes_dim_sum = 128;
|
|
|
|
static int64_t count_blocks(const String2TensorStorage& tensor_storage_map,
|
|
const std::string& prefix,
|
|
const std::string& block_prefix) {
|
|
int64_t count = 0;
|
|
for (const auto& [name, _] : tensor_storage_map) {
|
|
if (!starts_with(name, prefix)) {
|
|
continue;
|
|
}
|
|
size_t pos = name.find(block_prefix);
|
|
if (pos == std::string::npos) {
|
|
continue;
|
|
}
|
|
auto items = split_string(name.substr(pos), '.');
|
|
if (items.size() > 1) {
|
|
count = std::max<int64_t>(count, atoi(items[1].c_str()) + 1);
|
|
}
|
|
}
|
|
return count;
|
|
}
|
|
|
|
static LLaDAImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
|
|
LLaDAImageConfig config;
|
|
int64_t detected_q_dim = 0;
|
|
int64_t detected_kv_dim = 0;
|
|
|
|
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
|
if (!starts_with(name, prefix)) {
|
|
continue;
|
|
}
|
|
if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
|
|
int64_t patch_area = config.patch_size * config.patch_size;
|
|
config.in_channels = tensor_storage.ne[0] / patch_area;
|
|
config.hidden_size = tensor_storage.ne[1];
|
|
} else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
|
|
config.cap_feat_dim = tensor_storage.ne[0];
|
|
config.hidden_size = tensor_storage.ne[1];
|
|
} else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) {
|
|
config.semantic_feat_dim = tensor_storage.ne[0];
|
|
} else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) {
|
|
detected_q_dim = tensor_storage.ne[1];
|
|
} else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) {
|
|
detected_kv_dim = tensor_storage.ne[1];
|
|
} else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
|
|
int64_t patch_area = config.patch_size * config.patch_size;
|
|
config.out_channels = tensor_storage.ne[1] / patch_area;
|
|
}
|
|
}
|
|
|
|
int64_t detected_layers = count_blocks(tensor_storage_map, prefix, "layers.");
|
|
int64_t detected_refiner = std::max(count_blocks(tensor_storage_map, prefix, "noise_refiner."),
|
|
count_blocks(tensor_storage_map, prefix, "context_refiner."));
|
|
if (detected_layers > 0) {
|
|
config.num_layers = detected_layers;
|
|
}
|
|
if (detected_refiner > 0) {
|
|
config.num_refiner_layers = detected_refiner;
|
|
}
|
|
if (detected_q_dim > 0) {
|
|
config.num_heads = detected_q_dim / config.head_dim;
|
|
}
|
|
if (detected_kv_dim > 0) {
|
|
config.num_kv_heads = detected_kv_dim / config.head_dim;
|
|
} else if (detected_q_dim > 0) {
|
|
config.num_kv_heads = config.num_heads;
|
|
}
|
|
|
|
LOG_VERBOSE("llada_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64 ", cap_feat_dim = %" PRId64 ", semantic_feat_dim = %" PRId64,
|
|
config.num_layers,
|
|
config.num_refiner_layers,
|
|
config.hidden_size,
|
|
config.num_heads,
|
|
config.num_kv_heads,
|
|
config.in_channels,
|
|
config.out_channels,
|
|
config.cap_feat_dim,
|
|
config.semantic_feat_dim);
|
|
return config;
|
|
}
|
|
};
|
|
|
|
class LLaDAImageModel : public GGMLBlock {
|
|
protected:
|
|
LLaDAImageConfig config;
|
|
|
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
|
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
|
|
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
|
|
params["sigvq_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
|
|
}
|
|
|
|
std::shared_ptr<ZImage::JointTransformerBlock> make_block(bool modulation) {
|
|
return std::make_shared<ZImage::JointTransformerBlock>(0,
|
|
config.hidden_size,
|
|
config.head_dim,
|
|
config.num_heads,
|
|
config.num_kv_heads,
|
|
config.multiple_of,
|
|
config.ffn_dim_multiplier,
|
|
config.norm_eps,
|
|
config.qk_norm,
|
|
modulation,
|
|
false,
|
|
true);
|
|
}
|
|
|
|
public:
|
|
LLaDAImageModel() = default;
|
|
LLaDAImageModel(LLaDAImageConfig config)
|
|
: config(config) {
|
|
blocks["x_embedder"] = std::make_shared<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
|
|
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, ZImage::ADALN_EMBED_DIM);
|
|
|
|
blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps, false);
|
|
blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
|
|
|
|
blocks["semantic_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
|
|
blocks["semantic_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
|
|
blocks["sigvq_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
|
|
blocks["sigvq_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
|
|
|
|
for (int i = 0; i < config.num_refiner_layers; i++) {
|
|
blocks["noise_refiner." + std::to_string(i)] = make_block(true);
|
|
blocks["context_refiner." + std::to_string(i)] = make_block(false);
|
|
blocks["sigvq_refiner." + std::to_string(i)] = make_block(false);
|
|
}
|
|
for (int i = 0; i < config.num_layers; i++) {
|
|
blocks["layers." + std::to_string(i)] = make_block(true);
|
|
}
|
|
|
|
blocks["final_layer"] = std::make_shared<ZImage::FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
|
|
}
|
|
|
|
ggml_tensor* forward_core(GGMLRunnerContext* ctx,
|
|
ggml_tensor* x,
|
|
ggml_tensor* timestep,
|
|
ggml_tensor* context,
|
|
ggml_tensor* pe) {
|
|
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
|
|
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
|
|
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
|
|
auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
|
|
auto final_layer = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
|
|
|
|
auto txt_pad_token = params["cap_pad_token"];
|
|
auto img_pad_token = params["x_pad_token"];
|
|
|
|
int64_t N = x->ne[2];
|
|
int64_t n_img_token = x->ne[1];
|
|
int64_t n_txt_token = context->ne[1];
|
|
|
|
// sdcpp's flow denoiser already hands over sigma * 1000, which is the range the
|
|
// reference reaches via its own t_scale, so no further scaling here.
|
|
auto t_emb = t_embedder->forward(ctx, timestep);
|
|
|
|
auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size]
|
|
auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size]
|
|
sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.prelude", "txt");
|
|
sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.prelude", "img");
|
|
sd::ggml_graph_cut::mark_graph_cut(t_emb, "llada_image.prelude", "t_emb");
|
|
|
|
int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), ZImage::SEQ_MULTI_OF);
|
|
if (n_txt_pad_token > 0) {
|
|
auto txt_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, txt_pad_token, txt_pad_token->ne[0], n_txt_pad_token, N, 1);
|
|
txt = ggml_concat(ctx->ggml_ctx, txt, txt_pad_tokens, 1);
|
|
}
|
|
|
|
int64_t n_img_pad_token = Rope::bound_mod(static_cast<int>(n_img_token), ZImage::SEQ_MULTI_OF);
|
|
if (n_img_pad_token > 0) {
|
|
auto img_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_img_pad_token, N, 1);
|
|
img = ggml_concat(ctx->ggml_ctx, img, img_pad_tokens, 1);
|
|
}
|
|
|
|
GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
|
|
|
|
auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
|
|
auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
|
|
|
|
for (int i = 0; i < config.num_refiner_layers; i++) {
|
|
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
|
|
|
|
txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
|
|
sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.context_refiner." + std::to_string(i), "txt");
|
|
}
|
|
|
|
for (int i = 0; i < config.num_refiner_layers; i++) {
|
|
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
|
|
|
|
img = block->forward(ctx, img, img_pe, nullptr, t_emb);
|
|
sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.noise_refiner." + std::to_string(i), "img");
|
|
}
|
|
|
|
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);
|
|
sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.prelude", "txt_img");
|
|
|
|
for (int i = 0; i < config.num_layers; i++) {
|
|
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
|
|
|
|
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
|
|
sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.layers." + std::to_string(i), "txt_img");
|
|
}
|
|
|
|
txt_img = final_layer->forward(ctx, txt_img, t_emb);
|
|
|
|
return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
|
|
}
|
|
|
|
ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) {
|
|
int64_t n_pad = Rope::bound_mod(static_cast<int>(x->ne[1]), ZImage::SEQ_MULTI_OF);
|
|
if (n_pad == 0) {
|
|
return x;
|
|
}
|
|
auto pads = ggml_repeat_4d(ctx->ggml_ctx, pad_token, pad_token->ne[0], n_pad, x->ne[2], 1);
|
|
return ggml_concat(ctx->ggml_ctx, x, pads, 1);
|
|
}
|
|
|
|
// Editing runs one joint sequence carrying two timesteps: the caption and source latent
|
|
// are clean (t = 0) while the second caption copy and the target latent are noisy. adaLN
|
|
// is a linear map of the timestep embedding, so feeding a per-token embedding selects the
|
|
// right modulation exactly, without duplicating the modulation projections.
|
|
ggml_tensor* forward_editing(GGMLRunnerContext* ctx,
|
|
ggml_tensor* x,
|
|
ggml_tensor* timestep,
|
|
ggml_tensor* context,
|
|
ggml_tensor* semantic,
|
|
ggml_tensor* source_latent,
|
|
ggml_tensor* pe) {
|
|
ggml_context* gctx = ctx->ggml_ctx;
|
|
|
|
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
|
|
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
|
|
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
|
|
auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
|
|
auto sigvq_embed_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["sigvq_embedder.0"]);
|
|
auto sigvq_embed_1 = std::dynamic_pointer_cast<Linear>(blocks["sigvq_embedder.1"]);
|
|
auto final_layer = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
|
|
|
|
auto t_noisy = t_embedder->forward(ctx, timestep);
|
|
auto t_clean = t_embedder->forward(ctx, ggml_scale(gctx, timestep, 0.f));
|
|
|
|
auto per_token = [&](ggml_tensor* emb, int64_t n) {
|
|
return ggml_repeat_4d(gctx, emb, emb->ne[0], n, 1, 1);
|
|
};
|
|
|
|
auto cap = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context));
|
|
cap = pad_stream(ctx, cap, params["cap_pad_token"]);
|
|
int64_t cap_len = cap->ne[1];
|
|
cap = ggml_concat(gctx, cap, cap, 1);
|
|
|
|
auto src = pad_stream(ctx, x_embedder->forward(ctx, source_latent), params["x_pad_token"]);
|
|
auto tgt_embed = x_embedder->forward(ctx, x);
|
|
int64_t n_img_token = tgt_embed->ne[1];
|
|
auto tgt = pad_stream(ctx, tgt_embed, params["x_pad_token"]);
|
|
int64_t img_len = tgt->ne[1];
|
|
auto img = ggml_concat(gctx, src, tgt, 1);
|
|
|
|
ggml_tensor* sig = nullptr;
|
|
int64_t sig_len = 0;
|
|
if (semantic != nullptr) {
|
|
sig = sigvq_embed_1->forward(ctx, sigvq_embed_0->forward(ctx, semantic));
|
|
sig = pad_stream(ctx, sig, params["sigvq_pad_token"]);
|
|
sig_len = sig->ne[1];
|
|
}
|
|
|
|
GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]);
|
|
|
|
auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2);
|
|
auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
|
|
|
|
auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1);
|
|
|
|
for (int i = 0; i < config.num_refiner_layers; i++) {
|
|
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
|
|
cap = block->forward(ctx, cap, cap_pe, nullptr, nullptr);
|
|
}
|
|
for (int i = 0; i < config.num_refiner_layers; i++) {
|
|
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
|
|
img = block->forward(ctx, img, img_pe, nullptr, img_adaln);
|
|
}
|
|
if (sig != nullptr) {
|
|
auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
|
|
for (int i = 0; i < config.num_refiner_layers; i++) {
|
|
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["sigvq_refiner." + std::to_string(i)]);
|
|
sig = block->forward(ctx, sig, sig_pe, nullptr, nullptr);
|
|
}
|
|
}
|
|
|
|
auto seq = ggml_concat(gctx, cap, img, 1);
|
|
|
|
auto cap_adaln = ggml_concat(gctx, per_token(t_clean, cap_len), per_token(t_noisy, cap_len), 1);
|
|
auto seq_adaln = ggml_concat(gctx, cap_adaln, img_adaln, 1);
|
|
if (sig != nullptr) {
|
|
seq = ggml_concat(gctx, seq, sig, 1);
|
|
seq_adaln = ggml_concat(gctx, seq_adaln, per_token(t_clean, sig_len), 1);
|
|
}
|
|
|
|
for (int i = 0; i < config.num_layers; i++) {
|
|
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
|
|
seq = block->forward(ctx, seq, pe, nullptr, seq_adaln);
|
|
sd::ggml_graph_cut::mark_graph_cut(seq, "llada_image.layers." + std::to_string(i), "seq");
|
|
}
|
|
|
|
seq = final_layer->forward(ctx, seq, seq_adaln);
|
|
|
|
// Only the target latent is denoised; the source half of the image stream is context.
|
|
// The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of
|
|
// the latent grid that unpatchify reconstructs.
|
|
int64_t target_start = cap_len * 2 + img_len;
|
|
return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token);
|
|
}
|
|
|
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
|
ggml_tensor* x,
|
|
ggml_tensor* timestep,
|
|
ggml_tensor* context,
|
|
ggml_tensor* pe) {
|
|
// x: [N, C, H, W]
|
|
// timestep: [N,]
|
|
// context: [N, L, cap_feat_dim]
|
|
// pe: [L, d_head/2, 2, 2]
|
|
// return: [N, C, H, W]
|
|
int64_t W = x->ne[0];
|
|
int64_t H = x->ne[1];
|
|
|
|
int patch_size = config.patch_size;
|
|
|
|
auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
|
|
|
|
auto out = forward_core(ctx, img, timestep, context, pe);
|
|
|
|
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);
|
|
|
|
// The reference pipeline negates the model output before the scheduler step.
|
|
return ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
|
|
}
|
|
};
|
|
|
|
struct LLaDAImageRunner : public DiffusionModelRunner {
|
|
public:
|
|
LLaDAImageConfig config;
|
|
LLaDAImageModel llada_image;
|
|
std::vector<float> pe_vec;
|
|
|
|
LLaDAImageRunner(ggml_backend_t backend,
|
|
const String2TensorStorage& tensor_storage_map = {},
|
|
const std::string prefix = "",
|
|
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
|
: DiffusionModelRunner(backend, prefix, weight_manager),
|
|
config(LLaDAImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
|
|
llada_image = LLaDAImageModel(config);
|
|
llada_image.init(params_ctx, tensor_storage_map, prefix);
|
|
}
|
|
|
|
std::string get_desc() override {
|
|
return "llada_image";
|
|
}
|
|
|
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
|
|
llada_image.get_param_tensors(tensors, prefix);
|
|
}
|
|
|
|
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
|
const sd::Tensor<float>& timesteps_tensor,
|
|
const sd::Tensor<float>& context_tensor) {
|
|
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE);
|
|
ggml_tensor* x = make_input(x_tensor);
|
|
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
|
GGML_ASSERT(x->ne[3] == 1);
|
|
GGML_ASSERT(!context_tensor.empty());
|
|
ggml_tensor* context = make_input(context_tensor);
|
|
|
|
pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
|
|
static_cast<int>(x->ne[0]),
|
|
config.patch_size,
|
|
static_cast<int>(x->ne[3]),
|
|
static_cast<int>(context->ne[1]),
|
|
ZImage::SEQ_MULTI_OF,
|
|
config.theta,
|
|
config.axes_dim));
|
|
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 = llada_image.forward(&runner_ctx, x, timesteps, context, pe);
|
|
|
|
ggml_build_forward_expand(gf, out);
|
|
|
|
return gf;
|
|
}
|
|
|
|
sd::Tensor<float> compute(int n_threads,
|
|
const sd::Tensor<float>& x,
|
|
const sd::Tensor<float>& timesteps,
|
|
const sd::Tensor<float>& context) {
|
|
// x: [N, in_channels, h, w]
|
|
// timesteps: [N, ]
|
|
// context: [N, max_position, cap_feat_dim]
|
|
auto get_graph = [&]() -> ggml_cgraph* {
|
|
return build_graph(x, timesteps, context);
|
|
};
|
|
|
|
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
|
|
}
|
|
|
|
ggml_cgraph* build_edit_graph(const sd::Tensor<float>& x_tensor,
|
|
const sd::Tensor<float>& timesteps_tensor,
|
|
const sd::Tensor<float>& context_tensor,
|
|
const sd::Tensor<float>& semantic_tensor,
|
|
const sd::Tensor<float>& source_tensor) {
|
|
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE);
|
|
ggml_tensor* x = make_input(x_tensor);
|
|
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
|
ggml_tensor* context = make_input(context_tensor);
|
|
ggml_tensor* semantic = make_optional_input(semantic_tensor);
|
|
ggml_tensor* source = make_input(source_tensor);
|
|
GGML_ASSERT(x->ne[3] == 1);
|
|
|
|
pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
|
|
static_cast<int>(x->ne[0]),
|
|
config.patch_size,
|
|
static_cast<int>(context->ne[1]),
|
|
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
|
|
ZImage::SEQ_MULTI_OF,
|
|
config.theta,
|
|
config.axes_dim));
|
|
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();
|
|
|
|
int64_t W = x->ne[0];
|
|
int64_t H = x->ne[1];
|
|
auto target = DiT::pad_and_patchify(&runner_ctx, x, config.patch_size, config.patch_size, false);
|
|
auto src = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false);
|
|
|
|
auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe);
|
|
out = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false);
|
|
out = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f);
|
|
|
|
ggml_build_forward_expand(gf, out);
|
|
return gf;
|
|
}
|
|
|
|
sd::Tensor<float> compute(int n_threads,
|
|
const DiffusionParams& diffusion_params) override {
|
|
GGML_ASSERT(diffusion_params.x != nullptr);
|
|
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
|
|
|
const auto* extra = std::get_if<LLaDAImageDiffusionExtra>(&diffusion_params.extra);
|
|
bool has_semantic = extra != nullptr && extra->semantic != nullptr && !extra->semantic->empty();
|
|
bool has_ref_latent = diffusion_params.ref_latents != nullptr && !diffusion_params.ref_latents->empty();
|
|
if (has_semantic && !has_ref_latent) {
|
|
LOG_WARN("llada_image: SigVQ features without a reference latent are not supported; falling back to text to image");
|
|
}
|
|
if (has_ref_latent) {
|
|
const auto& source = diffusion_params.ref_latents->front();
|
|
if (source.shape() != diffusion_params.x->shape()) {
|
|
LOG_ERROR("llada_image: reference latent must match the target shape; use resize_vae_to_target=1");
|
|
return {};
|
|
}
|
|
auto get_graph = [&]() -> ggml_cgraph* {
|
|
return build_edit_graph(*diffusion_params.x,
|
|
*diffusion_params.timesteps,
|
|
tensor_or_empty(diffusion_params.context),
|
|
tensor_or_empty(extra != nullptr ? extra->semantic : nullptr),
|
|
source);
|
|
};
|
|
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
|
|
diffusion_params.x->dim());
|
|
}
|
|
|
|
return compute(n_threads,
|
|
*diffusion_params.x,
|
|
*diffusion_params.timesteps,
|
|
tensor_or_empty(diffusion_params.context));
|
|
}
|
|
};
|
|
|
|
} // namespace LLaDAImage
|
|
|
|
#endif // __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
|