DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-14 12:56:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support Sol-H3-VAE decoding for MiniMax-H3 (#2317)

Browse Source
This commit is contained in:
losewayy authored and GitHub committed 2026-11-10 23:29:17 +08:00
1 parent e1ca5e8bc7
commit 0bb132831f
5 files changed
+485 -72
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/vae/minimax_h3_sol_attention.hpp
+204
View File
@@ -0,0 +1,204 @@
#ifndef __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__
#define __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__
#include <algorithm>
#include <array>
#include <map>
#include <utility>
#include <vector>
#include "core/ggml_extend.h"
#include "core/tensor.hpp"
namespace MiniMaxH3VAE {
constexpr int64_t sol_window_size = 16;
constexpr int64_t sol_block_size = 5;
constexpr int64_t sol_global_rows = 5;
constexpr int64_t sol_spatial_shift = sol_window_size / 2;
constexpr int64_t sol_temporal_shift = 2;
struct SolAttentionPlan {
template <typename Tensor>
struct Batch {
Tensor output_rows;
Tensor queries;
Tensor keys;
Tensor previous_queries;
Tensor previous_keys;
};
template <typename Tensor>
struct Block {
Tensor rows;
std::vector<Batch<Tensor>> batches;
};
std::vector<Block<sd::Tensor<int32_t>>> blocks;
SolAttentionPlan(int64_t width, int64_t height, int64_t frames, bool shifted) {
const int64_t frame_size = width * height;
const int64_t pixels = frames * frame_size;
const int64_t block_rows = sol_block_size * frame_size;
const int64_t count = (frames + sol_block_size - 1) / sol_block_size;
const int64_t shift = shifted ? sol_spatial_shift : 0;
const int64_t time_shift = shifted ? sol_temporal_shift : 0;
auto pack = [](std::vector<int32_t> rows) {
const auto size = static_cast<int64_t>(rows.size());
return sd::Tensor<int32_t>({size}, std::move(rows));
};
for (int64_t block = 0; block < count; ++block) {
std::vector<int32_t> rows;
const int64_t start = block * block_rows;
const int64_t end = std::min(pixels, start + block_rows);
for (int64_t i = start; i < end; ++i) {
rows.push_back(static_cast<int32_t>(i));
}
std::vector<int32_t> globals;
std::vector<int32_t> queries;
for (int64_t i = 0; i < sol_global_rows; ++i) {
globals.push_back(static_cast<int32_t>(pixels + block * sol_global_rows + i));
queries.push_back(static_cast<int32_t>(rows.size()));
rows.push_back(globals.back());
}
blocks.push_back({pack(std::move(rows)), {}});
blocks.back().batches.push_back({pack(std::move(globals)), pack(std::move(queries)), {}, {}, {}});
}
struct Window {
std::vector<int32_t> output_rows;
std::vector<int32_t> queries;
std::vector<int32_t> keys;
std::vector<int32_t> previous_queries;
std::vector<int32_t> previous_keys;
};
for (int64_t t = -time_shift; t < frames; t += sol_block_size) {
const int64_t first = std::max<int64_t>(0, t);
const int64_t end = std::min(frames, t + sol_block_size);
const int64_t first_block = first / sol_block_size;
const int64_t last_block = (end - 1) / sol_block_size;
std::map<int64_t, std::vector<Window>> groups;
for (int64_t y = -shift; y < height; y += sol_window_size) {
for (int64_t x = -shift; x < width; x += sol_window_size) {
Window window;
for (int64_t i = first; i < end; ++i) {
for (int64_t j = std::max<int64_t>(0, y); j < std::min(height, y + sol_window_size); ++j) {
for (int64_t k = std::max<int64_t>(0, x); k < std::min(width, x + sol_window_size); ++k) {
const int64_t row = (i * height + j) * width + k;
window.output_rows.push_back(static_cast<int32_t>(row));
const int64_t block = i / sol_block_size;
auto& queries = block == last_block ? window.queries : window.previous_queries;
queries.push_back(static_cast<int32_t>(row - block * block_rows));
}
}
}
window.keys = window.queries;
window.previous_keys = window.previous_queries;
for (int64_t block = first_block; block <= last_block; ++block) {
const int64_t start = std::min(block_rows, pixels - block * block_rows);
auto& keys = block == last_block ? window.keys : window.previous_keys;
for (int64_t i = 0; i < sol_global_rows; ++i) {
keys.push_back(static_cast<int32_t>(start + i));
}
}
groups[window.output_rows.size()].push_back(std::move(window));
}
}
// Bound the workspace of backends that fall back to unfused attention.
constexpr size_t max_windows = 4;
for (const auto& [query_count, windows] : groups) {
for (size_t first_window = 0; first_window < windows.size(); first_window += max_windows) {
const auto n = static_cast<int64_t>(std::min(max_windows, windows.size() - first_window));
auto pack_windows = [&](auto member) -> sd::Tensor<int32_t> {
const auto rows = static_cast<int64_t>((windows[first_window].*member).size());
if (rows == 0) {
return {};
}
sd::Tensor<int32_t> indices({rows, n});
for (int64_t i = 0; i < n; ++i) {
const auto& source = windows[first_window + i].*member;
std::copy(source.begin(), source.end(), indices.data() + i * rows);
}
return indices;
};
blocks[last_block].batches.push_back({pack_windows(&Window::output_rows),
pack_windows(&Window::queries),
pack_windows(&Window::keys),
pack_windows(&Window::previous_queries),
pack_windows(&Window::previous_keys)});
}
}
}
}
struct Inputs {
std::vector<Block<ggml_tensor*>> blocks;
ggml_tensor* row_ids = nullptr;
template <typename Project>
ggml_tensor* forward(GGMLRunnerContext* runner_ctx,
ggml_tensor* x,
ggml_tensor* pe,
int64_t heads,
Project project) const {
auto ctx = runner_ctx->ggml_ctx;
const int64_t dim = x->ne[0];
auto output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, x->ne[1]);
auto positions = ggml_reshape_2d(ctx, pe, pe->ne[0] * pe->ne[1] * pe->ne[2], pe->ne[3]);
auto gather = [&](ggml_tensor* rows, ggml_tensor* indices) {
auto flat = ggml_reshape_1d(ctx, indices, ggml_nelements(indices));
return ggml_reshape_3d(ctx, ggml_get_rows(ctx, rows, flat), dim, indices->ne[0], indices->ne[1]);
};
std::array<ggml_tensor*, 3> previous{};
for (const auto& block : blocks) {
auto local_pe = ggml_reshape_4d(ctx, ggml_get_rows(ctx, positions, block.rows),
pe->ne[0], pe->ne[1], pe->ne[2], block.rows->ne[0]);
auto current = project(ggml_get_rows(ctx, x, block.rows), local_pe);
for (const auto& batch : block.batches) {
// Gather first so a crossing window never copies either full block.
auto window = [&](size_t i, ggml_tensor* now, ggml_tensor* before) {
// Global queries read the whole native block in its existing order.
auto rows = now == nullptr ? current[i] : gather(current[i], now);
if (before != nullptr) {
rows = ggml_concat(ctx, gather(previous[i], before), rows, 1);
}
return rows;
};
auto attended = ggml_ext_attention_ext(runner_ctx,
window(0, batch.queries, batch.previous_queries),
window(1, batch.keys, batch.previous_keys),
window(2, batch.keys, batch.previous_keys),
heads, nullptr, false, runner_ctx->flash_attn_enabled);
auto indices = ggml_reshape_1d(ctx, batch.output_rows, ggml_nelements(batch.output_rows));
output = ggml_set_rows(ctx, output, ggml_reshape_2d(ctx, attended, dim, indices->ne[0]), indices);
}
// Later windows can only read this block and its immediate successor.
previous = current;
}
return output;
}
};
template <typename BindInput>
Inputs bind(BindInput bind_input) const {
Inputs inputs;
for (const auto& block : blocks) {
Block<ggml_tensor*> bound{bind_input(block.rows), {}};
for (const auto& batch : block.batches) {
auto bind_optional = [&](const auto& indices) {
return indices.empty() ? nullptr : bind_input(indices);
};
bound.batches.push_back({bind_input(batch.output_rows), bind_input(batch.queries),
bind_optional(batch.keys), bind_optional(batch.previous_queries),
bind_optional(batch.previous_keys)});
}
inputs.blocks.push_back(std::move(bound));
}
return inputs;
}
};
} // namespace MiniMaxH3VAE
#endif // __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__
src/model/vae/minimax_h3_vae.hpp
+239 -61
View File
@@ -13,26 +13,34 @@
#include "core/rng_mt19937.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/vae/minimax_h3_sol_attention.hpp"
#include "model/vae/vae.hpp"
namespace MiniMaxH3VAE {
constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144;
static int64_t sol_output_frames(int64_t latent_frames, int patch_size_t) {
const int64_t blocks = (latent_frames + sol_block_size - 1) / sol_block_size;
return latent_frames * patch_size_t - blocks * (patch_size_t - 1);
}
struct MiniMaxH3VAEConfig {
int num_layers = 36;
int patch_size_t = 4;
bool sol = false;
static MiniMaxH3VAEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) {
MiniMaxH3VAEConfig config;
auto proj = tensors.find(prefix + ".decoder.proj_out.weight");
config.sol = tensors.count(prefix + ".decoder.blocks.0.scale1") != 0;
auto proj = tensors.find(prefix + ".decoder.proj_out.weight");
if (proj != tensors.end()) {
config.patch_size_t = static_cast<int>(proj->second.ne[1] / (3 * 16 * 16));
if (config.patch_size_t != 1 && config.patch_size_t != 4) {
throw std::runtime_error("unsupported MiniMax-H3 VAE temporal patch size");
}
int layers = 0;
const auto layer_prefix = prefix + ".decoder.transformer_blocks.";
const auto layer_prefix = prefix + (config.sol ? ".decoder.blocks." : ".decoder.transformer_blocks.");
for (const auto& [name, _] : tensors) {
if (starts_with(name, layer_prefix)) {
layers = std::max(layers, std::atoi(name.c_str() + layer_prefix.size()) + 1);
@@ -41,7 +49,12 @@ namespace MiniMaxH3VAE {
if (layers > 0) {
config.num_layers = layers;
}
LOG_VERBOSE("minimax_h3_vae: layers=%d, patch_size_t=%d", config.num_layers, config.patch_size_t);
}
if (config.sol || proj != tensors.end()) {
LOG_VERBOSE("minimax_h3_vae: decoder=%s, layers=%d, patch_size_t=%d",
config.sol ? "sol" : "original",
config.num_layers,
config.patch_size_t);
}
return config;
}
@@ -283,11 +296,10 @@ namespace MiniMaxH3VAE {
blocks["to_out"] = std::make_shared<Linear>(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
std::array<ggml_tensor*, 3> project(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];
@@ -320,6 +332,29 @@ namespace MiniMaxH3VAE {
k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
q = apply_partial_rope(ctx, q, pe);
k = apply_partial_rope(ctx, k, pe);
return {q, k, v};
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
RMSNorm& norm,
const SolAttentionPlan::Inputs* sol) {
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
if (sol != nullptr) {
auto project_rows = [&](ggml_tensor* rows, ggml_tensor* local_pe) {
auto [q, k, v] = project(ctx, norm.forward(ctx, rows), local_pe);
auto as_rows = [&](ggml_tensor* tensor) {
return ggml_reshape_2d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, tensor), dim, rows->ne[1]);
};
return std::array<ggml_tensor*, 3>{
as_rows(ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)),
as_rows(ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)),
as_rows(v)};
};
return to_out->forward(ctx, sol->forward(ctx, x, pe, num_head, project_rows));
}
auto [q, k, v] = project(ctx, norm.forward(ctx, x), pe);
auto out = ggml_ext_attention_ext(ctx,
q,
k,
@@ -370,21 +405,33 @@ namespace MiniMaxH3VAE {
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
ggml_tensor* pe,
const SolAttentionPlan::Inputs* sol,
int64_t rows_per_chunk) {
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"]));
x = ggml_add(ctx->ggml_ctx,
x,
ggml_mul(ctx->ggml_ctx,
attn->forward(ctx, x, pe, *norm1, sol),
params["scale1"]));
auto finish = [&](ggml_tensor* rows) {
auto update = ff->forward(ctx, norm2->forward(ctx, rows));
return ggml_add(ctx->ggml_ctx, rows, ggml_mul(ctx->ggml_ctx, update, params["scale2"]));
};
if (rows_per_chunk == 0 || x->ne[1] <= rows_per_chunk) {
return finish(x);
}
// Normalize and update each independent chunk before writing it to the layer output.
auto output = ggml_new_tensor_2d(ctx->ggml_ctx, x->type, x->ne[0], x->ne[1]);
for (int64_t row = 0; row < x->ne[1]; row += rows_per_chunk) {
auto input = ggml_ext_slice(ctx, x, 1, row, std::min(row + rows_per_chunk, x->ne[1]), false);
auto indices = ggml_ext_slice(ctx, sol->row_ids, 0, row, std::min(row + rows_per_chunk, x->ne[1]), false);
output = ggml_set_rows(ctx->ggml_ctx, output, finish(input), indices);
}
return output;
}
};
@@ -394,12 +441,13 @@ namespace MiniMaxH3VAE {
static constexpr int num_register_tokens = 4;
static constexpr int patch_size = 16;
int patch_size_t;
bool sol;
explicit Decoder(const MiniMaxH3VAEConfig& config)
: num_layers(config.num_layers), patch_size_t(config.patch_size_t) {
: num_layers(config.num_layers), patch_size_t(config.patch_size_t), sol(config.sol) {
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)] =
blocks[(sol ? "blocks." : "transformer_blocks.") + std::to_string(i)] =
std::make_shared<DecoderBlock>();
}
blocks["norm_out"] = std::make_shared<LayerNorm>(dim, 1e-5f, true, true);
@@ -418,14 +466,18 @@ namespace MiniMaxH3VAE {
GGML_TYPE_F32,
dim,
num_register_tokens);
if (patch_size_t != 1) {
if (patch_size_t != 1 && !sol) {
params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* z,
ggml_tensor* pe) {
ggml_tensor* pe,
ggml_tensor* shifted_pe,
const std::array<SolAttentionPlan::Inputs, 2>& sol_inputs,
ggml_tensor* pixel_mean,
ggml_tensor* pixel_std) {
int64_t width = z->ne[0];
int64_t height = z->ne[1];
int64_t num_frames = z->ne[2];
@@ -442,34 +494,68 @@ namespace MiniMaxH3VAE {
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_ext_slice(ctx, h, 1, 0, 1),
0.f);
h = ggml_concat(ctx->ggml_ctx, h, zero, 1);
auto registers = ggml_concat(ctx->ggml_ctx, params["register_tokens"], zero, 1);
if (sol) {
const int64_t num_blocks = (num_frames + sol_block_size - 1) / sol_block_size;
registers = ggml_repeat(ctx->ggml_ctx, registers,
ggml_new_tensor_2d(ctx->ggml_ctx, registers->type, dim, num_blocks * sol_global_rows));
}
h = ggml_concat(ctx->ggml_ctx, h, registers, 1);
const int64_t rows_per_chunk = sol ? width * height + sol_global_rows : 0;
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);
blocks[(sol ? "blocks." : "transformer_blocks.") + std::to_string(i)]);
h = block->forward(ctx, h, sol && i % 2 != 0 ? shifted_pe : pe, sol ? &sol_inputs[i % 2] : nullptr, rows_per_chunk);
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, h, 1, 0, num_patches);
return DiT::unpatchify_3d(ctx,
h,
num_frames,
height,
width,
patch_size_t,
patch_size,
patch_size,
true);
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto denormalize_and_clamp = [&](ggml_tensor* pixels) {
auto scaled = ggml_mul(ctx->ggml_ctx, pixels, pixel_std);
return ggml_clamp(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, scaled, pixel_mean), 0.f, 1.f);
};
if (sol) {
const int64_t frame_rows = width * height;
const int64_t output_frames = sol_output_frames(num_frames, patch_size_t);
const int64_t frame_size = width * height * patch_size * patch_size;
const int64_t channels = pixel_mean->ne[3];
auto output = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, frame_size, channels * output_frames);
int64_t frame_offset = 0;
for (int64_t t = 0; t < num_frames; t += sol_block_size) {
const int64_t end = std::min(t + sol_block_size, num_frames);
auto rows = ggml_ext_slice(ctx, h, 1, t * frame_rows, end * frame_rows, false);
auto patches = proj_out->forward(ctx, norm_out->forward(ctx, rows));
auto pixels = DiT::unpatchify_3d(ctx, patches, end - t, height, width,
patch_size_t, patch_size, patch_size, true);
// Keep only the last frame of the first latent patch in each native block.
pixels = ggml_ext_slice(ctx, pixels, 2, patch_size_t - 1, pixels->ne[2], false);
pixels = denormalize_and_clamp(pixels);
const int64_t count = pixels->ne[2];
for (int64_t c = 0; c < channels; ++c) {
auto plane = ggml_ext_slice(ctx, pixels, 3, c, c + 1, false);
auto frames = ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, plane), frame_size, count);
const int64_t start = c * output_frames + frame_offset;
auto indices = ggml_ext_slice(ctx, sol_inputs[0].row_ids, 0, start, start + count, false);
output = ggml_set_rows(ctx->ggml_ctx, output, frames, indices);
}
// Write each block before expanding the next block's pixel graph.
ctx->expand_graph(output);
frame_offset += count;
}
return ggml_reshape_4d(ctx->ggml_ctx, output, width * patch_size, height * patch_size, output_frames, channels);
}
h = proj_out->forward(ctx, norm_out->forward(ctx, h));
h = ggml_ext_slice(ctx, h, 1, 0, num_patches);
auto pixels = DiT::unpatchify_3d(ctx, h, num_frames, height, width,
patch_size_t, patch_size, patch_size, true);
return denormalize_and_clamp(pixels);
}
};
@@ -479,9 +565,7 @@ namespace MiniMaxH3VAE {
explicit MiniMaxH3VideoVAE(const MiniMaxH3VAEConfig& config)
: single_frame(config.patch_size_t == 1) {
blocks["encoder"] = std::make_shared<Encoder>();
blocks["quant_conv"] = std::make_shared<Conv3d>(48,
48,
std::tuple{1, 1, 1});
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});
@@ -504,15 +588,13 @@ namespace MiniMaxH3VAE {
ggml_tensor* decode(GGMLRunnerContext* ctx,
ggml_tensor* latent,
ggml_tensor* pe,
ggml_tensor* shifted_pe,
const std::array<SolAttentionPlan::Inputs, 2>& sol_inputs,
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);
return decoder->forward(ctx, post_quant->forward(ctx, latent), pe, shifted_pe, sol_inputs, pixel_mean, pixel_std);
}
};
@@ -553,6 +635,9 @@ namespace MiniMaxH3VAE {
}
std::string get_desc() override {
if (config.sol) {
return "sol_h3_vae";
}
return config.patch_size_t == 1 ? "p_h3_image_vae" : "minimax_h3_video_vae";
}
@@ -601,6 +686,10 @@ namespace MiniMaxH3VAE {
}
sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const override {
return config.sol ? params : resolve_original_tiling_params(params);
}
static sd_tiling_params_t resolve_original_tiling_params(sd_tiling_params_t params) {
if (!params.enabled) {
params.target_overlap = 0.25f;
}
@@ -655,7 +744,7 @@ namespace MiniMaxH3VAE {
bool circular_x = false,
bool circular_y = false) override {
auto input = ensure_video_shape(x);
auto tiling = resolve_tiling_params(tiling_params);
auto tiling = resolve_original_tiling_params(tiling_params);
if (config.patch_size_t == 1 && input.shape()[2] != 1) {
LOG_ERROR("P-H3-Image VAE encodes single-frame images only");
return {};
@@ -695,6 +784,46 @@ namespace MiniMaxH3VAE {
return result;
}
bool supports_temporal_tiling(VAETemporalDirection direction) const override {
return config.sol && direction == VAETemporalDirection::DECODE;
}
int get_default_temporal_tile_frames(VAETemporalDirection direction) const override {
SD_UNUSED(direction);
return 2 * sol_block_size;
}
int get_default_temporal_tile_overlap(VAETemporalDirection direction) const override {
SD_UNUSED(direction);
return sol_block_size;
}
sd::Tensor<float> _compute_temporal_tiled(int n_threads,
const sd::Tensor<float>& input,
VAETemporalDirection direction,
const VAETemporalTilingConfig& tiling) override {
GGML_ASSERT(supports_temporal_tiling(direction));
if (tiling.tile_frames % sol_block_size || tiling.overlap % sol_block_size ||
tiling.overlap >= tiling.tile_frames) {
LOG_ERROR("Sol-H3-VAE temporal tile size and overlap must be multiples of %d latent frames, with overlap smaller than the tile",
static_cast<int>(sol_block_size));
return {};
}
if (input.shape()[2] <= tiling.tile_frames) {
return _compute(n_threads, input, true);
}
auto plan = make_vae_temporal_tile_plan(input.shape()[2], tiling);
// With block-aligned starts, the same mapping gives tile lengths and output offsets.
auto frame_count = [&](int64_t frames) { return sol_output_frames(frames, config.patch_size_t); };
return process_vae_temporal_tiles_blended(input, plan, frame_count, frame_count,
[&](const sd::Tensor<float>& tile, const VAETemporalTile& window) {
LOG_VERBOSE("Sol-H3-VAE temporal tile %d/%d: latent frames [%lld, %lld)",
window.index + 1, static_cast<int>(plan.tiles.size()),
static_cast<long long>(window.start), static_cast<long long>(window.end));
return _compute(n_threads, tile, true);
});
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
@@ -702,7 +831,25 @@ namespace MiniMaxH3VAE {
bool circular_x = false,
bool circular_y = false,
bool silent = false) override {
auto input = ensure_video_shape(x);
auto input = ensure_video_shape(x);
if (config.sol) {
circular_x = circular_x || circular_x_enabled;
circular_y = circular_y || circular_y_enabled;
if (circular_x || circular_y) {
int tile_w = static_cast<int>(input.shape()[0]);
int tile_h = static_cast<int>(input.shape()[1]);
float overlap;
if (tiling_params.enabled &&
!get_tile_sizes(tile_w, tile_h, overlap, tiling_params, input.shape()[0], input.shape()[1])) {
return {};
}
if ((circular_x && tile_w >= input.shape()[0]) || (circular_y && tile_h >= input.shape()[1])) {
LOG_ERROR("Sol-H3-VAE circular decoding requires spatial tiles smaller than the input on every circular axis");
return {};
}
}
return VAE::decode(n_threads, input, tiling_params, decode_video, circular_x, circular_y, silent);
}
auto tiling = resolve_tiling_params(tiling_params);
if (config.patch_size_t == 1 && input.shape()[2] != 1) {
LOG_ERROR("P-H3-Image VAE decodes single-frame latents only");
@@ -801,23 +948,29 @@ namespace MiniMaxH3VAE {
sd::Tensor<float> build_rope(int64_t width,
int64_t height,
int64_t num_frames) {
int64_t num_frames,
bool shifted = false) {
std::vector<std::vector<float>> ids;
ids.reserve(static_cast<size_t>(width * height * num_frames + 5));
const int64_t global_rows = config.sol ? ((num_frames + sol_block_size - 1) / sol_block_size) * sol_global_rows : 5;
ids.reserve(static_cast<size_t>(width * height * num_frames + global_rows));
constexpr float two_pi = 6.28318530717958647692f;
const int64_t spatial_shift = shifted ? sol_spatial_shift : 0;
const int64_t time_shift = shifted ? sol_temporal_shift : 0;
for (int64_t t = 0; t < num_frames; ++t) {
float pt = (config.patch_size_t == 1 ? 2.f * (0.5f / 7.f) - 1.f
: 2.f * ((t + 0.5f) / num_frames) - 1.f) *
two_pi;
const float time = config.sol ? ((t + time_shift) % sol_block_size + 0.5f) / sol_block_size
: (config.patch_size_t == 1 ? 0.5f / 7.f : (t + 0.5f) / num_frames);
float pt = (2.f * time - 1.f) * two_pi;
for (int64_t h = 0; h < height; ++h) {
float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi;
const float y = config.sol ? ((h + spatial_shift) % sol_window_size + 0.5f) / sol_window_size : (h + 0.5f) / height;
float ph = (2.f * y - 1.f) * two_pi;
for (int64_t w = 0; w < width; ++w) {
float pw = (2.f * ((w + 0.5f) / width) - 1.f) * two_pi;
const float x = config.sol ? ((w + spatial_shift) % sol_window_size + 0.5f) / sol_window_size : (w + 0.5f) / width;
float pw = (2.f * x - 1.f) * two_pi;
ids.push_back({pt, ph, pw});
}
}
}
for (int i = 0; i < 5; ++i) {
for (int64_t i = 0; i < global_rows; ++i) {
ids.push_back({0.f, 0.f, 0.f});
}
auto values = Rope::embed_nd(ids,
@@ -835,24 +988,49 @@ namespace MiniMaxH3VAE {
const sd::Tensor<float>& z,
bool decode_graph) override {
auto input = ensure_video_shape(z);
std::unique_ptr<SolAttentionPlan> sol_plan;
std::unique_ptr<SolAttentionPlan> shifted_plan;
sd::Tensor<float> shifted_rope;
sd::Tensor<int32_t> sol_row_ids;
if (decode_graph) {
rope_cache = build_rope(input.shape()[0],
input.shape()[1],
input.shape()[2]);
if (config.sol) {
const auto& shape = input.shape();
sol_plan = std::make_unique<SolAttentionPlan>(shape[0], shape[1], shape[2], false);
shifted_plan = std::make_unique<SolAttentionPlan>(shape[0], shape[1], shape[2], true);
shifted_rope = build_rope(shape[0], shape[1], shape[2], true);
const int64_t rows = std::max(shifted_rope.shape()[3], pixel_mean.shape()[3] *
sol_output_frames(shape[2], config.patch_size_t));
sol_row_ids = sd::Tensor<int32_t>({rows});
for (int64_t i = 0; i < rows; ++i) {
sol_row_ids.data()[i] = static_cast<int32_t>(i);
}
}
}
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();
auto graph = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE);
auto runner_ctx = get_context(graph);
ggml_tensor* out = nullptr;
if (decode_graph) {
auto pe = make_input(rope_cache);
out = model.decode(&runner_ctx, value, pe, mean, std);
auto pe = make_input(rope_cache);
auto shifted_pe = config.sol ? make_input(shifted_rope) : nullptr;
std::array<SolAttentionPlan::Inputs, 2> sol_inputs;
if (sol_plan) {
auto bind_input = [&](const auto& tensor) { return make_input(tensor); };
sol_inputs = {sol_plan->bind(bind_input), shifted_plan->bind(bind_input)};
auto row_ids = make_input(sol_row_ids);
sol_inputs[0].row_ids = row_ids;
sol_inputs[1].row_ids = row_ids;
}
out = model.decode(&runner_ctx, value, pe, shifted_pe, sol_inputs, 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;
};
src/model/vae/vae_tiling.hpp
+18 -7
View File
@@ -116,14 +116,14 @@ inline sd::Tensor<float> process_vae_temporal_tiles(const sd::Tensor<float>& inp
return output;
}
template <typename Fn>
template <typename FrameCount, typename FrameOffset, typename Fn>
inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<float>& input,
const VAETemporalTilePlan& plan,
int output_scale,
FrameCount frame_count,
FrameOffset frame_offset,
Fn&& on_processing) {
GGML_ASSERT(output_scale >= 1);
const int64_t output_frames = 1 + (input.shape()[2] - 1) * output_scale;
const int overlap_frames = plan.overlap > 0 ? 1 + (plan.overlap - 1) * output_scale : 0;
const int64_t output_frames = frame_count(input.shape()[2]);
const int64_t overlap_frames = frame_count(plan.overlap);
std::vector<float> weights(static_cast<size_t>(output_frames), 0.f);
sd::Tensor<float> output;
@@ -138,7 +138,7 @@ inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<flo
return {};
}
const int64_t expected_tile_frames = 1 + (input_tile.shape()[2] - 1) * output_scale;
const int64_t expected_tile_frames = frame_count(input_tile.shape()[2]);
if (output_tile.dim() < 3 || output_tile.shape()[2] != expected_tile_frames) {
LOG_ERROR("unexpected temporal tile output shape: expected %lld frames, got %lld",
(long long)expected_tile_frames,
@@ -165,7 +165,7 @@ inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<flo
}
}
const int64_t output_start = tile.start * output_scale;
const int64_t output_start = frame_offset(tile.start);
const int64_t inner = output.shape()[0] * output.shape()[1];
const int64_t outer = output.numel() / (inner * output.shape()[2]);
const int64_t tile_frames = output_tile.shape()[2];
@@ -214,4 +214,15 @@ inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<flo
return output;
}
template <typename Fn>
inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<float>& input,
const VAETemporalTilePlan& plan,
int output_scale,
Fn&& on_processing) {
GGML_ASSERT(output_scale >= 1);
auto frame_count = [&](int64_t frames) { return frames > 0 ? 1 + (frames - 1) * output_scale : 0; };
auto frame_offset = [&](int64_t start) { return start * output_scale; };
return process_vae_temporal_tiles_blended(input, plan, frame_count, frame_offset, std::forward<Fn>(on_processing));
}
#endif // __SD_MODEL_VAE_VAE_TILING_HPP__