DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-08-11 06:36:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add hidream o1 image support

Browse Source
This commit is contained in:
leejet 2026-05-11 01:01:13 +08:00
parent 90e87bc846
commit 2f89058f24
11 changed files with 1301 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
src/conditioner.hpp
View File

@ -14,6 +14,13 @@ struct SDCondition {
sd::Tensor<float> c_concat; sd::Tensor<float> c_concat;
sd::Tensor<int32_t> c_t5_ids; sd::Tensor<int32_t> c_t5_ids;
sd::Tensor<float> c_t5_weights; sd::Tensor<float> c_t5_weights;
sd::Tensor<int32_t> c_input_ids;
sd::Tensor<int32_t> c_position_ids;
sd::Tensor<int32_t> c_token_types;
sd::Tensor<int32_t> c_image_embed_ranges;
sd::Tensor<int32_t> c_vinput_mask;
std::vector<sd::Tensor<float>> c_vlm_images;
std::vector<sd::Tensor<float>> c_ref_images;
std::vector<sd::Tensor<float>> extra_c_crossattns; std::vector<sd::Tensor<float>> extra_c_crossattns;
@ -26,10 +33,25 @@ struct SDCondition {
bool empty() const { bool empty() const {
if (!c_crossattn.empty() || !c_vector.empty() || !c_concat.empty() || if (!c_crossattn.empty() || !c_vector.empty() || !c_concat.empty() ||
!c_t5_ids.empty() || !c_t5_weights.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()) {
return false; return false;
} }
for (const auto& tensor : c_vlm_images) {
if (!tensor.empty()) {
return false;
}
}
for (const auto& tensor : c_ref_images) {
if (!tensor.empty()) {
return false;
}
}
for (const auto& tensor : extra_c_crossattns) { for (const auto& tensor : extra_c_crossattns) {
if (!tensor.empty()) { if (!tensor.empty()) {
return false; return false;

4
src/denoiser.hpp
View File

@ -1289,8 +1289,8 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
} }
sd::Tensor<float> denoised = std::move(denoised_opt); sd::Tensor<float> denoised = std::move(denoised_opt);
float sigma_from = sigmas[i]; float sigma_from = sigmas[i];
float sigma_to = sigmas[i + 1]; float sigma_to = sigmas[i + 1];
auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step(sigma_from, sigma_to, eta, is_flow_denoiser); auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step(sigma_from, sigma_to, eta, is_flow_denoiser);

84
src/diffusion_model.hpp
View File

@ -5,6 +5,7 @@
#include "anima.hpp" #include "anima.hpp"
#include "ernie_image.hpp" #include "ernie_image.hpp"
#include "flux.hpp" #include "flux.hpp"
#include "hidream_o1.hpp"
#include "mmdit.hpp" #include "mmdit.hpp"
#include "qwen_image.hpp" #include "qwen_image.hpp"
#include "tensor_ggml.hpp" #include "tensor_ggml.hpp"
@ -22,6 +23,12 @@ struct DiffusionParams {
const sd::Tensor<float>* t5_weights = nullptr; const sd::Tensor<float>* t5_weights = nullptr;
const sd::Tensor<float>* guidance = nullptr; const sd::Tensor<float>* guidance = nullptr;
const std::vector<sd::Tensor<float>>* ref_latents = nullptr; const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
const sd::Tensor<int32_t>* input_ids = nullptr;
const sd::Tensor<int32_t>* input_pos = nullptr;
const sd::Tensor<int32_t>* token_types = nullptr;
const sd::Tensor<int32_t>* image_embed_ranges = nullptr;
const sd::Tensor<int32_t>* vinput_mask = nullptr;
const std::vector<sd::Tensor<float>>* vlm_images = nullptr;
bool increase_ref_index = false; bool increase_ref_index = false;
int num_video_frames = -1; int num_video_frames = -1;
const std::vector<sd::Tensor<float>>* controls = nullptr; const std::vector<sd::Tensor<float>>* controls = nullptr;
@ -476,6 +483,83 @@ struct QwenImageModel : public DiffusionModel {
} }
}; };
struct HiDreamO1Model : public DiffusionModel {
std::string prefix;
HiDreamO1::HiDreamO1Runner hidream_o1;
HiDreamO1Model(ggml_backend_t backend,
bool offload_params_to_cpu,
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "model")
: prefix(prefix), hidream_o1(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
}
std::string get_desc() override {
return hidream_o1.get_desc();
}
void alloc_params_buffer() override {
hidream_o1.alloc_params_buffer();
}
void free_params_buffer() override {
hidream_o1.free_params_buffer();
}
void free_compute_buffer() override {
hidream_o1.free_compute_buffer();
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
hidream_o1.get_param_tensors(tensors, prefix);
}
size_t get_params_buffer_size() override {
return hidream_o1.get_params_buffer_size();
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
hidream_o1.set_weight_adapter(adapter);
}
int64_t get_adm_in_channels() override {
return 0;
}
void set_flash_attention_enabled(bool enabled) {
hidream_o1.set_flash_attention_enabled(enabled);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
hidream_o1.set_max_graph_vram_bytes(max_vram_bytes);
}
void set_circular_axes(bool circular_x, bool circular_y) override {
hidream_o1.set_circular_axes(circular_x, circular_y);
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
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<int32_t> empty_image_embed_ranges;
static const std::vector<sd::Tensor<float>> empty_images;
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,
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_images);
}
};
struct ZImageModel : public DiffusionModel { struct ZImageModel : public DiffusionModel {
std::string prefix; std::string prefix;
ZImage::ZImageRunner z_image; ZImage::ZImageRunner z_image;

95
src/ggml_extend.hpp
View File

@ -1698,13 +1698,41 @@ struct WeightAdapter {
}; };
struct GGMLRunnerContext { struct GGMLRunnerContext {
ggml_backend_t backend = nullptr; ggml_backend_t backend = nullptr;
ggml_context* ggml_ctx = nullptr; ggml_context* ggml_ctx = nullptr;
bool flash_attn_enabled = false; bool flash_attn_enabled = false;
bool conv2d_direct_enabled = false; bool conv2d_direct_enabled = false;
bool circular_x_enabled = false; bool circular_x_enabled = false;
bool circular_y_enabled = false; bool circular_y_enabled = false;
std::shared_ptr<WeightAdapter> weight_adapter = nullptr; std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
std::function<void(const std::string&, ggml_tensor*)> cache_tensor;
void capture_tensor(const std::string& name, ggml_tensor* tensor) {
if (debug_tensors == nullptr || tensor == nullptr) {
return;
}
ggml_tensor* snapshot = ggml_cont(ggml_ctx, tensor);
ggml_tensor* dst = ggml_dup_tensor(ggml_ctx, snapshot);
snapshot = ggml_cpy(ggml_ctx, snapshot, dst);
ggml_set_output(snapshot);
debug_tensors->push_back({snapshot, name});
}
ggml_tensor* load_cache_tensor(const std::string& name) const {
if (!get_cache_tensor) {
return nullptr;
}
return get_cache_tensor(name);
}
void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
if (!cache_tensor || tensor == nullptr) {
return;
}
cache_tensor(name, tensor);
}
}; };
struct GGMLRunner { struct GGMLRunner {
@ -1743,6 +1771,7 @@ protected:
std::map<ggml_tensor*, const void*> backend_tensor_data_map; std::map<ggml_tensor*, const void*> backend_tensor_data_map;
std::map<std::string, ggml_tensor*> cache_tensor_map; // name -> tensor std::map<std::string, ggml_tensor*> cache_tensor_map; // name -> tensor
std::vector<std::pair<ggml_tensor*, std::string>> debug_tensors;
const std::string final_result_name = "ggml_runner_final_result_tensor"; const std::string final_result_name = "ggml_runner_final_result_tensor";
bool flash_attn_enabled = false; bool flash_attn_enabled = false;
@ -1838,6 +1867,7 @@ protected:
} }
void free_compute_ctx() { void free_compute_ctx() {
debug_tensors.clear();
if (compute_ctx != nullptr) { if (compute_ctx != nullptr) {
ggml_free(compute_ctx); ggml_free(compute_ctx);
compute_ctx = nullptr; compute_ctx = nullptr;
@ -1884,6 +1914,16 @@ protected:
auto result = ggml_graph_node(gf, -1); auto result = ggml_graph_node(gf, -1);
ggml_set_name(result, final_result_name.c_str()); ggml_set_name(result, final_result_name.c_str());
} }
for (const auto& entry : debug_tensors) {
if (entry.first != nullptr) {
ggml_build_forward_expand(gf, entry.first);
}
}
for (const auto& entry : cache_tensor_map) {
if (entry.second != nullptr) {
ggml_build_forward_expand(gf, entry.second);
}
}
prepare_build_in_tensor_after(gf); prepare_build_in_tensor_after(gf);
return gf; return gf;
} }
@ -2031,6 +2071,21 @@ protected:
for (auto& kv : backend_tensor_data_map) { for (auto& kv : backend_tensor_data_map) {
auto tensor = kv.first; auto tensor = kv.first;
auto data = kv.second; auto data = kv.second;
if (tensor == nullptr || data == nullptr) {
continue;
}
const char* name = ggml_get_name(tensor);
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(),
name != nullptr ? name : "",
(long long)tensor->ne[0],
(long long)tensor->ne[1],
(long long)tensor->ne[2],
(long long)tensor->ne[3],
ggml_type_name(tensor->type));
continue;
}
if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
continue; continue;
@ -2421,6 +2476,22 @@ protected:
return std::nullopt; return std::nullopt;
} }
for (const auto& entry : debug_tensors) {
auto tensor = entry.first;
if (tensor == nullptr) {
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(),
entry.second.c_str(),
ggml_type_name(tensor->type));
continue;
}
auto debug_tensor = sd::make_sd_tensor_from_ggml<float>(tensor);
print_sd_tensor(debug_tensor, false, entry.second.c_str());
}
int64_t t_cache_begin = ggml_time_ms(); int64_t t_cache_begin = ggml_time_ms();
if (!copy_cache_tensors_to_cache_buffer(cache_keep_names)) { if (!copy_cache_tensors_to_cache_buffer(cache_keep_names)) {
if (free_compute_buffer_immediately) { if (free_compute_buffer_immediately) {
@ -2557,6 +2628,13 @@ public:
runner_ctx.circular_x_enabled = circular_x_enabled; runner_ctx.circular_x_enabled = circular_x_enabled;
runner_ctx.circular_y_enabled = circular_y_enabled; runner_ctx.circular_y_enabled = circular_y_enabled;
runner_ctx.weight_adapter = weight_adapter; runner_ctx.weight_adapter = weight_adapter;
runner_ctx.debug_tensors = &debug_tensors;
runner_ctx.get_cache_tensor = [this](const std::string& name) {
return this->get_cache_tensor_by_name(name);
};
runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
this->cache(name, tensor);
};
return runner_ctx; return runner_ctx;
} }
@ -2659,6 +2737,9 @@ public:
} }
void cache(const std::string name, ggml_tensor* tensor) { void cache(const std::string name, ggml_tensor* tensor) {
if (tensor != nullptr && tensor->view_src != nullptr) {
tensor = ggml_cont(compute_ctx, tensor);
}
cache_tensor_map[name] = tensor; cache_tensor_map[name] = tensor;
} }

