mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-24 20:20:37 +00:00
595 lines
34 KiB
C++
595 lines
34 KiB
C++
#ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
|
|
#define __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
|
|
|
|
#include "model/diffusion/qwen_image.hpp"
|
|
|
|
namespace Qwen {
|
|
|
|
struct QwenImage21Config {
|
|
int64_t in_channels = 64;
|
|
int64_t out_channels = 64;
|
|
int64_t hidden_size = 4096;
|
|
int64_t context_dim = 4096;
|
|
int64_t head_dim = 128;
|
|
int64_t intermediate_size = 12288;
|
|
int num_layers = 32;
|
|
bool fused_mlp = false;
|
|
std::vector<int> axes_dim = {16, 56, 56};
|
|
|
|
static QwenImage21Config detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
|
|
QwenImage21Config config;
|
|
auto find = [&](const std::string& suffix) -> const TensorStorage* {
|
|
auto it = weights.find(prefix + "." + suffix);
|
|
return it == weights.end() ? nullptr : &it->second;
|
|
};
|
|
if (auto w = find("img_in.weight")) {
|
|
config.in_channels = w->ne[0];
|
|
config.hidden_size = w->ne[1];
|
|
}
|
|
if (auto w = find("proj_out.weight")) {
|
|
config.out_channels = w->ne[1];
|
|
}
|
|
if (auto w = find("txt_in.in_layer.weight")) {
|
|
config.context_dim = w->ne[0];
|
|
}
|
|
if (auto w = find("transformer_blocks.0.attn.norm_q.weight")) {
|
|
config.head_dim = w->ne[0];
|
|
}
|
|
if (auto w = find("transformer_blocks.0.img_mlp.gate_up.weight")) {
|
|
config.intermediate_size = w->ne[1] / 2;
|
|
config.fused_mlp = true;
|
|
} else if (auto w = find("transformer_blocks.0.img_mlp.proj.weight")) {
|
|
config.intermediate_size = w->ne[1];
|
|
}
|
|
int layers = 0;
|
|
const std::string block_prefix = prefix + ".transformer_blocks.";
|
|
for (const auto& [name, _] : weights) {
|
|
if (starts_with(name, block_prefix)) {
|
|
layers = std::max(layers, atoi(name.substr(block_prefix.size()).c_str()) + 1);
|
|
}
|
|
}
|
|
if (layers > 0) {
|
|
config.num_layers = layers;
|
|
LOG_VERBOSE("qwen_image_2_1: layers = %d, hidden_size = %" PRId64 ", context_dim = %" PRId64,
|
|
layers, config.hidden_size, config.context_dim);
|
|
}
|
|
return config;
|
|
}
|
|
};
|
|
|
|
struct QwenImage21Segment {
|
|
int64_t start;
|
|
int64_t end;
|
|
int64_t context_start;
|
|
int image_index;
|
|
};
|
|
|
|
struct QwenImage21Layout {
|
|
std::vector<QwenImage21Segment> segments;
|
|
std::vector<std::vector<float>> positions;
|
|
int64_t prefix_length = 0;
|
|
Rope::PositionLayout rope_layout;
|
|
|
|
static QwenImage21Layout build(int64_t text_length,
|
|
const sd::Tensor<int32_t>& image_slots,
|
|
const std::vector<std::pair<int64_t, int64_t>>& image_shapes) {
|
|
if (image_shapes.empty() || (!image_slots.empty() && image_slots.numel() != text_length)) {
|
|
throw std::runtime_error("Qwen Image 2.1: invalid image token layout");
|
|
}
|
|
QwenImage21Layout layout;
|
|
int64_t position = 0;
|
|
int next_image = 0;
|
|
auto append_image = [&](int index, int64_t context_start) {
|
|
auto [height, width] = image_shapes[index];
|
|
int64_t start = static_cast<int64_t>(layout.positions.size());
|
|
layout.segments.push_back({start, start + height * width, context_start, index});
|
|
layout.rope_layout.append_image(static_cast<int>(height), static_cast<int>(width));
|
|
for (int64_t h = 0; h < height; ++h) {
|
|
for (int64_t w = 0; w < width; ++w) {
|
|
layout.positions.push_back({static_cast<float>(position),
|
|
static_cast<float>(h - (height - height / 2)),
|
|
static_cast<float>(w - (width - width / 2))});
|
|
}
|
|
}
|
|
position += std::max(height, width);
|
|
};
|
|
for (int64_t i = 0; i < text_length;) {
|
|
int tag = image_slots.empty() ? 0 : image_slots[i];
|
|
int64_t begin = i++;
|
|
while (i < text_length && (image_slots.empty() ? 0 : image_slots[i]) == tag) {
|
|
++i;
|
|
}
|
|
if (tag != 0) {
|
|
if (tag != next_image + 1 || next_image + 1 >= static_cast<int>(image_shapes.size()) ||
|
|
(i - begin) * 4 != image_shapes[next_image].first * image_shapes[next_image].second) {
|
|
throw std::runtime_error("Qwen Image 2.1: vision slots and reference latents must have matching sizes");
|
|
}
|
|
append_image(next_image++, begin);
|
|
} else {
|
|
int64_t start = static_cast<int64_t>(layout.positions.size());
|
|
layout.segments.push_back({start, start + i - begin, begin, -1});
|
|
layout.rope_layout.append_tokens(i - begin);
|
|
for (int64_t j = begin; j < i; ++j, ++position) {
|
|
float p = static_cast<float>(position);
|
|
layout.positions.push_back({p, p, p});
|
|
}
|
|
}
|
|
}
|
|
if (next_image + 1 != static_cast<int>(image_shapes.size())) {
|
|
throw std::runtime_error("Qwen Image 2.1: missing reference image slots");
|
|
}
|
|
layout.prefix_length = static_cast<int64_t>(layout.positions.size());
|
|
append_image(next_image, text_length);
|
|
return layout;
|
|
}
|
|
};
|
|
|
|
struct QwenImage21PrefixCache {
|
|
enum class Mode {
|
|
NONE,
|
|
STORE,
|
|
REUSE
|
|
};
|
|
Mode mode = Mode::NONE;
|
|
std::string name;
|
|
std::string cut_group;
|
|
int64_t prefix_length = 0;
|
|
ggml_type type = GGML_TYPE_F32;
|
|
bool* flash_attn_used = nullptr;
|
|
};
|
|
|
|
class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
|
|
public:
|
|
using RMSNorm::RMSNorm;
|
|
|
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
|
auto weight = params["weight"];
|
|
if (ctx->weight_adapter) {
|
|
weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
|
|
}
|
|
weight = ggml_scale_bias(ctx->ggml_ctx, weight, 1.f, 1.f);
|
|
return ggml_mul(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, x, eps), weight);
|
|
}
|
|
};
|
|
|
|
class QwenImage21TextProjection : public GGMLBlock {
|
|
public:
|
|
QwenImage21TextProjection(const QwenImage21Config& config) {
|
|
blocks["text_norm"] = std::make_shared<QwenImage21ZeroCenterRMSNorm>(config.context_dim, 1e-6f);
|
|
blocks["in_layer"] = std::make_shared<Linear>(config.context_dim, config.hidden_size, false);
|
|
blocks["out_layer"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, false);
|
|
}
|
|
|
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
|
x = std::dynamic_pointer_cast<QwenImage21ZeroCenterRMSNorm>(blocks["text_norm"])->forward(ctx, x);
|
|
x = std::dynamic_pointer_cast<Linear>(blocks["in_layer"])->forward(ctx, x);
|
|
x = ggml_ext_gelu(ctx->ggml_ctx, x);
|
|
return std::dynamic_pointer_cast<Linear>(blocks["out_layer"])->forward(ctx, x);
|
|
}
|
|
};
|
|
|
|
class QwenImage21Attention : public QwenImageAttention {
|
|
public:
|
|
QwenImage21Attention(const QwenImage21Config& config)
|
|
: QwenImageAttention(config.hidden_size, config.head_dim, config.hidden_size / config.head_dim, 0, 0, false, false) {
|
|
for (const auto* name : {"add_q_proj", "add_k_proj", "add_v_proj", "norm_added_q", "norm_added_k", "to_add_out"}) {
|
|
blocks.erase(name);
|
|
}
|
|
}
|
|
|
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
|
|
int64_t heads = x->ne[0] / dim_head;
|
|
auto project = [&](const char* name) {
|
|
auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
|
|
return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]);
|
|
};
|
|
auto q = project("to_q");
|
|
auto k = project("to_k");
|
|
auto v = project("to_v");
|
|
q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
|
|
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
|
|
q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
|
|
k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
|
|
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
|
|
// Preserve query-first attention evaluation while writing each layer's
|
|
// prefix before its full-sequence K/V can accumulate across layers.
|
|
ctx->expand_graph(q);
|
|
auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
|
|
auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
|
|
// Pack the contiguous data into wider rows so quantization blocks
|
|
// can exceed head_dim without padding or changing element order.
|
|
part = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
|
|
auto copy = ggml_cast(ctx->ggml_ctx, part, cache.type);
|
|
// Keep the copy in this layer's segment so graph cuts do not
|
|
// retain or recompute the full-sequence K/V in the final segment.
|
|
sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name);
|
|
ctx->persist_cache_tensor(cache.name + "." + name, copy);
|
|
};
|
|
persist(k, 1, "k");
|
|
persist(v, 2, "v");
|
|
}
|
|
auto attend = [&](ggml_tensor* aq, ggml_tensor* ak, ggml_tensor* av, ggml_tensor* mask) {
|
|
bool used_flash_attn = false;
|
|
auto out = ggml_ext_attention_ext(ctx, aq, ak, av, heads, mask, true, ctx->flash_attn_enabled, 1.f, &used_flash_attn);
|
|
if (cache.flash_attn_used != nullptr) {
|
|
*cache.flash_attn_used &= used_flash_attn;
|
|
}
|
|
return out;
|
|
};
|
|
ggml_tensor* result = nullptr;
|
|
if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
|
|
auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
|
|
auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
|
|
GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
|
|
if (prefix_k->type != k->type) {
|
|
prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type);
|
|
}
|
|
if (prefix_v->type != v->type) {
|
|
prefix_v = ggml_cast(ctx->ggml_ctx, prefix_v, v->type);
|
|
}
|
|
prefix_k = ggml_reshape_4d(ctx->ggml_ctx, prefix_k, dim_head, cache.prefix_length, heads, k->ne[3]);
|
|
prefix_v = ggml_reshape_4d(ctx->ggml_ctx, prefix_v, dim_head, heads, cache.prefix_length, v->ne[3]);
|
|
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
|
|
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
|
|
result = attend(q, k, v, nullptr);
|
|
} else {
|
|
for (size_t i = 0; i < segments.size(); ++i) {
|
|
const auto& segment = segments[i];
|
|
auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
|
|
auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
|
|
auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
|
|
auto out = attend(sq, sk, sv, masks[i]);
|
|
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
|
|
}
|
|
}
|
|
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
|
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
|
to_out->set_force_prec_f32(true);
|
|
}
|
|
return to_out->forward(ctx, result);
|
|
}
|
|
};
|
|
|
|
class QwenImage21TransformerBlock : public GGMLBlock {
|
|
public:
|
|
QwenImage21TransformerBlock(const QwenImage21Config& config) {
|
|
blocks["img_norm1"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
|
|
blocks["img_norm2"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
|
|
blocks["attn"] = std::make_shared<QwenImage21Attention>(config);
|
|
if (config.fused_mlp) {
|
|
blocks["img_mlp.gate_up"] = std::make_shared<Linear>(config.hidden_size, 2 * config.intermediate_size, false);
|
|
} else {
|
|
blocks["img_mlp.proj"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, false);
|
|
blocks["img_mlp.gate_layer"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, false);
|
|
}
|
|
blocks["img_mlp.out"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, false);
|
|
}
|
|
|
|
static ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
|
|
auto rows = ggml_ext_chunk(ctx, params, 2, 1);
|
|
auto apply = [&](ggml_tensor* part, ggml_tensor* row) {
|
|
row = gate ? ggml_tanh(ctx, row) : ggml_scale_bias(ctx, row, 1.f, 1.f);
|
|
return ggml_mul(ctx, part, row);
|
|
};
|
|
auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
|
|
if (prefix_length == 0) {
|
|
return target;
|
|
}
|
|
auto prefix = apply(ggml_ext_slice(ctx, x, 1, 0, prefix_length), rows[1]);
|
|
return ggml_concat(ctx, prefix, target, 1);
|
|
}
|
|
|
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
|
|
const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length;
|
|
auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
|
|
h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length);
|
|
h = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache);
|
|
x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true));
|
|
h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
|
|
h = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length);
|
|
ggml_tensor* gate;
|
|
auto fused = blocks.find("img_mlp.gate_up");
|
|
if (fused != blocks.end()) {
|
|
auto gate_up = std::dynamic_pointer_cast<Linear>(fused->second)->forward(ctx, h);
|
|
auto parts = ggml_ext_chunk(ctx->ggml_ctx, gate_up, 2, 0);
|
|
gate = parts[0];
|
|
h = parts[1];
|
|
} else {
|
|
gate = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.gate_layer"])->forward(ctx, h);
|
|
h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.proj"])->forward(ctx, h);
|
|
}
|
|
h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
|
|
h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
|
|
return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true));
|
|
}
|
|
};
|
|
|
|
class QwenImage21Model : public GGMLBlock {
|
|
QwenImage21Config config;
|
|
|
|
public:
|
|
QwenImage21Model(const QwenImage21Config& config)
|
|
: config(config) {
|
|
blocks["time_text_embed.timestep_embedder"] = std::make_shared<TimestepEmbedding>(256, config.hidden_size, 0, 0, false);
|
|
blocks["txt_in"] = std::make_shared<QwenImage21TextProjection>(config);
|
|
blocks["img_in"] = std::make_shared<Linear>(config.in_channels, config.hidden_size, false);
|
|
blocks["modulation.1"] = std::make_shared<Linear>(config.hidden_size, 4 * config.hidden_size, false);
|
|
blocks["norm_out.linear"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, false);
|
|
blocks["norm_out.norm"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
|
|
blocks["proj_out"] = std::make_shared<Linear>(config.hidden_size, config.out_channels, false);
|
|
for (int i = 0; i < config.num_layers; ++i) {
|
|
blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<QwenImage21TransformerBlock>(config);
|
|
}
|
|
}
|
|
|
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
|
|
auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
|
|
// Runtime flow timesteps already use the [0, 1000] scale.
|
|
time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
|
|
time = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["time_text_embed.timestep_embedder"])->forward(ctx, time);
|
|
time = ggml_silu(ctx->ggml_ctx, time);
|
|
auto modulation = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
|
|
auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
|
|
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
|
|
ggml_tensor* joint = nullptr;
|
|
if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
|
|
joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1));
|
|
} else {
|
|
auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
|
|
for (const auto& segment : layout.segments) {
|
|
ggml_tensor* h;
|
|
if (segment.image_index < 0) {
|
|
h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
|
|
segment.context_start + segment.end - segment.start);
|
|
} else {
|
|
auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
|
|
h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
|
|
}
|
|
joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
|
|
}
|
|
}
|
|
sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint");
|
|
for (int i = 0; i < config.num_layers; ++i) {
|
|
const std::string layer = "transformer_blocks." + std::to_string(i);
|
|
auto layer_cache = cache;
|
|
layer_cache.name = cache.name + "." + std::to_string(i);
|
|
layer_cache.cut_group = "qwen_image_2_1." + layer;
|
|
auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks[layer]);
|
|
joint = block->forward(ctx, joint, mod, pe, layout, masks, layer_cache);
|
|
sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint");
|
|
}
|
|
if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) {
|
|
joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
|
|
}
|
|
auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
|
|
joint = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
|
|
joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
|
|
joint = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, joint);
|
|
return DiT::unpatchify_and_crop(ctx->ggml_ctx, joint, x->ne[1], x->ne[0], 1, 1);
|
|
}
|
|
};
|
|
|
|
struct QwenImage21Runner : public DiffusionModelRunner {
|
|
QwenImage21Config config;
|
|
QwenImage21Model model;
|
|
std::vector<float> pe_data;
|
|
std::vector<sd::Tensor<float>> mask_data;
|
|
ggml_type prefix_cache_type = GGML_TYPE_COUNT;
|
|
bool prefix_cache_enabled = true;
|
|
bool prefix_cache_disabled = false;
|
|
bool prefix_cache_auto_f32 = false;
|
|
|
|
static bool supports_prefix_cache_type(ggml_type type) {
|
|
if (type == GGML_TYPE_F32) {
|
|
return true;
|
|
}
|
|
const auto* traits = ggml_get_type_traits(type);
|
|
if (traits->from_float_ref == nullptr || traits->to_float == nullptr) {
|
|
return false;
|
|
}
|
|
auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
|
|
if (cpu == nullptr) {
|
|
return false;
|
|
}
|
|
auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
|
|
ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free);
|
|
if (ctx == nullptr) {
|
|
return false;
|
|
}
|
|
// Some reference quantizers have no runtime copy support. Query the
|
|
// device through the registry so dynamically loaded CPU backends work.
|
|
auto source = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type));
|
|
auto encoded = ggml_cast(ctx.get(), source, type);
|
|
auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32);
|
|
return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded);
|
|
}
|
|
|
|
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
|
|
: DiffusionModelRunner(backend, prefix, weight_manager),
|
|
config(QwenImage21Config::detect_from_weights(weights, prefix)),
|
|
model(config) {
|
|
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
|
|
if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) {
|
|
LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str());
|
|
} else if (key == "qwen_image_2_1_prefix_cache_type") {
|
|
if (value == "auto") {
|
|
prefix_cache_type = GGML_TYPE_COUNT;
|
|
continue;
|
|
}
|
|
const auto type = sd_type_to_ggml_type(str_to_sd_type(value.c_str()));
|
|
if (type == GGML_TYPE_COUNT) {
|
|
LOG_WARN("ignoring unknown Qwen Image 2.1 cache type '%s'", value.c_str());
|
|
} else if (!supports_prefix_cache_type(type)) {
|
|
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': runtime conversion to and from F32 is unavailable", value.c_str());
|
|
} else if (config.hidden_size % ggml_blck_size(type) != 0) {
|
|
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': block size %" PRId64 " does not divide hidden size %" PRId64,
|
|
value.c_str(), ggml_blck_size(type), config.hidden_size);
|
|
} else {
|
|
prefix_cache_type = type;
|
|
}
|
|
}
|
|
}
|
|
model.init(params_ctx, weights, prefix);
|
|
}
|
|
|
|
std::string get_desc() override { return "qwen_image_2_1"; }
|
|
|
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
|
|
model.get_param_tensors(tensors, prefix);
|
|
}
|
|
|
|
bool has_prefix_cache(const QwenImage21PrefixCache& cache) {
|
|
for (int i = 0; i < config.num_layers; ++i) {
|
|
const auto name = cache.name + "." + std::to_string(i);
|
|
auto k = get_cache_tensor_by_name(name + ".k");
|
|
auto v = get_cache_tensor_by_name(name + ".v");
|
|
if (k == nullptr || v == nullptr || k->type != cache.type || v->type != cache.type ||
|
|
k->ne[0] != config.hidden_size || k->ne[1] != cache.prefix_length || k->ne[2] != 1 || k->ne[3] != 1 ||
|
|
v->ne[0] != config.hidden_size || v->ne[1] != cache.prefix_length || v->ne[2] != 1 || v->ne[3] != 1) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override {
|
|
const auto& x = tensor_or_empty(inputs.x);
|
|
const auto& context = tensor_or_empty(inputs.context);
|
|
if (x.empty() || context.empty() || context.dim() < 2 || context.shape()[0] != config.context_dim ||
|
|
tensor_or_empty(inputs.timesteps).numel() != 1 ||
|
|
x.dim() != 4 || x.shape()[3] != 1 || x.shape()[2] != config.in_channels) {
|
|
LOG_ERROR("Qwen Image 2.1 requires an image latent and text conditioning with batch size 1");
|
|
return {};
|
|
}
|
|
static const std::vector<sd::Tensor<float>> empty_refs;
|
|
const auto& refs = inputs.ref_latents && inputs.ref_image_params.pass_to_dit ? *inputs.ref_latents : empty_refs;
|
|
std::vector<std::pair<int64_t, int64_t>> shapes;
|
|
for (const auto& ref : refs) {
|
|
if (ref.dim() != 4 || ref.shape()[2] != config.in_channels || ref.shape()[3] != 1) {
|
|
LOG_ERROR("Qwen Image 2.1: invalid reference latent shape");
|
|
return {};
|
|
}
|
|
shapes.emplace_back(ref.shape()[1], ref.shape()[0]);
|
|
}
|
|
shapes.emplace_back(x.shape()[1], x.shape()[0]);
|
|
const auto* extra = std::get_if<QwenImage21DiffusionExtra>(&inputs.extra);
|
|
QwenImage21Layout layout;
|
|
try {
|
|
layout = QwenImage21Layout::build(context.shape()[1], tensor_or_empty(extra ? extra->image_slots : nullptr), shapes);
|
|
} catch (const std::exception& error) {
|
|
LOG_ERROR("%s", error.what());
|
|
return {};
|
|
}
|
|
if (!runner_started()) {
|
|
prefix_cache_disabled = false;
|
|
prefix_cache_auto_f32 = false;
|
|
}
|
|
QwenImage21PrefixCache cache;
|
|
if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
|
|
cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
|
|
".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
|
|
cache.prefix_length = layout.prefix_length;
|
|
if (prefix_cache_type != GGML_TYPE_COUNT) {
|
|
cache.type = prefix_cache_type;
|
|
} else if (!prefix_cache_auto_f32 && flash_attn_enabled && !sage_attn_enabled &&
|
|
(attn_scale <= 0.f || attn_scale == 1.f)) {
|
|
cache.type = GGML_TYPE_F16;
|
|
}
|
|
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
|
|
}
|
|
bool flash_attn_used = true;
|
|
auto run = [&](const QwenImage21PrefixCache& active_cache) {
|
|
flash_attn_used = true;
|
|
auto checked_cache = active_cache;
|
|
if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
|
|
checked_cache.flash_attn_used = &flash_attn_used;
|
|
}
|
|
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
|
|
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
|
|
Rope::Embedding embedding;
|
|
embedding.ids.assign(first_position, layout.positions.end());
|
|
const size_t offset = cached ? static_cast<size_t>(layout.prefix_length) : 0;
|
|
embedding.positions.token_count = embedding.ids.size();
|
|
for (auto region : layout.rope_layout.images) {
|
|
if (region.begin >= offset) {
|
|
region.begin -= offset;
|
|
embedding.positions.images.push_back(region);
|
|
}
|
|
}
|
|
embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
|
|
pe_data = finish_rope_pe(std::move(embedding));
|
|
mask_data.clear();
|
|
if (!cached) {
|
|
for (const auto& segment : layout.segments) {
|
|
sd::Tensor<float> mask;
|
|
if (segment.image_index < 0) {
|
|
mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
|
|
for (int64_t q = segment.start; q < segment.end; ++q) {
|
|
for (int64_t k = q + 1; k < segment.end; ++k) {
|
|
mask[k + segment.end * (q - segment.start)] = -INFINITY;
|
|
}
|
|
}
|
|
}
|
|
mask_data.push_back(std::move(mask));
|
|
}
|
|
}
|
|
auto build = [&]() {
|
|
auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
|
|
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
|
|
layout.positions.size() - (cached ? layout.prefix_length : 0));
|
|
set_backend_tensor_data(pe, pe_data.data());
|
|
std::vector<ggml_tensor*> masks, ref_inputs;
|
|
for (const auto& mask : mask_data) {
|
|
masks.push_back(mask.empty() ? nullptr : make_input(mask));
|
|
}
|
|
if (!cached) {
|
|
for (const auto& ref : refs) {
|
|
ref_inputs.push_back(make_input(ref));
|
|
}
|
|
}
|
|
auto ctx = get_context(graph);
|
|
auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
|
|
ref_inputs, pe, layout, masks, checked_cache);
|
|
if (!flash_attn_used) {
|
|
return static_cast<ggml_cgraph*>(nullptr);
|
|
}
|
|
ggml_build_forward_expand(graph, out);
|
|
return graph;
|
|
};
|
|
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
|
|
};
|
|
auto result = run(cache);
|
|
if (result.empty() && !flash_attn_used) {
|
|
// Casting an F16 cache back to F32 cannot recover its original values.
|
|
// Recompute the prefix before executing a graph that falls back from FA.
|
|
free_cache_ctx_and_buffer();
|
|
prefix_cache_auto_f32 = true;
|
|
cache.type = GGML_TYPE_F32;
|
|
cache.mode = QwenImage21PrefixCache::Mode::STORE;
|
|
LOG_DEBUG("Qwen Image 2.1: Flash Attention unavailable; using F32 prefix caching for this sampling run");
|
|
result = run(cache);
|
|
}
|
|
if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
|
|
(cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
|
|
// The failed graph has ended before persistent inputs are released.
|
|
free_cache_ctx_and_buffer();
|
|
prefix_cache_disabled = true;
|
|
LOG_WARN("Qwen Image 2.1: insufficient memory for prefix caching; retrying without it for this sampling run");
|
|
return run(QwenImage21PrefixCache{});
|
|
}
|
|
if (!result.empty() && cache.mode == QwenImage21PrefixCache::Mode::STORE) {
|
|
if (!has_prefix_cache(cache)) {
|
|
free_cache_ctx_and_buffer();
|
|
prefix_cache_disabled = true;
|
|
LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run");
|
|
} else {
|
|
LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens, %s)", extra->prefix_id, layout.prefix_length, ggml_type_name(cache.type));
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
};
|
|
}
|
|
|
|
#endif // __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
|