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

feat: add qwen image layered support (#1119)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-07-03 00:21:22 +08:00
1 parent 3590aa8d62
commit 556f04bb3f
20 files changed
+366 -165
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/common/rope.hpp
+83 -61
View File
@@ -12,6 +12,16 @@ namespace Rope {
ErnieImage,
};
enum class RefIndexMode {
FIXED,
INCREASE,
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);
@@ -346,7 +356,7 @@ namespace Rope {
int axes_dim_num,
int start_index,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
float ref_index_scale,
bool scale_rope,
int base_offset = 0) {
@@ -357,13 +367,15 @@ namespace Rope {
for (ggml_tensor* ref : ref_latents) {
int h_offset = 0;
int w_offset = 0;
if (!increase_ref_index) {
if (ref_index_mode == RefIndexMode::FIXED) {
if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
w_offset = curr_w_offset;
} else {
h_offset = curr_h_offset;
}
scale_rope = false;
} else if (ref_index_mode == RefIndexMode::DECREASE) {
index--;
}
auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
@@ -377,7 +389,7 @@ namespace Rope {
scale_rope);
ids = concat_ids(ids, ref_ids, bs);
if (increase_ref_index) {
if (ref_index_mode == RefIndexMode::INCREASE) {
index++;
}
@@ -395,7 +407,7 @@ namespace Rope {
int context_len,
std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
float ref_index_scale,
bool is_longcat) {
int x_index = is_longcat ? 1 : 0;
@@ -406,7 +418,7 @@ namespace Rope {
auto ids = concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
@@ -420,7 +432,7 @@ namespace Rope {
int context_len,
std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
float ref_index_scale,
int theta,
bool circular_h,
@@ -435,7 +447,7 @@ namespace Rope {
context_len,
txt_arange_dims,
ref_latents,
increase_ref_index,
ref_index_mode,
ref_index_scale,
is_longcat);
std::vector<std::vector<int>> wrap_dims;
@@ -481,17 +493,64 @@ namespace Rope {
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
int h,
int w,
int pt,
int ph,
int pw,
int bs,
int t_offset = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false) {
int t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw;
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
if (scale_rope) {
h_offset -= h_len / 2;
w_offset -= w_len / 2;
}
std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
for (int i = 0; i < t_len; ++i) {
for (int j = 0; j < h_len; ++j) {
for (int k = 0; k < w_len; ++k) {
int idx = i * h_len * w_len + j * w_len + k;
vid_ids[idx][0] = t_ids[i];
vid_ids[idx][1] = h_ids[j];
vid_ids[idx][2] = w_ids[k];
}
}
}
std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
for (int i = 0; i < bs; ++i) {
for (int j = 0; j < vid_ids.size(); ++j) {
vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
}
}
return vid_ids_repeated;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index) {
RefIndexMode ref_index_mode) {
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
int txt_id_start = std::max(h_len, w_len);
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
int txt_id_start = std::max(h_len, w_len) / 2;
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * txt_id_start + context_len - 1, context_len);
std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
for (int i = 0; i < bs; ++i) {
for (int j = 0; j < txt_ids.size(); ++j) {
@@ -499,28 +558,30 @@ namespace Rope {
}
}
int axes_dim_num = 3;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
ids = concat_ids(ids, refs_ids, bs);
int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
}
// Generate qwen_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int h,
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
@@ -533,7 +594,7 @@ namespace Rope {
// Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
size_t cursor = context_len; // ignore text tokens
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
@@ -684,46 +745,6 @@ namespace Rope {
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
int h,
int w,
int pt,
int ph,
int pw,
int bs,
int t_offset = 0,
int h_offset = 0,
int w_offset = 0) {
int t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw;
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
for (int i = 0; i < t_len; ++i) {
for (int j = 0; j < h_len; ++j) {
for (int k = 0; k < w_len; ++k) {
int idx = i * h_len * w_len + j * w_len + k;
vid_ids[idx][0] = t_ids[i];
vid_ids[idx][1] = h_ids[j];
vid_ids[idx][2] = w_ids[k];
}
}
}
std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
for (int i = 0; i < bs; ++i) {
for (int j = 0; j < vid_ids.size(); ++j) {
vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
}
}
return vid_ids_repeated;
}
// Generate wan positional embeddings
__STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
int h,
@@ -785,7 +806,8 @@ namespace Rope {
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index) {
RefIndexMode ref_index_mode) {
SD_UNUSED(ref_index_mode);
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
for (int i = 0; i < bs * padded_context_len; i++) {
@@ -816,12 +838,12 @@ namespace Rope {
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index);
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
src/model/diffusion/anima.hpp
+1 -1
View File
@@ -612,7 +612,7 @@ namespace Anima {
0,
{},
empty_ref_latents,
false,
Rope::RefIndexMode::FIXED,
1.0f,
false);
src/model/diffusion/dit.hpp
+12 -12
View File
@@ -135,23 +135,23 @@ namespace DiT {
return x;
}
inline ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t t_len,
int64_t h_len,
int64_t w_len,
int pt,
int ph,
int pw) {
// x: [N, t_len*h_len*w_len, pt*ph*pw*C]
inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
ggml_tensor* x,
int64_t t_len,
int64_t h_len,
int64_t w_len,
int pt,
int ph,
int pw) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
int64_t N = x->ne[3];
int64_t N = x->ne[2];
int64_t C = x->ne[0] / pt / ph / pw;
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw]
src/model/diffusion/flux.hpp
+8 -8
View File
@@ -1485,7 +1485,7 @@ namespace Flux {
const sd::Tensor<float>& y_tensor = {},
const sd::Tensor<float>& guidance_tensor = {},
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
bool increase_ref_index = false,
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
std::vector<int> skip_layers = {},
const sd::Tensor<float>& pulid_id_tensor = {},
float pulid_id_weight = 1.0f) {
@@ -1527,8 +1527,8 @@ namespace Flux {
}
std::set<int> txt_arange_dims;
if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
txt_arange_dims = {3};
increase_ref_index = true;
txt_arange_dims = {3};
ref_index_mode = Rope::RefIndexMode::INCREASE;
} else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2};
}
@@ -1539,7 +1539,7 @@ namespace Flux {
static_cast<int>(context->ne[1]),
txt_arange_dims,
ref_latents,
increase_ref_index,
ref_index_mode,
config.ref_index_scale,
config.theta,
circular_y_enabled,
@@ -1599,7 +1599,7 @@ namespace Flux {
const sd::Tensor<float>& y = {},
const sd::Tensor<float>& guidance = {},
const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false,
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED,
std::vector<int> skip_layers = std::vector<int>(),
const sd::Tensor<float>& pulid_id = {},
float pulid_id_weight = 1.0f) {
@@ -1610,7 +1610,7 @@ namespace Flux {
// guidance: [N, ]
// pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers, pulid_id, pulid_id_weight);
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -1632,7 +1632,7 @@ namespace Flux {
tensor_or_empty(diffusion_params.y),
tensor_or_empty(extra->guidance),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index,
diffusion_params.ref_index_mode,
extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
tensor_or_empty(extra->pulid_id),
extra->pulid_id_weight);
@@ -1683,7 +1683,7 @@ namespace Flux {
{},
guidance,
{},
false);
Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
src/model/diffusion/model.hpp
+2 -1
View File
@@ -7,6 +7,7 @@
#include "core/ggml_extend.hpp"
#include "core/tensor_ggml.hpp"
#include "model/common/rope.hpp"
#include "model_manager.h"
struct UNetDiffusionExtra {
@@ -73,7 +74,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;
bool increase_ref_index = false;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
DiffusionExtraParams extra = std::monostate{};
};
src/model/diffusion/qwen_image.hpp
+93 -27
View File
@@ -4,6 +4,7 @@
#include <memory>
#include "model/common/block.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
#include "model_loader.h"
@@ -23,6 +24,7 @@ namespace Qwen {
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
bool zero_cond_t = false;
bool use_additional_t_cond = false;
static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
QwenImageConfig config;
@@ -88,19 +90,33 @@ namespace Qwen {
};
struct QwenTimestepProjEmbeddings : public GGMLBlock {
protected:
bool use_additional_t_cond = false;
public:
QwenTimestepProjEmbeddings(int64_t embedding_dim) {
QwenTimestepProjEmbeddings(int64_t embedding_dim, bool use_additional_t_cond = false)
: use_additional_t_cond(use_additional_t_cond) {
blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
if (use_additional_t_cond) {
blocks["addition_t_embedding"] = std::shared_ptr<GGMLBlock>(new Embedding(2, embedding_dim));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps) {
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"]);
auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
auto timesteps_emb = timestep_embedder->forward(ctx, timesteps_proj);
if (use_additional_t_cond) {
GGML_ASSERT(addition_t_cond != nullptr);
auto addition_t_embedding = std::dynamic_pointer_cast<Embedding>(blocks["addition_t_embedding"]);
auto addition_t_emb = addition_t_embedding->forward(ctx, addition_t_cond);
timesteps_emb = ggml_add(ctx->ggml_ctx, timesteps_emb, addition_t_emb);
}
return timesteps_emb;
}
};
@@ -402,7 +418,7 @@ namespace Qwen {
QwenImageModel(QwenImageConfig config)
: 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));
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
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));
@@ -424,6 +440,7 @@ namespace Qwen {
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* addition_t_cond,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* modulate_index = nullptr) {
@@ -434,9 +451,9 @@ namespace Qwen {
auto norm_out = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto t_emb = time_text_embed->forward(ctx, timestep);
auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond);
if (config.zero_cond_t) {
auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond);
t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
}
auto img = img_in->forward(ctx, x);
@@ -469,33 +486,50 @@ namespace Qwen {
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* addition_t_cond,
ggml_tensor* context,
ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {},
ggml_tensor* modulate_index = nullptr) {
// Forward pass of DiT.
// x: [N, C, H, W]
// x: [N, C, H, W] or [N*C, T, H, W]
// timestep: [N,]
// context: [N, L, D]
// pe: [L, d_head/2, 2, 2]
// return: [N, C, H, W]
// return: [N, C, H, W] or [N*C, T, H, W]
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int64_t N = x->ne[3];
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = 1;
int64_t N = addition_t_cond != nullptr ? addition_t_cond->ne[0] : x->ne[3];
bool has_time_axis = false;
if (x->ne[3] != 1) {
T = x->ne[2];
has_time_axis = true;
}
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* {
input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size);
if (!has_time_axis) {
return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size);
}
if (input->ne[3] == 1) {
input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
}
return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
};
auto img = patchify_input(x);
int64_t img_tokens = img->ne[1];
if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) {
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
ref = patchify_input(ref);
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
}
}
auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index); // [N, h_len*w_len, ph*pw*C]
auto out = forward_orig(ctx, img, timestep, addition_t_cond, context, pe, modulate_index); // [N, h_len*w_len, ph*pw*C]
if (out->ne[1] > img_tokens) {
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [num_tokens, N, C * patch_size * patch_size]
@@ -503,7 +537,17 @@ namespace Qwen {
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [N, h*w, C * patch_size * patch_size]
}
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W]
if (has_time_axis) {
int pad_h = (config.patch_size - H % config.patch_size) % config.patch_size;
int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size;
int h_len = static_cast<int>((H + pad_h) / config.patch_size);
int w_len = static_cast<int>((W + pad_w) / config.patch_size);
out = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w]
out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W]
} else {
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W]
}
return out;
}
@@ -515,6 +559,7 @@ namespace Qwen {
QwenImageModel qwen_image;
std::vector<float> pe_vec;
std::vector<float> modulate_index_vec;
std::vector<int32_t> additional_t_cond_vec;
SDVersion version;
QwenImageRunner(ggml_backend_t backend,
@@ -524,9 +569,13 @@ namespace Qwen {
bool zero_cond_t = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
version(version) {
config.zero_cond_t = config.zero_cond_t || zero_cond_t;
qwen_image = QwenImageModel(config);
if (version == VERSION_QWEN_IMAGE_LAYERED) {
config.use_additional_t_cond = true;
}
qwen_image = QwenImageModel(config);
qwen_image.init(params_ctx, tensor_storage_map, prefix);
}
@@ -542,11 +591,11 @@ namespace Qwen {
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
bool increase_ref_index = false) {
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::INCREASE) {
ggml_cgraph* gf = new_graph_custom(QWEN_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(x->ne[3] == 1 || x_tensor.dim() == 5);
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
std::vector<ggml_tensor*> ref_latents;
@@ -555,13 +604,29 @@ namespace Qwen {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
int batch_size = static_cast<int>(x->ne[3]);
int time_len = 1;
if (x_tensor.dim() == 5) {
time_len = static_cast<int>(x_tensor.shape()[2]);
batch_size = static_cast<int>(x_tensor.shape()[4]);
}
ggml_tensor* addition_t_cond = nullptr;
if (version == VERSION_QWEN_IMAGE_LAYERED) {
additional_t_cond_vec.assign(static_cast<size_t>(batch_size), 0);
addition_t_cond = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, batch_size);
set_backend_tensor_data(addition_t_cond, additional_t_cond_vec.data());
ref_index_mode = Rope::RefIndexMode::DECREASE;
}
pe_vec = Rope::gen_qwen_image_pe(time_len,
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
increase_ref_index,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
@@ -604,6 +669,7 @@ namespace Qwen {
ggml_tensor* out = qwen_image.forward(&runner_ctx,
x,
timesteps,
addition_t_cond,
context,
pe,
ref_latents,
@@ -619,12 +685,12 @@ namespace Qwen {
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false) {
// x: [N, in_channels, h, w]
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::INCREASE) {
// x: [N, C, H, W] or [N*C, T, H, W]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -640,7 +706,7 @@ namespace Qwen {
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index);
diffusion_params.ref_index_mode);
}
void test() {
@@ -674,7 +740,7 @@ namespace Qwen {
timesteps,
context,
{},
false);
Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
src/model/diffusion/z_image.hpp
+6 -6
View File
@@ -575,7 +575,7 @@ namespace ZImage {
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
bool increase_ref_index = false) {
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -595,7 +595,7 @@ namespace ZImage {
static_cast<int>(context->ne[1]),
SEQ_MULTI_OF,
ref_latents,
increase_ref_index,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
@@ -626,12 +626,12 @@ namespace ZImage {
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false) {
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -647,7 +647,7 @@ namespace ZImage {
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index);
diffusion_params.ref_index_mode);
}
void test() {
@@ -681,7 +681,7 @@ namespace ZImage {
timesteps,
context,
{},
false);
Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
src/model/vae/wan_vae.hpp
+30 -23
View File
@@ -631,6 +631,7 @@ namespace WAN {
class Encoder3d : public GGMLBlock {
protected:
bool wan2_2;
int64_t in_channels;
int64_t dim;
int64_t z_dim;
std::vector<int> dim_mult;
@@ -641,12 +642,14 @@ namespace WAN {
public:
Encoder3d(int64_t dim = 128,
int64_t z_dim = 4,
int64_t in_channels = 3,
std::vector<int> dim_mult = {1, 2, 4, 4},
int num_res_blocks = 2,
std::vector<bool> temperal_downsample = {false, true, true},
bool wan2_2 = false,
bool is_2D = false)
: dim(dim),
: in_channels(in_channels),
dim(dim),
z_dim(z_dim),
dim_mult(dim_mult),
num_res_blocks(num_res_blocks),
@@ -659,11 +662,10 @@ namespace WAN {
dims.push_back(dim * u);
}
int64_t input_dim = wan2_2 ? 12 : 3;
if (is_2D) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(input_dim, dims[0], {3, 3}, {1, 1}, {1, 1}));
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_channels, dims[0], {3, 3}, {1, 1}, {1, 1}));
} else {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(input_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_channels, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
}
int index = 0;
@@ -812,6 +814,7 @@ namespace WAN {
class Decoder3d : public GGMLBlock {
protected:
bool wan2_2;
int64_t out_channels;
int64_t dim;
int64_t z_dim;
std::vector<int> dim_mult;
@@ -822,12 +825,14 @@ namespace WAN {
public:
Decoder3d(int64_t dim = 128,
int64_t z_dim = 4,
int64_t out_channels = 3,
std::vector<int> dim_mult = {1, 2, 4, 4},
int num_res_blocks = 2,
std::vector<bool> temperal_upsample = {true, true, false},
bool wan2_2 = false,
bool is_2D = false)
: dim(dim),
: out_channels(out_channels),
dim(dim),
z_dim(z_dim),
dim_mult(dim_mult),
num_res_blocks(num_res_blocks),
@@ -889,7 +894,7 @@ namespace WAN {
// output blocks
blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
int64_t final_dim = wan2_2 ? 12 : 3;
int64_t final_dim = out_channels;
// head.1 is nn.SiLU()
if (is_2D) {
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, final_dim, {3, 3}, {1, 1}, {1, 1}));
@@ -1002,6 +1007,8 @@ namespace WAN {
public:
bool wan2_2 = false;
bool decode_only = true;
int64_t input_channels = 3;
int patch_size = 1;
int64_t dim = 96;
int64_t dec_dim = 96;
int64_t z_dim = 16;
@@ -1026,16 +1033,22 @@ namespace WAN {
}
public:
WanVAE(bool decode_only = true, bool wan2_2 = false, bool is_2D = false)
: decode_only(decode_only), wan2_2(wan2_2), is_2D(is_2D) {
WanVAE(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_2D = false)
: decode_only(decode_only),
wan2_2(version == VERSION_WAN2_2_TI2V),
is_2D(is_2D) {
// attn_scales is always []
if (wan2_2) {
dim = 160;
dec_dim = 256;
z_dim = 48;
dim = 160;
dec_dim = 256;
z_dim = 48;
input_channels = 12;
patch_size = 2;
_conv_num = 34;
_enc_conv_num = 26;
} else if (version == VERSION_QWEN_IMAGE_LAYERED) {
input_channels = 4;
}
if (is_2D) {
@@ -1044,14 +1057,14 @@ namespace WAN {
}
if (!decode_only) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, temperal_downsample, wan2_2, is_2D));
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, input_channels, dim_mult, num_res_blocks, temperal_downsample, wan2_2, is_2D));
if (is_2D) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1}));
} else {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1}));
}
}
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, dim_mult, num_res_blocks, temperal_upsample, wan2_2, is_2D));
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, input_channels, dim_mult, num_res_blocks, temperal_upsample, wan2_2, is_2D));
if (is_2D) {
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, z_dim, {1, 1}));
} else {
@@ -1125,9 +1138,7 @@ namespace WAN {
clear_cache();
if (wan2_2) {
x = patchify(ctx->ggml_ctx, x, 2, b);
}
x = patchify(ctx->ggml_ctx, x, patch_size, b);
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
@@ -1202,9 +1213,7 @@ namespace WAN {
}
}
}
if (wan2_2) {
out = unpatchify(ctx->ggml_ctx, out, 2, b);
}
out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
clear_cache();
return out;
@@ -1225,9 +1234,7 @@ namespace WAN {
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w]
_conv_idx = 0;
auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
if (wan2_2) {
out = unpatchify(ctx->ggml_ctx, out, 2, b);
}
out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
return out;
}
@@ -1257,7 +1264,7 @@ namespace WAN {
if (is_2D) {
LOG_DEBUG("USING 2D VAE");
}
ae = WanVAE(decode_only, version == VERSION_WAN2_2_TI2V, is_2D);
ae = WanVAE(decode_only, version, is_2D);
ae.init(params_ctx, tensor_storage_map, prefix);
}