upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-07-30 23:00:41 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add qwen image layered support (#1119)

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

3
examples/cli/main.cpp
View File

@@ -481,7 +481,8 @@ bool save_results(const SDCliParams& cli_params,
if (!img.data)
return false;
const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx;
int images_per_batch = gen_params.batch_count > 0 ? std::max(1, num_results / gen_params.batch_count) : 1;
const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx / images_per_batch;
std::string params = gen_params.embed_image_metadata
? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
: "";

12
examples/common/common.cpp
View File

@@ -1007,6 +1007,10 @@ ArgOptions SDGenerationParams::get_options() {
"--batch-count",
"batch count",
&batch_count},
{"",
"--qwen-image-layers",
"number of Qwen Image Layered layers; latent/output count is layers + 1 (default: 3)",
&qwen_image_layers},
{"",
"--video-frames",
"video frames (default: 1)",
@@ -1827,6 +1831,7 @@ bool SDGenerationParams::from_json_str(
load_if_exists("width", width);
load_if_exists("height", height);
load_if_exists("batch_count", batch_count);
load_if_exists("qwen_image_layers", qwen_image_layers);
load_if_exists("video_frames", video_frames);
load_if_exists("fps", fps);
load_if_exists("upscale_repeats", upscale_repeats);
@@ -2251,6 +2256,11 @@ bool SDGenerationParams::validate(SDMode mode) {
return false;
}
if (qwen_image_layers < 0) {
LOG_ERROR("error: qwen_image_layers must be non-negative");
return false;
}
if (sample_params.sample_steps <= 0) {
LOG_ERROR("error: the sample_steps must be greater than 0\n");
return false;
@@ -2417,6 +2427,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.strength = strength;
params.seed = seed;
params.batch_count = batch_count;
params.qwen_image_layers = qwen_image_layers;
params.control_image = control_image.get();
params.control_strength = control_strength;
params.pm_params = pm_params;
@@ -2542,6 +2553,7 @@ std::string SDGenerationParams::to_string() const {
<< " width: " << width << ",\n"
<< " height: " << height << ",\n"
<< " batch_count: " << batch_count << ",\n"
<< " qwen_image_layers: " << qwen_image_layers << ",\n"
<< " init_image_path: \"" << init_image_path << "\",\n"
<< " end_image_path: \"" << end_image_path << "\",\n"
<< " mask_image_path: \"" << mask_image_path << "\",\n"

1
examples/common/common.h
View File

@@ -197,6 +197,7 @@ struct SDGenerationParams {
int width = -1;
int height = -1;
int batch_count = 1;
int qwen_image_layers = 3;
int64_t seed = 42;
float strength = 0.75f;
float control_strength = 0.9f;

5
examples/server/async_jobs.cpp
View File

@@ -2,6 +2,7 @@
#include "async_jobs.h"
#include <algorithm>
#include <iomanip>
#include <sstream>
@@ -195,6 +196,8 @@ bool execute_img_gen_job(ServerRuntime& runtime,
encoded_format = EncodedImageFormat::WEBP;
}
int batch_count = job.img_gen.gen_params.batch_count;
int images_per_batch = batch_count > 0 ? std::max(1, num_results / batch_count) : 1;
for (int i = 0; i < num_results; ++i) {
if (results[i].data == nullptr) {
continue;
@@ -203,7 +206,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
const std::string metadata = job.img_gen.gen_params.embed_image_metadata
? get_image_params(*runtime.ctx_params,
job.img_gen.gen_params,
job.img_gen.gen_params.seed + i)
job.img_gen.gen_params.seed + i / images_per_batch)
: "";
auto image_bytes = encode_image_to_vector(encoded_format,
results[i].data,

12
examples/server/routes_openai.cpp
View File

@@ -284,14 +284,16 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
out["data"] = json::array();
out["output_format"] = request.output_format;
for (int i = 0; i < request.gen_params.batch_count; ++i) {
int result_count = results.count();
int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1;
for (int i = 0; i < result_count; ++i) {
if (results[i].data == nullptr) {
continue;
}
std::string params = request.gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params,
request.gen_params,
request.gen_params.seed + i)
request.gen_params.seed + i / images_per_batch)
: "";
auto image_bytes = encode_image_to_vector(request.output_format == "jpeg"
? EncodedImageFormat::JPEG
@@ -356,14 +358,16 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
out["data"] = json::array();
out["output_format"] = request.output_format;
for (int i = 0; i < request.gen_params.batch_count; ++i) {
int result_count = results.count();
int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1;
for (int i = 0; i < result_count; ++i) {
if (results[i].data == nullptr) {
continue;
}
std::string params = request.gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params,
request.gen_params,
request.gen_params.seed + i)
request.gen_params.seed + i / images_per_batch)
: "";
auto image_bytes = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : EncodedImageFormat::PNG,
results[i].data,

3
examples/server/routes_sdapi.cpp
View File

@@ -311,6 +311,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
out["parameters"] = j;
out["info"] = "";
int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, num_results / request.gen_params.batch_count) : 1;
for (int i = 0; i < num_results; ++i) {
if (results[i].data == nullptr) {
continue;
@@ -319,7 +320,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
std::string params = request.gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params,
request.gen_params,
request.gen_params.seed + i)
request.gen_params.seed + i / images_per_batch)
: "";
auto image_bytes = encode_image_to_vector(EncodedImageFormat::PNG,
results[i].data,

1
examples/server/routes_sdcpp.cpp
View File

@@ -126,6 +126,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
{"strength", defaults.strength},
{"seed", defaults.seed},
{"batch_count", defaults.batch_count},
{"qwen_image_layers", defaults.qwen_image_layers},
{"auto_resize_ref_image", defaults.auto_resize_ref_image},
{"increase_ref_index", defaults.increase_ref_index},
{"control_strength", defaults.control_strength},

1
include/stable-diffusion.h
View File

@@ -380,6 +380,7 @@ typedef struct {
sd_tiling_params_t vae_tiling_params;
sd_cache_params_t cache;
sd_hires_params_t hires;
int qwen_image_layers;
} sd_img_gen_params_t;
typedef struct {

2
src/core/ggml_extend.hpp
View File

@@ -391,7 +391,7 @@ __STATIC_INLINE__ uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t*
int64_t width = input->ne[0];
int64_t height = input->ne[1];
int64_t channels = input->ne[2];
GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
GGML_ASSERT(input->type == GGML_TYPE_F32);
if (image_data == nullptr) {
image_data = (uint8_t*)malloc(width * height * channels);
}

3
src/model.h
View File

@@ -36,6 +36,7 @@ enum SDVersion {
VERSION_WAN2_2_I2V,
VERSION_WAN2_2_TI2V,
VERSION_QWEN_IMAGE,
VERSION_QWEN_IMAGE_LAYERED,
VERSION_ANIMA,
VERSION_FLUX2,
VERSION_FLUX2_KLEIN,
@@ -127,7 +128,7 @@ static inline bool sd_version_is_wan(SDVersion version) {
}
static inline bool sd_version_is_qwen_image(SDVersion version) {
if (version == VERSION_QWEN_IMAGE) {
if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
return true;
}
return false;

144
src/model/common/rope.hpp
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;

2
src/model/diffusion/anima.hpp
View File

@@ -612,7 +612,7 @@ namespace Anima {
0,
{},
empty_ref_latents,
false,
Rope::RefIndexMode::FIXED,
1.0f,
false);

24
src/model/diffusion/dit.hpp
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]

16
src/model/diffusion/flux.hpp
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());

3
src/model/diffusion/model.hpp
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{};
};

120
src/model/diffusion/qwen_image.hpp
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());

12
src/model/diffusion/z_image.hpp
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());

53
src/model/vae/wan_vae.hpp
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);
}

3
src/model_loader.cpp
View File

@@ -476,6 +476,9 @@ SDVersion ModelLoader::get_sd_version() {
return VERSION_MINIT2I;
}
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
return VERSION_QWEN_IMAGE_LAYERED;
}
return VERSION_QWEN_IMAGE;
}
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {

111
src/stable-diffusion.cpp
View File

@@ -84,6 +84,7 @@ const char* model_version_to_str[] = {
"Wan 2.2 I2V",
"Wan 2.2 TI2V",
"Qwen Image",
"Qwen Image Layered",
"Anima",
"Flux.2",
"Flux.2 klein",
@@ -754,13 +755,14 @@ public:
}
}
} else if (sd_version_is_qwen_image(version)) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
enable_vision,
model_manager);
diffusion_model = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
diffusion_model = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
@@ -2118,9 +2120,9 @@ public:
sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
std::vector<sd::Tensor<float>> controls;
DiffusionParams diffusion_params;
diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.increase_ref_index = increase_ref_index;
diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
sd::guidance::GuidanceInput step_guidance_input;
step_guidance_input.step = step;
step_guidance_input.schedule_size = sigmas.size();
@@ -2371,6 +2373,10 @@ public:
return latent_channel;
}
int get_image_channels() const {
return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3;
}
int get_image_seq_len(int h, int w) {
int vae_scale_factor = get_vae_scale_factor();
return (h / vae_scale_factor) * (w / vae_scale_factor);
@@ -2958,6 +2964,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
sd_img_gen_params->seed = -1;
sd_img_gen_params->batch_count = 1;
sd_img_gen_params->control_strength = 0.9f;
sd_img_gen_params->qwen_image_layers = 3;
sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f};
sd_img_gen_params->pulid_params = {nullptr, 1.0f};
sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
@@ -2984,6 +2991,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
"seed: %" PRId64
"\n"
"batch_count: %d\n"
"qwen_image_layers: %d\n"
"ref_images_count: %d\n"
"auto_resize_ref_image: %s\n"
"increase_ref_index: %s\n"
@@ -3000,6 +3008,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
sd_img_gen_params->strength,
sd_img_gen_params->seed,
sd_img_gen_params->batch_count,
sd_img_gen_params->qwen_image_layers,
sd_img_gen_params->ref_images_count,
BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
BOOL_STR(sd_img_gen_params->increase_ref_index),
@@ -3267,6 +3276,7 @@ struct GenerationRequest {
bool has_ref_images = false;
const sd_cache_params_t* cache_params = nullptr;
int batch_count = 1;
int qwen_image_layers = 3;
int shifted_timestep = 0;
float strength = 1.f;
float control_strength = 0.f;
@@ -3292,6 +3302,7 @@ struct GenerationRequest {
diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor();
seed = sd_img_gen_params->seed;
batch_count = sd_img_gen_params->batch_count;
qwen_image_layers = std::max(0, sd_img_gen_params->qwen_image_layers);
clip_skip = sd_img_gen_params->clip_skip;
shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep;
strength = sd_img_gen_params->strength;
@@ -3996,6 +4007,34 @@ public:
ImageVaeAxesGuard& operator=(const ImageVaeAxesGuard&) = delete;
};
static sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, int channels) {
if (image.empty()) {
return image;
}
GGML_ASSERT(image.dim() == 4);
int64_t current_channels = image.shape()[2];
if (current_channels == channels) {
return image;
}
if (channels == 4) {
sd::Tensor<float> alpha = sd::full<float>({image.shape()[0], image.shape()[1], 1, image.shape()[3]}, 1.f);
if (current_channels == 3) {
return sd::ops::concat(image, alpha, 2);
}
if (current_channels == 1) {
sd::Tensor<float> rgb = sd::ops::concat(image, image, 2);
rgb = sd::ops::concat(rgb, image, 2);
return sd::ops::concat(rgb, alpha, 2);
}
}
if (channels == 3 && current_channels >= 3) {
return sd::ops::slice(image, 2, 0, 3);
}
GGML_ABORT("cannot convert image tensor from %lld to %d channels",
(long long)current_channels,
channels);
}
static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request,
@@ -4005,6 +4044,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
sd::Tensor<float> init_image_tensor;
sd::Tensor<float> control_image_tensor;
sd::Tensor<float> mask_image_tensor;
int image_channels = sd_ctx->sd->get_image_channels();
if (sd_img_gen_params->init_image.data != nullptr) {
LOG_INFO("IMG2IMG");
@@ -4021,7 +4061,8 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
plan->sample_steps = static_cast<int>(plan->sigmas.size() - 1);
}
init_image_tensor = sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height);
init_image_tensor = ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height),
image_channels);
}
if (sd_img_gen_params->mask_image.data != nullptr) {
@@ -4053,7 +4094,11 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
sd::Tensor<float> init_latent;
sd::Tensor<float> control_latent;
if (init_image_tensor.empty()) {
init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->qwen_image_layers + 1, true);
} else {
init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
}
} else {
init_latent = sd_ctx->sd->encode_first_stage(init_image_tensor);
if (init_latent.empty()) {
@@ -4072,12 +4117,13 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
std::vector<sd::Tensor<float>> ref_images;
for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) {
ref_images.push_back(sd_image_to_tensor(sd_img_gen_params->ref_images[i]));
ref_images.push_back(ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->ref_images[i]),
image_channels));
}
if (ref_images.empty() && sd_version_is_unet_edit(sd_ctx->sd->version)) {
LOG_WARN("This model needs at least one reference image; using an empty reference");
ref_images.push_back(sd::zeros<float>({request->width, request->height, 3, 1}));
ref_images.push_back(sd::zeros<float>({request->width, request->height, image_channels, 1}));
request->guidance.img_cfg = request->guidance.txt_cfg;
request->use_img_uncond = false;
}
@@ -4103,7 +4149,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
vae_width = round(vae_width / factor) * factor;
auto resized_ref_img = sd::ops::interpolate(ref_images[i],
{static_cast<int>(vae_width), static_cast<int>(vae_height), 3, 1});
{static_cast<int>(vae_width),
static_cast<int>(vae_height),
ref_images[i].shape()[2],
ref_images[i].shape()[3]});
LOG_DEBUG("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64,
static_cast<int>(i),
@@ -4333,13 +4382,41 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
cancelled = true;
break;
}
int64_t t1 = ggml_time_ms();
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
return nullptr;
int64_t t1 = ggml_time_ms();
if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
int qwen_image_latent_layers = request.qwen_image_layers + 1;
if (final_latents[i].dim() < 5 || final_latents[i].shape()[2] < qwen_image_latent_layers) {
LOG_ERROR("qwen image layered expected at least %d latent layers, got shape dim=%d",
qwen_image_latent_layers,
final_latents[i].dim());
return nullptr;
}
for (int layer_index = 0; layer_index < qwen_image_latent_layers; layer_index++) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling latent decodings");
cancelled = true;
break;
}
sd::Tensor<float> layer_latent = sd::ops::slice(final_latents[i], 2, layer_index, layer_index + 1);
layer_latent.squeeze_(2);
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(layer_latent);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for latent %zu layer %d", i + 1, layer_index + 1);
return nullptr;
}
decoded_images.push_back(std::move(image));
}
if (cancelled) {
break;
}
} else {
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
return nullptr;
}
decoded_images.push_back(std::move(image));
}
decoded_images.push_back(std::move(image));
int64_t t2 = ggml_time_ms();
LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000);
}