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

feat: add Mage-Flow support (#1808)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-07-23 01:23:25 +08:00
1 parent 35fb21f3e0
commit 8a51eb9284
18 files changed
+861 -26
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/common/rope.hpp
+37
View File
@@ -654,6 +654,43 @@ namespace Rope {
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
}
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
int w,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
const int axes_dim_num = static_cast<int>(axes_dim.size());
auto make_image_ids = [=](int image_h, int image_w, int image_index) {
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
std::vector<float>(axes_dim_num, 0.f));
int h_start = -(image_h - image_h / 2);
int w_start = -(image_w - image_w / 2);
for (int b = 0; b < bs; ++b) {
for (int y = 0; y < image_h; ++y) {
for (int x = 0; x < image_w; ++x) {
auto& id = image_ids[static_cast<size_t>(b) * image_h * image_w + y * image_w + x];
id[0] = static_cast<float>(image_index);
id[1] = static_cast<float>(h_start + y);
id[2] = static_cast<float>(w_start + x);
}
}
}
return image_ids;
};
auto ids = gen_flux_txt_ids(bs, context_len, axes_dim_num, {});
auto img_ids = make_image_ids(h, w, 0);
ids = concat_ids(ids, img_ids, bs);
for (size_t i = 0; i < ref_latents.size(); ++i) {
auto ref_ids = make_image_ids(static_cast<int>(ref_latents[i]->ne[1]),
static_cast<int>(ref_latents[i]->ne[0]),
static_cast<int>(i + 1));
ids = concat_ids(ids, ref_ids, bs);
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
int w,
int bs,
src/model/diffusion/mage_flow.hpp
+162
View File
@@ -0,0 +1,162 @@
#ifndef __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
#define __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
#include <cmath>
#include <memory>
#include "model/diffusion/qwen_image.hpp"
namespace MageFlow {
constexpr int MAGE_FLOW_GRAPH_SIZE = 20480;
// Mage-Flow was trained with BF16-rounded timestep frequencies; using Qwen's F32 projection degrades generation quality.
struct MageFlowTimestepProjEmbeddings : public Qwen::QwenTimestepProjEmbeddings {
static constexpr int TIMESTEP_DIM = 256;
static constexpr int HALF_DIM = TIMESTEP_DIM / 2;
std::vector<float> frequencies;
std::vector<float> timesteps_proj;
explicit MageFlowTimestepProjEmbeddings(int64_t embedding_dim)
: QwenTimestepProjEmbeddings(embedding_dim), frequencies(HALF_DIM) {
for (int i = 0; i < HALF_DIM; ++i) {
float frequency = std::exp(-std::log(10000.f) * static_cast<float>(i) / HALF_DIM);
frequencies[i] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(frequency));
}
}
void prepare(const sd::Tensor<float>& timesteps) {
size_t num_timesteps = static_cast<size_t>(timesteps.numel());
timesteps_proj.resize(static_cast<size_t>(TIMESTEP_DIM) * num_timesteps);
for (size_t b = 0; b < num_timesteps; ++b) {
float sigma = ggml_bf16_to_fp32(ggml_fp32_to_bf16(timesteps.values()[b] / 1000.f));
for (int i = 0; i < HALF_DIM; ++i) {
float argument = sigma * frequencies[i] * 1000.f;
timesteps_proj[b * TIMESTEP_DIM + i] =
ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::cos(argument)));
timesteps_proj[b * TIMESTEP_DIM + HALF_DIM + i] =
ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::sin(argument)));
}
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps,
ggml_tensor* addition_t_cond = nullptr) override {
GGML_ASSERT(addition_t_cond == nullptr);
GGML_ASSERT(timesteps_proj.size() ==
static_cast<size_t>(TIMESTEP_DIM * ggml_nelements(timesteps)));
auto projection = ggml_new_tensor_2d(ctx->ggml_ctx,
GGML_TYPE_F32,
TIMESTEP_DIM,
ggml_nelements(timesteps));
ctx->bind_backend_tensor_data(projection, timesteps_proj.data());
auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
return timestep_embedder->forward(ctx, projection);
}
};
struct MageFlowRunner : public DiffusionModelRunner {
public:
Qwen::QwenImageConfig config;
Qwen::QwenImageModel mage_flow;
std::shared_ptr<MageFlowTimestepProjEmbeddings> time_text_embed;
std::vector<float> pe_vec;
MageFlowRunner(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.patch_size = 1;
config.in_channels = 128;
config.out_channels = 128;
config.num_layers = 12;
config.attention_head_dim = 128;
config.num_attention_heads = 24;
config.joint_attention_dim = 2560;
config.theta = 10000;
config.axes_dim = {16, 56, 56};
config.axes_dim_sum = 128;
time_text_embed = std::make_shared<MageFlowTimestepProjEmbeddings>(
config.num_attention_heads * config.attention_head_dim);
mage_flow = Qwen::QwenImageModel(config, time_text_embed);
mage_flow.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "mage_flow";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
mage_flow.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,
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
ggml_cgraph* gf = new_graph_custom(MAGE_FLOW_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);
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));
}
int batch_size = static_cast<int>(x->ne[3]);
pe_vec = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
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());
time_text_embed->prepare(timesteps_tensor);
auto runner_ctx = get_context();
auto out = mage_flow.forward(&runner_ctx,
x,
timesteps,
nullptr,
context,
pe,
ref_latents);
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,
const std::vector<sd::Tensor<float>>& ref_latents = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
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),
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
}
};
} // namespace MageFlow
#endif // __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
src/model/diffusion/model.hpp
+2
View File
@@ -26,6 +26,7 @@ struct RefImageParams {
RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
int vlm_min_size = -1;
int vlm_max_size = -1;
bool resize_vae_to_target = false;
};
const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
@@ -34,6 +35,7 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
{"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}},
{"mage_flow", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, -1, 384, true}},
{"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}},
src/model/diffusion/qwen_image.hpp
+10 -6
View File
@@ -103,9 +103,9 @@ namespace Qwen {
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps,
ggml_tensor* addition_t_cond = nullptr) {
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps,
ggml_tensor* addition_t_cond = nullptr) {
// timesteps: [N,]
// return: [N, embedding_dim]
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
@@ -416,10 +416,14 @@ namespace Qwen {
public:
QwenImageModel() {}
QwenImageModel(QwenImageConfig config)
QwenImageModel(QwenImageConfig config,
std::shared_ptr<QwenTimestepProjEmbeddings> time_text_embed = nullptr)
: config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
if (time_text_embed == nullptr) {
time_text_embed = std::make_shared<QwenTimestepProjEmbeddings>(inner_dim, config.use_additional_t_cond);
}
blocks["time_text_embed"] = std::move(time_text_embed);
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
src/model/vae/mage_vae.hpp
+521
View File
@@ -0,0 +1,521 @@
#ifndef __SD_MODEL_VAE_MAGE_VAE_HPP__
#define __SD_MODEL_VAE_MAGE_VAE_HPP__
#include "model/diffusion/dit.hpp"
#include "model/vae/vae.hpp"
namespace MageVAE {
constexpr int MAGE_VAE_GRAPH_SIZE = 327680;
constexpr int HIDDEN_SIZE = 384;
constexpr int LATENT_CHANNELS = 128;
constexpr int PATCH_SIZE = 16;
struct LayerNorm2d : public UnaryBlock {
int64_t channels;
bool affine;
std::string prefix;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
this->prefix = prefix;
if (affine) {
params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
}
}
LayerNorm2d(int64_t channels, bool affine = true)
: channels(channels), affine(affine) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* weight = affine ? params["weight"] : nullptr;
ggml_tensor* bias = affine ? params["bias"] : nullptr;
if (affine && ctx->weight_adapter) {
weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
bias = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
}
// [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels.
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f);
return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
}
};
inline ggml_tensor* modulate_2d(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift);
}
inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx,
ggml_tensor* x,
Conv2d* projection) {
auto pooled = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
pooled = ggml_mean(ctx->ggml_ctx, pooled);
pooled = ggml_reshape_4d(ctx->ggml_ctx, pooled, 1, 1, x->ne[2], x->ne[3]);
pooled = ggml_sigmoid(ctx->ggml_ctx, projection->forward(ctx, pooled));
return ggml_mul(ctx->ggml_ctx, x, pooled);
}
struct TimestepEmbedder : public GGMLBlock {
TimestepEmbedder() {
blocks["mlp.0"] = std::make_shared<Linear>(256, HIDDEN_SIZE);
blocks["mlp.2"] = std::make_shared<Linear>(HIDDEN_SIZE, HIDDEN_SIZE);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) {
auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto x = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f);
x = linear_0->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
return linear_2->forward(ctx, x);
}
};
struct EncoderDiCoBlock : public UnaryBlock {
explicit EncoderDiCoBlock(int64_t channels) {
blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv2"] = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["conv3"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["ca.1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv4"] = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
blocks["conv5"] = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
blocks["norm1"] = std::make_shared<LayerNorm2d>(channels);
blocks["norm2"] = std::make_shared<LayerNorm2d>(channels);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
auto ca = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
auto x = norm1->forward(ctx, input);
x = conv1->forward(ctx, x);
x = conv2->forward(ctx, x);
x = ggml_gelu(ctx->ggml_ctx, x);
x = channel_attention(ctx, x, ca.get());
x = conv3->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, input, x);
auto h = norm2->forward(ctx, x);
h = conv4->forward(ctx, h);
h = ggml_gelu(ctx->ggml_ctx, h);
h = conv5->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct DiCoBlock : public GGMLBlock {
explicit DiCoBlock(int64_t channels) {
blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv2"] = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["conv3"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["ca.1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv4"] = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
blocks["conv5"] = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
blocks["norm1"] = std::make_shared<LayerNorm2d>(channels, false);
blocks["norm2"] = std::make_shared<LayerNorm2d>(channels, false);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(channels, channels * 6);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* condition) {
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
auto ca = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
auto ada = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0);
auto x = norm1->forward(ctx, input);
x = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]);
x = conv1->forward(ctx, x);
x = conv2->forward(ctx, x);
x = ggml_gelu(ctx->ggml_ctx, x);
x = channel_attention(ctx, x, ca.get());
x = conv3->forward(ctx, x);
auto gate_1 = ggml_reshape_4d(ctx->ggml_ctx, chunks[2], 1, 1, chunks[2]->ne[0], chunks[2]->ne[1]);
x = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1));
auto h = norm2->forward(ctx, x);
h = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]);
h = conv4->forward(ctx, h);
h = ggml_gelu(ctx->ggml_ctx, h);
h = conv5->forward(ctx, h);
auto gate_2 = ggml_reshape_4d(ctx->ggml_ctx, chunks[5], 1, 1, chunks[5]->ne[0], chunks[5]->ne[1]);
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, h, gate_2));
}
};
struct MageResnetBlock : public UnaryBlock {
explicit MageResnetBlock(int64_t channels) {
blocks["norm1"] = std::make_shared<GroupNorm32>(channels);
blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["norm2"] = std::make_shared<GroupNorm32>(channels);
blocks["conv2"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto norm2 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm2"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
auto x = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, input)));
x = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, x)));
return ggml_add(ctx->ggml_ctx, input, x);
}
};
inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx,
ggml_tensor* x,
int pad_w,
int pad_h) {
if (pad_w > 0) {
auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
edge = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]);
x = ggml_concat(ctx, x, edge, 0);
}
if (pad_h > 0) {
auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
edge = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]);
x = ggml_concat(ctx, x, edge, 1);
}
return x;
}
struct MageAttnBlock : public UnaryBlock {
int64_t channels;
int patch_size;
MageAttnBlock(int64_t channels, int patch_size = 32)
: channels(channels), patch_size(patch_size) {
blocks["norm"] = std::make_shared<GroupNorm32>(channels);
blocks["q"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["k"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["v"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["proj_out"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
}
ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) {
x = DiT::patchify(ctx, x, patch_size, patch_size);
x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
// [N, np, C, P] -> [N, np, P, C] for attention over P pixels.
x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]);
}
ggml_tensor* from_patches(ggml_context* ctx,
ggml_tensor* x,
int64_t patch_count,
int64_t batch_size,
int64_t h_patches,
int64_t w_patches) {
x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
// [N, np, P, C] -> [N, np, C, P] before spatial unpatchify.
x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size);
return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
auto q_proj = std::dynamic_pointer_cast<Conv2d>(blocks["q"]);
auto k_proj = std::dynamic_pointer_cast<Conv2d>(blocks["k"]);
auto v_proj = std::dynamic_pointer_cast<Conv2d>(blocks["v"]);
auto proj_out = std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"]);
int64_t width = input->ne[0];
int64_t height = input->ne[1];
int64_t batch = input->ne[3];
int pad_w = (patch_size - static_cast<int>(width % patch_size)) % patch_size;
int pad_h = (patch_size - static_cast<int>(height % patch_size)) % patch_size;
int64_t wp = (width + pad_w) / patch_size;
int64_t hp = (height + pad_h) / patch_size;
int64_t np = wp * hp;
auto h = norm->forward(ctx, input);
auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h);
auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h);
auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h);
q = to_patches(ctx->ggml_ctx, q);
k = to_patches(ctx->ggml_ctx, k);
v = to_patches(ctx->ggml_ctx, v);
h = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
h = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
if (pad_h > 0) {
h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);
}
if (pad_w > 0) {
h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width);
}
return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h));
}
};
struct Decoder : public UnaryBlock {
Decoder() {
blocks["conv_in"] = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["block.0"] = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
blocks["block.1"] = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
blocks["block.2"] = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
blocks["block.3"] = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
blocks["block.4"] = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
blocks["norm_out"] = std::make_shared<GroupNorm32>(HIDDEN_SIZE);
blocks["conv_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
x = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"])->forward(ctx, x);
for (int i = 0; i < 5; ++i) {
x = std::dynamic_pointer_cast<UnaryBlock>(blocks["block." + std::to_string(i)])->forward(ctx, x);
}
x = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm_out"])->forward(ctx, x);
x = ggml_silu(ctx->ggml_ctx, x);
return std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"])->forward(ctx, x);
}
};
struct DConvEncoder : public UnaryBlock {
DConvEncoder() {
blocks["patch_cond_embed"] = std::make_shared<Conv2d>(3, 768, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE});
for (int i = 0; i < 2; ++i) {
blocks["head_blocks." + std::to_string(i)] = std::make_shared<EncoderDiCoBlock>(768);
}
blocks["proj_down"] = std::make_shared<Conv2d>(768, HIDDEN_SIZE, std::pair{1, 1});
blocks["z_proj"] = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{1, 1});
blocks["fuse_proj"] = std::make_shared<Conv2d>(HIDDEN_SIZE * 2, HIDDEN_SIZE, std::pair{1, 1});
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>();
for (int i = 0; i < 21; ++i) {
blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
}
blocks["norm_out"] = std::make_shared<LayerNorm2d>(HIDDEN_SIZE);
blocks["proj_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, LATENT_CHANNELS * 2, std::pair{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image) override {
auto cond = std::dynamic_pointer_cast<Conv2d>(blocks["patch_cond_embed"])->forward(ctx, image);
for (int i = 0; i < 2; ++i) {
cond = std::dynamic_pointer_cast<EncoderDiCoBlock>(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond);
}
cond = std::dynamic_pointer_cast<Conv2d>(blocks["proj_down"])->forward(ctx, cond);
auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
z = std::dynamic_pointer_cast<Conv2d>(blocks["z_proj"])->forward(ctx, z);
z = ggml_concat(ctx->ggml_ctx, cond, z, 2);
z = std::dynamic_pointer_cast<Conv2d>(blocks["fuse_proj"])->forward(ctx, z);
auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1);
t = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]);
auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
for (int i = 0; i < 21; ++i) {
z = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, z, c);
}
z = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm_out"])->forward(ctx, z);
return std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"])->forward(ctx, z);
}
};
struct MLPResBlock : public GGMLBlock {
MLPResBlock() {
blocks["in_ln"] = std::make_shared<LayerNorm>(32, 1e-6f);
blocks["mlp.0"] = std::make_shared<Linear>(32, 32);
blocks["mlp.2"] = std::make_shared<Linear>(32, 32);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(32, 96);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) {
auto params = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0);
auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["in_ln"])->forward(ctx, x);
h = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]);
h = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, h);
h = ggml_silu(ctx->ggml_ctx, h);
h = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, chunks[2], h));
}
};
struct DConvDenoiser : public GGMLBlock {
DConvDenoiser() {
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>();
blocks["y_embedder_x"] = std::make_shared<Conv2d>(HIDDEN_SIZE, 32 * PATCH_SIZE * PATCH_SIZE, std::pair{1, 1});
blocks["x_embedder.embedder.0"] = std::make_shared<Linear>(3 + 32 + 64, 32);
blocks["s_embedder.proj1"] = std::make_shared<Conv2d>(3, LATENT_CHANNELS, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{0, 0}, std::pair{1, 1}, false);
blocks["s_embedder.proj2"] = std::make_shared<Conv2d>(LATENT_CHANNELS + HIDDEN_SIZE, HIDDEN_SIZE, std::pair{1, 1});
for (int i = 0; i < 21; ++i) {
blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
}
blocks["dec_net.cond_embed"] = std::make_shared<Linear>(HIDDEN_SIZE, PATCH_SIZE * PATCH_SIZE * 32);
blocks["dec_net.input_proj"] = std::make_shared<Linear>(32, 32);
for (int i = 0; i < 3; ++i) {
blocks["dec_net.res_blocks." + std::to_string(i)] = std::make_shared<MLPResBlock>();
}
blocks["final_layer.norm"] = std::make_shared<RMSNorm>(32);
blocks["final_layer.linear"] = std::make_shared<Linear>(32, 3);
blocks["y_embedder.decoder"] = std::make_shared<Decoder>();
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
auto cond = std::dynamic_pointer_cast<Decoder>(blocks["y_embedder.decoder"])->forward(ctx, latent);
int64_t w = cond->ne[0];
int64_t h = cond->ne[1];
int64_t n = cond->ne[3];
int64_t length = w * h;
auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
auto t = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1);
t = ggml_reshape_1d(ctx->ggml_ctx, t, n);
auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
auto s0 = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj1"])->forward(ctx, image);
s0 = ggml_concat(ctx->ggml_ctx, s0, cond, 2);
auto s = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj2"])->forward(ctx, s0);
for (int i = 0; i < 21; ++i) {
s = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c);
}
// [N, C, H, W] -> [N*H*W, C].
s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n);
auto y = std::dynamic_pointer_cast<Conv2d>(blocks["y_embedder_x"])->forward(ctx, cond);
// Split 32*P channels as [32, P], then produce [N*L, P, 32].
y = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n);
y = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
y = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n);
auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
dct = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1);
auto x = ggml_concat(ctx->ggml_ctx, zeros, y, 0);
x = ggml_concat(ctx->ggml_ctx, x, dct, 0);
x = std::dynamic_pointer_cast<Linear>(blocks["x_embedder.embedder.0"])->forward(ctx, x);
x = std::dynamic_pointer_cast<Linear>(blocks["dec_net.input_proj"])->forward(ctx, x);
auto dec_cond = std::dynamic_pointer_cast<Linear>(blocks["dec_net.cond_embed"])->forward(ctx, s);
dec_cond = ggml_reshape_3d(ctx->ggml_ctx, dec_cond, 32, PATCH_SIZE * PATCH_SIZE, length * n);
for (int i = 0; i < 3; ++i) {
x = std::dynamic_pointer_cast<MLPResBlock>(blocks["dec_net.res_blocks." + std::to_string(i)])->forward(ctx, x, dec_cond);
}
x = std::dynamic_pointer_cast<RMSNorm>(blocks["final_layer.norm"])->forward(ctx, x);
x = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, x);
// [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify.
x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n);
return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
}
};
struct MageVAEModel : public GGMLBlock {
MageVAEModel() {
blocks["student.dconv_encoder"] = std::make_shared<DConvEncoder>();
blocks["pipeline"] = std::make_shared<DConvDenoiser>();
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* image) {
return std::dynamic_pointer_cast<DConvEncoder>(blocks["student.dconv_encoder"])->forward(ctx, image);
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
return std::dynamic_pointer_cast<DConvDenoiser>(blocks["pipeline"])->forward(ctx, latent, dct);
}
};
struct MageVAERunner : public VAE {
MageVAEModel model;
std::vector<float> dct_vec;
MageVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(VERSION_MAGE_FLOW, backend, prefix, weight_manager) {
model = MageVAEModel();
model.init(params_ctx, tensor_storage_map, prefix);
dct_vec.resize(64 * PATCH_SIZE * PATCH_SIZE);
constexpr float pi = 3.14159265358979323846f;
for (int py = 0; py < PATCH_SIZE; ++py) {
float y = static_cast<float>(py) / static_cast<float>(PATCH_SIZE - 1);
for (int px = 0; px < PATCH_SIZE; ++px) {
float x = static_cast<float>(px) / static_cast<float>(PATCH_SIZE - 1);
int pos = py * PATCH_SIZE + px;
for (int fy = 0; fy < 8; ++fy) {
for (int fx = 0; fx < 8; ++fx) {
int freq = fx * 8 + fy;
float freq_x = static_cast<float>(fx) * 8.f / 7.f;
float freq_y = static_cast<float>(fy) * 8.f / 7.f;
float coeff = 1.f / (1.f + freq_x * freq_y);
dct_vec[freq + 64 * pos] = std::cos(x * freq_x * pi) *
std::cos(y * freq_y * pi) * coeff;
}
}
}
}
}
std::string get_desc() override {
return "mage_vae";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& input_tensor, bool decode_graph) {
ggml_cgraph* gf = new_graph_custom(MAGE_VAE_GRAPH_SIZE);
auto input = make_input(input_tensor);
auto runner_ctx = get_context();
ggml_tensor* dct = nullptr;
if (decode_graph) {
dct = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, 64, PATCH_SIZE * PATCH_SIZE, 1);
set_backend_tensor_data(dct, dct_vec.data());
}
auto out = decode_graph ? model.decode(&runner_ctx, input, dct) : model.encode(&runner_ctx, input);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& input,
bool decode_graph) override {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), input.dim());
}
int get_encoder_output_channels(int input_channels) override {
SD_UNUSED(input_channels);
return LATENT_CHANNELS * 2;
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
const auto chunks = sd::ops::chunk(vae_output, 2, 2);
const auto& mean = chunks[0];
const auto& logvar = chunks[1];
sd::Tensor<float> stddev = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -20.0f, 10.0f));
sd::Tensor<float> noise = sd::Tensor<float>::randn_like(mean, rng);
sd::Tensor<float> latents = mean + stddev * noise;
return latents;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return latents;
}
};
} // namespace MageVAE
#endif // __SD_MODEL_VAE_MAGE_VAE_HPP__
src/model/vae/vae.hpp
+1 -1
View File
@@ -74,7 +74,7 @@ public:
int scale_factor = 8;
if (version == VERSION_LTXAV) {
scale_factor = 32;
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version)) {
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16;