upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-10-01 18:37:53 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add minimax-h3 support (#1854)

Browse Source
This commit is contained in:
leejet
2026-08-04 23:14:18 +08:00
committed by GitHub
parent b4e67d1221
commit ea7f0c87cf
26 changed files with 4293 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/model/diffusion/minimax_h3.hpp
+1178
View File
File diff suppressed because it is too large Load Diff
src/model/diffusion/model.hpp
+24
View File
@@ -87,6 +87,29 @@ struct LTXAVDiffusionExtra {
const sd::Tensor<float>* video_positions = nullptr;
};
enum class MiniMaxH3ReferenceKind : int32_t {
IMAGE,
VIDEO,
AUDIO,
VIDEO_AUDIO,
};
struct MiniMaxH3ReferenceBlock {
MiniMaxH3ReferenceKind kind = MiniMaxH3ReferenceKind::IMAGE;
int32_t video_index = -1;
int32_t audio_index = -1;
};
struct MiniMaxH3DiffusionExtra {
const sd::Tensor<int32_t>* text_token_tags = nullptr;
const sd::Tensor<int32_t>* keyframe_indices = nullptr;
const std::vector<sd::Tensor<float>>* reference_audio_latents = nullptr;
const std::vector<MiniMaxH3ReferenceBlock>* reference_blocks = nullptr;
int audio_length = 0;
float video_sigma_shift = 12.f;
float audio_sigma_shift = 3.f;
};
struct MiniT2IDiffusionExtra {
const sd::Tensor<float>* mask = nullptr;
};
@@ -106,6 +129,7 @@ using DiffusionExtraParams = std::variant<std::monostate,
WanDiffusionExtra,
HiDreamO1DiffusionExtra,
LTXAVDiffusionExtra,
MiniMaxH3DiffusionExtra,
MiniT2IDiffusionExtra,
HunyuanVideoDiffusionExtra>;
src/model/te/llm.hpp
+369 -26
View File
@@ -79,9 +79,20 @@ namespace LLM {
int window_size = 112;
int num_position_embeddings = 0;
std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
bool split_patch_embed = false;
std::vector<int> deepstack_visual_indexes;
bool split_patch_embed = false;
};
struct ImageGrid {
int index = 0;
int size = 0;
int grid_h = 0;
int grid_w = 0;
};
using ImageEmbeds = std::vector<std::pair<int, sd::Tensor<float>>>;
using DeepStackImageEmbeds = std::vector<ImageEmbeds>;
struct LLMConfig {
LLMArch arch = LLMArch::QWEN2_5_VL;
int64_t num_layers = 28;
@@ -93,6 +104,7 @@ namespace LLM {
bool qkv_bias = true;
bool attention_out_bias = false;
bool qk_norm = false;
bool final_norm = true;
bool rms_norm_add = false;
bool normalize_input = false;
int64_t vocab_size = 152064;
@@ -257,9 +269,20 @@ namespace LLM {
if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
config.num_heads = 16;
}
if (arch == LLMArch::QWEN3_VL && config.num_layers == 50 && config.hidden_size == 5120) {
config.num_heads = 64;
config.final_norm = false;
}
if (detected_vision_layers > 0) {
config.vision.num_layers = detected_vision_layers;
}
if (arch == LLMArch::QWEN3_VL) {
if (config.vision.num_layers == 24) {
config.vision.deepstack_visual_indexes = {5, 11, 17};
} else if (config.vision.num_layers == 27) {
config.vision.deepstack_visual_indexes = {8, 16, 24};
}
}
LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
config.num_layers,
config.vocab_size,
@@ -541,6 +564,37 @@ namespace LLM {
return input_embed;
}
static ggml_tensor* add_deepstack_image_embeds(GGMLRunnerContext* ctx,
ggml_tensor* x,
const std::vector<std::pair<int, ggml_tensor*>>& image_embeds) {
if (image_embeds.empty()) {
return x;
}
GGML_ASSERT(x->ne[2] == 1);
auto raw_x = ggml_cast(ctx->ggml_ctx, x, image_embeds[0].second->type);
int64_t token_start = 0;
ggml_tensor* output = nullptr;
for (const auto& [index, image_embed] : image_embeds) {
GGML_ASSERT(index >= token_start);
GGML_ASSERT(index + image_embed->ne[1] <= raw_x->ne[1]);
if (index > token_start) {
auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, index);
output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
}
auto visual_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, index, index + image_embed->ne[1]);
visual_embed = ggml_add(ctx->ggml_ctx, visual_embed, image_embed);
output = output == nullptr ? visual_embed : ggml_concat(ctx->ggml_ctx, output, visual_embed, 1);
token_start = index + image_embed->ne[1];
}
if (token_start < raw_x->ne[1]) {
auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, raw_x->ne[1]);
output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
}
GGML_ASSERT(output != nullptr && output->ne[1] == raw_x->ne[1]);
return output;
}
struct VisionMLP : public GGMLBlock {
protected:
LLMVisionArch arch_;
@@ -723,6 +777,33 @@ namespace LLM {
}
};
struct Qwen3VLDeepStackMerger : public GGMLBlock {
protected:
int64_t merge_dim;
public:
Qwen3VLDeepStackMerger(int64_t dim,
int64_t context_dim,
int64_t spatial_merge_size)
: merge_dim(context_dim * spatial_merge_size * spatial_merge_size) {
blocks["norm"] = std::make_shared<LayerNorm>(merge_dim, 1e-6f);
blocks["linear_fc1"] = std::make_shared<Linear>(merge_dim, merge_dim, true);
blocks["linear_fc2"] = std::make_shared<Linear>(merge_dim, 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 = ggml_reshape_2d(ctx->ggml_ctx, x, merge_dim, ggml_nelements(x) / merge_dim);
x = norm->forward(ctx, x);
x = linear_fc1->forward(ctx, x);
x = ggml_gelu_erf(ctx->ggml_ctx, x);
return linear_fc2->forward(ctx, x);
}
};
struct VisionAttention : public GGMLBlock {
protected:
bool llama_cpp_style;
@@ -844,6 +925,7 @@ namespace LLM {
int spatial_merge_size;
int num_grid_per_side;
std::set<int> fullatt_block_indexes;
std::vector<int> deepstack_visual_indexes;
public:
VisionModel(bool llama_cpp_style,
@@ -853,7 +935,8 @@ namespace LLM {
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<int>(std::sqrt(vision_params.num_position_embeddings)) : 0),
fullatt_block_indexes(vision_params.fullatt_block_indexes) {
fullatt_block_indexes(vision_params.fullatt_block_indexes),
deepstack_visual_indexes(vision_params.deepstack_visual_indexes) {
blocks["patch_embed"] = std::shared_ptr<GGMLBlock>(new VisionPatchEmbed(vision_params.split_patch_embed,
arch_,
vision_params.patch_size,
@@ -875,6 +958,11 @@ namespace LLM {
vision_params.out_hidden_size,
vision_params.hidden_size,
spatial_merge_size));
for (size_t i = 0; i < deepstack_visual_indexes.size(); ++i) {
blocks["deepstack_merger_list." + std::to_string(i)] = std::make_shared<Qwen3VLDeepStackMerger>(vision_params.out_hidden_size,
vision_params.hidden_size,
spatial_merge_size);
}
}
std::shared_ptr<Embedding> pos_embedder() {
@@ -893,13 +981,13 @@ namespace LLM {
return spatial_merge_size;
}
ggml_tensor* 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* pos_embeds = nullptr) {
std::vector<ggml_tensor*> forward_outputs(GGMLRunnerContext* ctx,
ggml_tensor* pixel_values,
ggml_tensor* pe,
ggml_tensor* window_index,
ggml_tensor* window_inverse_index,
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)]
@@ -919,6 +1007,7 @@ namespace LLM {
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]);
}
std::vector<ggml_tensor*> deepstack_outputs;
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<VisionBlock>(blocks["blocks." + std::to_string(i)]);
@@ -926,8 +1015,12 @@ namespace LLM {
if (fullatt_block_indexes.find(i) != fullatt_block_indexes.end()) {
mask = nullptr;
}
x = block->forward(ctx, x, pe, mask);
if (i == 0) {
x = block->forward(ctx, x, pe, mask);
auto deepstack_it = std::find(deepstack_visual_indexes.begin(), deepstack_visual_indexes.end(), i);
if (deepstack_it != deepstack_visual_indexes.end()) {
size_t deepstack_index = static_cast<size_t>(std::distance(deepstack_visual_indexes.begin(), deepstack_it));
auto deepstack_merger = std::dynamic_pointer_cast<Qwen3VLDeepStackMerger>(blocks["deepstack_merger_list." + std::to_string(deepstack_index)]);
deepstack_outputs.push_back(deepstack_merger->forward(ctx, x));
}
sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x");
}
@@ -939,7 +1032,19 @@ namespace LLM {
x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index);
}
return x;
std::vector<ggml_tensor*> outputs = {x};
outputs.insert(outputs.end(), deepstack_outputs.begin(), deepstack_outputs.end());
return outputs;
}
ggml_tensor* 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* pos_embeds = nullptr) {
return forward_outputs(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask, pos_embeds)[0];
}
};
@@ -1263,7 +1368,9 @@ namespace LLM {
for (int i = 0; i < num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, i));
}
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
if (config.final_norm) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
}
}
ggml_tensor* embed(GGMLRunnerContext* ctx,
@@ -1278,9 +1385,11 @@ namespace LLM {
ggml_tensor* input_pos,
ggml_tensor* attention_mask,
std::set<int> out_layers,
ggml_tensor* sliding_attention_mask = nullptr,
bool return_all_hidden_states = false) {
auto norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"]);
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds = {},
ggml_tensor* sliding_attention_mask = nullptr,
bool return_all_hidden_states = false) {
auto norm = config.final_norm ? std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])
: nullptr;
std::vector<ggml_tensor*> intermediate_outputs;
if (config.normalize_input) {
@@ -1295,6 +1404,9 @@ namespace LLM {
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask);
if (i < static_cast<int>(deepstack_image_embeds.size())) {
x = add_deepstack_image_embeds(ctx, x, deepstack_image_embeds[static_cast<size_t>(i)]);
}
if (return_all_hidden_states || out_layers.size() > 1) {
x = ggml_cont(ctx->ggml_ctx, x);
}
@@ -1308,7 +1420,7 @@ namespace LLM {
}
}
auto normed_x = norm->forward(ctx, x);
auto normed_x = norm == nullptr ? x : norm->forward(ctx, x);
if (return_all_hidden_states) {
intermediate_outputs.push_back(normed_x);
x = intermediate_outputs[0];
@@ -1336,6 +1448,7 @@ namespace LLM {
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
// input_ids: [N, n_token]
@@ -1347,6 +1460,7 @@ namespace LLM {
input_pos,
attention_mask,
std::move(out_layers),
deepstack_image_embeds,
sliding_attention_mask,
return_all_hidden_states);
}
@@ -1372,6 +1486,7 @@ namespace LLM {
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
// input_ids: [N, n_token]
@@ -1383,6 +1498,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states);
return x;
@@ -1524,7 +1640,8 @@ namespace LLM {
std::vector<float>& window_mask_vec,
std::vector<float>& pe_vec,
std::array<std::vector<int32_t>, 4>& pos_embed_idx_data,
std::array<std::vector<float>, 4>& pos_embed_weight_data) {
std::array<std::vector<float>, 4>& pos_embed_weight_data,
std::vector<ggml_tensor*>* output_tensors = nullptr) {
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);
@@ -1556,7 +1673,11 @@ namespace LLM {
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);
runner->set_backend_tensor_data(pe, pe_vec.data());
return vision_model->forward(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds);
auto outputs = vision_model->forward_outputs(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds);
if (output_tensors != nullptr) {
*output_tensors = outputs;
}
return outputs[0];
}
int llm_grid_h = grid_h / vision_params.spatial_merge_size;
@@ -1622,7 +1743,11 @@ namespace LLM {
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);
auto output = vision_model->forward(runner_ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
if (output_tensors != nullptr) {
*output_tensors = {output};
}
return output;
}
public:
@@ -1667,6 +1792,7 @@ namespace LLM {
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
auto hidden_states = model.forward(ctx,
@@ -1675,6 +1801,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states); // [N, n_token, hidden_size]
return hidden_states;
@@ -1693,7 +1820,9 @@ namespace LLM {
ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
const sd::Tensor<float>& attention_mask_tensor,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
const ImageEmbeds& image_embeds_tensor,
const DeepStackImageEmbeds& deepstack_image_embeds_tensor,
const std::vector<ImageGrid>& image_grids,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
@@ -1704,6 +1833,13 @@ namespace LLM {
ggml_tensor* embed = make_input(embed_tensor);
image_embeds.emplace_back(idx, embed);
}
std::vector<std::vector<std::pair<int, ggml_tensor*>>> deepstack_image_embeds(deepstack_image_embeds_tensor.size());
for (size_t layer = 0; layer < deepstack_image_embeds_tensor.size(); ++layer) {
deepstack_image_embeds[layer].reserve(deepstack_image_embeds_tensor[layer].size());
for (const auto& [idx, embed_tensor] : deepstack_image_embeds_tensor[layer]) {
deepstack_image_embeds[layer].emplace_back(idx, make_input(embed_tensor));
}
}
int64_t n_tokens = input_ids->ne[0];
if (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
@@ -1724,6 +1860,30 @@ namespace LLM {
input_pos_vec[2 * n_tokens + i] = i;
input_pos_vec[3 * n_tokens + i] = 0;
}
if (config.arch == LLMArch::QWEN3_VL && !image_grids.empty()) {
int offset = 0;
for (const auto& grid : image_grids) {
int end = grid.index + grid.size;
int grid_h = grid.grid_h / config.vision.spatial_merge_size;
int grid_w = grid.grid_w / config.vision.spatial_merge_size;
int len_max = std::max(grid_h, grid_w);
int next_pos = grid.index + len_max + offset;
GGML_ASSERT(grid.index >= 0 && end <= n_tokens);
GGML_ASSERT(grid_h > 0 && grid_w > 0 && grid.size == grid_h * grid_w);
for (int token = end; token < n_tokens; ++token) {
int pos = next_pos + token - end;
input_pos_vec[token] = pos;
input_pos_vec[n_tokens + token] = pos;
input_pos_vec[2 * n_tokens + token] = pos;
}
for (int token = 0; token < grid.size; ++token) {
input_pos_vec[grid.index + token] = grid.index + offset;
input_pos_vec[n_tokens + grid.index + token] = grid.index + offset + token / grid_w;
input_pos_vec[2 * n_tokens + grid.index + token] = grid.index + offset + token % grid_w;
}
offset += len_max - grid.size;
}
}
}
auto input_pos = ggml_new_tensor_1d(compute_ctx,
@@ -1781,6 +1941,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states);
@@ -1792,16 +1953,20 @@ namespace LLM {
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<float>& attention_mask,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
const ImageEmbeds& image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true,
const DeepStackImageEmbeds& deepstack_image_embeds = {},
const std::vector<ImageGrid>& image_grids = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids,
attention_mask,
image_embeds,
deepstack_image_embeds,
image_grids,
out_layers,
return_all_hidden_states);
};
@@ -1851,6 +2016,24 @@ namespace LLM {
pos_embed_weight_data_);
}
std::vector<ggml_tensor*> encode_image_outputs(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
std::vector<ggml_tensor*> outputs;
encode_image_common(this,
compute_ctx,
runner_ctx,
image,
config.vision,
model.vision_model(),
window_index_vec,
window_inverse_index_vec,
window_mask_vec,
pe_vec,
pos_embed_idx_data_,
pos_embed_weight_data_,
&outputs);
return outputs;
}
ggml_cgraph* build_encode_image_graph(const sd::Tensor<float>& image_tensor) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
@@ -1875,6 +2058,166 @@ namespace LLM {
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
}
ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
auto runner_ctx = get_context();
auto outputs = encode_image_outputs(&runner_ctx, image);
GGML_ASSERT(!outputs.empty());
auto combined = outputs[0];
for (size_t i = 1; i < outputs.size(); ++i) {
combined = ggml_concat(compute_ctx, combined, outputs[i], 0);
}
ggml_build_forward_expand(gf, combined);
return gf;
}
static sd::Tensor<float> process_video_block_tensor(const sd::Tensor<float>& frames,
const LLMVisionConfig& vision_params) {
GGML_ASSERT(frames.dim() == 5);
GGML_ASSERT(frames.shape()[2] == vision_params.temporal_patch_size);
GGML_ASSERT(frames.shape()[3] == vision_params.in_channels);
GGML_ASSERT(frames.shape()[4] == 1);
int64_t width = frames.shape()[0];
int64_t height = frames.shape()[1];
int64_t temporal = frames.shape()[2];
int64_t channels = frames.shape()[3];
int64_t patch = vision_params.patch_size;
int64_t merge = vision_params.spatial_merge_size;
int64_t grid_w = width / patch;
int64_t grid_h = height / patch;
int64_t feature = channels * temporal * patch * patch;
int64_t token_count = grid_h * grid_w;
sd::Tensor<float> output({feature, token_count});
int64_t token = 0;
for (int64_t block_h = 0; block_h < grid_h / merge; ++block_h) {
for (int64_t block_w = 0; block_w < grid_w / merge; ++block_w) {
for (int64_t inner_h = 0; inner_h < merge; ++inner_h) {
for (int64_t inner_w = 0; inner_w < merge; ++inner_w) {
int64_t patch_h = block_h * merge + inner_h;
int64_t patch_w = block_w * merge + inner_w;
int64_t offset = 0;
for (int64_t c = 0; c < channels; ++c) {
for (int64_t t = 0; t < temporal; ++t) {
for (int64_t y = 0; y < patch; ++y) {
for (int64_t x = 0; x < patch; ++x) {
output.index(offset++, token) =
frames.index(patch_w * patch + x,
patch_h * patch + y,
t,
c,
0);
}
}
}
}
++token;
}
}
}
}
return output;
}
ggml_cgraph* build_encode_video_block_outputs_graph(const sd::Tensor<float>& pixel_values_tensor,
int grid_h,
int grid_w) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
auto pixel_values = make_input(pixel_values_tensor);
auto runner_ctx = get_context();
auto vision = model.vision_model();
int head_dim = static_cast<int>(config.vision.hidden_size / config.vision.num_heads);
auto pos_embeds = build_patch_pos_embeds(&runner_ctx, vision, grid_h, grid_w);
window_index_vec.resize(static_cast<size_t>((grid_h / config.vision.spatial_merge_size) *
(grid_w / config.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,
config.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());
auto outputs = vision->forward_outputs(&runner_ctx,
pixel_values,
pe,
nullptr,
nullptr,
nullptr,
pos_embeds);
GGML_ASSERT(!outputs.empty());
auto combined = outputs[0];
for (size_t i = 1; i < outputs.size(); ++i) {
combined = ggml_concat(compute_ctx, combined, outputs[i], 0);
}
ggml_build_forward_expand(gf, combined);
return gf;
}
std::vector<sd::Tensor<float>> encode_image_outputs(const int n_threads,
const sd::Tensor<float>& image,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_image_outputs_graph(image);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
if (combined.empty()) {
return {};
}
size_t output_count = config.vision.deepstack_visual_indexes.size() + 1;
GGML_ASSERT(combined.shape()[0] == config.hidden_size * static_cast<int64_t>(output_count));
std::vector<sd::Tensor<float>> outputs;
outputs.reserve(output_count);
for (size_t i = 0; i < output_count; ++i) {
outputs.push_back(sd::ops::slice(combined,
0,
static_cast<int64_t>(i) * config.hidden_size,
static_cast<int64_t>(i + 1) * config.hidden_size));
}
return outputs;
}
std::vector<sd::Tensor<float>> encode_video_block_outputs(const int n_threads,
const sd::Tensor<float>& frames,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
int grid_h = static_cast<int>(frames.shape()[1] / config.vision.patch_size);
int grid_w = static_cast<int>(frames.shape()[0] / config.vision.patch_size);
auto pixel_values = process_video_block_tensor(frames, config.vision);
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph,
n_threads,
auto_free,
free_compute_buffer,
free_compute_params));
if (combined.empty()) {
return {};
}
size_t output_count = config.vision.deepstack_visual_indexes.size() + 1;
GGML_ASSERT(combined.shape()[0] == config.hidden_size * static_cast<int64_t>(output_count));
std::vector<sd::Tensor<float>> outputs;
outputs.reserve(output_count);
for (size_t i = 0; i < output_count; ++i) {
outputs.push_back(sd::ops::slice(combined,
0,
static_cast<int64_t>(i) * config.hidden_size,
static_cast<int64_t>(i + 1) * config.hidden_size));
}
return outputs;
}
};
struct LLMEmbedder {
src/model/vae/audio_vae.hpp
+28
View File
@@ -0,0 +1,28 @@
#ifndef __SD_MODEL_VAE_AUDIO_VAE_HPP__
#define __SD_MODEL_VAE_AUDIO_VAE_HPP__
#include "core/ggml_extend.hpp"
struct AudioVAERunner : public GGMLRunner {
AudioVAERunner(ggml_backend_t backend,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {}
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual size_t get_params_mem_size() = 0;
virtual std::string get_desc() = 0;
virtual sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& waveform) {
SD_UNUSED(n_threads);
SD_UNUSED(waveform);
return {};
}
virtual sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) = 0;
virtual int input_sample_rate() const {
return output_sample_rate();
}
virtual int output_sample_rate() const = 0;
};
#endif // __SD_MODEL_VAE_AUDIO_VAE_HPP__
src/model/vae/ltx_audio_vae.hpp
+11 -6
View File
@@ -8,6 +8,7 @@
#include <vector>
#include "core/ggml_extend.hpp"
#include "model/vae/audio_vae.hpp"
#include "model_loader.h"
#include "model_manager.h"
@@ -996,7 +997,7 @@ namespace LTXV {
}
};
struct LTXAudioVAERunner : public GGMLRunner {
struct LTXAudioVAERunner : public AudioVAERunner {
LTXAudioVAEConfig config;
LTXAudioVAE model;
std::string weight_prefix;
@@ -1006,7 +1007,7 @@ namespace LTXV {
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
: AudioVAERunner(backend, weight_manager),
weight_prefix(prefix),
config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
model(config) {
@@ -1017,20 +1018,20 @@ namespace LTXV {
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
size_t get_params_mem_size() {
size_t get_params_mem_size() override {
return model.get_params_mem_size();
}
std::string get_desc() {
std::string get_desc() override {
return "ltx_audio_vae";
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) {
const sd::Tensor<float>& latent_tensor) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto latent = make_input(latent_tensor);
@@ -1047,6 +1048,10 @@ namespace LTXV {
return result;
}
int output_sample_rate() const override {
return config.output_sample_rate();
}
void test(const std::string& input_path) {
auto z = sd::load_tensor_from_file_as_tensor<float>(input_path);
GGML_ASSERT(!z.empty());
src/model/vae/minimax_h3_audio_vae.hpp
+497
View File
@@ -0,0 +1,497 @@
#ifndef __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
#define __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
#include <array>
#include <string>
#include <vector>
#include "model/vae/audio_vae.hpp"
#include "model/vae/ltx_audio_vae.hpp"
namespace MiniMaxH3 {
struct AudioSnake1D : public UnaryBlock {
int64_t channels;
explicit AudioSnake1D(int64_t channels)
: channels(channels) {}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["alpha"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, channels, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto alpha = params["alpha"];
auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha));
oscillation = ggml_mul(ctx->ggml_ctx, oscillation, oscillation);
auto eps = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), 1e-9f);
return ggml_add(ctx->ggml_ctx,
x,
ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, alpha, eps)));
}
};
struct AudioEncoderResidualUnit : public GGMLBlock {
int64_t channels;
AudioEncoderResidualUnit(int64_t channels, int dilation)
: channels(channels) {
blocks["block.0"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.1"] = std::make_shared<LTXV::Conv1D>(channels,
channels,
7,
1,
3 * dilation,
dilation);
blocks["block.2"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.3"] = std::make_shared<LTXV::Conv1D>(channels, channels, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto act1 = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.0"]);
auto conv1 = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.1"]);
auto act2 = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.2"]);
auto conv2 = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.3"]);
auto h = conv2->forward(ctx, act2->forward(ctx, conv1->forward(ctx, act1->forward(ctx, x))));
if (x->ne[0] != h->ne[0]) {
int64_t pad = (x->ne[0] - h->ne[0]) / 2;
x = ggml_ext_slice(ctx->ggml_ctx, x, 0, pad, x->ne[0] - pad);
}
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct AudioEncoderBlock : public GGMLBlock {
int64_t out_channels;
AudioEncoderBlock(int64_t out_channels, int stride)
: out_channels(out_channels) {
int64_t in_channels = out_channels / 2;
blocks["block.0"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 1);
blocks["block.1"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 3);
blocks["block.2"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 9);
blocks["block.3"] = std::make_shared<AudioSnake1D>(in_channels);
blocks["block.4"] = std::make_shared<LTXV::Conv1D>(in_channels,
out_channels,
2 * stride,
stride,
static_cast<int>(std::ceil(stride / 2.f)));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
for (int i = 0; i < 3; ++i) {
auto unit = std::dynamic_pointer_cast<AudioEncoderResidualUnit>(blocks["block." + std::to_string(i)]);
x = unit->forward(ctx, x);
}
auto act = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.3"]);
auto conv = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.4"]);
return conv->forward(ctx, act->forward(ctx, x));
}
};
struct AudioEncoder : public GGMLBlock {
static constexpr std::array<int, 5> strides = {2, 4, 4, 5, 5};
AudioEncoder() {
int64_t channels = 64;
blocks["block.0"] = std::make_shared<LTXV::Conv1D>(1, channels, 7, 1, 3);
for (size_t i = 0; i < strides.size(); ++i) {
channels *= 2;
blocks["block." + std::to_string(i + 1)] = std::make_shared<AudioEncoderBlock>(channels, strides[i]);
}
blocks["block.6"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.7"] = std::make_shared<LTXV::Conv1D>(channels, 2048, 3, 1, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto input = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.0"]);
x = input->forward(ctx, x);
for (size_t i = 0; i < strides.size(); ++i) {
auto block = std::dynamic_pointer_cast<AudioEncoderBlock>(blocks["block." + std::to_string(i + 1)]);
x = block->forward(ctx, x);
}
auto act = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.6"]);
auto out = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.7"]);
return out->forward(ctx, act->forward(ctx, x));
}
};
struct AudioGeGLUMLP : public GGMLBlock {
AudioGeGLUMLP(int64_t hidden_size, int64_t intermediate_size) {
blocks["norm"] = std::make_shared<LayerNorm>(hidden_size);
blocks["w0"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
blocks["w1"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
blocks["w2"] = std::make_shared<Linear>(intermediate_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto w0 = std::dynamic_pointer_cast<Linear>(blocks["w0"]);
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
x = norm->forward(ctx, x);
auto gate = ggml_ext_gelu(ctx->ggml_ctx, w0->forward(ctx, x), true);
return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, w1->forward(ctx, x)));
}
};
struct AudioCausalAttention : public GGMLBlock {
static constexpr int64_t in_channels = 2048;
static constexpr int64_t out_channels = 32;
static constexpr int64_t num_head = 8;
static constexpr int64_t head_dim = in_channels / num_head;
AudioCausalAttention() {
blocks["qkv"] = std::make_shared<Linear>(in_channels, in_channels * 3, false);
blocks["proj"] = std::make_shared<Linear>(out_channels, out_channels, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
params["q_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
params["v_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto qkv_layer = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_layer->forward(ctx, x), 3, 0);
auto bias_shape = [&](ggml_tensor* bias) {
return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1);
};
auto q = ggml_add(ctx->ggml_ctx, qkv[0], bias_shape(params["q_bias"]));
auto k = qkv[1];
auto v = ggml_add(ctx->ggml_ctx, qkv[2], bias_shape(params["v_bias"]));
int64_t sequence = x->ne[1];
auto mask = ggml_diag_mask_inf(ctx->ggml_ctx,
ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1),
0);
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
ctx->backend,
q,
k,
v,
num_head,
mask,
false,
ctx->flash_attn_enabled);
int64_t batch = attn_out->ne[2] * attn_out->ne[3];
attn_out = ggml_reshape_4d(ctx->ggml_ctx, attn_out, head_dim, num_head, sequence, batch);
attn_out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, attn_out, 1, 0, 2, 3));
attn_out = ggml_mean(ctx->ggml_ctx, attn_out);
attn_out = ggml_reshape_3d(ctx->ggml_ctx, attn_out, head_dim, sequence, batch);
constexpr int64_t pool = head_dim / out_channels;
attn_out = ggml_reshape_4d(ctx->ggml_ctx, attn_out, pool, out_channels, sequence, batch);
attn_out = ggml_mean(ctx->ggml_ctx, attn_out);
attn_out = ggml_reshape_3d(ctx->ggml_ctx, attn_out, out_channels, sequence, batch);
return proj->forward(ctx, attn_out);
}
};
struct AudioAttentionProjection : public GGMLBlock {
AudioAttentionProjection() {
blocks["norm1"] = std::make_shared<LayerNorm>(2048);
blocks["attn"] = std::make_shared<AudioCausalAttention>();
blocks["proj"] = std::make_shared<Linear>(2048, 32, true);
blocks["norm3"] = std::make_shared<LayerNorm>(2048);
blocks["norm2"] = std::make_shared<LayerNorm>(32);
blocks["mlp"] = std::make_shared<AudioGeGLUMLP>(32, 64);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<AudioCausalAttention>(blocks["attn"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto norm3 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm3"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto mlp = std::dynamic_pointer_cast<AudioGeGLUMLP>(blocks["mlp"]);
x = ggml_add(ctx->ggml_ctx,
proj->forward(ctx, norm3->forward(ctx, x)),
attn->forward(ctx, norm1->forward(ctx, x)));
return ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, norm2->forward(ctx, x)));
}
};
struct AudioAMPBlock : public GGMLBlock {
int channels;
AudioAMPBlock(int channels,
int kernel_size,
const std::array<int, 3>& dilations)
: channels(channels) {
for (int i = 0; i < 3; ++i) {
blocks["activations." + std::to_string(i * 2)] =
std::make_shared<LTXV::Activation1D>(channels);
blocks["activations." + std::to_string(i * 2 + 1)] =
std::make_shared<LTXV::Activation1D>(channels);
blocks["convs1." + std::to_string(i)] =
std::make_shared<LTXV::Conv1D>(channels,
channels,
kernel_size,
1,
(kernel_size * dilations[i] - dilations[i]) / 2,
dilations[i]);
blocks["convs2." + std::to_string(i)] =
std::make_shared<LTXV::Conv1D>(channels,
channels,
kernel_size,
1,
kernel_size / 2);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
for (int i = 0; i < 3; ++i) {
auto act1 = std::dynamic_pointer_cast<LTXV::Activation1D>(
blocks["activations." + std::to_string(i * 2)]);
auto act2 = std::dynamic_pointer_cast<LTXV::Activation1D>(
blocks["activations." + std::to_string(i * 2 + 1)]);
auto conv1 = std::dynamic_pointer_cast<LTXV::Conv1D>(
blocks["convs1." + std::to_string(i)]);
auto conv2 = std::dynamic_pointer_cast<LTXV::Conv1D>(
blocks["convs2." + std::to_string(i)]);
auto h = conv1->forward(ctx, act1->forward(ctx, x));
h = conv2->forward(ctx, act2->forward(ctx, h));
x = ggml_add(ctx->ggml_ctx, x, h);
}
return x;
}
};
struct BigVGAN : public GGMLBlock {
static constexpr int initial_channels = 1024;
static constexpr int num_kernels = 3;
static constexpr int num_upsamples = 7;
static constexpr std::array<int, num_upsamples> rates = {5, 5, 2, 2, 2, 2, 2};
static constexpr std::array<int, num_upsamples> kernels = {9, 9, 4, 4, 4, 4, 4};
static constexpr std::array<int, num_kernels> res_kernels = {3, 7, 11};
BigVGAN() {
blocks["conv_pre"] = std::make_shared<LTXV::Conv1D>(2048,
initial_channels,
7,
1,
3);
int channels = initial_channels;
for (int i = 0; i < num_upsamples; ++i) {
int next_channels = initial_channels / (1 << (i + 1));
blocks["ups." + std::to_string(i) + ".0"] =
std::make_shared<LTXV::ConvTranspose1D>(channels,
next_channels,
kernels[i],
rates[i],
(kernels[i] - rates[i]) / 2);
for (int j = 0; j < num_kernels; ++j) {
blocks["resblocks." + std::to_string(i * num_kernels + j)] =
std::make_shared<AudioAMPBlock>(next_channels,
res_kernels[j],
std::array<int, 3>{1, 3, 5});
}
channels = next_channels;
}
blocks["activation_post"] = std::make_shared<LTXV::Activation1D>(channels);
blocks["conv_post"] = std::make_shared<LTXV::Conv1D>(channels,
1,
7,
1,
3,
1,
false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv_pre = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_pre"]);
x = conv_pre->forward(ctx, x);
for (int i = 0; i < num_upsamples; ++i) {
auto up = std::dynamic_pointer_cast<LTXV::ConvTranspose1D>(
blocks["ups." + std::to_string(i) + ".0"]);
x = up->forward(ctx, x);
ggml_tensor* sum = nullptr;
for (int j = 0; j < num_kernels; ++j) {
auto block = std::dynamic_pointer_cast<AudioAMPBlock>(
blocks["resblocks." + std::to_string(i * num_kernels + j)]);
auto value = block->forward(ctx, x);
sum = sum == nullptr ? value : ggml_add(ctx->ggml_ctx, sum, value);
}
x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.f / num_kernels);
}
auto activation = std::dynamic_pointer_cast<LTXV::Activation1D>(blocks["activation_post"]);
auto conv_post = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_post"]);
return ggml_clamp(ctx->ggml_ctx,
conv_post->forward(ctx, activation->forward(ctx, x)),
-1.f,
1.f);
}
};
struct AudioVAE : public GGMLBlock {
static constexpr int kLatentChannels = 32;
AudioVAE() {
blocks["encoder"] = std::make_shared<AudioEncoder>();
blocks["pre_block"] = std::make_shared<AudioAttentionProjection>();
blocks["mean_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels, kLatentChannels, 1);
blocks["logs_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels, kLatentChannels, 1);
blocks["dec_in_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels,
2048,
1);
blocks["decoder"] = std::make_shared<BigVGAN>();
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["latents_mean"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels);
params["latents_std"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels);
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* waveform) {
GGML_ASSERT(waveform->ne[1] == 2);
auto encoder = std::dynamic_pointer_cast<AudioEncoder>(blocks["encoder"]);
auto pre = std::dynamic_pointer_cast<AudioAttentionProjection>(blocks["pre_block"]);
auto mean_proj = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["mean_proj"]);
waveform = ggml_reshape_3d(ctx->ggml_ctx, waveform, waveform->ne[0], 1, waveform->ne[1]);
auto x = encoder->forward(ctx, waveform); // [B*S, 2048, T]
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = pre->forward(ctx, x);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
auto z = mean_proj->forward(ctx, x);
auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
auto std = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
z = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent) {
GGML_ASSERT(latent->ne[1] == 2 && latent->ne[2] == kLatentChannels);
latent = ggml_cont(ctx->ggml_ctx,
ggml_permute(ctx->ggml_ctx, latent, 0, 2, 1, 3));
auto mean = ggml_reshape_4d(ctx->ggml_ctx,
params["latents_mean"],
1,
kLatentChannels,
1,
1);
auto std = ggml_reshape_4d(ctx->ggml_ctx,
params["latents_std"],
1,
kLatentChannels,
1,
1);
latent = ggml_add(ctx->ggml_ctx,
ggml_mul(ctx->ggml_ctx, latent, std),
mean);
auto dec_in = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["dec_in_proj"]);
auto decoder = std::dynamic_pointer_cast<BigVGAN>(blocks["decoder"]);
int64_t streams = latent->ne[2] * latent->ne[3];
latent = ggml_reshape_3d(ctx->ggml_ctx,
latent,
latent->ne[0],
latent->ne[1],
streams);
ggml_tensor* waveform = nullptr;
for (int64_t stream = 0; stream < streams; ++stream) {
auto stream_latent = ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1);
auto stream_waveform = decoder->forward(ctx, dec_in->forward(ctx, stream_latent));
waveform = waveform == nullptr
? stream_waveform
: ggml_concat(ctx->ggml_ctx, waveform, stream_waveform, 2);
}
return ggml_reshape_4d(ctx->ggml_ctx,
waveform,
waveform->ne[0],
streams,
1,
1);
}
};
struct AudioVAERunner : public ::AudioVAERunner {
AudioVAE model;
std::string weight_prefix;
AudioVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: ::AudioVAERunner(backend, weight_manager),
weight_prefix(prefix) {
model.init(params_ctx, tensor_storage_map, prefix);
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
size_t get_params_mem_size() override {
return model.get_params_mem_size();
}
std::string get_desc() override {
return "minimax_h3_audio_vae";
}
int output_sample_rate() const override {
return 32000;
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& waveform) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto input = make_input(waveform);
auto runner_ctx = get_context();
auto latent = model.encode(&runner_ctx, input);
auto graph = new_graph_custom(655360);
ggml_build_forward_expand(graph, latent);
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
(t1 - t0) / 1000.f);
return result;
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto latent = make_input(latent_tensor);
auto runner_ctx = get_context();
auto waveform = model.decode(&runner_ctx, latent);
auto graph = new_graph_custom(655360);
ggml_build_forward_expand(graph, waveform);
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",
(t1 - t0) / 1000.f);
return result;
}
};
} // namespace MiniMaxH3
#endif // __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
src/model/vae/minimax_h3_vae.hpp
+805
View File
@@ -0,0 +1,805 @@
#ifndef __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
#define __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
#include <algorithm>
#include <array>
#include <cmath>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/vae/vae.hpp"
namespace MiniMaxH3VAE {
constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144;
struct CausalConv3d : public Conv3d {
std::tuple<int, int, int> temporal_padding;
CausalConv3d(int64_t in_channels,
int64_t out_channels,
std::tuple<int, int, int> kernel_size,
std::tuple<int, int, int> stride = {1, 1, 1},
std::tuple<int, int, int> padding = {0, 0, 0})
: Conv3d(in_channels,
out_channels,
kernel_size,
stride,
{0, 0, 0}),
temporal_padding(padding) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto reflect_pad = [&](ggml_tensor* value, int dim, int amount) {
for (int i = 0; i < amount; ++i) {
GGML_ASSERT(value->ne[dim] > 1);
auto left = ggml_ext_slice(ctx->ggml_ctx, value, dim, 1, 2);
auto right = ggml_ext_slice(ctx->ggml_ctx,
value,
dim,
value->ne[dim] - 2,
value->ne[dim] - 1);
value = ggml_concat(ctx->ggml_ctx, left, value, dim);
value = ggml_concat(ctx->ggml_ctx, value, right, dim);
}
return value;
};
x = reflect_pad(x, 0, std::get<2>(temporal_padding));
x = reflect_pad(x, 1, std::get<1>(temporal_padding));
int temporal_pad = std::get<0>(temporal_padding) * 2;
if (temporal_pad > 0) {
x = ggml_ext_pad_ext(ctx->ggml_ctx,
ctx->backend,
x,
0,
0,
0,
0,
temporal_pad,
0,
0,
0);
}
return Conv3d::forward(ctx, x);
}
};
struct TemporalGroupNorm : public GroupNorm {
explicit TemporalGroupNorm(int64_t channels)
: GroupNorm(32, channels, 1e-6f, true) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* result = nullptr;
for (int64_t t = 0; t < x->ne[2]; ++t) {
auto frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, t, t + 1);
GGML_ASSERT(frame->ne[3] % num_channels == 0);
int64_t batch_size = frame->ne[3] / num_channels;
frame = ggml_cont(ctx->ggml_ctx, frame);
frame = ggml_reshape_4d(ctx->ggml_ctx,
frame,
frame->ne[0],
frame->ne[1],
num_channels,
batch_size);
frame = GroupNorm::forward(ctx, frame);
frame = ggml_reshape_4d(ctx->ggml_ctx,
frame,
frame->ne[0],
frame->ne[1],
1,
num_channels * batch_size);
result = result == nullptr ? frame : ggml_concat(ctx->ggml_ctx, result, frame, 2);
}
return result;
}
};
struct Downsample3D : public GGMLBlock {
int spatial_stride;
Downsample3D(int64_t in_channels,
int64_t out_channels,
int temporal_stride,
int spatial_stride)
: spatial_stride(spatial_stride) {
blocks["conv"] = std::make_shared<CausalConv3d>(in_channels,
out_channels,
std::tuple{3, 3, 3},
std::tuple{temporal_stride, spatial_stride, spatial_stride},
std::tuple{1, 0, 0});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
if (spatial_stride == 2) {
GGML_ASSERT(x->ne[0] > 1 && x->ne[1] > 1);
auto right = ggml_ext_slice(ctx->ggml_ctx, x, 0, x->ne[0] - 2, x->ne[0] - 1);
x = ggml_concat(ctx->ggml_ctx, x, right, 0);
auto bottom = ggml_ext_slice(ctx->ggml_ctx, x, 1, x->ne[1] - 2, x->ne[1] - 1);
x = ggml_concat(ctx->ggml_ctx, x, bottom, 1);
}
return std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"])->forward(ctx, x);
}
};
struct ResnetBlock3D : public GGMLBlock {
int64_t in_channels;
int64_t out_channels;
ResnetBlock3D(int64_t in_channels,
int64_t out_channels)
: in_channels(in_channels), out_channels(out_channels) {
blocks["norm1"] = std::make_shared<TemporalGroupNorm>(in_channels);
blocks["norm2"] = std::make_shared<TemporalGroupNorm>(out_channels);
blocks["conv1"] = std::make_shared<CausalConv3d>(in_channels,
out_channels,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
blocks["conv2"] = std::make_shared<CausalConv3d>(out_channels,
out_channels,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
if (in_channels != out_channels) {
blocks["nin_shortcut"] = std::make_shared<CausalConv3d>(in_channels,
out_channels,
std::tuple{1, 1, 1});
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm2"]);
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
auto h = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, x)));
h = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, h)));
if (in_channels != out_channels) {
x = std::dynamic_pointer_cast<CausalConv3d>(blocks["nin_shortcut"])->forward(ctx, x);
}
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct Encoder : public GGMLBlock {
static constexpr int levels = 6;
static constexpr std::array<int, levels> multipliers = {1, 2, 2, 4, 4, 8};
static constexpr std::array<int, levels> spatial_down = {2, 2, 2, 2, 1, 1};
static constexpr std::array<int, levels> temporal_down = {1, 2, 2, 1, 1, 1};
Encoder() {
constexpr int ch = 128;
blocks["conv_in"] = std::make_shared<CausalConv3d>(3,
ch,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
int64_t previous = ch;
for (int level = 0; level < levels; ++level) {
int64_t current = ch * multipliers[level];
for (int block = 0; block < 2; ++block) {
blocks["down." + std::to_string(level) + ".block." + std::to_string(block)] =
std::make_shared<ResnetBlock3D>(block == 0 ? previous : current,
current);
}
if (spatial_down[level] * temporal_down[level] > 1) {
blocks["down." + std::to_string(level) + ".downsample"] =
std::make_shared<Downsample3D>(current,
current,
temporal_down[level],
spatial_down[level]);
}
previous = current;
}
blocks["norm_out"] = std::make_shared<TemporalGroupNorm>(previous);
blocks["conv_out"] = std::make_shared<CausalConv3d>(previous,
48,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
x = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"])->forward(ctx, x);
for (int level = 0; level < levels; ++level) {
for (int block = 0; block < 2; ++block) {
x = std::dynamic_pointer_cast<ResnetBlock3D>(
blocks["down." + std::to_string(level) + ".block." + std::to_string(block)])
->forward(ctx, x);
}
auto downsample = blocks.find("down." + std::to_string(level) + ".downsample");
if (downsample != blocks.end()) {
x = std::dynamic_pointer_cast<Downsample3D>(downsample->second)->forward(ctx, x);
}
}
auto norm = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm_out"]);
auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
return conv->forward(ctx, ggml_silu(ctx->ggml_ctx, norm->forward(ctx, x)));
}
};
static ggml_tensor* attention_layout(ggml_context* ctx, ggml_tensor* x) {
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));
return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
}
static ggml_tensor* apply_partial_rope(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
int64_t rot_dim = pe->ne[2] * 2;
auto rotated = Rope::apply_rope(ctx,
ggml_ext_slice(ctx, x, 0, 0, rot_dim),
pe,
false);
if (rot_dim == x->ne[0]) {
return rotated;
}
auto tail = attention_layout(ctx,
ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0]));
return ggml_concat(ctx, rotated, tail, 0);
}
struct DecoderAttention : public GGMLBlock {
static constexpr int num_head = 32;
static constexpr int head_dim = 64;
static constexpr int dim = num_head * head_dim;
DecoderAttention() {
blocks["to_qkv"] = std::make_shared<Linear>(dim, dim * 3, true);
blocks["to_out"] = std::make_shared<Linear>(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto to_qkv = std::dynamic_pointer_cast<Linear>(blocks["to_qkv"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
auto qkv_projection = to_qkv->forward(ctx, x);
int64_t sequence = x->ne[1];
int64_t batch_size = x->ne[2] * x->ne[3];
qkv_projection = ggml_reshape_4d(ctx->ggml_ctx,
qkv_projection,
3 * head_dim,
num_head,
sequence,
batch_size);
auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0);
auto q = ggml_reshape_4d(ctx->ggml_ctx,
qkv[0],
head_dim,
num_head,
sequence,
batch_size);
auto k = ggml_reshape_4d(ctx->ggml_ctx,
qkv[1],
head_dim,
num_head,
sequence,
batch_size);
auto v = ggml_reshape_4d(ctx->ggml_ctx,
qkv[2],
head_dim,
num_head,
sequence,
batch_size);
q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f);
k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
q = apply_partial_rope(ctx->ggml_ctx, q, pe);
k = apply_partial_rope(ctx->ggml_ctx, k, pe);
auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
ctx->backend,
q,
k,
v,
num_head,
nullptr,
true,
ctx->flash_attn_enabled);
return to_out->forward(ctx, out);
}
};
struct DecoderFeedForward : public GGMLBlock {
static constexpr int dim = 2048;
static constexpr int kInnerDim = dim * 4;
DecoderFeedForward() {
blocks["w1"] = std::make_shared<Linear>(dim, kInnerDim * 2, true);
blocks["w2"] = std::make_shared<Linear>(kInnerDim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto gate = ggml_ext_chunk(ctx->ggml_ctx, w1->forward(ctx, x), 2, 0);
return w2->forward(ctx,
ggml_mul(ctx->ggml_ctx,
ggml_silu(ctx->ggml_ctx, gate[0]),
gate[1]));
}
};
struct DecoderBlock : public GGMLBlock {
static constexpr int dim = 2048;
DecoderBlock() {
blocks["norm1"] = std::make_shared<RMSNorm>(dim, 1e-5f);
blocks["attn"] = std::make_shared<DecoderAttention>();
blocks["norm2"] = std::make_shared<RMSNorm>(dim, 1e-5f);
blocks["ff"] = std::make_shared<DecoderFeedForward>();
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["scale1"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
params["scale2"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<DecoderAttention>(blocks["attn"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto ff = std::dynamic_pointer_cast<DecoderFeedForward>(blocks["ff"]);
x = ggml_add(ctx->ggml_ctx,
x,
ggml_mul(ctx->ggml_ctx,
attn->forward(ctx, norm1->forward(ctx, x), pe),
params["scale1"]));
return ggml_add(ctx->ggml_ctx,
x,
ggml_mul(ctx->ggml_ctx,
ff->forward(ctx, norm2->forward(ctx, x)),
params["scale2"]));
}
};
struct Decoder : public GGMLBlock {
static constexpr int dim = 2048;
static constexpr int num_layers = 36;
static constexpr int num_register_tokens = 4;
static constexpr int patch_size = 16;
static constexpr int patch_size_t = 4;
Decoder() {
blocks["x_embedder"] = std::make_shared<Linear>(24, dim, true);
for (int i = 0; i < num_layers; ++i) {
blocks["transformer_blocks." + std::to_string(i)] =
std::make_shared<DecoderBlock>();
}
blocks["norm_out"] = std::make_shared<LayerNorm>(dim, 1e-5f, true, true);
blocks["proj_out"] = std::make_shared<Linear>(dim,
3 * patch_size_t * patch_size * patch_size,
true,
true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["register_tokens"] = ggml_new_tensor_2d(ctx,
GGML_TYPE_F32,
dim,
num_register_tokens);
params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* z,
ggml_tensor* pe) {
int64_t width = z->ne[0];
int64_t height = z->ne[1];
int64_t num_frames = z->ne[2];
int64_t batch_size = z->ne[3] / 24;
GGML_ASSERT(batch_size == 1);
z = ggml_cont(ctx->ggml_ctx,
ggml_ext_torch_permute(ctx->ggml_ctx, z, 3, 0, 1, 2));
z = ggml_reshape_3d(ctx->ggml_ctx,
z,
24,
width * height * num_frames,
batch_size);
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto h = x_embedder->forward(ctx, z);
int64_t num_patches = h->ne[1];
h = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1);
auto zero = ggml_ext_scale(ctx->ggml_ctx,
ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, 1),
0.f);
h = ggml_concat(ctx->ggml_ctx, h, zero, 1);
for (int i = 0; i < num_layers; ++i) {
auto block = std::dynamic_pointer_cast<DecoderBlock>(
blocks["transformer_blocks." + std::to_string(i)]);
h = block->forward(ctx, h, pe);
sd::ggml_graph_cut::mark_graph_cut(h,
"minimax_h3_vae.decoder.blocks." + std::to_string(i),
"hidden_states");
}
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
h = proj_out->forward(ctx, norm_out->forward(ctx, h));
h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, num_patches);
return DiT::unpatchify_3d(ctx->ggml_ctx,
h,
num_frames,
height,
width,
patch_size_t,
patch_size,
patch_size,
true);
}
};
struct MiniMaxH3VideoVAE : public GGMLBlock {
MiniMaxH3VideoVAE() {
blocks["encoder"] = std::make_shared<Encoder>();
blocks["quant_conv"] = std::make_shared<Conv3d>(48,
48,
std::tuple{1, 1, 1});
blocks["post_quant_conv"] = std::make_shared<Conv3d>(24,
24,
std::tuple{1, 1, 1});
blocks["decoder"] = std::make_shared<Decoder>();
}
ggml_tensor* encode(GGMLRunnerContext* ctx,
ggml_tensor* pixels,
ggml_tensor* pixel_mean,
ggml_tensor* pixel_std) {
pixels = ggml_div(ctx->ggml_ctx,
ggml_sub(ctx->ggml_ctx, pixels, pixel_mean),
pixel_std);
auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
auto quant = std::dynamic_pointer_cast<Conv3d>(blocks["quant_conv"]);
auto moments = quant->forward(ctx, encoder->forward(ctx, pixels));
return ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24);
}
ggml_tensor* decode(GGMLRunnerContext* ctx,
ggml_tensor* latent,
ggml_tensor* pe,
ggml_tensor* pixel_mean,
ggml_tensor* pixel_std) {
auto post_quant = std::dynamic_pointer_cast<Conv3d>(blocks["post_quant_conv"]);
auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
auto pixels = decoder->forward(ctx, post_quant->forward(ctx, latent), pe);
pixels = ggml_add(ctx->ggml_ctx,
ggml_mul(ctx->ggml_ctx, pixels, pixel_std),
pixel_mean);
return ggml_clamp(ctx->ggml_ctx, pixels, 0.f, 1.f);
}
};
struct MiniMaxH3VideoVAERunner : public VAE {
MiniMaxH3VideoVAE model;
sd::Tensor<float> pixel_mean;
sd::Tensor<float> pixel_std;
sd::Tensor<float> latents_mean;
sd::Tensor<float> latents_std;
sd::Tensor<float> rope_cache;
MiniMaxH3VideoVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "first_stage_model",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(VERSION_MINIMAX_H3, backend, prefix, weight_manager),
pixel_mean({1, 1, 1, 3}, {0.485f, 0.456f, 0.406f}),
pixel_std({1, 1, 1, 3}, {0.229f, 0.224f, 0.225f}),
latents_mean({1, 1, 1, 24},
{0.858090341091156f, -0.960659146308899f, 1.066164016723633f, -0.509032547473907f,
-0.272758185863495f, -1.367541432380676f, -0.255325496196747f, -0.269075542688370f,
-0.537684082984924f, -0.046409729868174f, 0.665737032890320f, 0.196901276707649f,
-0.546060800552368f, -0.403534203767776f, -0.236830249428749f, 0.259284526109695f,
-0.301339447498322f, 0.211341992020607f, -1.120684862136841f, 0.358193337917328f,
-0.042251437902451f, 0.260482996702194f, 0.228640928864479f, 0.705603182315826f}),
latents_std({1, 1, 1, 24},
{1.222377419471741f, 1.276726365089417f, 1.683177471160889f, 1.754945516586304f,
1.563621640205383f, 2.194143533706665f, 0.965313792228699f, 1.056988596916199f,
0.841948926448822f, 0.772995293140411f, 1.895593762397766f, 0.946841835975647f,
0.799680948257446f, 0.449889004230499f, 0.719739973545075f, 0.693629324436188f,
2.961095094680786f, 2.769419908523560f, 3.049618482589722f, 2.108805418014527f,
3.276226282119751f, 3.162735700607300f, 2.281681299209595f, 2.612784385681153f}) {
scale_input = false;
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "minimax_h3_video_vae";
}
int get_encoder_output_channels(int input_channels) override {
SD_UNUSED(input_channels);
return 24;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output,
std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents * latents_std + latents_mean;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return (latents - latents_mean) / latents_std;
}
static sd::Tensor<float> ensure_video_shape(const sd::Tensor<float>& tensor) {
if (tensor.dim() == 5) {
return tensor;
}
GGML_ASSERT(tensor.dim() == 4);
return tensor.reshape({tensor.shape()[0],
tensor.shape()[1],
1,
tensor.shape()[2],
tensor.shape()[3]});
}
static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) {
params.enabled = true;
params.tile_size_x = 16;
params.tile_size_y = 16;
params.target_overlap = 0.25f;
return params;
}
static sd::Tensor<float> repeat_last_frame(const sd::Tensor<float>& input,
int64_t count) {
auto result = input;
auto last = sd::ops::slice(input, 2, input.shape()[2] - 1, input.shape()[2]);
for (int64_t i = 0; i < count; ++i) {
result = sd::ops::concat(result, last, 2);
}
return result;
}
static sd::Tensor<float> blend_temporal(const sd::Tensor<float>& previous,
const sd::Tensor<float>& current,
int64_t extent) {
auto output = current;
extent = std::min({extent, previous.shape()[2], current.shape()[2]});
int64_t previous_start = previous.shape()[2] - extent;
for (int64_t b = 0; b < current.shape()[4]; ++b) {
for (int64_t c = 0; c < current.shape()[3]; ++c) {
for (int64_t t = 0; t < extent; ++t) {
float wb = static_cast<float>(t) / extent;
float wa = 1.f - wb;
for (int64_t h = 0; h < current.shape()[1]; ++h) {
for (int64_t w = 0; w < current.shape()[0]; ++w) {
output.index(w, h, t, c, b) =
previous.index(w, h, previous_start + t, c, b) * wa +
current.index(w, h, t, c, b) * wb;
}
}
}
}
}
return output;
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) override {
auto input = ensure_video_shape(x);
auto tiling = h3_tiling(tiling_params);
if (input.shape()[2] == 1) {
auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y);
if (!encoded.empty() && encoded.shape()[2] > 1) {
encoded = sd::ops::slice(encoded,
2,
encoded.shape()[2] - 1,
encoded.shape()[2]);
}
return encoded;
}
int64_t pad = (-input.shape()[2]) % 17;
if (pad < 0) {
pad += 17;
}
if (pad > 0) {
input = repeat_last_frame(input, pad);
}
sd::Tensor<float> result;
for (int64_t start = 0; start < input.shape()[2]; start += 17) {
auto chunk = sd::ops::slice(input, 2, start, start + 17);
auto encoded = VAE::encode(n_threads, chunk, tiling, circular_x, circular_y);
if (encoded.empty()) {
return {};
}
result = result.empty() ? std::move(encoded)
: sd::ops::concat(result, encoded, 2);
}
if (result.shape()[2] > 3) {
result = sd::ops::slice(result, 2, 0, result.shape()[2] - 3);
}
return result;
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) override {
auto input = ensure_video_shape(x);
auto tiling = h3_tiling(tiling_params);
if (input.shape()[2] == 1) {
auto decoded = VAE::decode(n_threads,
input,
tiling,
decode_video,
circular_x,
circular_y,
silent);
if (!decoded.empty() && decoded.shape()[2] > 1) {
decoded = sd::ops::slice(decoded,
2,
decoded.shape()[2] - 1,
decoded.shape()[2]);
}
return decoded;
}
constexpr int64_t tokens_per_chunk = 5;
constexpr int64_t token_drop = 3;
constexpr int64_t token_overlap = 2;
constexpr int64_t frames_per_chunk = 20;
constexpr int64_t frame_pre_padding = 3;
constexpr int64_t frame_overlap = 5;
int64_t pseudo_tokens = input.shape()[2] + token_drop;
int64_t pad_tokens = (tokens_per_chunk - pseudo_tokens % tokens_per_chunk) % tokens_per_chunk;
pseudo_tokens += pad_tokens;
int64_t num_chunks = pseudo_tokens / tokens_per_chunk - 1;
if (num_chunks < 1) {
pad_tokens += tokens_per_chunk;
num_chunks += 1;
}
if (pad_tokens > 0) {
input = repeat_last_frame(input, pad_tokens);
}
sd::Tensor<float> result;
sd::Tensor<float> overlap;
for (int64_t i = 0; i < num_chunks; ++i) {
int64_t start = i * tokens_per_chunk;
int64_t end = std::min(start + tokens_per_chunk + token_overlap,
input.shape()[2]);
auto chunk = sd::ops::slice(input, 2, start, end);
auto decoded = VAE::decode(n_threads,
chunk,
tiling,
true,
circular_x,
circular_y,
silent);
if (decoded.empty()) {
return {};
}
int64_t first_end = std::min<int64_t>(frames_per_chunk, decoded.shape()[2]);
auto first = sd::ops::slice(decoded,
2,
std::min<int64_t>(frame_pre_padding, first_end),
first_end);
if (!overlap.empty()) {
first = blend_temporal(overlap, first, frame_overlap);
overlap = {};
}
result = result.empty() ? std::move(first)
: sd::ops::concat(result, first, 2);
if (decoded.shape()[2] > frames_per_chunk + frame_pre_padding) {
overlap = sd::ops::slice(decoded,
2,
frames_per_chunk + frame_pre_padding,
decoded.shape()[2]);
}
if (i == num_chunks - 1 && !overlap.empty()) {
result = sd::ops::concat(result, overlap, 2);
overlap = {};
}
}
int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5;
expected_frames = std::max<int64_t>(1, expected_frames);
if (result.shape()[2] > expected_frames) {
result = sd::ops::slice(result, 2, 0, expected_frames);
}
return result;
}
sd::Tensor<float> build_rope(int64_t width,
int64_t height,
int64_t num_frames) {
std::vector<std::vector<float>> ids;
ids.reserve(static_cast<size_t>(width * height * num_frames + 5));
constexpr float two_pi = 6.28318530717958647692f;
for (int64_t t = 0; t < num_frames; ++t) {
float pt = (2.f * ((t + 0.5f) / num_frames) - 1.f) * two_pi;
for (int64_t h = 0; h < height; ++h) {
float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi;
for (int64_t w = 0; w < width; ++w) {
float pw = (2.f * ((w + 0.5f) / width) - 1.f) * two_pi;
ids.push_back({pt, ph, pw});
}
}
}
for (int i = 0; i < 5; ++i) {
ids.push_back({0.f, 0.f, 0.f});
}
auto values = Rope::embed_nd(ids,
1,
100.f,
std::vector<int>{16, 16, 16});
return sd::Tensor<float>({2,
2,
24,
static_cast<int64_t>(ids.size())},
std::move(values));
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z,
bool decode_graph) override {
auto input = ensure_video_shape(z);
if (decode_graph) {
rope_cache = build_rope(input.shape()[0],
input.shape()[1],
input.shape()[2]);
}
auto get_graph = [&]() -> ggml_cgraph* {
auto value = make_input(input);
auto mean = make_input(pixel_mean);
auto std = make_input(pixel_std);
auto runner_ctx = get_context();
ggml_tensor* out = nullptr;
if (decode_graph) {
auto pe = make_input(rope_cache);
out = model.decode(&runner_ctx, value, pe, mean, std);
} else {
out = model.encode(&runner_ctx, value, mean, std);
}
auto graph = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE);
ggml_build_forward_expand(graph, out);
return graph;
};
return restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph,
n_threads,
false,
false,
false),
5);
}
};
} // namespace MiniMaxH3VAE
#endif // __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
src/model/vae/vae.hpp
+15 -14
View File
@@ -74,7 +74,7 @@ public:
int scale_factor = 8;
if (version == VERSION_LTXAV) {
scale_factor = 32;
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16;
@@ -115,11 +115,11 @@ public:
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) {
virtual sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) {
int64_t t0 = ggml_time_ms();
sd::Tensor<float> input = x;
sd::Tensor<float> output;
@@ -136,7 +136,8 @@ public:
// Image VAE encode is more sensitive to tile boundary context than decode.
// Keep the smaller legacy factor for video VAEs, but default image encode
// tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f
: 2.0f;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
output = tiled_compute(input,
@@ -166,13 +167,13 @@ public:
return std::move(output);
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) {
virtual sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) {
int64_t t0 = ggml_time_ms();
sd::Tensor<float> input = x;
sd::Tensor<float> output;