DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-08-10 06:07:58 +00: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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
20 changed files with 366 additions and 165 deletions
Show Stats 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) if (!img.data)
return false; 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 std::string params = gen_params.embed_image_metadata
? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode) ? 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", "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",
"video frames (default: 1)", "video frames (default: 1)",
@ -1827,6 +1831,7 @@ bool SDGenerationParams::from_json_str(
load_if_exists("width", width); load_if_exists("width", width);
load_if_exists("height", height); load_if_exists("height", height);
load_if_exists("batch_count", batch_count); 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("video_frames", video_frames);
load_if_exists("fps", fps); load_if_exists("fps", fps);
load_if_exists("upscale_repeats", upscale_repeats); load_if_exists("upscale_repeats", upscale_repeats);
@ -2251,6 +2256,11 @@ bool SDGenerationParams::validate(SDMode mode) {
return false; 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) { if (sample_params.sample_steps <= 0) {
LOG_ERROR("error: the sample_steps must be greater than 0\n"); LOG_ERROR("error: the sample_steps must be greater than 0\n");
return false; return false;
@ -2417,6 +2427,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.strength = strength; params.strength = strength;
params.seed = seed; params.seed = seed;
params.batch_count = batch_count; params.batch_count = batch_count;
params.qwen_image_layers = qwen_image_layers;
params.control_image = control_image.get(); params.control_image = control_image.get();
params.control_strength = control_strength; params.control_strength = control_strength;
params.pm_params = pm_params; params.pm_params = pm_params;
@ -2542,6 +2553,7 @@ std::string SDGenerationParams::to_string() const {
<< " width: " << width << ",\n" << " width: " << width << ",\n"
<< " height: " << height << ",\n" << " height: " << height << ",\n"
<< " batch_count: " << batch_count << ",\n" << " batch_count: " << batch_count << ",\n"
<< " qwen_image_layers: " << qwen_image_layers << ",\n"
<< " init_image_path: \"" << init_image_path << "\",\n" << " init_image_path: \"" << init_image_path << "\",\n"
<< " end_image_path: \"" << end_image_path << "\",\n" << " end_image_path: \"" << end_image_path << "\",\n"
<< " mask_image_path: \"" << mask_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 width = -1;
int height = -1; int height = -1;
int batch_count = 1; int batch_count = 1;
int qwen_image_layers = 3;
int64_t seed = 42; int64_t seed = 42;
float strength = 0.75f; float strength = 0.75f;
float control_strength = 0.9f; float control_strength = 0.9f;

5
examples/server/async_jobs.cpp
View File

@ -2,6 +2,7 @@
#include "async_jobs.h" #include "async_jobs.h"
#include <algorithm>
#include <iomanip> #include <iomanip>
#include <sstream> #include <sstream>
@ -195,6 +196,8 @@ bool execute_img_gen_job(ServerRuntime& runtime,
encoded_format = EncodedImageFormat::WEBP; 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) { for (int i = 0; i < num_results; ++i) {
if (results[i].data == nullptr) { if (results[i].data == nullptr) {
continue; continue;
@ -203,7 +206,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
const std::string metadata = job.img_gen.gen_params.embed_image_metadata const std::string metadata = job.img_gen.gen_params.embed_image_metadata
? get_image_params(*runtime.ctx_params, ? get_image_params(*runtime.ctx_params,
job.img_gen.gen_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, auto image_bytes = encode_image_to_vector(encoded_format,
results[i].data, 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["data"] = json::array();
out["output_format"] = request.output_format; 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) { if (results[i].data == nullptr) {
continue; continue;
} }
std::string params = request.gen_params.embed_image_metadata std::string params = request.gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params, ? get_image_params(*runtime->ctx_params,
request.gen_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" auto image_bytes = encode_image_to_vector(request.output_format == "jpeg"
? EncodedImageFormat::JPEG ? EncodedImageFormat::JPEG
@ -356,14 +358,16 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
out["data"] = json::array(); out["data"] = json::array();
out["output_format"] = request.output_format; 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) { if (results[i].data == nullptr) {
continue; continue;
} }
std::string params = request.gen_params.embed_image_metadata std::string params = request.gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params, ? get_image_params(*runtime->ctx_params,
request.gen_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, auto image_bytes = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : EncodedImageFormat::PNG,
results[i].data, 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["parameters"] = j;
out["info"] = ""; 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) { for (int i = 0; i < num_results; ++i) {
if (results[i].data == nullptr) { if (results[i].data == nullptr) {
continue; continue;
@ -319,7 +320,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
std::string params = request.gen_params.embed_image_metadata std::string params = request.gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params, ? get_image_params(*runtime->ctx_params,
request.gen_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, auto image_bytes = encode_image_to_vector(EncodedImageFormat::PNG,
results[i].data, 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}, {"strength", defaults.strength},
{"seed", defaults.seed}, {"seed", defaults.seed},
{"batch_count", defaults.batch_count}, {"batch_count", defaults.batch_count},
{"qwen_image_layers", defaults.qwen_image_layers},
{"auto_resize_ref_image", defaults.auto_resize_ref_image}, {"auto_resize_ref_image", defaults.auto_resize_ref_image},
{"increase_ref_index", defaults.increase_ref_index}, {"increase_ref_index", defaults.increase_ref_index},
{"control_strength", defaults.control_strength}, {"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_tiling_params_t vae_tiling_params;
sd_cache_params_t cache; sd_cache_params_t cache;
sd_hires_params_t hires; sd_hires_params_t hires;
int qwen_image_layers;
} sd_img_gen_params_t; } sd_img_gen_params_t;
typedef struct { 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 width = input->ne[0];
int64_t height = input->ne[1]; int64_t height = input->ne[1];
int64_t channels = input->ne[2]; 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) { if (image_data == nullptr) {
image_data = (uint8_t*)malloc(width * height * channels); 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_I2V,
VERSION_WAN2_2_TI2V, VERSION_WAN2_2_TI2V,
VERSION_QWEN_IMAGE, VERSION_QWEN_IMAGE,
VERSION_QWEN_IMAGE_LAYERED,
VERSION_ANIMA, VERSION_ANIMA,
VERSION_FLUX2, VERSION_FLUX2,
VERSION_FLUX2_KLEIN, 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) { 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 true;
} }
return false; return false;

144
src/model/common/rope.hpp
View File

@ -12,6 +12,16 @@ namespace Rope {
ErnieImage, 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> template <class T>
__STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) { __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
std::vector<T> result(num); std::vector<T> result(num);
@ -346,7 +356,7 @@ namespace Rope {
int axes_dim_num, int axes_dim_num,
int start_index, int start_index,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
bool scale_rope, bool scale_rope,
int base_offset = 0) { int base_offset = 0) {
@ -357,13 +367,15 @@ namespace Rope {
for (ggml_tensor* ref : ref_latents) { for (ggml_tensor* ref : ref_latents) {
int h_offset = 0; int h_offset = 0;
int w_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) { if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
w_offset = curr_w_offset; w_offset = curr_w_offset;
} else { } else {
h_offset = curr_h_offset; h_offset = curr_h_offset;
} }
scale_rope = false; scale_rope = false;
} else if (ref_index_mode == RefIndexMode::DECREASE) {
index--;
} }
auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]), auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
@ -377,7 +389,7 @@ namespace Rope {
scale_rope); scale_rope);
ids = concat_ids(ids, ref_ids, bs); ids = concat_ids(ids, ref_ids, bs);
if (increase_ref_index) { if (ref_index_mode == RefIndexMode::INCREASE) {
index++; index++;
} }
@ -395,7 +407,7 @@ namespace Rope {
int context_len, int context_len,
std::set<int> txt_arange_dims, std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
bool is_longcat) { bool is_longcat) {
int x_index = is_longcat ? 1 : 0; int x_index = is_longcat ? 1 : 0;
@ -406,7 +418,7 @@ namespace Rope {
auto ids = concat_ids(txt_ids, img_ids, bs); auto ids = concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) { 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); ids = concat_ids(ids, refs_ids, bs);
} }
return ids; return ids;
@ -420,7 +432,7 @@ namespace Rope {
int context_len, int context_len,
std::set<int> txt_arange_dims, std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
int theta, int theta,
bool circular_h, bool circular_h,
@ -435,7 +447,7 @@ namespace Rope {
context_len, context_len,
txt_arange_dims, txt_arange_dims,
ref_latents, ref_latents,
increase_ref_index, ref_index_mode,
ref_index_scale, ref_index_scale,
is_longcat); is_longcat);
std::vector<std::vector<int>> wrap_dims; 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); 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 w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, 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 h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size; int w_len = (w + (patch_size / 2)) / patch_size;
int txt_id_start = std::max(h_len, w_len); int txt_id_start = std::max(h_len, w_len) / 2;
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len); 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)); std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
for (int i = 0; i < bs; ++i) { for (int i = 0; i < bs; ++i) {
for (int j = 0; j < txt_ids.size(); ++j) { for (int j = 0; j < txt_ids.size(); ++j) {
@ -499,28 +558,30 @@ namespace Rope {
} }
} }
int axes_dim_num = 3; 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); auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
if (ref_latents.size() > 0) { 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); int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
ids = concat_ids(ids, refs_ids, bs); 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; return ids;
} }
// Generate qwen_image positional embeddings // 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 w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index, RefIndexMode ref_index_mode,
int theta, int theta,
bool circular_h, bool circular_h,
bool circular_w, bool circular_w,
const std::vector<int>& axes_dim) { 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; 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 // 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) { 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. // 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)); wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
size_t cursor = context_len; // ignore text tokens 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) { for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) { if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len; 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); 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 // Generate wan positional embeddings
__STATIC_INLINE__ std::vector<float> gen_wan_pe(int t, __STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
int h, int h,
@ -785,7 +806,8 @@ namespace Rope {
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents, 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); 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)); 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++) { for (int i = 0; i < bs * padded_context_len; i++) {
@ -816,12 +838,12 @@ namespace Rope {
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
bool increase_ref_index, RefIndexMode ref_index_mode,
int theta, int theta,
bool circular_h, bool circular_h,
bool circular_w, bool circular_w,
const std::vector<int>& axes_dim) { 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; std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size; 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, 0,
{}, {},
empty_ref_latents, empty_ref_latents,
false, Rope::RefIndexMode::FIXED,
1.0f, 1.0f,
false); false);

24
src/model/diffusion/dit.hpp
View File

@ -135,23 +135,23 @@ namespace DiT {
return x; return x;
} }
inline ggml_tensor* unpatchify(ggml_context* ctx, inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
ggml_tensor* x, ggml_tensor* x,
int64_t t_len, int64_t t_len,
int64_t h_len, int64_t h_len,
int64_t w_len, int64_t w_len,
int pt, int pt,
int ph, int ph,
int pw) { int pw) {
// x: [N, t_len*h_len*w_len, pt*ph*pw*C] // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
// return: [N*C, t_len*pt, h_len*ph, w_len*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; int64_t C = x->ne[0] / pt / ph / pw;
GGML_ASSERT(C * pt * ph * pw == x->ne[0]); 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_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, 1, 2, 0, 3)); // [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*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_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_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] 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>& y_tensor = {},
const sd::Tensor<float>& guidance_tensor = {}, const sd::Tensor<float>& guidance_tensor = {},
const std::vector<sd::Tensor<float>>& ref_latents_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 = {}, std::vector<int> skip_layers = {},
const sd::Tensor<float>& pulid_id_tensor = {}, const sd::Tensor<float>& pulid_id_tensor = {},
float pulid_id_weight = 1.0f) { float pulid_id_weight = 1.0f) {
@ -1527,8 +1527,8 @@ namespace Flux {
} }
std::set<int> txt_arange_dims; std::set<int> txt_arange_dims;
if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
txt_arange_dims = {3}; txt_arange_dims = {3};
increase_ref_index = true; ref_index_mode = Rope::RefIndexMode::INCREASE;
} else if (version == VERSION_OVIS_IMAGE) { } else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2}; txt_arange_dims = {1, 2};
} }
@ -1539,7 +1539,7 @@ namespace Flux {
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
txt_arange_dims, txt_arange_dims,
ref_latents, ref_latents,
increase_ref_index, ref_index_mode,
config.ref_index_scale, config.ref_index_scale,
config.theta, config.theta,
circular_y_enabled, circular_y_enabled,
@ -1599,7 +1599,7 @@ namespace Flux {
const sd::Tensor<float>& y = {}, const sd::Tensor<float>& y = {},
const sd::Tensor<float>& guidance = {}, const sd::Tensor<float>& guidance = {},
const std::vector<sd::Tensor<float>>& ref_latents = {}, 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>(), std::vector<int> skip_layers = std::vector<int>(),
const sd::Tensor<float>& pulid_id = {}, const sd::Tensor<float>& pulid_id = {},
float pulid_id_weight = 1.0f) { float pulid_id_weight = 1.0f) {
@ -1610,7 +1610,7 @@ namespace Flux {
// guidance: [N, ] // guidance: [N, ]
// pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048] // pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048]
auto get_graph = [&]() -> ggml_cgraph* { 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()); 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(diffusion_params.y),
tensor_or_empty(extra->guidance), tensor_or_empty(extra->guidance),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, 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, extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
tensor_or_empty(extra->pulid_id), tensor_or_empty(extra->pulid_id),
extra->pulid_id_weight); extra->pulid_id_weight);
@ -1683,7 +1683,7 @@ namespace Flux {
{}, {},
guidance, guidance,
{}, {},
false); Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms(); int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty()); GGML_ASSERT(!out_opt.empty());

