diff --git a/docs/qwen_image_2.1.md b/docs/qwen_image_2.1.md index 31008395..3eaae207 100644 --- a/docs/qwen_image_2.1.md +++ b/docs/qwen_image_2.1.md @@ -40,6 +40,14 @@ Pass the reference image with `-r` and describe the edit in `-p`. Vision weights For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`. +### Prefix cache + +By default, the first denoising call for each fixed condition saves the text and reference-image keys and values from every transformer layer. Later calls only compute the target-image tokens. Positive and negative conditions use separate caches, which are released when sampling ends. + +The cache uses FP32 on all attention backends. For the default 32-layer model, a prefix of 4096 tokens takes about 4 GiB per condition, in addition to weights and working buffers. The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path. + +Disable this optimization with `--model-args qwen_image_2_1_prefix_cache=false`. It reuses step-independent activations; numerical results can still differ slightly because the matrix sizes change. + ### Alpha channel This model supports alpha channel output. As the model determines whether to output a regular image or with transparency through the prompt, according to [official recommendation](https://github.com/QwenLM/Qwen-Image-2.1#transparent-image-generation-rgba), use the following prompt format for better results: diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 81345199..24dc937b 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -518,7 +518,7 @@ ArgOptions SDContextParams::get_options() { {"", "--model-args", "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, " - "chroma_t5_mask_pad, qwen_image_zero_cond_t", + "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache", (int)',', &model_args}, {"", diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 182fd098..59573210 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -644,6 +644,10 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, std::optional> output; try { output = execute_graph(graph, n_threads, no_return, read_outputs); + } catch (const std::bad_alloc&) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + LOG_ERROR("%s graph allocation failed", get_desc().c_str()); + return std::nullopt; } catch (const std::exception& error) { last_compute_status_ = GGML_STATUS_FAILED; LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), @@ -964,10 +968,16 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), index + 1, plan.segments.size(), segment.group_name.c_str()); - if (!execute_segment(segment_graph, n_threads) || - !cache_.capture(segment_graph) || - !cut_cache_.capture(graph, segment, get_desc().c_str())) { - return fail_segment("execution or output caching"); + if (!execute_segment(segment_graph, n_threads)) { + return fail_segment("execution"); + } + auto cache_status = cache_.capture(segment_graph); + if (cache_status == GGML_STATUS_SUCCESS) { + cache_status = cut_cache_.capture(graph, segment, get_desc().c_str()); + } + if (cache_status != GGML_STATUS_SUCCESS) { + last_compute_status_ = cache_status; + return fail_segment("output caching"); } sync_runtime_residency(); if (last) { diff --git a/src/core/runner_cache.cpp b/src/core/runner_cache.cpp index 09562106..b640bb45 100644 --- a/src/core/runner_cache.cpp +++ b/src/core/runner_cache.cpp @@ -26,10 +26,13 @@ namespace sd { std::unique_ptr CachedTensor::copy(ggml_backend_t backend, const std::string& name, - ggml_tensor* source) { + ggml_tensor* source, + ggml_status& status) { + status = GGML_STATUS_FAILED; if (ggml_graph_cut::tensor_buffer(source) == nullptr) { return nullptr; } + status = GGML_STATUS_ALLOC_FAILED; auto entry = std::make_unique(); entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); if (entry->context == nullptr) { @@ -50,6 +53,7 @@ namespace sd { } else { ggml_backend_tensor_copy(source, entry->tensor); } + status = GGML_STATUS_SUCCESS; return entry; } @@ -106,9 +110,9 @@ namespace sd { return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending; } - bool RunnerCache::capture(ggml_cgraph* graph) { + ggml_status RunnerCache::capture(ggml_cgraph* graph) { if (outputs_.empty()) { - return true; + return GGML_STATUS_SUCCESS; } const auto tensors = cache_graph_tensors(graph); for (const auto& output : outputs_) { @@ -116,14 +120,15 @@ namespace sd { continue; } GGML_ASSERT(ggml_is_contiguous(output.second)); - auto entry = CachedTensor::copy(backend_, output.first, output.second); + ggml_status status; + auto entry = CachedTensor::copy(backend_, output.first, output.second, status); if (entry == nullptr) { - return false; + return status; } pending_[output.first] = std::move(entry); } ggml_backend_synchronize(backend_); - return true; + return GGML_STATUS_SUCCESS; } void RunnerCache::graph_end(bool success) { @@ -180,9 +185,9 @@ namespace sd { } } - bool GraphCutTensorCache::capture(ggml_cgraph* graph, - const ggml_graph_cut::Segment& segment, - const char* log_desc) { + ggml_status GraphCutTensorCache::capture(ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment, + const char* log_desc) { size_t copied_bytes = 0; size_t copied_count = 0; for (int index : segment.output_node_indices) { @@ -191,10 +196,11 @@ namespace sd { !segment.future_cut_names.count(output->name)) { continue; } - auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output)); + ggml_status status; + auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status); if (entry == nullptr) { LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name); - return false; + return status; } const size_t size = ggml_backend_buffer_get_size(entry->buffer); copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size; @@ -206,6 +212,6 @@ namespace sd { LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)", log_desc, copied_bytes / (1024.f * 1024.f), copied_count); } - return true; + return GGML_STATUS_SUCCESS; } } diff --git a/src/core/runner_cache.h b/src/core/runner_cache.h index bd331a4b..f1b5cf91 100644 --- a/src/core/runner_cache.h +++ b/src/core/runner_cache.h @@ -20,7 +20,8 @@ namespace sd { ~CachedTensor(); static std::unique_ptr copy(ggml_backend_t backend, const std::string& name, - ggml_tensor* source); + ggml_tensor* source, + ggml_status& status); }; using CachedTensors = std::map>; @@ -41,7 +42,8 @@ namespace sd { const std::map& outputs() const { return outputs_; } size_t pending_bytes(ggml_cgraph* graph) const; size_t resident_bytes(ggml_backend_dev_t device) const; - bool capture(ggml_cgraph* graph); + bool empty() const { return committed_.empty(); } + ggml_status capture(ggml_cgraph* graph); void graph_end(bool success); void clear(); }; @@ -57,7 +59,7 @@ namespace sd { size_t resident_bytes(ggml_backend_dev_t device) const; size_t estimate_output_bytes(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment) const; - bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); + ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); void prune(const std::unordered_set& keep_names); void clear() { tensors_.clear(); } }; diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index fb4e362c..f82fc36f 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -71,6 +71,8 @@ struct AnimaDiffusionExtra { struct QwenImage21DiffusionExtra { const sd::Tensor* image_slots = nullptr; + // Nonzero IDs identify immutable prefix inputs within one sampling run. + uint64_t prefix_id = 0; }; struct WanDiffusionExtra { diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index cb044053..1364de13 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -121,6 +121,18 @@ namespace Qwen { } }; + struct QwenImage21PrefixCache { + enum class Mode { + NONE, + STORE, + REUSE + }; + Mode mode = Mode::NONE; + std::string name; + std::string cut_group; + int64_t prefix_length = 0; + }; + class QwenImage21ZeroCenterRMSNorm : public RMSNorm { public: using RMSNorm::RMSNorm; @@ -160,7 +172,7 @@ namespace Qwen { } } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector& segments, const std::vector& masks) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector& segments, const std::vector& masks, const QwenImage21PrefixCache& cache) { int64_t heads = x->ne[0] / dim_head; auto project = [&](const char* name) { auto h = std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); @@ -173,14 +185,36 @@ namespace Qwen { k = std::dynamic_pointer_cast(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) { + auto persist = [&](ggml_tensor* tensor, int axis, const char* name) { + auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length); + auto copy = ggml_new_tensor(ctx->ggml_ctx, GGML_TYPE_F32, 4, part->ne); + copy = ggml_cpy(ctx->ggml_ctx, part, copy); + // 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"); + } ggml_tensor* result = nullptr; - 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 = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled); - result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); + 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); + k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1); + v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2); + result = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, true, ctx->flash_attn_enabled); + } 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 = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled); + result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); + } } auto to_out = std::dynamic_pointer_cast(blocks["to_out.0"]); if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { @@ -219,13 +253,14 @@ namespace Qwen { return ggml_concat(ctx, prefix, target, 1); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks) { - auto h = std::dynamic_pointer_cast(blocks["img_norm1"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length); - h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks); - x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true)); - h = std::dynamic_pointer_cast(blocks["img_norm2"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length); + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks, const QwenImage21PrefixCache& cache) { + const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length; + auto h = std::dynamic_pointer_cast(blocks["img_norm1"])->forward(ctx, x); + h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length); + h = std::dynamic_pointer_cast(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(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()) { @@ -239,7 +274,7 @@ namespace Qwen { } h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate)); h = std::dynamic_pointer_cast(blocks["img_mlp.out"])->forward(ctx, h); - return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], layout.prefix_length, true)); + return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true)); } }; @@ -261,7 +296,7 @@ namespace Qwen { } } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& 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); @@ -269,27 +304,37 @@ namespace Qwen { time = ggml_silu(ctx->ggml_ctx, time); auto modulation = std::dynamic_pointer_cast(blocks["modulation.1"])->forward(ctx, time); auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0); - auto text = std::dynamic_pointer_cast(blocks["txt_in"])->forward(ctx, context); auto img_in = std::dynamic_pointer_cast(blocks["img_in"]); ggml_tensor* joint = nullptr; - 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(refs.size()) ? x : refs[segment.image_index]; - h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1)); + 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(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(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); } - 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) { - auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); - joint = block->forward(ctx, joint, mod, pe, layout, masks); - sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint"); + 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(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]); } - joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]); auto scale = std::dynamic_pointer_cast(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]); joint = std::dynamic_pointer_cast(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)); @@ -303,11 +348,18 @@ namespace Qwen { QwenImage21Model model; std::vector pe_data; std::vector> mask_data; + bool prefix_cache_enabled = true; + bool prefix_cache_disabled = false; - QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr 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()); + } + } model.init(params_ctx, weights, prefix); } @@ -317,6 +369,22 @@ namespace Qwen { 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 != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 || + k->ne[0] != config.head_dim || k->ne[1] != cache.prefix_length || + k->ne[2] != config.hidden_size / config.head_dim || k->ne[3] != 1 || + v->ne[0] != config.head_dim || v->ne[1] != config.hidden_size / config.head_dim || + v->ne[2] != cache.prefix_length || v->ne[3] != 1) { + return false; + } + } + return true; + } + sd::Tensor compute(int n_threads, const DiffusionParams& inputs) override { const auto& x = tensor_or_empty(inputs.x); const auto& context = tensor_or_empty(inputs.context); @@ -345,38 +413,75 @@ namespace Qwen { LOG_ERROR("%s", error.what()); return {}; } - pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim); - mask_data.clear(); - for (const auto& segment : layout.segments) { - sd::Tensor mask; - if (segment.image_index < 0) { - mask = sd::Tensor::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; + if (!runner_started()) { + prefix_cache_disabled = 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); + cache.prefix_length = layout.prefix_length; + cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE; + } + auto run = [&](const QwenImage21PrefixCache& active_cache) { + const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE; + const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0); + pe_data = Rope::embed_nd(std::vector>(first_position, layout.positions.end()), 1, 10000.f, config.axes_dim); + mask_data.clear(); + if (!cached) { + for (const auto& segment : layout.segments) { + sd::Tensor mask; + if (segment.image_index < 0) { + mask = sd::Tensor::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)); } } - 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()); - set_backend_tensor_data(pe, pe_data.data()); - std::vector masks, ref_inputs; - for (const auto& mask : mask_data) { - masks.push_back(mask.empty() ? nullptr : make_input(mask)); - } - for (const auto& ref : refs) { - ref_inputs.push_back(make_input(ref)); - } - auto ctx = get_context(); - auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context), - ref_inputs, pe, layout, masks); - ggml_build_forward_expand(graph, out); - return graph; + 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 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(); + auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context), + ref_inputs, pe, layout, masks, active_cache); + ggml_build_forward_expand(graph, out); + return graph; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim()); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim()); + auto 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)", extra->prefix_id, layout.prefix_length); + } + } + return result; } }; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 78426f08..02209395 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -2255,6 +2256,15 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptris_enabled(); }); + using QwenPrefixInputs = std::tuple*, const sd::Tensor*, + const std::vector>*>; + std::vector qwen_prefix_inputs; + RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; const bool apply_denoise_mask = !denoise_mask.empty() && @@ -2524,6 +2534,18 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptrbefore_diffusion(diffusion_params, step); } + if (cache_qwen_prefix) { + auto* extra = std::get_if(&diffusion_params.extra); + if (extra != nullptr) { + auto key = std::make_tuple(diffusion_params.context, extra->image_slots, + diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr); + auto entry = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key); + extra->prefix_id = static_cast(entry - qwen_prefix_inputs.begin()) + 1; + if (entry == qwen_prefix_inputs.end()) { + qwen_prefix_inputs.push_back(key); + } + } + } auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params); if (output_opt.empty()) { LOG_ERROR("diffusion model compute failed"); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 626bdfee..9b8819a0 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -291,7 +291,8 @@ namespace sd::model_builders { result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", - weight_manager); + weight_manager, + sd_ctx_params->model_args); } else { result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), tensor_storage_map,