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/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));