3
src/model/diffusion/model.hpp
View File

@ -7,6 +7,7 @@
#include "core/ggml_extend.hpp" #include "core/ggml_extend.hpp"
#include "core/tensor_ggml.hpp" #include "core/tensor_ggml.hpp"
#include "model/common/rope.hpp"
#include "model_manager.h" #include "model_manager.h"
struct UNetDiffusionExtra { struct UNetDiffusionExtra {
@ -73,7 +74,7 @@ struct DiffusionParams {
const sd::Tensor<float>* c_concat = nullptr; const sd::Tensor<float>* c_concat = nullptr;
const sd::Tensor<float>* y = nullptr; const sd::Tensor<float>* y = nullptr;
const std::vector<sd::Tensor<float>>* ref_latents = 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{}; DiffusionExtraParams extra = std::monostate{};
}; };

120
src/model/diffusion/qwen_image.hpp
View File

@ -4,6 +4,7 @@
#include <memory> #include <memory>
#include "model/common/block.hpp" #include "model/common/block.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/flux.hpp" #include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp" #include "model/diffusion/model.hpp"
#include "model_loader.h" #include "model_loader.h"
@ -23,6 +24,7 @@ namespace Qwen {
std::vector<int> axes_dim = {16, 56, 56}; std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128; int axes_dim_sum = 128;
bool zero_cond_t = false; 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) { static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
QwenImageConfig config; QwenImageConfig config;
@ -88,19 +90,33 @@ namespace Qwen {
}; };
struct QwenTimestepProjEmbeddings : public GGMLBlock { struct QwenTimestepProjEmbeddings : public GGMLBlock {
protected:
bool use_additional_t_cond = false;
public: 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)); 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* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps) { ggml_tensor* timesteps,
ggml_tensor* addition_t_cond = nullptr) {
// timesteps: [N,] // timesteps: [N,]
// return: [N, embedding_dim] // return: [N, embedding_dim]
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]); 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_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
auto timesteps_emb = timestep_embedder->forward(ctx, timesteps_proj); 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; return timesteps_emb;
} }
}; };
@ -402,7 +418,7 @@ namespace Qwen {
QwenImageModel(QwenImageConfig config) QwenImageModel(QwenImageConfig config)
: config(config) { : config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim; 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["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["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)); 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* forward_orig(GGMLRunnerContext* ctx,
ggml_tensor* x, ggml_tensor* x,
ggml_tensor* timestep, ggml_tensor* timestep,
ggml_tensor* addition_t_cond,
ggml_tensor* context, ggml_tensor* context,
ggml_tensor* pe, ggml_tensor* pe,
ggml_tensor* modulate_index = nullptr) { ggml_tensor* modulate_index = nullptr) {
@ -434,9 +451,9 @@ namespace Qwen {
auto norm_out = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]); auto norm_out = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_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) { 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); t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
} }
auto img = img_in->forward(ctx, x); auto img = img_in->forward(ctx, x);
@ -469,33 +486,50 @@ namespace Qwen {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x, ggml_tensor* x,
ggml_tensor* timestep, ggml_tensor* timestep,
ggml_tensor* addition_t_cond,
ggml_tensor* context, ggml_tensor* context,
ggml_tensor* pe, ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {}, std::vector<ggml_tensor*> ref_latents = {},
ggml_tensor* modulate_index = nullptr) { ggml_tensor* modulate_index = nullptr) {
// Forward pass of DiT. // Forward pass of DiT.
// x: [N, C, H, W] // x: [N, C, H, W] or [N*C, T, H, W]
// timestep: [N,] // timestep: [N,]
// context: [N, L, D] // context: [N, L, D]
// pe: [L, d_head/2, 2, 2] // 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 W = x->ne[0];
int64_t H = x->ne[1]; int64_t H = x->ne[1];
int64_t C = x->ne[2]; int64_t T = 1;
int64_t N = x->ne[3]; 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]; int64_t img_tokens = img->ne[1];
if (ref_latents.size() > 0) { if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) { 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); 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) { 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] 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 = 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; return out;
} }
@ -515,6 +559,7 @@ namespace Qwen {
QwenImageModel qwen_image; QwenImageModel qwen_image;
std::vector<float> pe_vec; std::vector<float> pe_vec;
std::vector<float> modulate_index_vec; std::vector<float> modulate_index_vec;
std::vector<int32_t> additional_t_cond_vec;
SDVersion version; SDVersion version;
QwenImageRunner(ggml_backend_t backend, QwenImageRunner(ggml_backend_t backend,
@ -524,9 +569,13 @@ namespace Qwen {
bool zero_cond_t = false, bool zero_cond_t = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager), : 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; 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); 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>& timesteps_tensor,
const sd::Tensor<float>& context_tensor, const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_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_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor); ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_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_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor); ggml_tensor* context = make_input(context_tensor);
std::vector<ggml_tensor*> ref_latents; std::vector<ggml_tensor*> ref_latents;
@ -555,13 +604,29 @@ namespace Qwen {
ref_latents.push_back(make_input(ref_latent_tensor)); 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]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), batch_size,
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ref_latents, ref_latents,
increase_ref_index, ref_index_mode,
config.theta, config.theta,
circular_y_enabled, circular_y_enabled,
circular_x_enabled, circular_x_enabled,
@ -604,6 +669,7 @@ namespace Qwen {
ggml_tensor* out = qwen_image.forward(&runner_ctx, ggml_tensor* out = qwen_image.forward(&runner_ctx,
x, x,
timesteps, timesteps,
addition_t_cond,
context, context,
pe, pe,
ref_latents, ref_latents,
@ -619,12 +685,12 @@ namespace Qwen {
const sd::Tensor<float>& timesteps, const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context, const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {}, const std::vector<sd::Tensor<float>>& ref_latents = {},
bool increase_ref_index = false) { Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::INCREASE) {
// x: [N, in_channels, h, w] // x: [N, C, H, W] or [N*C, T, H, W]
// timesteps: [N, ] // timesteps: [N, ]
// context: [N, max_position, hidden_size] // context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* { 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()); 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, *diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context), tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index); diffusion_params.ref_index_mode);
} }
void test() { void test() {
@ -674,7 +740,7 @@ namespace Qwen {
timesteps, timesteps,
context, context,
{}, {},
false); Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms(); int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty()); 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>& timesteps_tensor,
const sd::Tensor<float>& context_tensor, const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_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_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor); ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor);
@ -595,7 +595,7 @@ namespace ZImage {
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
SEQ_MULTI_OF, SEQ_MULTI_OF,
ref_latents, ref_latents,
increase_ref_index, ref_index_mode,
config.theta, config.theta,
circular_y_enabled, circular_y_enabled,
circular_x_enabled, circular_x_enabled,
@ -626,12 +626,12 @@ namespace ZImage {
const sd::Tensor<float>& timesteps, const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context, const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {}, 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] // x: [N, in_channels, h, w]
// timesteps: [N, ] // timesteps: [N, ]
// context: [N, max_position, hidden_size] // context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* { 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()); 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, *diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context), tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.increase_ref_index); diffusion_params.ref_index_mode);
} }
void test() { void test() {
@ -681,7 +681,7 @@ namespace ZImage {
timesteps, timesteps,
context, context,
{}, {},
false); Rope::RefIndexMode::FIXED);
int64_t t1 = ggml_time_ms(); int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty()); GGML_ASSERT(!out_opt.empty());

