From dacdbaa5f1fad4d30350b6da28f7634dbc94136d Mon Sep 17 00:00:00 2001 From: leejet Date: Thu, 14 May 2026 23:18:18 +0800 Subject: [PATCH] refactor --- src/conditioner.hpp | 10 +- src/denoiser.hpp | 24 +- src/diffusion_model.hpp | 51 ++- src/ggml_extend.hpp | 83 ++++- src/ggml_graph_cut.cpp | 30 +- src/hidream_o1.hpp | 662 +++++++++-------------------------- src/llm.hpp | 720 +++++++++++++++++++++++++-------------- src/stable-diffusion.cpp | 30 +- 8 files changed, 808 insertions(+), 802 deletions(-) diff --git a/src/conditioner.hpp b/src/conditioner.hpp index cc51718a..5050eeff 100644 --- a/src/conditioner.hpp +++ b/src/conditioner.hpp @@ -17,9 +17,8 @@ struct SDCondition { sd::Tensor c_input_ids; sd::Tensor c_position_ids; sd::Tensor c_token_types; - sd::Tensor c_image_embed_ranges; sd::Tensor c_vinput_mask; - std::vector> c_vlm_images; + std::vector>> c_image_embeds; std::vector> c_ref_images; std::vector> extra_c_crossattns; @@ -35,13 +34,12 @@ struct SDCondition { if (!c_crossattn.empty() || !c_vector.empty() || !c_concat.empty() || !c_t5_ids.empty() || !c_t5_weights.empty() || !c_input_ids.empty() || !c_position_ids.empty() || - !c_token_types.empty() || !c_image_embed_ranges.empty() || - !c_vinput_mask.empty()) { + !c_token_types.empty() || !c_vinput_mask.empty()) { return false; } - for (const auto& tensor : c_vlm_images) { - if (!tensor.empty()) { + for (const auto& image_embed : c_image_embeds) { + if (!image_embed.second.empty()) { return false; } } diff --git a/src/denoiser.hpp b/src/denoiser.hpp index e7a25a6d..444cc45b 100644 --- a/src/denoiser.hpp +++ b/src/denoiser.hpp @@ -872,8 +872,9 @@ static sd::Tensor sample_euler_flow_flash(denoise_cb_t model, const std::vector& sigmas, std::shared_ptr rng, float eta) { - float s_noise = eta; - int steps = static_cast(sigmas.size()) - 1; + constexpr float noise_clip_std = 2.5f; + float s_noise = eta; + int steps = static_cast(sigmas.size()) - 1; for (int i = 0; i < steps; i++) { float sigma = sigmas[i]; float sigma_next = sigmas[i + 1]; @@ -887,7 +888,24 @@ static sd::Tensor sample_euler_flow_flash(denoise_cb_t model, continue; } auto noise = sd::Tensor::randn_like(x, rng); - x = sigma_next * noise * s_noise + (1.0f - sigma_next) * denoised; + if (noise_clip_std > 0.0f && noise.numel() > 0) { + double mean = 0.0; + for (int64_t j = 0; j < noise.numel(); ++j) { + mean += static_cast(noise[j]); + } + mean /= static_cast(noise.numel()); + + double variance = 0.0; + for (int64_t j = 0; j < noise.numel(); ++j) { + double centered = static_cast(noise[j]) - mean; + variance += centered * centered; + } + variance /= static_cast(noise.numel()); + + float clip_val = noise_clip_std * static_cast(std::sqrt(variance)); + noise = sd::ops::clamp(noise, -clip_val, clip_val); + } + x = sigma_next * noise * s_noise + (1.0f - sigma_next) * denoised; } return x; } diff --git a/src/diffusion_model.hpp b/src/diffusion_model.hpp index d0204baf..26021ef2 100644 --- a/src/diffusion_model.hpp +++ b/src/diffusion_model.hpp @@ -14,28 +14,28 @@ #include "z_image.hpp" struct DiffusionParams { - const sd::Tensor* x = nullptr; - const sd::Tensor* timesteps = nullptr; - const sd::Tensor* context = nullptr; - const sd::Tensor* c_concat = nullptr; - const sd::Tensor* y = nullptr; - const sd::Tensor* t5_ids = nullptr; - const sd::Tensor* t5_weights = nullptr; - const sd::Tensor* guidance = nullptr; - const std::vector>* ref_latents = nullptr; - const sd::Tensor* input_ids = nullptr; - const sd::Tensor* input_pos = nullptr; - const sd::Tensor* token_types = nullptr; - const sd::Tensor* image_embed_ranges = nullptr; - const sd::Tensor* vinput_mask = nullptr; - const std::vector>* vlm_images = nullptr; - bool increase_ref_index = false; - int num_video_frames = -1; - const std::vector>* controls = nullptr; - float control_strength = 0.f; - const sd::Tensor* vace_context = nullptr; - float vace_strength = 1.f; - const std::vector* skip_layers = nullptr; + const sd::Tensor* x = nullptr; + const sd::Tensor* timesteps = nullptr; + const sd::Tensor* context = nullptr; + const sd::Tensor* c_concat = nullptr; + const sd::Tensor* y = nullptr; + const sd::Tensor* t5_ids = nullptr; + const sd::Tensor* t5_weights = nullptr; + const sd::Tensor* guidance = nullptr; + const std::vector>* ref_latents = nullptr; + const sd::Tensor* input_ids = nullptr; + const sd::Tensor* input_pos = nullptr; + const sd::Tensor* token_types = nullptr; + const sd::Tensor* vinput_mask = nullptr; + const std::vector>* vlm_images = nullptr; + const std::vector>>* image_embeds = nullptr; + bool increase_ref_index = false; + int num_video_frames = -1; + const std::vector>* controls = nullptr; + float control_strength = 0.f; + const sd::Tensor* vace_context = nullptr; + float vace_strength = 1.f; + const std::vector* skip_layers = nullptr; }; template @@ -545,17 +545,16 @@ struct HiDreamO1Model : public DiffusionModel { GGML_ASSERT(diffusion_params.input_ids != nullptr); GGML_ASSERT(diffusion_params.input_pos != nullptr); GGML_ASSERT(diffusion_params.token_types != nullptr); - static const sd::Tensor empty_image_embed_ranges; static const std::vector> empty_images; + static const std::vector>> empty_image_embeds; return hidream_o1.compute(n_threads, *diffusion_params.x, *diffusion_params.timesteps, *diffusion_params.input_ids, *diffusion_params.input_pos, *diffusion_params.token_types, - diffusion_params.image_embed_ranges ? *diffusion_params.image_embed_ranges : empty_image_embed_ranges, - diffusion_params.vinput_mask ? *diffusion_params.vinput_mask : empty_image_embed_ranges, - diffusion_params.vlm_images ? *diffusion_params.vlm_images : empty_images, + tensor_or_empty(diffusion_params.vinput_mask), + diffusion_params.image_embeds ? *diffusion_params.image_embeds : empty_image_embeds, diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_images); } }; diff --git a/src/ggml_extend.hpp b/src/ggml_extend.hpp index c1939bcc..77e2774f 100644 --- a/src/ggml_extend.hpp +++ b/src/ggml_extend.hpp @@ -280,6 +280,9 @@ __STATIC_INLINE__ void print_sd_tensor(const sd::Tensor& tensor, bool shape_o if (shape_only) { return; } + if (tensor.empty()) { + return; + } int range = 3; std::vector shape = tensor.shape(); while (shape.size() < 4) { @@ -2021,9 +2024,13 @@ protected: ggml_backend_buffer_t src_buf = sd::ggml_graph_cut::tensor_buffer(src); ggml_backend_buffer_t dst_buf = sd::ggml_graph_cut::tensor_buffer(dst); if (src_buf == nullptr || dst_buf == nullptr) { - LOG_ERROR("%s cache copy tensor buffer missing: name=%s src_buffer=%p src_view_src=%p src_view_src_buffer=%p dst_buffer=%p", + LOG_ERROR("%s cache copy tensor buffer missing: name=%s op=%s src0=%p src0_name=%s src0_buffer=%p src_buffer=%p src_view_src=%p src_view_src_buffer=%p dst_buffer=%p", get_desc().c_str(), src && src->name[0] != '\0' ? src->name : "", + src ? ggml_op_name(src->op) : "", + src ? src->src[0] : nullptr, + (src && src->src[0] && src->src[0]->name[0] != '\0') ? src->src[0]->name : "", + (src && src->src[0]) ? sd::ggml_graph_cut::tensor_buffer(src->src[0]) : nullptr, src ? src->buffer : nullptr, src ? src->view_src : nullptr, (src && src->view_src) ? src->view_src->buffer : nullptr, @@ -2055,6 +2062,42 @@ protected: return true; } + template + std::optional> read_graph_tensor(ggml_tensor* tensor, const char* label) { + if (tensor == nullptr) { + LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label); + return std::nullopt; + } + if (tensor->type != sd::GGMLTypeTraits::type) { + LOG_ERROR("%s %s tensor type mismatch: got %s", + get_desc().c_str(), + label, + ggml_type_name(tensor->type)); + return std::nullopt; + } + ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor); + if (buf == nullptr) { + LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p", + get_desc().c_str(), + label, + tensor->name[0] != '\0' ? tensor->name : "", + ggml_op_name(tensor->op), + tensor->buffer, + tensor->view_src, + tensor->view_src ? tensor->view_src->buffer : nullptr, + tensor->data); + return std::nullopt; + } + + sd::Tensor result(sd::shape_from_ggml(tensor)); + if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor) || tensor->buffer == nullptr) { + ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); + } else { + ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); + } + return result; + } + void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true) { GGML_ASSERT(gf != nullptr); std::unordered_set graph_tensor_set; @@ -2075,6 +2118,9 @@ protected: continue; } const char* name = ggml_get_name(tensor); + if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { + continue; + } if (tensor->buffer == nullptr) { LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s", get_desc().c_str(), @@ -2087,10 +2133,6 @@ protected: continue; } - if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { - continue; - } - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; if (buf == nullptr) { LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", @@ -2476,11 +2518,32 @@ protected: return std::nullopt; } + std::unordered_set debug_graph_tensor_set; + const int n_debug_leafs = sd::ggml_graph_cut::leaf_count(gf); + const int n_debug_nodes = ggml_graph_n_nodes(gf); + debug_graph_tensor_set.reserve(static_cast(n_debug_leafs + n_debug_nodes)); + for (int i = 0; i < n_debug_leafs; ++i) { + debug_graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); + } + for (int i = 0; i < n_debug_nodes; ++i) { + debug_graph_tensor_set.insert(ggml_graph_node(gf, i)); + } + for (const auto& entry : debug_tensors) { auto tensor = entry.first; if (tensor == nullptr) { continue; } + if (debug_graph_tensor_set.find(tensor) == debug_graph_tensor_set.end()) { + continue; + } + ggml_backend_buffer_t tensor_buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (tensor_buf == nullptr) { + LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", + get_desc().c_str(), + entry.second.c_str()); + continue; + } if (tensor->type != GGML_TYPE_F32) { LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", get_desc().c_str(), @@ -2505,7 +2568,15 @@ protected: auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); std::optional> output; if (!no_return) { - output = sd::make_sd_tensor_from_ggml(result); + output = read_graph_tensor(result, "output"); + if (!output.has_value()) { + if (free_compute_buffer_immediately) { + free_compute_buffer(); + } else if (use_partial_param_offload) { + restore_partial_params(); + } + return std::nullopt; + } } else { output = sd::Tensor(); } diff --git a/src/ggml_graph_cut.cpp b/src/ggml_graph_cut.cpp index f206f2d2..0958d888 100644 --- a/src/ggml_graph_cut.cpp +++ b/src/ggml_graph_cut.cpp @@ -45,6 +45,21 @@ namespace sd::ggml_graph_cut { return params_tensor_set.find(tensor) != params_tensor_set.end(); } + static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) { + GGML_ASSERT(gf != nullptr); + if (name == nullptr || name[0] == '\0') { + return -1; + } + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (node != nullptr && std::strcmp(node->name, name) == 0) { + return i; + } + } + return -1; + } + static Plan::InputShape input_shape(const ggml_tensor* tensor) { Plan::InputShape shape; if (tensor == nullptr) { @@ -244,6 +259,11 @@ namespace sd::ggml_graph_cut { if (tensor == nullptr) { return nullptr; } + if (tensor_buffer(tensor) == nullptr && tensor->src[0] != nullptr && + ggml_nelements(tensor->src[0]) == ggml_nelements(tensor) && + ggml_nbytes(tensor->src[0]) == ggml_nbytes(tensor)) { + return cache_source_tensor(tensor->src[0]); + } return tensor->view_src ? tensor->view_src : tensor; } @@ -503,11 +523,15 @@ namespace sd::ggml_graph_cut { log_desc); } - ggml_tensor* final_output = ggml_graph_node(gf, -1); - if (final_output != nullptr && available_cut_output_node_indices.find(n_nodes - 1) == available_cut_output_node_indices.end()) { + int final_output_index = graph_node_index_by_name(gf, "ggml_runner_final_result_tensor"); + if (final_output_index < 0) { + final_output_index = n_nodes - 1; + } + ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr; + if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) { Segment final_segment; final_segment.group_name = "ggml_runner.final"; - final_segment.output_node_indices.push_back(n_nodes - 1); + final_segment.output_node_indices.push_back(final_output_index); build_segment(gf, plan, final_segment, diff --git a/src/hidream_o1.hpp b/src/hidream_o1.hpp index 76ad629f..908f2de3 100644 --- a/src/hidream_o1.hpp +++ b/src/hidream_o1.hpp @@ -70,12 +70,16 @@ namespace HiDreamO1 { sd::Tensor resized; if (s1 < s2) { int64_t resized_h = static_cast(std::llround(height / s1)); - resized = sd::ops::interpolate(image, {new_size.first, resized_h, image.shape()[2], image.shape()[3]}); + resized = sd::ops::interpolate(image, + {new_size.first, resized_h, image.shape()[2], image.shape()[3]}, + sd::ops::InterpolateMode::Bicubic); int64_t top = (resized_h - new_size.second) / 2; resized = sd::ops::slice(resized, 1, top, top + new_size.second); } else { int64_t resized_w = static_cast(std::llround(width / s2)); - resized = sd::ops::interpolate(image, {resized_w, new_size.second, image.shape()[2], image.shape()[3]}); + resized = sd::ops::interpolate(image, + {resized_w, new_size.second, image.shape()[2], image.shape()[3]}, + sd::ops::InterpolateMode::Bicubic); int64_t left = (resized_w - new_size.first) / 2; resized = sd::ops::slice(resized, 0, left, left + new_size.first); } @@ -202,298 +206,34 @@ namespace HiDreamO1 { struct HiDreamO1Params { LLM::LLMParams llm; - int patch_size = PATCH_SIZE; - int num_position_embeddings = 2304; - std::vector deepstack_visual_indexes; + int patch_size = PATCH_SIZE; }; - struct VisionMLP : public GGMLBlock { - VisionMLP(int64_t hidden_size, int64_t intermediate_size) { - blocks["linear_fc1"] = std::make_shared(hidden_size, intermediate_size, true); - blocks["linear_fc2"] = std::make_shared(intermediate_size, hidden_size, true); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto linear_fc1 = std::dynamic_pointer_cast(blocks["linear_fc1"]); - auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); - - x = linear_fc1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x); - x = linear_fc2->forward(ctx, x); - return x; - } - }; - - struct VisionPatchEmbed : public GGMLBlock { - int patch_size; - int temporal_patch_size; - int64_t in_channels; - int64_t embed_dim; - - VisionPatchEmbed(int patch_size, - int temporal_patch_size, - int64_t in_channels, - int64_t embed_dim) - : patch_size(patch_size), - temporal_patch_size(temporal_patch_size), - in_channels(in_channels), - embed_dim(embed_dim) { - blocks["proj"] = std::make_shared(in_channels, - embed_dim, - std::tuple{temporal_patch_size, patch_size, patch_size}, - std::tuple{temporal_patch_size, patch_size, patch_size}, - std::tuple{0, 0, 0}, - std::tuple{1, 1, 1}, - true); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto proj = std::dynamic_pointer_cast(blocks["proj"]); - x = ggml_reshape_4d(ctx->ggml_ctx, - x, - patch_size, - patch_size, - temporal_patch_size, - ggml_nelements(x) / (temporal_patch_size * patch_size * patch_size)); - x = proj->forward(ctx, x); - x = ggml_reshape_2d(ctx->ggml_ctx, x, embed_dim, ggml_nelements(x) / embed_dim); - return x; - } - }; - - struct VisionPatchMerger : public GGMLBlock { - int64_t hidden_size; - bool use_postshuffle_norm; - - VisionPatchMerger(int64_t dim, - int64_t context_dim, - int spatial_merge_size, - bool use_postshuffle_norm) - : hidden_size(context_dim * spatial_merge_size * spatial_merge_size), - use_postshuffle_norm(use_postshuffle_norm) { - blocks["norm"] = std::make_shared(use_postshuffle_norm ? hidden_size : context_dim, 1e-6f); - blocks["linear_fc1"] = std::make_shared(hidden_size, hidden_size, true); - blocks["linear_fc2"] = std::make_shared(hidden_size, dim, true); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto norm = std::dynamic_pointer_cast(blocks["norm"]); - auto linear_fc1 = std::dynamic_pointer_cast(blocks["linear_fc1"]); - auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); - - x = norm->forward(ctx, x); - x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size); - x = linear_fc1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x); - x = linear_fc2->forward(ctx, x); - return x; - } - }; - - struct VisionAttention : public GGMLBlock { - int head_dim; - int num_heads; - - VisionAttention(int64_t hidden_size, int num_heads) - : num_heads(num_heads) { - head_dim = static_cast(hidden_size / num_heads); - GGML_ASSERT(num_heads * head_dim == hidden_size); - blocks["qkv"] = std::make_shared(hidden_size, hidden_size * 3, true); - blocks["proj"] = std::make_shared(hidden_size, hidden_size, true); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* pe) { - auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); - auto proj = std::dynamic_pointer_cast(blocks["proj"]); - - auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); - - auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]); - auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]); - auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, qkv_vec[2]->ne[1], qkv_vec[2]->ne[2]); - - x = Rope::attention(ctx, q, k, v, pe, nullptr, 1.f, false); - x = proj->forward(ctx, x); - return x; - } - }; - - struct VisionBlock : public GGMLBlock { - VisionBlock(int64_t hidden_size, - int64_t intermediate_size, - int num_heads) { - blocks["norm1"] = std::make_shared(hidden_size, 1e-6f); - blocks["norm2"] = std::make_shared(hidden_size, 1e-6f); - blocks["attn"] = std::make_shared(hidden_size, num_heads); - blocks["mlp"] = std::make_shared(hidden_size, intermediate_size); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* pe) { - auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); - auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); - auto attn = std::dynamic_pointer_cast(blocks["attn"]); - auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); - - auto residual = x; - x = norm1->forward(ctx, x); - x = attn->forward(ctx, x, pe); - x = ggml_add_inplace(ctx->ggml_ctx, x, residual); - - residual = x; - x = norm2->forward(ctx, x); - x = mlp->forward(ctx, x); - x = ggml_add_inplace(ctx->ggml_ctx, x, residual); - return x; - } - }; - - struct VisionOutput { - ggml_tensor* hidden_states = nullptr; - std::vector deepstack_hidden_states; - }; - - struct VisionModel : public GGMLBlock { - int num_layers; - int spatial_merge_size; - int num_grid_per_side; - std::vector deepstack_visual_indexes; - - VisionModel(int num_layers, - int64_t in_channels, - int64_t hidden_size, - int64_t out_hidden_size, - int64_t intermediate_size, - int num_heads, - int spatial_merge_size, - int patch_size, - int temporal_patch_size, - int num_position_embeddings, - std::vector deepstack_visual_indexes) - : num_layers(num_layers), - spatial_merge_size(spatial_merge_size), - num_grid_per_side(static_cast(std::sqrt(num_position_embeddings))), - deepstack_visual_indexes(std::move(deepstack_visual_indexes)) { - blocks["patch_embed"] = std::make_shared(patch_size, - temporal_patch_size, - in_channels, - hidden_size); - blocks["pos_embed"] = std::make_shared(num_position_embeddings, hidden_size); - for (int i = 0; i < num_layers; ++i) { - blocks["blocks." + std::to_string(i)] = std::make_shared(hidden_size, - intermediate_size, - num_heads); - } - blocks["merger"] = std::make_shared(out_hidden_size, - hidden_size, - spatial_merge_size, - false); - for (int i = 0; i < static_cast(this->deepstack_visual_indexes.size()); ++i) { - blocks["deepstack_merger_list." + std::to_string(i)] = std::make_shared(out_hidden_size, - hidden_size, - spatial_merge_size, - true); - } - } - - ggml_tensor* fast_pos_embed_interpolate(GGMLRunnerContext* ctx, - int grid_h, - int grid_w) { - auto pos_embed = std::dynamic_pointer_cast(blocks["pos_embed"]); - std::vector idx_list[4]; - std::vector weight_list[4]; - idx_list[0].reserve(static_cast(grid_h * grid_w)); - idx_list[1].reserve(static_cast(grid_h * grid_w)); - idx_list[2].reserve(static_cast(grid_h * grid_w)); - idx_list[3].reserve(static_cast(grid_h * grid_w)); - weight_list[0].reserve(static_cast(grid_h * grid_w)); - weight_list[1].reserve(static_cast(grid_h * grid_w)); - weight_list[2].reserve(static_cast(grid_h * grid_w)); - weight_list[3].reserve(static_cast(grid_h * grid_w)); - - double max_index = static_cast(num_grid_per_side - 1); - for (int h = 0; h < grid_h; ++h) { - double h_pos = grid_h == 1 ? 0.0 : max_index * h / static_cast(grid_h - 1); - int h_floor = static_cast(std::floor(h_pos)); - int h_ceil = std::min(h_floor + 1, num_grid_per_side - 1); - double dh = h_pos - h_floor; - for (int w = 0; w < grid_w; ++w) { - double w_pos = grid_w == 1 ? 0.0 : max_index * w / static_cast(grid_w - 1); - int w_floor = static_cast(std::floor(w_pos)); - int w_ceil = std::min(w_floor + 1, num_grid_per_side - 1); - double dw = w_pos - w_floor; - - idx_list[0].push_back(h_floor * num_grid_per_side + w_floor); - idx_list[1].push_back(h_floor * num_grid_per_side + w_ceil); - idx_list[2].push_back(h_ceil * num_grid_per_side + w_floor); - idx_list[3].push_back(h_ceil * num_grid_per_side + w_ceil); - - weight_list[0].push_back(static_cast((1.0 - dh) * (1.0 - dw))); - weight_list[1].push_back(static_cast((1.0 - dh) * dw)); - weight_list[2].push_back(static_cast(dh * (1.0 - dw))); - weight_list[3].push_back(static_cast(dh * dw)); - } - } - - ggml_tensor* patch_pos_embeds = nullptr; - for (int i = 0; i < 4; ++i) { - auto idx_tensor = ggml_new_tensor_1d(ctx->ggml_ctx, GGML_TYPE_I32, static_cast(idx_list[i].size())); - std::memcpy(idx_tensor->data, idx_list[i].data(), idx_list[i].size() * sizeof(int32_t)); - auto embed = pos_embed->forward(ctx, idx_tensor); - auto weight_tensor = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, 1, static_cast(weight_list[i].size())); - std::memcpy(weight_tensor->data, weight_list[i].data(), weight_list[i].size() * sizeof(float)); - embed = ggml_mul(ctx->ggml_ctx, embed, weight_tensor); - patch_pos_embeds = patch_pos_embeds == nullptr ? embed : ggml_add(ctx->ggml_ctx, patch_pos_embeds, embed); - } - - patch_pos_embeds = ggml_reshape_4d(ctx->ggml_ctx, - patch_pos_embeds, - patch_pos_embeds->ne[0], - spatial_merge_size, - grid_w / spatial_merge_size, - grid_h * spatial_merge_size); - patch_pos_embeds = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, patch_pos_embeds, 0, 1, 3, 2)); - patch_pos_embeds = ggml_reshape_2d(ctx->ggml_ctx, - patch_pos_embeds, - patch_pos_embeds->ne[0], - ggml_nelements(patch_pos_embeds) / patch_pos_embeds->ne[0]); - return patch_pos_embeds; - } - - VisionOutput forward(GGMLRunnerContext* ctx, - ggml_tensor* pixel_values, - ggml_tensor* pe, - int grid_h, - int grid_w) { - auto patch_embed = std::dynamic_pointer_cast(blocks["patch_embed"]); - auto merger = std::dynamic_pointer_cast(blocks["merger"]); - - auto x = patch_embed->forward(ctx, pixel_values); - auto pos_embeds = fast_pos_embed_interpolate(ctx, grid_h, grid_w); - x = ggml_add(ctx->ggml_ctx, x, pos_embeds); - x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], 1); - - VisionOutput out; - for (int i = 0; i < num_layers; ++i) { - auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); - x = block->forward(ctx, x, pe); - for (int j = 0; j < static_cast(deepstack_visual_indexes.size()); ++j) { - if (deepstack_visual_indexes[j] == i) { - auto deepstack_merger = std::dynamic_pointer_cast(blocks["deepstack_merger_list." + std::to_string(j)]); - out.deepstack_hidden_states.push_back(deepstack_merger->forward(ctx, x)); - break; - } - } - } - - out.hidden_states = merger->forward(ctx, x); - return out; - } - }; + static inline HiDreamO1Params make_hidream_o1_params() { + HiDreamO1Params params; + params.llm.arch = LLM::LLMArch::QWEN3_VL; + params.llm.hidden_size = 4096; + params.llm.intermediate_size = 12288; + params.llm.num_layers = 36; + params.llm.num_heads = 32; + params.llm.num_kv_heads = 8; + params.llm.head_dim = 128; + params.llm.qkv_bias = false; + params.llm.qk_norm = true; + params.llm.vocab_size = 151936; + params.llm.rms_norm_eps = 1e-6f; + params.llm.vision.arch = LLM::LLMVisionArch::QWEN3_VL; + params.llm.vision.num_layers = 27; + params.llm.vision.hidden_size = 1152; + params.llm.vision.intermediate_size = 4304; + params.llm.vision.num_heads = 16; + params.llm.vision.out_hidden_size = 4096; + params.llm.vision.patch_size = 16; + params.llm.vision.spatial_merge_size = 2; + params.llm.vision.temporal_patch_size = 2; + params.llm.vision.num_position_embeddings = 2304; + return params; + } struct HiDreamO1Model : public GGMLBlock { HiDreamO1Params params; @@ -502,17 +242,6 @@ namespace HiDreamO1 { explicit HiDreamO1Model(HiDreamO1Params params) : params(std::move(params)) { blocks["language_model"] = std::make_shared(this->params.llm); - blocks["visual"] = std::make_shared(this->params.llm.vision.num_layers, - this->params.llm.vision.in_channels, - this->params.llm.vision.hidden_size, - this->params.llm.vision.out_hidden_size, - this->params.llm.vision.intermediate_size, - this->params.llm.vision.num_heads, - this->params.llm.vision.spatial_merge_size, - this->params.llm.vision.patch_size, - this->params.llm.vision.temporal_patch_size, - this->params.num_position_embeddings, - this->params.deepstack_visual_indexes); blocks["t_embedder1"] = std::make_shared(this->params.llm.hidden_size); blocks["x_embedder"] = std::make_shared(this->params.patch_size * this->params.patch_size * 3, this->params.llm.hidden_size / 4, @@ -525,10 +254,6 @@ namespace HiDreamO1 { return std::dynamic_pointer_cast(blocks["language_model"]); } - std::shared_ptr vision_model() { - return std::dynamic_pointer_cast(blocks["visual"]); - } - std::shared_ptr timestep_embedder() { return std::dynamic_pointer_cast(blocks["t_embedder1"]); } @@ -542,43 +267,80 @@ namespace HiDreamO1 { } }; - struct HiDreamO1Runner : public GGMLRunner { + struct HiDreamO1VisionRunner : public GGMLRunner { HiDreamO1Params params; - HiDreamO1Model model; + std::shared_ptr model; std::vector window_index_vec; std::vector window_inverse_index_vec; std::vector window_mask_vec; std::vector pe_vec; + std::array, 4> pos_embed_idx_data_; + std::array, 4> pos_embed_weight_data_; + + HiDreamO1VisionRunner(ggml_backend_t backend, + bool offload_params_to_cpu, + const String2TensorStorage& tensor_storage_map = {}, + const std::string& prefix = "model.visual") + : GGMLRunner(backend, offload_params_to_cpu), + params(make_hidream_o1_params()), + model(std::make_shared(false, params.llm.vision)) { + model->init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "hidream_o1_vision"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix = "model.visual") { + model->get_param_tensors(tensors, prefix); + } + + ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) { + return LLM::LLMRunner::encode_image_common(this, + compute_ctx, + runner_ctx, + image, + params.llm.vision, + model, + window_index_vec, + window_inverse_index_vec, + window_mask_vec, + pe_vec, + pos_embed_idx_data_, + pos_embed_weight_data_); + } + + ggml_cgraph* build_graph(const sd::Tensor& image_tensor) { + ggml_cgraph* gf = new_graph_custom(HIDREAM_O1_GRAPH_SIZE); + ggml_tensor* image = make_input(image_tensor); + auto runner_ctx = get_context(); + auto image_embeds = encode_image(&runner_ctx, image); + ggml_build_forward_expand(gf, image_embeds); + return gf; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& image) { + auto get_graph = [&]() { + return build_graph(image); + }; + auto output = GGMLRunner::compute(get_graph, n_threads, false); + return output.has_value() ? std::move(output.value()) : sd::Tensor(); + } + }; + + struct HiDreamO1Runner : public GGMLRunner { + HiDreamO1Params params; + HiDreamO1Model model; + std::vector attention_mask_vec; HiDreamO1Runner(ggml_backend_t backend, bool offload_params_to_cpu, const String2TensorStorage& tensor_storage_map = {}, const std::string& prefix = "model") - : GGMLRunner(backend, offload_params_to_cpu) { - params.llm.arch = LLM::LLMArch::QWEN3_VL; - params.llm.hidden_size = 4096; - params.llm.intermediate_size = 12288; - params.llm.num_layers = 36; - params.llm.num_heads = 32; - params.llm.num_kv_heads = 8; - params.llm.head_dim = 128; - params.llm.qkv_bias = false; - params.llm.qk_norm = true; - params.llm.vocab_size = 151936; - params.llm.rms_norm_eps = 1e-6f; - params.llm.vision.num_layers = 27; - params.llm.vision.hidden_size = 1152; - params.llm.vision.intermediate_size = 4304; - params.llm.vision.num_heads = 16; - params.llm.vision.out_hidden_size = 4096; - params.llm.vision.patch_size = 16; - params.llm.vision.spatial_merge_size = 2; - params.llm.vision.temporal_patch_size = 2; - params.num_position_embeddings = 2304; - params.deepstack_visual_indexes = {8, 16, 24}; - + : GGMLRunner(backend, offload_params_to_cpu), + params(make_hidream_o1_params()) { model = HiDreamO1Model(params); model.init(params_ctx, tensor_storage_map, prefix); } @@ -591,105 +353,13 @@ namespace HiDreamO1 { model.get_param_tensors(tensors, prefix); } - ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) { - int64_t C = image->ne[2]; - int64_t H = image->ne[1]; - int64_t W = image->ne[0]; - int64_t mh = params.llm.vision.spatial_merge_size; - int64_t mw = params.llm.vision.spatial_merge_size; - int64_t pt = params.llm.vision.temporal_patch_size; - int64_t ph = params.llm.vision.patch_size; - int64_t pw = params.llm.vision.patch_size; - - image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C); - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); - image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw); - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); - image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph)); - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); - image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph)); - image = ggml_concat(ctx, image, image, 0); - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); - image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph)); - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); - image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph)); - return image; - } - - ggml_tensor* concat_seq(GGMLRunnerContext* ctx, ggml_tensor* a, ggml_tensor* b) { - if (a == nullptr) { - return b; - } - if (b == nullptr) { - return a; - } - return ggml_concat(ctx->ggml_ctx, a, b, 1); - } - - ggml_tensor* scatter_visual_embeds(GGMLRunnerContext* ctx, - ggml_tensor* inputs_embeds, - const sd::Tensor& image_embed_ranges_tensor, - ggml_tensor* visual_embeds) { - if (visual_embeds == nullptr || image_embed_ranges_tensor.empty()) { - return inputs_embeds; - } - - ggml_tensor* output = nullptr; - int prev_end = 0; - int n_ranges = static_cast(image_embed_ranges_tensor.shape()[1]); - int visual_offset = 0; - for (int i = 0; i < n_ranges; ++i) { - int start = image_embed_ranges_tensor.values()[i * 2]; - int len = image_embed_ranges_tensor.values()[i * 2 + 1]; - - if (start > prev_end) { - output = concat_seq(ctx, output, ggml_ext_slice(ctx->ggml_ctx, inputs_embeds, 1, prev_end, start)); - } - - output = concat_seq(ctx, - output, - ggml_ext_slice(ctx->ggml_ctx, visual_embeds, 1, visual_offset, visual_offset + len)); - prev_end = start + len; - visual_offset += len; - } - - if (prev_end < inputs_embeds->ne[1]) { - output = concat_seq(ctx, output, ggml_ext_slice(ctx->ggml_ctx, inputs_embeds, 1, prev_end, inputs_embeds->ne[1])); - } - return output == nullptr ? inputs_embeds : output; - } - - VisionOutput encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) { - auto vision = model.vision_model(); - GGML_ASSERT(image->ne[1] % (params.llm.vision.patch_size * params.llm.vision.spatial_merge_size) == 0); - GGML_ASSERT(image->ne[0] % (params.llm.vision.patch_size * params.llm.vision.spatial_merge_size) == 0); - - int grid_h = static_cast(image->ne[1]) / params.llm.vision.patch_size; - int grid_w = static_cast(image->ne[0]) / params.llm.vision.patch_size; - - auto pixel_values = process_image(compute_ctx, image); - - int head_dim = static_cast(params.llm.vision.hidden_size / params.llm.vision.num_heads); - std::vector window_index_vec(static_cast((grid_h / params.llm.vision.spatial_merge_size) * (grid_w / params.llm.vision.spatial_merge_size))); - for (int i = 0; i < static_cast(window_index_vec.size()); ++i) { - window_index_vec[static_cast(i)] = i; - } - pe_vec = Rope::gen_qwen2vl_pe(grid_h, grid_w, params.llm.vision.spatial_merge_size, window_index_vec, 10000, {head_dim / 2, head_dim / 2}); - int pos_len = static_cast(pe_vec.size() / head_dim / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); - - return vision->forward(runner_ctx, pixel_values, pe, grid_h, grid_w); - } - ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timestep_tensor, const sd::Tensor& input_ids_tensor, const sd::Tensor& input_pos_tensor, const sd::Tensor& token_types_tensor, - const sd::Tensor& image_embed_ranges_tensor, const sd::Tensor& vinput_mask_tensor, - const std::vector>& vlm_images, + const std::vector>>& image_embeds_tensor, const std::vector>& ref_images) { ggml_cgraph* gf = new_graph_custom(HIDREAM_O1_GRAPH_SIZE); ggml_tensor* x = make_input(x_tensor); @@ -702,11 +372,6 @@ namespace HiDreamO1 { auto x_embedder = model.patch_embedder(); auto final_layer2 = model.final_layer(); - std::vector vlm_image_tensors; - for (const auto& image : vlm_images) { - vlm_image_tensors.push_back(make_input(image)); - } - std::vector ref_image_tensors; for (const auto& image : ref_images) { ref_image_tensors.push_back(make_input(image)); @@ -725,15 +390,14 @@ namespace HiDreamO1 { auto attention_mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, total_seq_len, total_seq_len); set_backend_tensor_data(attention_mask, attention_mask_vec.data()); - auto runner_ctx = get_context(); - ggml_tensor* visual_embeds = nullptr; - for (size_t i = 0; i < vlm_image_tensors.size(); ++i) { - auto image_output = encode_image(&runner_ctx, vlm_image_tensors[i]); - visual_embeds = visual_embeds == nullptr ? image_output.hidden_states : ggml_concat(compute_ctx, visual_embeds, image_output.hidden_states, 1); + auto runner_ctx = get_context(); + auto txt = text_model->embed(&runner_ctx, input_ids); + std::vector> image_embeds; + image_embeds.reserve(image_embeds_tensor.size()); + for (const auto& image_embed : image_embeds_tensor) { + image_embeds.emplace_back(image_embed.first, make_input(image_embed.second)); } - - auto txt = text_model->embed(&runner_ctx, input_ids); - txt = scatter_visual_embeds(&runner_ctx, txt, image_embed_ranges_tensor, visual_embeds); + txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds); auto t_emb = t_embedder1->forward(&runner_ctx, timestep); int64_t txt_seq_len = input_ids->ne[0]; @@ -765,15 +429,7 @@ namespace HiDreamO1 { } x_pred_start = first_vinput; } - auto x_pred = ggml_view_3d(compute_ctx, - x_pred_all, - x_pred_all->ne[0], - target_tokens, - x_pred_all->ne[2], - x_pred_all->nb[1], - x_pred_all->nb[2], - x_pred_start * x_pred_all->nb[1]); - x_pred = ggml_cont(compute_ctx, x_pred); + auto x_pred = ggml_ext_slice(compute_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens); x_pred = DiT::unpatchify_and_crop(compute_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE); float sigma = 1.0f - timestep_tensor.values()[0]; @@ -790,12 +446,11 @@ namespace HiDreamO1 { const sd::Tensor& input_ids, const sd::Tensor& input_pos, const sd::Tensor& token_types, - const sd::Tensor& image_embed_ranges, const sd::Tensor& vinput_mask, - const std::vector>& vlm_images, + const std::vector>>& image_embeds, const std::vector>& ref_images) { auto get_graph = [&]() { - return build_graph(x, timestep, input_ids, input_pos, token_types, image_embed_ranges, vinput_mask, vlm_images, ref_images); + return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } @@ -803,19 +458,43 @@ namespace HiDreamO1 { struct HiDreamO1Conditioner : public Conditioner { Qwen2Tokenizer tokenizer; + std::shared_ptr vision_runner; + + HiDreamO1Conditioner(ggml_backend_t backend, + bool offload_params_to_cpu, + const String2TensorStorage& tensor_storage_map = {}) + : vision_runner(std::make_shared(backend, offload_params_to_cpu, tensor_storage_map)) {} void get_param_tensors(std::map& tensors) override { - SD_UNUSED(tensors); + vision_runner->get_param_tensors(tensors); } - void alloc_params_buffer() override {} - void free_params_buffer() override {} - size_t get_params_buffer_size() override { return 0; } - void set_flash_attention_enabled(bool enabled) override { SD_UNUSED(enabled); } + void alloc_params_buffer() override { + vision_runner->alloc_params_buffer(); + } + + void free_params_buffer() override { + vision_runner->free_params_buffer(); + } + + size_t get_params_buffer_size() override { + return vision_runner->get_params_buffer_size(); + } + + void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override { + vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes); + } + + void set_flash_attention_enabled(bool enabled) override { + vision_runner->set_flash_attention_enabled(enabled); + } + + void set_weight_adapter(const std::shared_ptr& adapter) override { + vision_runner->set_weight_adapter(adapter); + } SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) override { - SD_UNUSED(n_threads); SDCondition result; int width = conditioner_params.width; @@ -827,12 +506,11 @@ namespace HiDreamO1 { ref_images = *conditioner_params.ref_images; } - std::vector> vlm_images; + std::vector>> vlm_images; std::vector> image_grids; std::vector skip_vision_start; std::string prompt = "<|im_start|>user\n"; - std::vector image_ranges; if (ref_images.empty()) { prompt += conditioner_params.text; @@ -849,12 +527,9 @@ namespace HiDreamO1 { std::vector token_types(input_ids_pad.size(), 0); int txt_seq_len = static_cast(input_ids.size()); int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; - for (int i = bgn; i < bgn + target_image_len + TIMESTEP_TOKEN_NUM; ++i) { + for (int i = bgn; i < static_cast(token_types.size()); ++i) { token_types[i] = 1; } - for (int i = txt_seq_len - TIMESTEP_TOKEN_NUM; i < txt_seq_len; ++i) { - token_types[i] = 3; - } auto position_ids = build_position_ids(input_ids_pad, image_grids, skip_vision_start); @@ -867,11 +542,10 @@ namespace HiDreamO1 { } std::vector vinput_mask_shape{static_cast(vinput_mask.size())}; - result.c_input_ids = sd::Tensor(input_shape, std::move(input_ids)); - result.c_position_ids = sd::Tensor(position_shape, position_ids); - result.c_token_types = sd::Tensor(token_type_shape, std::move(token_types)); - result.c_vinput_mask = sd::Tensor(vinput_mask_shape, std::move(vinput_mask)); - result.c_image_embed_ranges = sd::Tensor(); + result.c_input_ids = sd::Tensor(input_shape, std::move(input_ids)); + result.c_position_ids = sd::Tensor(position_shape, position_ids); + result.c_token_types = sd::Tensor(token_type_shape, std::move(token_types)); + result.c_vinput_mask = sd::Tensor(vinput_mask_shape, std::move(vinput_mask)); return result; } @@ -899,21 +573,27 @@ namespace HiDreamO1 { } for (const auto& ref_image : ref_images) { - auto patch_img = resize_to_area(ref_image, max_size); - patch_img = sd::ops::clamp(patch_img, 0.0f, 1.0f); - patch_img = patch_img * 2.0f - 1.0f; - result.c_ref_images.push_back(std::move(patch_img)); + auto resized_ref = resize_to_area(ref_image, max_size); + resized_ref = sd::ops::clamp(resized_ref, 0.0f, 1.0f); - auto dims = calculate_dimensions(cond_img_size, static_cast(ref_image.shape()[0]) / static_cast(ref_image.shape()[1])); - auto vlm_image = clip_preprocess(ref_image, dims.first, dims.second); + // VLM image: Qwen3-VL expects mean=[0.5]/std=[0.5] (i.e. range [-1,1]), + // not CLIP normalization. Resize the already-resized ref directly to + // (cond_w, cond_h) to match the Python pipeline's pil_r.resize(). + auto dims = calculate_dimensions(cond_img_size, + static_cast(resized_ref.shape()[0]) / static_cast(resized_ref.shape()[1])); + sd::Tensor vlm_image = sd::ops::interpolate( + resized_ref, + {dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]}); + vlm_image = vlm_image * 2.0f - 1.0f; int64_t image_tokens = static_cast(dims.first / PATCH_SIZE) * static_cast(dims.second / PATCH_SIZE); + + auto patch_img = resized_ref * 2.0f - 1.0f; + result.c_ref_images.push_back(std::move(patch_img)); int64_t prompt_start = static_cast(tokenizer.encode(prompt + "<|vision_start|>", nullptr).size()); prompt += "<|vision_start|>"; prompt += repeat_special_token("<|image_pad|>", image_tokens); prompt += "<|vision_end|>"; - image_ranges.push_back(static_cast(prompt_start)); - image_ranges.push_back(static_cast(image_tokens)); - result.c_vlm_images.push_back(std::move(vlm_image)); + vlm_images.emplace_back(static_cast(prompt_start), std::move(vlm_image)); image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE}); skip_vision_start.push_back(0); } @@ -928,10 +608,8 @@ namespace HiDreamO1 { image_grids.push_back({1, static_cast(height / PATCH_SIZE), static_cast(width / PATCH_SIZE)}); skip_vision_start.push_back(1); - int64_t total_ref_len = 0; for (const auto& ref_image : result.c_ref_images) { int64_t ref_len = static_cast(ref_image.shape()[0] / PATCH_SIZE) * static_cast(ref_image.shape()[1] / PATCH_SIZE); - total_ref_len += ref_len; input_ids_pad.push_back(VISION_START_TOKEN_ID); input_ids_pad.insert(input_ids_pad.end(), ref_len - 1, IMAGE_TOKEN_ID); image_grids.push_back({1, static_cast(ref_image.shape()[1] / PATCH_SIZE), static_cast(ref_image.shape()[0] / PATCH_SIZE)}); @@ -941,32 +619,32 @@ namespace HiDreamO1 { std::vector token_types(input_ids_pad.size(), 0); int txt_seq_len = static_cast(input_ids.size()); int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; - int end = bgn + static_cast(target_image_len) + TIMESTEP_TOKEN_NUM; - for (int i = bgn; i < end; ++i) { + for (int i = bgn; i < static_cast(token_types.size()); ++i) { token_types[i] = 1; } - for (int i = end; i < end + total_ref_len; ++i) { - token_types[i] = 2; - } - for (int i = txt_seq_len - TIMESTEP_TOKEN_NUM; i < txt_seq_len; ++i) { - token_types[i] = 3; - } std::vector input_shape{static_cast(input_ids.size())}; std::vector position_shape{static_cast(input_ids_pad.size() * 4)}; std::vector token_type_shape{static_cast(token_types.size())}; - std::vector image_range_shape{2, static_cast(image_ranges.size() / 2)}; std::vector vinput_mask(token_types.size(), 0); for (int i = txt_seq_len; i < static_cast(vinput_mask.size()); ++i) { vinput_mask[static_cast(i)] = 1; } std::vector vinput_mask_shape{static_cast(vinput_mask.size())}; - result.c_input_ids = sd::Tensor(input_shape, std::move(input_ids)); - result.c_position_ids = sd::Tensor(position_shape, build_position_ids(input_ids_pad, image_grids, skip_vision_start)); - result.c_token_types = sd::Tensor(token_type_shape, std::move(token_types)); - result.c_image_embed_ranges = sd::Tensor(image_range_shape, std::move(image_ranges)); - result.c_vinput_mask = sd::Tensor(vinput_mask_shape, std::move(vinput_mask)); + result.c_input_ids = sd::Tensor(input_shape, std::move(input_ids)); + result.c_position_ids = sd::Tensor(position_shape, build_position_ids(input_ids_pad, image_grids, skip_vision_start)); + result.c_token_types = sd::Tensor(token_type_shape, std::move(token_types)); + result.c_vinput_mask = sd::Tensor(vinput_mask_shape, std::move(vinput_mask)); + result.c_image_embeds.reserve(vlm_images.size()); + for (const auto& vlm_image : vlm_images) { + auto image_embed = vision_runner->compute(n_threads, vlm_image.second); + if (image_embed.empty()) { + LOG_ERROR("hidream_o1 conditioner: encode VLM image failed"); + return SDCondition(); + } + result.c_image_embeds.emplace_back(vlm_image.first, std::move(image_embed)); + } return result; } }; diff --git a/src/llm.hpp b/src/llm.hpp index cc6bf417..8509b6b7 100644 --- a/src/llm.hpp +++ b/src/llm.hpp @@ -2,7 +2,10 @@ #define __LLM_HPP__ #include +#include +#include #include +#include #include #include #include @@ -41,7 +44,13 @@ namespace LLM { "ministral3.3b", }; + enum class LLMVisionArch { + QWEN2_5_VL, + QWEN3_VL, + }; + struct LLMVisionParams { + LLMVisionArch arch = LLMVisionArch::QWEN2_5_VL; int num_layers = 32; int64_t hidden_size = 1280; int64_t intermediate_size = 3420; @@ -52,6 +61,7 @@ namespace LLM { int patch_size = 14; int spatial_merge_size = 2; int window_size = 112; + int num_position_embeddings = 0; std::set fullatt_block_indexes = {7, 15, 23, 31}; }; @@ -92,6 +102,84 @@ namespace LLM { } }; + static ggml_tensor* splice_image_embeds(GGMLRunnerContext* ctx, + ggml_tensor* x, + const std::vector>& image_embeds) { + if (image_embeds.empty()) { + return x; + } + + GGML_ASSERT(x->ne[2] == 1); // N == 1 + + auto raw_x = ggml_cast(ctx->ggml_ctx, x, image_embeds[0].second->type); + int64_t txt_token_start = 0; + int64_t txt_token_end = 0; + ggml_tensor* input_embed = nullptr; + + for (int i = 0; i < image_embeds.size(); i++) { + if (i == 0) { + txt_token_start = 0; + } else { + txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1]; + } + txt_token_end = image_embeds[i].first; + + auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + if (input_embed == nullptr) { + input_embed = txt_embed; + } else { + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1); + } + + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1); + } + + txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1]; + txt_token_end = raw_x->ne[1]; + + auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); + GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]); + return input_embed; + } + + struct VisionMLP : public GGMLBlock { + protected: + LLMVisionArch arch_; + + public: + VisionMLP(LLMVisionArch arch, int64_t hidden_size, int64_t intermediate_size) + : arch_(arch) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + blocks["linear_fc1"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["linear_fc2"] = std::make_shared(intermediate_size, hidden_size, true); + } else { + blocks["gate_proj"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["up_proj"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["down_proj"] = std::make_shared(intermediate_size, hidden_size, true); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + auto linear_fc1 = std::dynamic_pointer_cast(blocks["linear_fc1"]); + auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); + x = linear_fc1->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x); + x = linear_fc2->forward(ctx, x); + } else { + auto gate_proj = std::dynamic_pointer_cast(blocks["gate_proj"]); + auto up_proj = std::dynamic_pointer_cast(blocks["up_proj"]); + auto down_proj = std::dynamic_pointer_cast(blocks["down_proj"]); + auto h = gate_proj->forward(ctx, x); + h = ggml_silu_inplace(ctx->ggml_ctx, h); + h = ggml_mul_inplace(ctx->ggml_ctx, h, up_proj->forward(ctx, x)); + x = down_proj->forward(ctx, h); + } + return x; + } + }; + struct VisionPatchEmbed : public GGMLBlock { protected: bool llama_cpp_style; @@ -102,6 +190,7 @@ namespace LLM { public: VisionPatchEmbed(bool llama_cpp_style, + LLMVisionArch arch, int patch_size = 14, int temporal_patch_size = 2, int64_t in_channels = 3, @@ -111,36 +200,35 @@ namespace LLM { temporal_patch_size(temporal_patch_size), in_channels(in_channels), embed_dim(embed_dim) { + bool bias = arch == LLMVisionArch::QWEN3_VL; if (llama_cpp_style) { blocks["proj.0"] = std::shared_ptr(new Conv2d(in_channels, embed_dim, {patch_size, patch_size}, - {patch_size, patch_size}, // stride - {0, 0}, // padding - {1, 1}, // dilation - false)); + {patch_size, patch_size}, + {0, 0}, + {1, 1}, + bias)); blocks["proj.1"] = std::shared_ptr(new Conv2d(in_channels, embed_dim, {patch_size, patch_size}, - {patch_size, patch_size}, // stride - {0, 0}, // padding - {1, 1}, // dilation - false)); + {patch_size, patch_size}, + {0, 0}, + {1, 1}, + bias)); } else { std::tuple kernel_size = {(int)temporal_patch_size, (int)patch_size, (int)patch_size}; blocks["proj"] = std::shared_ptr(new Conv3d(in_channels, embed_dim, kernel_size, - kernel_size, // stride - {0, 0, 0}, // padding - {1, 1, 1}, // dilation - false)); + kernel_size, + {0, 0, 0}, + {1, 1, 1}, + bias)); } } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [N*grid_t*grid_h*grid_w, in_channels, temporal_patch_size*patch_size*patch_size] - // return: [N*grid_t*grid_h*grid_w, embed_dim] x = ggml_reshape_4d(ctx->ggml_ctx, x, patch_size, @@ -172,22 +260,43 @@ namespace LLM { } }; - struct PatchMerger : public GGMLBlock { + struct VisionPatchMerger : public GGMLBlock { protected: + LLMVisionArch arch_; int64_t hidden_size; public: - PatchMerger(int64_t dim, - int64_t context_dim, - int64_t spatial_merge_size) { - hidden_size = context_dim * spatial_merge_size * spatial_merge_size; - blocks["ln_q"] = std::shared_ptr(new RMSNorm(context_dim, 1e-6f)); - blocks["mlp.0"] = std::shared_ptr(new Linear(hidden_size, hidden_size)); - // mlp.1 is nn.GELU() - blocks["mlp.2"] = std::shared_ptr(new Linear(hidden_size, dim)); + VisionPatchMerger(LLMVisionArch arch, + int64_t dim, + int64_t context_dim, + int64_t spatial_merge_size) + : arch_(arch), + hidden_size(context_dim * spatial_merge_size * spatial_merge_size) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + blocks["norm"] = std::make_shared(context_dim, 1e-6f); + blocks["linear_fc1"] = std::make_shared(hidden_size, hidden_size, true); + blocks["linear_fc2"] = std::make_shared(hidden_size, dim, true); + } else { + blocks["ln_q"] = std::make_shared(context_dim, 1e-6f); + blocks["mlp.0"] = std::make_shared(hidden_size, hidden_size); + blocks["mlp.2"] = std::make_shared(hidden_size, dim); + } } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto linear_fc1 = std::dynamic_pointer_cast(blocks["linear_fc1"]); + auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); + + x = norm->forward(ctx, x); + x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size); + x = linear_fc1->forward(ctx, x); + x = ggml_gelu_erf(ctx->ggml_ctx, x); + x = linear_fc2->forward(ctx, x); + return x; + } + auto ln_q = std::dynamic_pointer_cast(blocks["ln_q"]); auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); @@ -262,16 +371,35 @@ namespace LLM { }; struct VisionBlock : public GGMLBlock { + protected: + LLMVisionArch arch_; + + ggml_tensor* forward_norm(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + auto norm = std::dynamic_pointer_cast(blocks[name]); + return norm->forward(ctx, x); + } + auto norm = std::dynamic_pointer_cast(blocks[name]); + return norm->forward(ctx, x); + } + public: VisionBlock(bool llama_cpp_style, + LLMVisionArch arch, int64_t hidden_size, int64_t intermediate_size, int num_heads, - float eps = 1e-6f) { - blocks["attn"] = std::shared_ptr(new VisionAttention(llama_cpp_style, hidden_size, num_heads)); - blocks["mlp"] = std::shared_ptr(new MLP(hidden_size, intermediate_size, true)); - blocks["norm1"] = std::shared_ptr(new RMSNorm(hidden_size, eps)); - blocks["norm2"] = std::shared_ptr(new RMSNorm(hidden_size, eps)); + float eps = 1e-6f) + : arch_(arch) { + blocks["attn"] = std::shared_ptr(new VisionAttention(llama_cpp_style, hidden_size, num_heads)); + blocks["mlp"] = std::shared_ptr(new VisionMLP(arch_, hidden_size, intermediate_size)); + if (arch_ == LLMVisionArch::QWEN3_VL) { + blocks["norm1"] = std::shared_ptr(new LayerNorm(hidden_size, eps)); + blocks["norm2"] = std::shared_ptr(new LayerNorm(hidden_size, eps)); + } else { + blocks["norm1"] = std::shared_ptr(new RMSNorm(hidden_size, eps)); + blocks["norm2"] = std::shared_ptr(new RMSNorm(hidden_size, eps)); + } } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -279,18 +407,16 @@ namespace LLM { ggml_tensor* pe, ggml_tensor* mask = nullptr) { // x: [N, n_token, hidden_size] - auto attn = std::dynamic_pointer_cast(blocks["attn"]); - auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); - auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); - auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); auto residual = x; - x = norm1->forward(ctx, x); + x = forward_norm(ctx, "norm1", x); x = attn->forward(ctx, x, pe, mask); x = ggml_add_inplace(ctx->ggml_ctx, x, residual); residual = x; - x = norm2->forward(ctx, x); + x = forward_norm(ctx, "norm2", x); x = mlp->forward(ctx, x); x = ggml_add_inplace(ctx->ggml_ctx, x, residual); @@ -300,38 +426,58 @@ namespace LLM { struct VisionModel : public GGMLBlock { protected: + LLMVisionArch arch_; int num_layers; int spatial_merge_size; + int num_grid_per_side; std::set fullatt_block_indexes; public: VisionModel(bool llama_cpp_style, - int num_layers, - int64_t in_channels, - int64_t hidden_size, - int64_t out_hidden_size, - int64_t intermediate_size, - int num_heads, - int spatial_merge_size, - int patch_size, - int temporal_patch_size, - int window_size, - std::set fullatt_block_indexes = {7, 15, 23, 31}, - float eps = 1e-6f) - : num_layers(num_layers), fullatt_block_indexes(std::move(fullatt_block_indexes)), spatial_merge_size(spatial_merge_size) { + const LLMVisionParams& vision_params, + float eps = 1e-6f) + : arch_(vision_params.arch), + num_layers(vision_params.num_layers), + spatial_merge_size(vision_params.spatial_merge_size), + num_grid_per_side(vision_params.num_position_embeddings > 0 ? static_cast(std::sqrt(vision_params.num_position_embeddings)) : 0), + fullatt_block_indexes(vision_params.fullatt_block_indexes) { blocks["patch_embed"] = std::shared_ptr(new VisionPatchEmbed(llama_cpp_style, - patch_size, - temporal_patch_size, - in_channels, - hidden_size)); + arch_, + vision_params.patch_size, + vision_params.temporal_patch_size, + vision_params.in_channels, + vision_params.hidden_size)); + if (vision_params.num_position_embeddings > 0) { + blocks["pos_embed"] = std::make_shared(vision_params.num_position_embeddings, vision_params.hidden_size); + } for (int i = 0; i < num_layers; i++) { blocks["blocks." + std::to_string(i)] = std::shared_ptr(new VisionBlock(llama_cpp_style, - hidden_size, - intermediate_size, - num_heads, + arch_, + vision_params.hidden_size, + vision_params.intermediate_size, + vision_params.num_heads, eps)); } - blocks["merger"] = std::shared_ptr(new PatchMerger(out_hidden_size, hidden_size, spatial_merge_size)); + blocks["merger"] = std::shared_ptr(new VisionPatchMerger(arch_, + vision_params.out_hidden_size, + vision_params.hidden_size, + spatial_merge_size)); + } + + std::shared_ptr pos_embedder() { + auto it = blocks.find("pos_embed"); + if (it == blocks.end()) { + return nullptr; + } + return std::dynamic_pointer_cast(it->second); + } + + int get_num_grid_per_side() const { + return num_grid_per_side; + } + + int get_spatial_merge_size() const { + return spatial_merge_size; } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -339,20 +485,26 @@ namespace LLM { ggml_tensor* pe, ggml_tensor* window_index, ggml_tensor* window_inverse_index, - ggml_tensor* window_mask) { + ggml_tensor* window_mask, + ggml_tensor* pos_embeds = nullptr) { // pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw] // window_index: [grid_t*(H/mh/ph)*(W/mw/pw)] // window_inverse_index: [grid_t*(H/mh/ph)*(W/mw/pw)] // window_mask: [grid_h*grid_w, grid_h*grid_w] auto patch_embed = std::dynamic_pointer_cast(blocks["patch_embed"]); - auto merger = std::dynamic_pointer_cast(blocks["merger"]); + auto merger = std::dynamic_pointer_cast(blocks["merger"]); auto x = patch_embed->forward(ctx, pixel_values); sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.prelude", "x"); + if (pos_embeds != nullptr) { + x = ggml_add(ctx->ggml_ctx, x, pos_embeds); + } - x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]); - x = ggml_get_rows(ctx->ggml_ctx, x, window_index); - x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] / spatial_merge_size / spatial_merge_size, x->ne[1] * spatial_merge_size * spatial_merge_size, x->ne[2], x->ne[3]); + if (window_index != nullptr) { + x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]); + x = ggml_get_rows(ctx->ggml_ctx, x, window_index); + x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] / spatial_merge_size / spatial_merge_size, x->ne[1] * spatial_merge_size * spatial_merge_size, x->ne[2], x->ne[3]); + } for (int i = 0; i < num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); @@ -362,13 +514,17 @@ namespace LLM { mask = nullptr; } x = block->forward(ctx, x, pe, mask); + if (i == 0) { + } sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x"); } x = merger->forward(ctx, x); sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.final", "x"); - x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index); + if (window_inverse_index != nullptr) { + x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index); + } return x; } @@ -510,47 +666,6 @@ namespace LLM { return x; } - ggml_tensor* splice_image_embeds(GGMLRunnerContext* ctx, - ggml_tensor* x, - std::vector> image_embeds) { - if (image_embeds.empty()) { - return x; - } - - GGML_ASSERT(x->ne[2] == 1); // N == 1 - - auto raw_x = ggml_cast(ctx->ggml_ctx, x, image_embeds[0].second->type); - int64_t txt_token_start = 0; - int64_t txt_token_end = 0; - ggml_tensor* input_embed = nullptr; - - for (int i = 0; i < image_embeds.size(); i++) { - if (i == 0) { - txt_token_start = 0; - } else { - txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1]; - } - txt_token_end = image_embeds[i].first; - - auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); - if (input_embed == nullptr) { - input_embed = txt_embed; - } else { - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1); - } - - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1); - } - - txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1]; - txt_token_end = raw_x->ne[1]; - - auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); - GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]); - return input_embed; - } - ggml_tensor* forward_embeds(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* input_pos, @@ -593,7 +708,7 @@ namespace LLM { // input_ids: [N, n_token] // return: [N, n_token, hidden_size] auto x = embed(ctx, input_ids); - x = splice_image_embeds(ctx, x, std::move(image_embeds)); + x = splice_image_embeds(ctx, x, image_embeds); return forward_embeds(ctx, x, input_pos, attention_mask, std::move(out_layers)); } }; @@ -608,18 +723,7 @@ namespace LLM { : enable_vision(enable_vision), params(params) { blocks["model"] = std::shared_ptr(new TextModel(params)); if (enable_vision) { - blocks["visual"] = std::shared_ptr(new VisionModel(llama_cpp_style, - params.vision.num_layers, - params.vision.in_channels, - params.vision.hidden_size, - params.vision.out_hidden_size, - params.vision.intermediate_size, - params.vision.num_heads, - params.vision.spatial_merge_size, - params.vision.patch_size, - params.vision.temporal_patch_size, - params.vision.window_size, - params.vision.fullatt_block_indexes)); + blocks["visual"] = std::shared_ptr(new VisionModel(llama_cpp_style, params.vision)); } } @@ -636,15 +740,20 @@ namespace LLM { return x; } + std::shared_ptr vision_model() { + GGML_ASSERT(enable_vision); + return std::dynamic_pointer_cast(blocks["visual"]); + } + ggml_tensor* vision_forward(GGMLRunnerContext* ctx, ggml_tensor* pixel_values, ggml_tensor* pe, ggml_tensor* window_index, ggml_tensor* window_inverse_index, - ggml_tensor* window_mask) { + ggml_tensor* window_mask, + ggml_tensor* pos_embeds = nullptr) { GGML_ASSERT(enable_vision); - auto vision_model = std::dynamic_pointer_cast(blocks["visual"]); - return vision_model->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask); + return vision_model()->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask, pos_embeds); } }; @@ -659,7 +768,215 @@ namespace LLM { std::vector window_index_vec; std::vector window_inverse_index_vec; std::vector pe_vec; + std::array, 4> pos_embed_idx_data_; + std::array, 4> pos_embed_weight_data_; + static ggml_tensor* process_image_common(ggml_context* ctx, + ggml_tensor* image, + const LLMVisionParams& vision_params) { + // image: [C, H, W] + // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1 + int64_t C = image->ne[2]; + int64_t H = image->ne[1]; + int64_t W = image->ne[0]; + int64_t mh = vision_params.spatial_merge_size; + int64_t mw = vision_params.spatial_merge_size; + int64_t pt = vision_params.temporal_patch_size; + int64_t ph = vision_params.patch_size; + int64_t pw = vision_params.patch_size; + + image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C); // [C*H, (W/mw/pw), mw, pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); // [mw, C*H, (W/mw/pw), pw] + image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw); // [mw, C, H, (W/mw/pw)*pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); // [H, mw, C, (W/mw/pw)*pw] + image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph)); // [(H/mh/ph)*mh, ph, mw*C*(W/mw/pw), pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw] + image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph)); // [(H/mh/ph)*mh*mw, C, (W/mw/pw), ph*pw] + image = ggml_concat(ctx, image, image, 0); // [(H/mh/ph)*mh*mw, C, (W/mw/pw), pt*ph*pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw] + image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph)); // [(H/mh/ph), mh*mw, (W/mw/pw), C*pt*ph*pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw] + image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph)); // [(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw] + return image; + } + + static ggml_tensor* build_patch_pos_embeds_common(GGMLRunner* runner, + ggml_context* compute_ctx, + GGMLRunnerContext* runner_ctx, + std::shared_ptr vision, + int grid_h, + int grid_w, + std::array, 4>& pos_embed_idx_data, + std::array, 4>& pos_embed_weight_data) { + auto pos_embed = vision->pos_embedder(); + GGML_ASSERT(pos_embed != nullptr); + for (int i = 0; i < 4; ++i) { + pos_embed_idx_data[i].clear(); + pos_embed_weight_data[i].clear(); + pos_embed_idx_data[i].reserve(static_cast(grid_h * grid_w)); + pos_embed_weight_data[i].reserve(static_cast(grid_h * grid_w)); + } + + int num_grid_per_side = vision->get_num_grid_per_side(); + double max_index = static_cast(num_grid_per_side - 1); + int merge_size = vision->get_spatial_merge_size(); + GGML_ASSERT(grid_h % merge_size == 0); + GGML_ASSERT(grid_w % merge_size == 0); + for (int bh = 0; bh < grid_h / merge_size; ++bh) { + for (int bw = 0; bw < grid_w / merge_size; ++bw) { + for (int ih = 0; ih < merge_size; ++ih) { + int h = bh * merge_size + ih; + double h_pos = grid_h == 1 ? 0.0 : max_index * h / static_cast(grid_h - 1); + int h_floor = static_cast(std::floor(h_pos)); + int h_ceil = std::min(h_floor + 1, num_grid_per_side - 1); + double dh = h_pos - h_floor; + for (int iw = 0; iw < merge_size; ++iw) { + int w = bw * merge_size + iw; + double w_pos = grid_w == 1 ? 0.0 : max_index * w / static_cast(grid_w - 1); + int w_floor = static_cast(std::floor(w_pos)); + int w_ceil = std::min(w_floor + 1, num_grid_per_side - 1); + double dw = w_pos - w_floor; + + pos_embed_idx_data[0].push_back(h_floor * num_grid_per_side + w_floor); + pos_embed_idx_data[1].push_back(h_floor * num_grid_per_side + w_ceil); + pos_embed_idx_data[2].push_back(h_ceil * num_grid_per_side + w_floor); + pos_embed_idx_data[3].push_back(h_ceil * num_grid_per_side + w_ceil); + + pos_embed_weight_data[0].push_back(static_cast((1.0 - dh) * (1.0 - dw))); + pos_embed_weight_data[1].push_back(static_cast((1.0 - dh) * dw)); + pos_embed_weight_data[2].push_back(static_cast(dh * (1.0 - dw))); + pos_embed_weight_data[3].push_back(static_cast(dh * dw)); + } + } + } + } + + ggml_tensor* patch_pos_embeds = nullptr; + for (int i = 0; i < 4; ++i) { + auto idx_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast(pos_embed_idx_data[i].size())); + runner->set_backend_tensor_data(idx_tensor, pos_embed_idx_data[i].data()); + auto embed = pos_embed->forward(runner_ctx, idx_tensor); + auto weight_tensor = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, 1, static_cast(pos_embed_weight_data[i].size())); + runner->set_backend_tensor_data(weight_tensor, pos_embed_weight_data[i].data()); + embed = ggml_mul(compute_ctx, embed, weight_tensor); + patch_pos_embeds = patch_pos_embeds == nullptr ? embed : ggml_add(compute_ctx, patch_pos_embeds, embed); + } + return patch_pos_embeds; + } + + static ggml_tensor* encode_image_common(GGMLRunner* runner, + ggml_context* compute_ctx, + GGMLRunnerContext* runner_ctx, + ggml_tensor* image, + const LLMVisionParams& vision_params, + std::shared_ptr vision_model, + std::vector& window_index_vec, + std::vector& window_inverse_index_vec, + std::vector& window_mask_vec, + std::vector& pe_vec, + std::array, 4>& pos_embed_idx_data, + std::array, 4>& pos_embed_weight_data) { + GGML_ASSERT(image->ne[1] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0); + GGML_ASSERT(image->ne[0] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0); + + int grid_h = static_cast(image->ne[1]) / vision_params.patch_size; + int grid_w = static_cast(image->ne[0]) / vision_params.patch_size; + + auto pixel_values = process_image_common(compute_ctx, image, vision_params); + int head_dim = static_cast(vision_params.hidden_size / vision_params.num_heads); + + if (vision_params.arch == LLMVisionArch::QWEN3_VL) { + auto pos_embeds = build_patch_pos_embeds_common(runner, + compute_ctx, + runner_ctx, + vision_model, + grid_h, + grid_w, + pos_embed_idx_data, + pos_embed_weight_data); + window_index_vec.resize(static_cast((grid_h / vision_params.spatial_merge_size) * (grid_w / vision_params.spatial_merge_size))); + for (int i = 0; i < static_cast(window_index_vec.size()); ++i) { + window_index_vec[static_cast(i)] = i; + } + pe_vec = Rope::gen_qwen2vl_pe(grid_h, + grid_w, + vision_params.spatial_merge_size, + window_index_vec, + 10000, + {head_dim / 2, head_dim / 2}); + int pos_len = static_cast(pe_vec.size() / head_dim / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); + runner->set_backend_tensor_data(pe, pe_vec.data()); + return vision_model->forward(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds); + } + + int llm_grid_h = grid_h / vision_params.spatial_merge_size; + int llm_grid_w = grid_w / vision_params.spatial_merge_size; + int vit_merger_window_size = vision_params.window_size / vision_params.patch_size / vision_params.spatial_merge_size; + + int inverse_index = 0; + window_index_vec.resize(llm_grid_h * llm_grid_w); + window_inverse_index_vec.resize(llm_grid_h * llm_grid_w); + std::vector seqlens; + for (int ih = 0; ih < llm_grid_h; ih += vit_merger_window_size) { + for (int iw = 0; iw < llm_grid_w; iw += vit_merger_window_size) { + int win_h = std::min(vit_merger_window_size, llm_grid_h - ih); + int win_w = std::min(vit_merger_window_size, llm_grid_w - iw); + for (int iy = 0; iy < win_h; iy++) { + for (int ix = 0; ix < win_w; ix++) { + int index = (ih + iy) * llm_grid_w + iw + ix; + window_index_vec[inverse_index] = index; + window_inverse_index_vec[index] = inverse_index; + inverse_index++; + } + } + seqlens.push_back(win_h * win_w * vision_params.spatial_merge_size * vision_params.spatial_merge_size); + } + } + auto window_index = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, llm_grid_h * llm_grid_w); + auto window_inverse_index = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, llm_grid_h * llm_grid_w); + runner->set_backend_tensor_data(window_index, window_index_vec.data()); + runner->set_backend_tensor_data(window_inverse_index, window_inverse_index_vec.data()); + + window_mask_vec.resize((grid_h * grid_w) * (grid_h * grid_w)); + int window_start_index = 0; + for (int seq_index = 0; seq_index < seqlens.size(); seq_index++) { + int window_end_index = window_start_index + seqlens[seq_index]; + GGML_ASSERT(window_end_index <= grid_h * grid_w); + for (int i = window_start_index; i < window_end_index; i++) { + for (int j = 0; j < grid_h * grid_w; j++) { + float mask_value = -INFINITY; + if (j >= window_start_index && j < window_end_index) { + mask_value = 0; + } + GGML_ASSERT((i * (grid_h * grid_w) + j) < window_mask_vec.size()); + window_mask_vec[i * (grid_h * grid_w) + j] = mask_value; + } + } + window_start_index = window_end_index; + } + + auto window_mask = ggml_new_tensor_2d(compute_ctx, + GGML_TYPE_F32, + grid_h * grid_w, + grid_h * grid_w); + runner->set_backend_tensor_data(window_mask, window_mask_vec.data()); + + pe_vec = Rope::gen_qwen2vl_pe(grid_h, + grid_w, + vision_params.spatial_merge_size, + window_inverse_index_vec, + 10000, + {head_dim / 2, head_dim / 2}); + int pos_len = static_cast(pe_vec.size() / head_dim / 2); + + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); + runner->set_backend_tensor_data(pe, pe_vec.data()); + + return vision_model->forward(runner_ctx, pixel_values, pe, window_index, window_inverse_index, window_mask); + } + + public: LLMRunner(LLMArch arch, ggml_backend_t backend, bool offload_params_to_cpu, @@ -761,8 +1078,9 @@ namespace LLM { ggml_tensor* input_pos, ggml_tensor* window_index, ggml_tensor* window_inverse_index, - ggml_tensor* window_mask) { - auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask); + ggml_tensor* window_mask, + ggml_tensor* pos_embeds = nullptr) { + auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask, pos_embeds); return hidden_states; } @@ -848,30 +1166,36 @@ namespace LLM { } ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) { - // image: [C, H, W] - // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1 - int64_t C = image->ne[2]; - int64_t H = image->ne[1]; - int64_t W = image->ne[0]; - int64_t mh = params.vision.spatial_merge_size; - int64_t mw = params.vision.spatial_merge_size; - int64_t pt = params.vision.temporal_patch_size; - int64_t ph = params.vision.patch_size; - int64_t pw = params.vision.patch_size; + return process_image_common(ctx, image, params.vision); + } - image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C); // [C*H, (W/mw/pw), mw, pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); // [mw, C*H, (W/mw/pw), pw] - image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw); // [mw, C, H, (W/mw/pw)*pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); // [H, mw, C, (W/mw/pw)*pw] - image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph)); // [(H/mh/ph)*mh, ph, mw*C*(W/mw/pw), pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw] - image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph)); // [(H/mh/ph)*mh*mw, C, (W/mw/pw), ph*pw] - image = ggml_concat(ctx, image, image, 0); // [(H/mh/ph)*mh*mw, C, (W/mw/pw), pt*ph*pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw] - image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph)); // [(H/mh/ph), mh*mw, (W/mw/pw), C*pt*ph*pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw] - image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph)); // [(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw] - return image; + ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx, + std::shared_ptr vision, + int grid_h, + int grid_w) { + return build_patch_pos_embeds_common(this, + compute_ctx, + runner_ctx, + vision, + grid_h, + grid_w, + pos_embed_idx_data_, + pos_embed_weight_data_); + } + + ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) { + return encode_image_common(this, + compute_ctx, + runner_ctx, + image, + params.vision, + model.vision_model(), + window_index_vec, + window_inverse_index_vec, + window_mask_vec, + pe_vec, + pos_embed_idx_data_, + pos_embed_weight_data_); } ggml_cgraph* build_encode_image_graph(const sd::Tensor& image_tensor) { @@ -881,116 +1205,8 @@ namespace LLM { GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0); GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0); - int grid_t = 1; - int grid_h = static_cast(image->ne[1]) / params.vision.patch_size; - int grid_w = static_cast(image->ne[0]) / params.vision.patch_size; - int llm_grid_h = grid_h / params.vision.spatial_merge_size; - int llm_grid_w = grid_w / params.vision.spatial_merge_size; - int vit_merger_window_size = params.vision.window_size / params.vision.patch_size / params.vision.spatial_merge_size; - - auto pixel_values = process_image(compute_ctx, image); - - // window index - int inverse_index = 0; - window_index_vec.resize(llm_grid_h * llm_grid_w); - window_inverse_index_vec.resize(llm_grid_h * llm_grid_w); - std::vector seqlens; - for (int ih = 0; ih < llm_grid_h; ih += vit_merger_window_size) { - for (int iw = 0; iw < llm_grid_w; iw += vit_merger_window_size) { - int win_h = std::min(vit_merger_window_size, llm_grid_h - ih); - int win_w = std::min(vit_merger_window_size, llm_grid_w - iw); - for (int iy = 0; iy < win_h; iy++) { - for (int ix = 0; ix < win_w; ix++) { - int index = (ih + iy) * llm_grid_w + iw + ix; - window_index_vec[inverse_index] = index; - window_inverse_index_vec[index] = inverse_index; - inverse_index++; - } - } - seqlens.push_back(win_h * win_w * params.vision.spatial_merge_size * params.vision.spatial_merge_size); - } - } - // printf("window_index: "); - // for (int i : window_index_vec) { - // printf("%d ", i); - // } - // printf("\n"); - // printf("window_inverse_index: "); - // for (int i : window_inverse_index_vec) { - // printf("%d ", i); - // } - // printf("\n"); - // printf("seqlens: "); - // for (int i : seqlens) { - // printf("%d ", i); - // } - // printf("\n"); - auto window_index = ggml_new_tensor_1d(compute_ctx, - GGML_TYPE_I32, - llm_grid_h * llm_grid_w); - auto window_inverse_index = ggml_new_tensor_1d(compute_ctx, - GGML_TYPE_I32, - llm_grid_h * llm_grid_w); - set_backend_tensor_data(window_index, window_index_vec.data()); - set_backend_tensor_data(window_inverse_index, window_inverse_index_vec.data()); - - // window mask - int seq_window_size = (vit_merger_window_size * params.vision.spatial_merge_size) * (vit_merger_window_size * params.vision.spatial_merge_size); - window_mask_vec.resize((grid_h * grid_w) * (grid_h * grid_w)); - int window_start_index = 0; - for (int seq_index = 0; seq_index < seqlens.size(); seq_index++) { - int window_end_index = window_start_index + seqlens[seq_index]; - // LOG_DEBUG("%d %d", window_start_index, window_end_index); - GGML_ASSERT(window_end_index <= grid_h * grid_w); - for (int i = window_start_index; i < window_end_index; i++) { - for (int j = 0; j < grid_h * grid_w; j++) { - float mask_value = -INFINITY; - if (j >= window_start_index && j < window_end_index) { - mask_value = 0; - } - GGML_ASSERT((i * (grid_h * grid_w) + j) < window_mask_vec.size()); - window_mask_vec[i * (grid_h * grid_w) + j] = mask_value; - } - } - window_start_index = window_end_index; - // printf("\n"); - } - // printf("window_mask: \n"); - // for (int i = 0; i < grid_h*grid_w; i++) { - // for (int j = 0; j < grid_h*grid_w; j++) { - // printf("%f ", window_mask_vec[i * (grid_h * grid_w) + j]); - // } - // printf("\n"); - // } - auto window_mask = ggml_new_tensor_2d(compute_ctx, - GGML_TYPE_F32, - grid_h * grid_w, - grid_h * grid_w); - set_backend_tensor_data(window_mask, window_mask_vec.data()); - - // pe - int head_dim = static_cast(params.vision.hidden_size / params.vision.num_heads); - pe_vec = Rope::gen_qwen2vl_pe(grid_h, - grid_w, - params.vision.spatial_merge_size, - window_inverse_index_vec, - 10000, - {head_dim / 2, head_dim / 2}); - int pos_len = static_cast(pe_vec.size() / head_dim / 2); - // LOG_DEBUG("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); - auto runnter_ctx = get_context(); - ggml_tensor* hidden_states = vision_forward(&runnter_ctx, - pixel_values, - pe, - window_index, - window_inverse_index, - window_mask); + ggml_tensor* hidden_states = encode_image(&runnter_ctx, image); ggml_build_forward_expand(gf, hidden_states); return gf; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index ba2b3ea6..b97d1f47 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -494,7 +494,9 @@ public: version, sd_ctx_params->qwen_image_zero_cond_t); } else if (version == VERSION_HIDREAM_O1) { - cond_stage_model = std::make_shared(); + cond_stage_model = std::make_shared(clip_backend, + offload_params_to_cpu, + tensor_storage_map); diffusion_model = std::make_shared(backend, offload_params_to_cpu, tensor_storage_map, @@ -806,6 +808,7 @@ public: } if (version == VERSION_HIDREAM_O1) { ignore_tensors.insert("lm_head."); + ignore_tensors.insert("model.visual.deepstack_merger_list."); } bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap); if (!success) { @@ -1698,19 +1701,18 @@ public: auto run_condition = [&](const SDCondition& condition, const sd::Tensor* c_concat_override = nullptr, const std::vector* local_skip_layers = nullptr) -> sd::Tensor { - diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; - diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); - diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; - diffusion_params.t5_ids = condition.c_t5_ids.empty() ? nullptr : &condition.c_t5_ids; - diffusion_params.t5_weights = condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights; - diffusion_params.input_ids = condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids; - diffusion_params.input_pos = condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids; - diffusion_params.token_types = condition.c_token_types.empty() ? nullptr : &condition.c_token_types; - diffusion_params.image_embed_ranges = condition.c_image_embed_ranges.empty() ? nullptr : &condition.c_image_embed_ranges; - diffusion_params.vinput_mask = condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask; - diffusion_params.vlm_images = condition.c_vlm_images.empty() ? nullptr : &condition.c_vlm_images; - diffusion_params.ref_latents = condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images; - diffusion_params.skip_layers = local_skip_layers; + diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; + diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); + diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; + diffusion_params.t5_ids = condition.c_t5_ids.empty() ? nullptr : &condition.c_t5_ids; + diffusion_params.t5_weights = condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights; + diffusion_params.input_ids = condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids; + diffusion_params.input_pos = condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids; + diffusion_params.token_types = condition.c_token_types.empty() ? nullptr : &condition.c_token_types; + diffusion_params.vinput_mask = condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask; + diffusion_params.image_embeds = condition.c_image_embeds.empty() ? nullptr : &condition.c_image_embeds; + diffusion_params.ref_latents = condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images; + diffusion_params.skip_layers = local_skip_layers; sd::Tensor cached_output; if (step_cache.before_condition(&condition, noised_input, &cached_output)) {