mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-08-10 14:16:50 +00:00
806 lines
39 KiB
C++
806 lines
39 KiB
C++
#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__
|