975
src/hidream_o1.hpp Normal file
View File

@ -0,0 +1,975 @@
#ifndef __SD_HIDREAM_O1_H__
#define __SD_HIDREAM_O1_H__
#include <algorithm>
#include <array>
#include <cmath>
#include <cstring>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "common_dit.hpp"
#include "conditioner.hpp"
#include "llm.hpp"
#include "util.h"
namespace HiDreamO1 {
constexpr int HIDREAM_O1_GRAPH_SIZE = 32768;
constexpr int PATCH_SIZE = 32;
constexpr int TIMESTEP_TOKEN_NUM = 1;
constexpr int IMAGE_TOKEN_ID = 151655;
constexpr int VISION_START_TOKEN_ID = 151652;
static inline std::string repeat_special_token(const std::string& token, int64_t count) {
std::string out;
out.reserve(static_cast<size_t>(count) * token.size());
for (int64_t i = 0; i < count; ++i) {
out += token;
}
return out;
}
static inline std::pair<int, int> calculate_dimensions(int max_size, double ratio) {
int width = static_cast<int>(std::sqrt(max_size * max_size * ratio));
int height = static_cast<int>(width / ratio);
width = (width / PATCH_SIZE) * PATCH_SIZE;
height = (height / PATCH_SIZE) * PATCH_SIZE;
width = std::max(width, PATCH_SIZE);
height = std::max(height, PATCH_SIZE);
return {width, height};
}
static inline sd::Tensor<float> resize_to_area(const sd::Tensor<float>& image, int image_size) {
int64_t width = image.shape()[0];
int64_t height = image.shape()[1];
int64_t s_max = static_cast<int64_t>(image_size) * image_size;
double scale = std::sqrt(static_cast<double>(s_max) / static_cast<double>(width * height));
std::vector<std::pair<int64_t, int64_t>> sizes = {
{(static_cast<int64_t>(std::llround(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::llround(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
{(static_cast<int64_t>(std::llround(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::floor(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
{(static_cast<int64_t>(std::floor(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::llround(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
{(static_cast<int64_t>(std::floor(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::floor(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
};
std::sort(sizes.begin(), sizes.end(), [](const auto& a, const auto& b) {
return a.first * a.second > b.first * b.second;
});
std::pair<int64_t, int64_t> new_size = sizes.back();
for (const auto& size : sizes) {
if (size.first > 0 && size.second > 0 && size.first * size.second <= s_max) {
new_size = size;
break;
}
}
double s1 = static_cast<double>(width) / static_cast<double>(new_size.first);
double s2 = static_cast<double>(height) / static_cast<double>(new_size.second);
sd::Tensor<float> resized;
if (s1 < s2) {
int64_t resized_h = static_cast<int64_t>(std::llround(height / s1));
resized = sd::ops::interpolate(image, {new_size.first, resized_h, image.shape()[2], image.shape()[3]});
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<int64_t>(std::llround(width / s2));
resized = sd::ops::interpolate(image, {resized_w, new_size.second, image.shape()[2], image.shape()[3]});
int64_t left = (resized_w - new_size.first) / 2;
resized = sd::ops::slice(resized, 0, left, left + new_size.first);
}
return resized;
}
static inline std::vector<int32_t> build_position_ids(const std::vector<int32_t>& input_ids,
const std::vector<std::array<int32_t, 3>>& image_grids,
const std::vector<int32_t>& skip_vision_start_token) {
std::vector<int32_t> position_ids(4 * input_ids.size(), 0);
int image_index = 0;
int st = 0;
int fix_point = 4096;
std::vector<int32_t> out_t;
std::vector<int32_t> out_h;
std::vector<int32_t> out_w;
while (st < static_cast<int>(input_ids.size())) {
int ed = st;
while (ed < static_cast<int>(input_ids.size()) && input_ids[ed] != IMAGE_TOKEN_ID) {
ed++;
}
if (ed >= static_cast<int>(input_ids.size())) {
int st_idx = out_t.empty() ? 0 : (*std::max_element(out_t.begin(), out_t.end()) + 1);
for (int i = 0; i < static_cast<int>(input_ids.size()) - st; ++i) {
out_t.push_back(st_idx + i);
out_h.push_back(st_idx + i);
out_w.push_back(st_idx + i);
}
break;
}
int text_len = std::max(0, ed - st - skip_vision_start_token[image_index]);
int st_idx = out_t.empty() ? 0 : (*std::max_element(out_t.begin(), out_t.end()) + 1);
for (int i = 0; i < text_len; ++i) {
out_t.push_back(st_idx + i);
out_h.push_back(st_idx + i);
out_w.push_back(st_idx + i);
}
auto grid = image_grids[image_index];
int base;
if (skip_vision_start_token[image_index]) {
if (fix_point > 0) {
base = fix_point;
fix_point = 0;
} else {
base = st_idx;
}
} else {
base = text_len + st_idx;
}
for (int32_t ti = 0; ti < grid[0]; ++ti) {
for (int32_t hi = 0; hi < grid[1]; ++hi) {
for (int32_t wi = 0; wi < grid[2]; ++wi) {
out_t.push_back(base + ti);
out_h.push_back(base + hi);
out_w.push_back(base + wi);
}
}
}
st = ed + grid[0] * grid[1] * grid[2];
image_index++;
}
GGML_ASSERT(out_t.size() == input_ids.size());
for (size_t i = 0; i < input_ids.size(); ++i) {
// ggml IMROPE consumes 4 flattened position streams:
// [t, h, w, e]
// llama.cpp's generic Qwen-VL fallback expands text positions as
// [pos, pos, pos, 0]. Keep the extra stream zeroed here too.
position_ids[i] = out_t[i];
position_ids[input_ids.size() + i] = out_h[i];
position_ids[input_ids.size() * 2 + i] = out_w[i];
position_ids[input_ids.size() * 3 + i] = 0;
}
return position_ids;
}
struct TimestepEmbedder : public GGMLBlock {
int frequency_embedding_size = 256;
TimestepEmbedder(int64_t hidden_size) {
blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1000.0f);
emb = mlp_0->forward(ctx, emb);
emb = ggml_silu_inplace(ctx->ggml_ctx, emb);
emb = mlp_2->forward(ctx, emb);
return emb;
}
};
struct BottleneckPatchEmbed : public GGMLBlock {
BottleneckPatchEmbed(int64_t in_dim, int64_t pca_dim, int64_t embed_dim) {
blocks["proj1"] = std::make_shared<Linear>(in_dim, pca_dim, false);
blocks["proj2"] = std::make_shared<Linear>(pca_dim, embed_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj1 = std::dynamic_pointer_cast<Linear>(blocks["proj1"]);
auto proj2 = std::dynamic_pointer_cast<Linear>(blocks["proj2"]);
return proj2->forward(ctx, proj1->forward(ctx, x));
}
};
struct FinalLayer : public GGMLBlock {
FinalLayer(int64_t hidden_size, int64_t out_dim) {
blocks["linear"] = std::make_shared<Linear>(hidden_size, out_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
return linear->forward(ctx, x);
}
};
struct HiDreamO1Params {
LLM::LLMParams llm;
int patch_size = PATCH_SIZE;
int num_position_embeddings = 2304;
std::vector<int> deepstack_visual_indexes;
};
struct VisionMLP : public GGMLBlock {
VisionMLP(int64_t hidden_size, int64_t intermediate_size) {
blocks["linear_fc1"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
blocks["linear_fc2"] = std::make_shared<Linear>(intermediate_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
auto linear_fc2 = std::dynamic_pointer_cast<Linear>(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<Conv3d>(in_channels,
embed_dim,
std::tuple<int, int, int>{temporal_patch_size, patch_size, patch_size},
std::tuple<int, int, int>{temporal_patch_size, patch_size, patch_size},
std::tuple<int, int, int>{0, 0, 0},
std::tuple<int, int, int>{1, 1, 1},
true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj = std::dynamic_pointer_cast<Conv3d>(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<LayerNorm>(use_postshuffle_norm ? hidden_size : context_dim, 1e-6f);
blocks["linear_fc1"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
blocks["linear_fc2"] = std::make_shared<Linear>(hidden_size, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
auto linear_fc2 = std::dynamic_pointer_cast<Linear>(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<int>(hidden_size / num_heads);
GGML_ASSERT(num_heads * head_dim == hidden_size);
blocks["qkv"] = std::make_shared<Linear>(hidden_size, hidden_size * 3, true);
blocks["proj"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto proj = std::dynamic_pointer_cast<Linear>(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<LayerNorm>(hidden_size, 1e-6f);
blocks["norm2"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f);
blocks["attn"] = std::make_shared<VisionAttention>(hidden_size, num_heads);
blocks["mlp"] = std::make_shared<VisionMLP>(hidden_size, intermediate_size);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto attn = std::dynamic_pointer_cast<VisionAttention>(blocks["attn"]);
auto mlp = std::dynamic_pointer_cast<VisionMLP>(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<ggml_tensor*> deepstack_hidden_states;
};
struct VisionModel : public GGMLBlock {
int num_layers;
int spatial_merge_size;
int num_grid_per_side;
std::vector<int> 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<int> deepstack_visual_indexes)
: num_layers(num_layers),
spatial_merge_size(spatial_merge_size),
num_grid_per_side(static_cast<int>(std::sqrt(num_position_embeddings))),
deepstack_visual_indexes(std::move(deepstack_visual_indexes)) {
blocks["patch_embed"] = std::make_shared<VisionPatchEmbed>(patch_size,
temporal_patch_size,
in_channels,
hidden_size);
blocks["pos_embed"] = std::make_shared<Embedding>(num_position_embeddings, hidden_size);
for (int i = 0; i < num_layers; ++i) {
blocks["blocks." + std::to_string(i)] = std::make_shared<VisionBlock>(hidden_size,
intermediate_size,
num_heads);
}
blocks["merger"] = std::make_shared<VisionPatchMerger>(out_hidden_size,
hidden_size,
spatial_merge_size,
false);
for (int i = 0; i < static_cast<int>(this->deepstack_visual_indexes.size()); ++i) {
blocks["deepstack_merger_list." + std::to_string(i)] = std::make_shared<VisionPatchMerger>(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<Embedding>(blocks["pos_embed"]);
std::vector<int32_t> idx_list[4];
std::vector<float> weight_list[4];
idx_list[0].reserve(static_cast<size_t>(grid_h * grid_w));
idx_list[1].reserve(static_cast<size_t>(grid_h * grid_w));
idx_list[2].reserve(static_cast<size_t>(grid_h * grid_w));
idx_list[3].reserve(static_cast<size_t>(grid_h * grid_w));
weight_list[0].reserve(static_cast<size_t>(grid_h * grid_w));
weight_list[1].reserve(static_cast<size_t>(grid_h * grid_w));
weight_list[2].reserve(static_cast<size_t>(grid_h * grid_w));
weight_list[3].reserve(static_cast<size_t>(grid_h * grid_w));
double max_index = static_cast<double>(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<double>(grid_h - 1);
int h_floor = static_cast<int>(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<double>(grid_w - 1);
int w_floor = static_cast<int>(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<float>((1.0 - dh) * (1.0 - dw)));
weight_list[1].push_back(static_cast<float>((1.0 - dh) * dw));
weight_list[2].push_back(static_cast<float>(dh * (1.0 - dw)));
weight_list[3].push_back(static_cast<float>(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<int64_t>(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<int64_t>(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<VisionPatchEmbed>(blocks["patch_embed"]);
auto merger = std::dynamic_pointer_cast<VisionPatchMerger>(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<VisionBlock>(blocks["blocks." + std::to_string(i)]);
x = block->forward(ctx, x, pe);
for (int j = 0; j < static_cast<int>(deepstack_visual_indexes.size()); ++j) {
if (deepstack_visual_indexes[j] == i) {
auto deepstack_merger = std::dynamic_pointer_cast<VisionPatchMerger>(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;
}
};
struct HiDreamO1Model : public GGMLBlock {
HiDreamO1Params params;
HiDreamO1Model() = default;
explicit HiDreamO1Model(HiDreamO1Params params)
: params(std::move(params)) {
blocks["language_model"] = std::make_shared<LLM::TextModel>(this->params.llm);
blocks["visual"] = std::make_shared<VisionModel>(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<TimestepEmbedder>(this->params.llm.hidden_size);
blocks["x_embedder"] = std::make_shared<BottleneckPatchEmbed>(this->params.patch_size * this->params.patch_size * 3,
this->params.llm.hidden_size / 4,
this->params.llm.hidden_size);
blocks["final_layer2"] = std::make_shared<FinalLayer>(this->params.llm.hidden_size,
this->params.patch_size * this->params.patch_size * 3);
}
std::shared_ptr<LLM::TextModel> text_model() {
return std::dynamic_pointer_cast<LLM::TextModel>(blocks["language_model"]);
}
std::shared_ptr<VisionModel> vision_model() {
return std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
}
std::shared_ptr<TimestepEmbedder> timestep_embedder() {
return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder1"]);
}
std::shared_ptr<BottleneckPatchEmbed> patch_embedder() {
return std::dynamic_pointer_cast<BottleneckPatchEmbed>(blocks["x_embedder"]);
}
std::shared_ptr<FinalLayer> final_layer() {
return std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer2"]);
}
};
struct HiDreamO1Runner : public GGMLRunner {
HiDreamO1Params params;
HiDreamO1Model model;
std::vector<int> window_index_vec;
std::vector<int> window_inverse_index_vec;
std::vector<float> window_mask_vec;
std::vector<float> pe_vec;
std::vector<float> 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};
model = HiDreamO1Model(params);
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "hidream_o1";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
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<int32_t>& 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<int>(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<int>(image->ne[1]) / params.llm.vision.patch_size;
int grid_w = static_cast<int>(image->ne[0]) / params.llm.vision.patch_size;
auto pixel_values = process_image(compute_ctx, image);
int head_dim = static_cast<int>(params.llm.vision.hidden_size / params.llm.vision.num_heads);
std::vector<int> window_index_vec(static_cast<size_t>((grid_h / params.llm.vision.spatial_merge_size) * (grid_w / params.llm.vision.spatial_merge_size)));
for (int i = 0; i < static_cast<int>(window_index_vec.size()); ++i) {
window_index_vec[static_cast<size_t>(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<int>(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<float>& x_tensor,
const sd::Tensor<float>& timestep_tensor,
const sd::Tensor<int32_t>& input_ids_tensor,
const sd::Tensor<int32_t>& input_pos_tensor,
const sd::Tensor<int32_t>& token_types_tensor,
const sd::Tensor<int32_t>& image_embed_ranges_tensor,
const sd::Tensor<int32_t>& vinput_mask_tensor,
const std::vector<sd::Tensor<float>>& vlm_images,
const std::vector<sd::Tensor<float>>& ref_images) {
ggml_cgraph* gf = new_graph_custom(HIDREAM_O1_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timestep = make_input(timestep_tensor);
ggml_tensor* input_ids = make_input(input_ids_tensor);
ggml_tensor* input_pos = make_input(input_pos_tensor);
auto text_model = model.text_model();
auto t_embedder1 = model.timestep_embedder();
auto x_embedder = model.patch_embedder();
auto final_layer2 = model.final_layer();
std::vector<ggml_tensor*> vlm_image_tensors;
for (const auto& image : vlm_images) {
vlm_image_tensors.push_back(make_input(image));
}
std::vector<ggml_tensor*> ref_image_tensors;
for (const auto& image : ref_images) {
ref_image_tensors.push_back(make_input(image));
}
attention_mask_vec = std::vector<float>(static_cast<size_t>(token_types_tensor.shape()[0] * token_types_tensor.shape()[0]), 0.0f);
int64_t total_seq_len = token_types_tensor.shape()[0];
for (int64_t query = 0; query < total_seq_len; ++query) {
bool is_gen = token_types_tensor.values()[static_cast<size_t>(query)] > 0;
for (int64_t key = 0; key < total_seq_len; ++key) {
if (!is_gen && key > query) {
attention_mask_vec[static_cast<size_t>(query * total_seq_len + key)] = -INFINITY;
}
}
}
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 txt = text_model->embed(&runner_ctx, input_ids);
txt = scatter_visual_embeds(&runner_ctx, txt, image_embed_ranges_tensor, visual_embeds);
auto t_emb = t_embedder1->forward(&runner_ctx, timestep);
int64_t txt_seq_len = input_ids->ne[0];
if (txt_seq_len > 1) {
auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1);
txt = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1);
} else {
txt = ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1);
}
auto vinputs = DiT::pad_and_patchify(&runner_ctx, x, PATCH_SIZE, PATCH_SIZE);
int64_t target_tokens = vinputs->ne[1];
for (ggml_tensor* ref_image : ref_image_tensors) {
auto ref = DiT::pad_and_patchify(&runner_ctx, ref_image, PATCH_SIZE, PATCH_SIZE);
vinputs = ggml_concat(compute_ctx, vinputs, ref, 1);
}
auto vis = x_embedder->forward(&runner_ctx, vinputs);
auto inputs_embeds = ggml_concat(compute_ctx, txt, vis, 1);
auto hidden_states = text_model->forward_embeds(&runner_ctx, inputs_embeds, input_pos, attention_mask, {});
auto x_pred_all = final_layer2->forward(&runner_ctx, hidden_states);
int64_t x_pred_start = txt_seq_len;
if (!vinput_mask_tensor.empty()) {
int64_t seq_len = static_cast<int64_t>(vinput_mask_tensor.shape()[0]);
int64_t first_vinput = 0;
while (first_vinput < seq_len && vinput_mask_tensor.values()[static_cast<size_t>(first_vinput)] == 0) {
first_vinput++;
}
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);
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];
sigma = std::max(1e-6f, sigma);
auto out = ggml_scale(compute_ctx, ggml_sub(compute_ctx, x, x_pred), 1.0f / sigma);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timestep,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<int32_t>& input_pos,
const sd::Tensor<int32_t>& token_types,
const sd::Tensor<int32_t>& image_embed_ranges,
const sd::Tensor<int32_t>& vinput_mask,
const std::vector<sd::Tensor<float>>& vlm_images,
const std::vector<sd::Tensor<float>>& 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 restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
};
struct HiDreamO1Conditioner : public Conditioner {
Qwen2Tokenizer tokenizer;
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
SD_UNUSED(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); }
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
SD_UNUSED(n_threads);
SDCondition result;
int width = conditioner_params.width;
int height = conditioner_params.height;
int64_t target_image_len = static_cast<int64_t>(width / PATCH_SIZE) * static_cast<int64_t>(height / PATCH_SIZE);
std::vector<sd::Tensor<float>> ref_images;
if (conditioner_params.ref_images != nullptr) {
ref_images = *conditioner_params.ref_images;
}
std::vector<sd::Tensor<float>> vlm_images;
std::vector<std::array<int32_t, 3>> image_grids;
std::vector<int32_t> skip_vision_start;
std::string prompt = "<|im_start|>user\n";
std::vector<int32_t> image_ranges;
if (ref_images.empty()) {
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
auto input_ids = tokenizer.encode(prompt, nullptr);
std::vector<int32_t> input_ids_pad = input_ids;
input_ids_pad.push_back(VISION_START_TOKEN_ID);
input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID);
image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(width / PATCH_SIZE)});
skip_vision_start.push_back(1);
std::vector<int32_t> token_types(input_ids_pad.size(), 0);
int txt_seq_len = static_cast<int>(input_ids.size());
int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM;
for (int i = bgn; i < bgn + target_image_len + TIMESTEP_TOKEN_NUM; ++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);
std::vector<int64_t> input_shape{static_cast<int64_t>(input_ids.size())};
std::vector<int64_t> position_shape{static_cast<int64_t>(input_ids_pad.size() * 4)};
std::vector<int64_t> token_type_shape{static_cast<int64_t>(token_types.size())};
std::vector<int32_t> vinput_mask(token_types.size(), 0);
for (int64_t i = txt_seq_len; i < static_cast<int64_t>(vinput_mask.size()); ++i) {
vinput_mask[static_cast<size_t>(i)] = 1;
}
std::vector<int64_t> vinput_mask_shape{static_cast<int64_t>(vinput_mask.size())};
result.c_input_ids = sd::Tensor<int32_t>(input_shape, std::move(input_ids));
result.c_position_ids = sd::Tensor<int32_t>(position_shape, position_ids);
result.c_token_types = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
result.c_vinput_mask = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
result.c_image_embed_ranges = sd::Tensor<int32_t>();
return result;
}
int K = static_cast<int>(ref_images.size());
int max_size;
if (K == 1) {
max_size = std::max(height, width);
} else if (K == 2) {
max_size = std::max(height, width) * 48 / 64;
} else if (K <= 4) {
max_size = std::max(height, width) / 2;
} else if (K <= 8) {
max_size = std::max(height, width) * 24 / 64;
} else {
max_size = std::max(height, width) / 4;
}
int cond_img_size;
if (K <= 4) {
cond_img_size = 384;
} else if (K <= 8) {
cond_img_size = 384 * 48 / 64;
} else {
cond_img_size = 384 / 2;
}
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 dims = calculate_dimensions(cond_img_size, static_cast<double>(ref_image.shape()[0]) / static_cast<double>(ref_image.shape()[1]));
auto vlm_image = clip_preprocess(ref_image, dims.first, dims.second);
int64_t image_tokens = static_cast<int64_t>(dims.first / PATCH_SIZE) * static_cast<int64_t>(dims.second / PATCH_SIZE);
int64_t prompt_start = static_cast<int64_t>(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<int32_t>(prompt_start));
image_ranges.push_back(static_cast<int32_t>(image_tokens));
result.c_vlm_images.push_back(std::move(vlm_image));
image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE});
skip_vision_start.push_back(0);
}
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
auto input_ids = tokenizer.encode(prompt, nullptr);
std::vector<int32_t> input_ids_pad = input_ids;
input_ids_pad.push_back(VISION_START_TOKEN_ID);
input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID);
image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(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<int64_t>(ref_image.shape()[0] / PATCH_SIZE) * static_cast<int64_t>(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<int32_t>(ref_image.shape()[1] / PATCH_SIZE), static_cast<int32_t>(ref_image.shape()[0] / PATCH_SIZE)});
skip_vision_start.push_back(1);
}
std::vector<int32_t> token_types(input_ids_pad.size(), 0);
int txt_seq_len = static_cast<int>(input_ids.size());
int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM;
int end = bgn + static_cast<int>(target_image_len) + TIMESTEP_TOKEN_NUM;
for (int i = bgn; i < end; ++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<int64_t> input_shape{static_cast<int64_t>(input_ids.size())};
std::vector<int64_t> position_shape{static_cast<int64_t>(input_ids_pad.size() * 4)};
std::vector<int64_t> token_type_shape{static_cast<int64_t>(token_types.size())};
std::vector<int64_t> image_range_shape{2, static_cast<int64_t>(image_ranges.size() / 2)};
std::vector<int32_t> vinput_mask(token_types.size(), 0);
for (int i = txt_seq_len; i < static_cast<int>(vinput_mask.size()); ++i) {
vinput_mask[static_cast<size_t>(i)] = 1;
}
std::vector<int64_t> vinput_mask_shape{static_cast<int64_t>(vinput_mask.size())};
result.c_input_ids = sd::Tensor<int32_t>(input_shape, std::move(input_ids));
result.c_position_ids = sd::Tensor<int32_t>(position_shape, build_position_ids(input_ids_pad, image_grids, skip_vision_start));
result.c_token_types = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
result.c_image_embed_ranges = sd::Tensor<int32_t>(image_range_shape, std::move(image_ranges));
result.c_vinput_mask = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
return result;
}
};
} // namespace HiDreamO1
#endif // __SD_HIDREAM_O1_H__

133
src/llm.hpp
View File

@ -27,6 +27,7 @@ namespace LLM {
enum class LLMArch { enum class LLMArch {
QWEN2_5_VL, QWEN2_5_VL,
QWEN3, QWEN3,
QWEN3_VL,
MISTRAL_SMALL_3_2, MISTRAL_SMALL_3_2,
MINISTRAL_3_3B, MINISTRAL_3_3B,
ARCH_COUNT, ARCH_COUNT,
@ -35,6 +36,7 @@ namespace LLM {
static const char* llm_arch_to_str[] = { static const char* llm_arch_to_str[] = {
"qwen2.5vl", "qwen2.5vl",
"qwen3", "qwen3",
"qwen3vl",
"mistral_small3.2", "mistral_small3.2",
"ministral3.3b", "ministral3.3b",
}; };
@ -430,6 +432,10 @@ namespace LLM {
} else if (arch == LLMArch::QWEN3) { } else if (arch == LLMArch::QWEN3) {
q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
} else if (arch == LLMArch::QWEN3_VL) {
int sections[4] = {24, 20, 20, 0};
q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
k = ggml_rope_multi(ctx->ggml_ctx, k, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
} else { } else {
int sections[4] = {16, 24, 24, 0}; int sections[4] = {16, 24, 24, 0};
q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_MROPE, 128000, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_MROPE, 128000, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@ -485,10 +491,11 @@ namespace LLM {
struct TextModel : public GGMLBlock { struct TextModel : public GGMLBlock {
protected: protected:
int64_t num_layers; int64_t num_layers;
LLMParams params;
public: public:
TextModel(const LLMParams& params) TextModel(const LLMParams& params)
: num_layers(params.num_layers) { : num_layers(params.num_layers), params(params) {
blocks["embed_tokens"] = std::shared_ptr<GGMLBlock>(new Embedding(params.vocab_size, params.hidden_size)); blocks["embed_tokens"] = std::shared_ptr<GGMLBlock>(new Embedding(params.vocab_size, params.hidden_size));
for (int i = 0; i < num_layers; i++) { for (int i = 0; i < num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(params)); blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(params));
@ -496,62 +503,63 @@ namespace LLM {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(params.hidden_size, params.rms_norm_eps)); blocks["norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(params.hidden_size, params.rms_norm_eps));
} }
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* embed(GGMLRunnerContext* ctx,
ggml_tensor* input_ids, ggml_tensor* input_ids) {
ggml_tensor* input_pos,
ggml_tensor* attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
std::set<int> out_layers) {
// input_ids: [N, n_token]
// return: [N, n_token, hidden_size]
auto embed_tokens = std::dynamic_pointer_cast<Embedding>(blocks["embed_tokens"]); auto embed_tokens = std::dynamic_pointer_cast<Embedding>(blocks["embed_tokens"]);
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]); auto x = embed_tokens->forward(ctx, input_ids);
return x;
}
auto x = embed_tokens->forward(ctx, input_ids); ggml_tensor* splice_image_embeds(GGMLRunnerContext* ctx,
sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.prelude", "x"); ggml_tensor* x,
std::vector<std::pair<int, ggml_tensor*>> image_embeds) {
std::vector<ggml_tensor*> intermediate_outputs; if (image_embeds.empty()) {
return x;
if (image_embeds.size() > 0) {
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);
}
auto image_embed = image_embeds[i].second;
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embed, 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]);
x = input_embed;
} }
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,
ggml_tensor* attention_mask,
std::set<int> out_layers) {
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
std::vector<ggml_tensor*> intermediate_outputs;
sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.prelude", "x");
for (int i = 0; i < num_layers; i++) { for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]); auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
@ -570,10 +578,23 @@ namespace LLM {
for (int i = 1; i < intermediate_outputs.size(); i++) { for (int i = 1; i < intermediate_outputs.size(); i++) {
x = ggml_concat(ctx->ggml_ctx, x, intermediate_outputs[i], 0); x = ggml_concat(ctx->ggml_ctx, x, intermediate_outputs[i], 0);
} }
} else { return x;
x = norm->forward(ctx, x);
} }
return x;
return norm->forward(ctx, x);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* input_pos,
ggml_tensor* attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
std::set<int> out_layers) {
// 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));
return forward_embeds(ctx, x, input_pos, attention_mask, std::move(out_layers));
} }
}; };

4
src/model.cpp
View File

@ -437,6 +437,10 @@ SDVersion ModelLoader::get_sd_version() {
if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) { if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) {
return VERSION_SD3; return VERSION_SD3;
} }
if (tensor_storage.name.find("model.x_embedder.proj1.weight") != std::string::npos &&
tensor_storage_map.find("model.language_model.layers.0.self_attn.q_proj.weight") != tensor_storage_map.end()) {
return VERSION_HIDREAM_O1;
}
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) { if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
return VERSION_QWEN_IMAGE; return VERSION_QWEN_IMAGE;
} }

2
src/model.h
View File

@ -42,6 +42,7 @@ enum SDVersion {
VERSION_ANIMA, VERSION_ANIMA,
VERSION_FLUX2, VERSION_FLUX2,
VERSION_FLUX2_KLEIN, VERSION_FLUX2_KLEIN,
VERSION_HIDREAM_O1,
VERSION_Z_IMAGE, VERSION_Z_IMAGE,
VERSION_OVIS_IMAGE, VERSION_OVIS_IMAGE,
VERSION_ERNIE_IMAGE, VERSION_ERNIE_IMAGE,
@ -163,6 +164,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
sd_version_is_sd3(version) || sd_version_is_sd3(version) ||
sd_version_is_wan(version) || sd_version_is_wan(version) ||
sd_version_is_qwen_image(version) || sd_version_is_qwen_image(version) ||
version == VERSION_HIDREAM_O1 ||
sd_version_is_anima(version) || sd_version_is_anima(version) ||
sd_version_is_z_image(version) || sd_version_is_z_image(version) ||
sd_version_is_ernie_image(version)) { sd_version_is_ernie_image(version)) {

47
src/stable-diffusion.cpp
View File

@ -52,6 +52,7 @@ const char* model_version_to_str[] = {
"Anima", "Anima",
"Flux.2", "Flux.2",
"Flux.2 klein", "Flux.2 klein",
"HiDream O1",
"Z-Image", "Z-Image",
"Ovis Image", "Ovis Image",
"Ernie Image", "Ernie Image",
@ -491,6 +492,12 @@ public:
"model.diffusion_model", "model.diffusion_model",
version, version,
sd_ctx_params->qwen_image_zero_cond_t); sd_ctx_params->qwen_image_zero_cond_t);
} else if (version == VERSION_HIDREAM_O1) {
cond_stage_model = std::make_shared<HiDreamO1::HiDreamO1Conditioner>();
diffusion_model = std::make_shared<HiDreamO1Model>(backend,
offload_params_to_cpu,
tensor_storage_map,
"model");
} else if (sd_version_is_anima(version)) { } else if (sd_version_is_anima(version)) {
cond_stage_model = std::make_shared<AnimaConditioner>(clip_backend, cond_stage_model = std::make_shared<AnimaConditioner>(clip_backend,
offload_params_to_cpu, offload_params_to_cpu,
@ -625,7 +632,7 @@ public:
} }
}; };
if (version == VERSION_CHROMA_RADIANCE) { if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
LOG_INFO("using FakeVAE"); LOG_INFO("using FakeVAE");
first_stage_model = std::make_shared<FakeVAE>(version, first_stage_model = std::make_shared<FakeVAE>(version,
vae_backend, vae_backend,
@ -796,6 +803,9 @@ public:
ignore_tensors.insert("text_encoders.llm.vision_tower."); ignore_tensors.insert("text_encoders.llm.vision_tower.");
ignore_tensors.insert("text_encoders.llm.multi_modal_projector."); ignore_tensors.insert("text_encoders.llm.multi_modal_projector.");
} }
if (version == VERSION_HIDREAM_O1) {
ignore_tensors.insert("lm_head.");
}
bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap); bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap);
if (!success) { if (!success) {
LOG_ERROR("load tensors from model loader failed"); LOG_ERROR("load tensors from model loader failed");
@ -898,6 +908,7 @@ public:
} else if (sd_version_is_sd3(version) || } else if (sd_version_is_sd3(version) ||
sd_version_is_wan(version) || sd_version_is_wan(version) ||
sd_version_is_qwen_image(version) || sd_version_is_qwen_image(version) ||
version == VERSION_HIDREAM_O1 ||
sd_version_is_anima(version) || sd_version_is_anima(version) ||
sd_version_is_ernie_image(version) || sd_version_is_ernie_image(version) ||
sd_version_is_z_image(version)) { sd_version_is_z_image(version)) {
@ -1495,6 +1506,9 @@ public:
if (sd_version_is_anima(version)) { if (sd_version_is_anima(version)) {
return std::vector<float>{t / static_cast<float>(TIMESTEPS)}; return std::vector<float>{t / static_cast<float>(TIMESTEPS)};
} }
if (version == VERSION_HIDREAM_O1) {
return std::vector<float>{1.0f - (t / static_cast<float>(TIMESTEPS))};
}
if (sd_version_is_z_image(version)) { if (sd_version_is_z_image(version)) {
return std::vector<float>{1000.f - t}; return std::vector<float>{1000.f - t};
} }
@ -1607,6 +1621,10 @@ public:
LOG_WARN("SLG is incompatible with this model type"); LOG_WARN("SLG is incompatible with this model type");
} }
if (version == VERSION_HIDREAM_O1 && !noise.empty()) {
noise *= eta;
}
int64_t t0 = ggml_time_us(); int64_t t0 = ggml_time_us();
sd::Tensor<float> x_t = !noise.empty() sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, init_latent) ? denoiser->noise_scaling(sigmas[0], noise, init_latent)
@ -1679,12 +1697,19 @@ public:
auto run_condition = [&](const SDCondition& condition, auto run_condition = [&](const SDCondition& condition,
const sd::Tensor<float>* c_concat_override = nullptr, const sd::Tensor<float>* c_concat_override = nullptr,
const std::vector<int>* local_skip_layers = nullptr) -> sd::Tensor<float> { const std::vector<int>* local_skip_layers = nullptr) -> sd::Tensor<float> {
diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; 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.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.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_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.t5_weights = condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights;
diffusion_params.skip_layers = local_skip_layers; 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;
sd::Tensor<float> cached_output; sd::Tensor<float> cached_output;
if (step_cache.before_condition(&condition, noised_input, &cached_output)) { if (step_cache.before_condition(&condition, noised_input, &cached_output)) {
@ -1831,6 +1856,8 @@ public:
if (sd_version_is_dit(version)) { if (sd_version_is_dit(version)) {
if (version == VERSION_WAN2_2_TI2V) { if (version == VERSION_WAN2_2_TI2V) {
latent_channel = 48; latent_channel = 48;
} else if (version == VERSION_HIDREAM_O1) {
latent_channel = 3;
} else if (version == VERSION_CHROMA_RADIANCE) { } else if (version == VERSION_CHROMA_RADIANCE) {
latent_channel = 3; latent_channel = 3;
} else if (sd_version_uses_flux2_vae(version)) { } else if (sd_version_uses_flux2_vae(version)) {
@ -2518,6 +2545,9 @@ static float resolve_eta(sd_ctx_t* sd_ctx,
float eta, float eta,
enum sample_method_t sample_method) { enum sample_method_t sample_method) {
if (eta == INFINITY) { if (eta == INFINITY) {
if (sd_ctx->sd->version == VERSION_HIDREAM_O1) {
return 8.f;
}
switch (sample_method) { switch (sample_method) {
case DDIM_TRAILING_SAMPLE_METHOD: case DDIM_TRAILING_SAMPLE_METHOD:
case TCD_SAMPLE_METHOD: case TCD_SAMPLE_METHOD:
@ -3009,6 +3039,9 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
std::vector<sd::Tensor<float>> ref_latents; std::vector<sd::Tensor<float>> ref_latents;
for (size_t i = 0; i < ref_images.size(); i++) { for (size_t i = 0; i < ref_images.size(); i++) {
if (sd_ctx->sd->version == VERSION_HIDREAM_O1) {
continue;
}
sd::Tensor<float> ref_latent; sd::Tensor<float> ref_latent;
if (request->auto_resize_ref_image) { if (request->auto_resize_ref_image) {
LOG_DEBUG("auto resize ref images"); LOG_DEBUG("auto resize ref images");

5
src/tokenizers/qwen2_tokenizer.cpp
View File

@ -81,6 +81,11 @@ Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) {
"</tool_response>", "</tool_response>",
"<think>", "<think>",
"</think>", "</think>",
"<|boi_token|>",
"<|bor_token|>",
"<|eor_token|>",
"<|bot_token|>",
"<|tms_token|>",
}; };
if (merges_utf8_str.size() > 0) { if (merges_utf8_str.size() > 0) {

2
src/vae.hpp
View File

@ -71,7 +71,7 @@ public:
scale_factor = 16; scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) { } else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16; scale_factor = 16;
} else if (version == VERSION_CHROMA_RADIANCE) { } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
scale_factor = 1; scale_factor = 1;
} }
return scale_factor; return scale_factor;