53
src/model/vae/wan_vae.hpp
View File

@ -631,6 +631,7 @@ namespace WAN {
class Encoder3d : public GGMLBlock { class Encoder3d : public GGMLBlock {
protected: protected:
bool wan2_2; bool wan2_2;
int64_t in_channels;
int64_t dim; int64_t dim;
int64_t z_dim; int64_t z_dim;
std::vector<int> dim_mult; std::vector<int> dim_mult;
@ -641,12 +642,14 @@ namespace WAN {
public: public:
Encoder3d(int64_t dim = 128, Encoder3d(int64_t dim = 128,
int64_t z_dim = 4, int64_t z_dim = 4,
int64_t in_channels = 3,
std::vector<int> dim_mult = {1, 2, 4, 4}, std::vector<int> dim_mult = {1, 2, 4, 4},
int num_res_blocks = 2, int num_res_blocks = 2,
std::vector<bool> temperal_downsample = {false, true, true}, std::vector<bool> temperal_downsample = {false, true, true},
bool wan2_2 = false, bool wan2_2 = false,
bool is_2D = false) bool is_2D = false)
: dim(dim), : in_channels(in_channels),
dim(dim),
z_dim(z_dim), z_dim(z_dim),
dim_mult(dim_mult), dim_mult(dim_mult),
num_res_blocks(num_res_blocks), num_res_blocks(num_res_blocks),
@ -659,11 +662,10 @@ namespace WAN {
dims.push_back(dim * u); dims.push_back(dim * u);
} }
int64_t input_dim = wan2_2 ? 12 : 3;
if (is_2D) { 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 { } 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; int index = 0;
@ -812,6 +814,7 @@ namespace WAN {
class Decoder3d : public GGMLBlock { class Decoder3d : public GGMLBlock {
protected: protected:
bool wan2_2; bool wan2_2;
int64_t out_channels;
int64_t dim; int64_t dim;
int64_t z_dim; int64_t z_dim;
std::vector<int> dim_mult; std::vector<int> dim_mult;
@ -822,12 +825,14 @@ namespace WAN {
public: public:
Decoder3d(int64_t dim = 128, Decoder3d(int64_t dim = 128,
int64_t z_dim = 4, int64_t z_dim = 4,
int64_t out_channels = 3,
std::vector<int> dim_mult = {1, 2, 4, 4}, std::vector<int> dim_mult = {1, 2, 4, 4},
int num_res_blocks = 2, int num_res_blocks = 2,
std::vector<bool> temperal_upsample = {true, true, false}, std::vector<bool> temperal_upsample = {true, true, false},
bool wan2_2 = false, bool wan2_2 = false,
bool is_2D = false) bool is_2D = false)
: dim(dim), : out_channels(out_channels),
dim(dim),
z_dim(z_dim), z_dim(z_dim),
dim_mult(dim_mult), dim_mult(dim_mult),
num_res_blocks(num_res_blocks), num_res_blocks(num_res_blocks),
@ -889,7 +894,7 @@ namespace WAN {
// output blocks // output blocks
blocks["head.0"] = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim)); 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() // head.1 is nn.SiLU()
if (is_2D) { if (is_2D) {
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, final_dim, {3, 3}, {1, 1}, {1, 1})); 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: public:
bool wan2_2 = false; bool wan2_2 = false;
bool decode_only = true; bool decode_only = true;
int64_t input_channels = 3;
int patch_size = 1;
int64_t dim = 96; int64_t dim = 96;
int64_t dec_dim = 96; int64_t dec_dim = 96;
int64_t z_dim = 16; int64_t z_dim = 16;
@ -1026,16 +1033,22 @@ namespace WAN {
} }
public: public:
WanVAE(bool decode_only = true, bool wan2_2 = false, bool is_2D = false) WanVAE(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_2D = false)
: decode_only(decode_only), wan2_2(wan2_2), is_2D(is_2D) { : decode_only(decode_only),
wan2_2(version == VERSION_WAN2_2_TI2V),
is_2D(is_2D) {
// attn_scales is always [] // attn_scales is always []
if (wan2_2) { if (wan2_2) {
dim = 160; dim = 160;
dec_dim = 256; dec_dim = 256;
z_dim = 48; z_dim = 48;
input_channels = 12;
patch_size = 2;
_conv_num = 34; _conv_num = 34;
_enc_conv_num = 26; _enc_conv_num = 26;
} else if (version == VERSION_QWEN_IMAGE_LAYERED) {
input_channels = 4;
} }
if (is_2D) { if (is_2D) {
@ -1044,14 +1057,14 @@ namespace WAN {
} }
if (!decode_only) { 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) { if (is_2D) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1})); blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1}));
} else { } else {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1})); 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) { if (is_2D) {
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, z_dim, {1, 1})); blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, z_dim, {1, 1}));
} else { } else {
@ -1125,9 +1138,7 @@ namespace WAN {
clear_cache(); clear_cache();
if (wan2_2) { x = patchify(ctx->ggml_ctx, x, patch_size, b);
x = patchify(ctx->ggml_ctx, x, 2, b);
}
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x"); // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]); auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
@ -1202,9 +1213,7 @@ namespace WAN {
} }
} }
} }
if (wan2_2) { out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
out = unpatchify(ctx->ggml_ctx, out, 2, b);
}
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out"); // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
clear_cache(); clear_cache();
return out; 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] auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w]
_conv_idx = 0; _conv_idx = 0;
auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i); auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
if (wan2_2) { out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
out = unpatchify(ctx->ggml_ctx, out, 2, b);
}
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out"); // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
return out; return out;
} }
@ -1257,7 +1264,7 @@ namespace WAN {
if (is_2D) { if (is_2D) {
LOG_DEBUG("USING 2D VAE"); 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); 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; return VERSION_MINIT2I;
} }
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) { 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; return VERSION_QWEN_IMAGE;
} }
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) { 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 I2V",
"Wan 2.2 TI2V", "Wan 2.2 TI2V",
"Qwen Image", "Qwen Image",
"Qwen Image Layered",
"Anima", "Anima",
"Flux.2", "Flux.2",
"Flux.2 klein", "Flux.2 klein",
@ -754,13 +755,14 @@ public:
} }
} }
} else if (sd_version_is_qwen_image(version)) { } 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, tensor_storage_map,
version, version,
"", "",
true, enable_vision,
model_manager); 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, tensor_storage_map,
"model.diffusion_model", "model.diffusion_model",
version, version,
@ -2118,9 +2120,9 @@ public:
sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma); sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
std::vector<sd::Tensor<float>> controls; std::vector<sd::Tensor<float>> controls;
DiffusionParams diffusion_params; DiffusionParams diffusion_params;
diffusion_params.x = &noised_input; diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor; diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.increase_ref_index = increase_ref_index; diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
sd::guidance::GuidanceInput step_guidance_input; sd::guidance::GuidanceInput step_guidance_input;
step_guidance_input.step = step; step_guidance_input.step = step;
step_guidance_input.schedule_size = sigmas.size(); step_guidance_input.schedule_size = sigmas.size();
@ -2371,6 +2373,10 @@ public:
return latent_channel; 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 get_image_seq_len(int h, int w) {
int vae_scale_factor = get_vae_scale_factor(); int vae_scale_factor = get_vae_scale_factor();
return (h / vae_scale_factor) * (w / 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->seed = -1;
sd_img_gen_params->batch_count = 1; sd_img_gen_params->batch_count = 1;
sd_img_gen_params->control_strength = 0.9f; 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->pm_params = {nullptr, 0, nullptr, 20.f};
sd_img_gen_params->pulid_params = {nullptr, 1.0f}; 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}; 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 "seed: %" PRId64
"\n" "\n"
"batch_count: %d\n" "batch_count: %d\n"
"qwen_image_layers: %d\n"
"ref_images_count: %d\n" "ref_images_count: %d\n"
"auto_resize_ref_image: %s\n" "auto_resize_ref_image: %s\n"
"increase_ref_index: %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->strength,
sd_img_gen_params->seed, sd_img_gen_params->seed,
sd_img_gen_params->batch_count, sd_img_gen_params->batch_count,
sd_img_gen_params->qwen_image_layers,
sd_img_gen_params->ref_images_count, sd_img_gen_params->ref_images_count,
BOOL_STR(sd_img_gen_params->auto_resize_ref_image), BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
BOOL_STR(sd_img_gen_params->increase_ref_index), BOOL_STR(sd_img_gen_params->increase_ref_index),
@ -3267,6 +3276,7 @@ struct GenerationRequest {
bool has_ref_images = false; bool has_ref_images = false;
const sd_cache_params_t* cache_params = nullptr; const sd_cache_params_t* cache_params = nullptr;
int batch_count = 1; int batch_count = 1;
int qwen_image_layers = 3;
int shifted_timestep = 0; int shifted_timestep = 0;
float strength = 1.f; float strength = 1.f;
float control_strength = 0.f; float control_strength = 0.f;
@ -3292,6 +3302,7 @@ struct GenerationRequest {
diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor(); diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor();
seed = sd_img_gen_params->seed; seed = sd_img_gen_params->seed;
batch_count = sd_img_gen_params->batch_count; 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; clip_skip = sd_img_gen_params->clip_skip;
shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep; shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep;
strength = sd_img_gen_params->strength; strength = sd_img_gen_params->strength;
@ -3996,6 +4007,34 @@ public:
ImageVaeAxesGuard& operator=(const ImageVaeAxesGuard&) = delete; 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, static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params, const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request, GenerationRequest* request,
@ -4005,6 +4044,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
sd::Tensor<float> init_image_tensor; sd::Tensor<float> init_image_tensor;
sd::Tensor<float> control_image_tensor; sd::Tensor<float> control_image_tensor;
sd::Tensor<float> mask_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) { if (sd_img_gen_params->init_image.data != nullptr) {
LOG_INFO("IMG2IMG"); 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); 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) { 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> init_latent;
sd::Tensor<float> control_latent; sd::Tensor<float> control_latent;
if (init_image_tensor.empty()) { 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 { } else {
init_latent = sd_ctx->sd->encode_first_stage(init_image_tensor); init_latent = sd_ctx->sd->encode_first_stage(init_image_tensor);
if (init_latent.empty()) { if (init_latent.empty()) {
@ -4072,12 +4117,13 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
std::vector<sd::Tensor<float>> ref_images; std::vector<sd::Tensor<float>> ref_images;
for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) { 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)) { 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"); 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->guidance.img_cfg = request->guidance.txt_cfg;
request->use_img_uncond = false; 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; vae_width = round(vae_width / factor) * factor;
auto resized_ref_img = sd::ops::interpolate(ref_images[i], 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, LOG_DEBUG("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64,
static_cast<int>(i), static_cast<int>(i),
@ -4333,13 +4382,41 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
cancelled = true; cancelled = true;
break; break;
} }
int64_t t1 = ggml_time_ms(); int64_t t1 = ggml_time_ms();
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]); if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
if (image.empty()) { int qwen_image_latent_layers = request.qwen_image_layers + 1;
LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1); if (final_latents[i].dim() < 5 || final_latents[i].shape()[2] < qwen_image_latent_layers) {
return nullptr; 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(); int64_t t2 = ggml_time_ms();
LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000); LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000);
} }