Skip to content

Commit 6100d83

Browse files
stduhpfstduhpf
authored
feat: add taeh3 support (#1874)
1 parent de298c2 commit 6100d83

2 files changed

Lines changed: 115 additions & 12 deletions

File tree

‎src/model/vae/tae.hpp‎

Lines changed: 109 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -426,10 +426,10 @@ class TinyVideoDecoder : public UnaryBlock {
426426
static const int num_layers = 3;
427427
int channels[num_layers + 1] = {256, 128, 64, 64};
428428
int patch_size = 1;
429-
int t_upscale = 1;
430429
bool is_wide = false;
431-
430+
432431
public:
432+
int t_upscale = 1;
433433
TinyVideoDecoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_upscale = {false, true, true}, bool is_wide = false)
434434
: z_channels(z_channels), patch_size(patch_size), is_wide(is_wide) {
435435
t_upscale = 1;
@@ -536,6 +536,10 @@ class TAEHV : public GGMLBlock {
536536
patch = 4;
537537
time_downscale = {true, true, true};
538538
time_upscale = {true, true, true};
539+
} else if (sd_version_is_minimax_h3(version)) {
540+
z_channels = 24;
541+
patch = 2;
542+
time_downscale = {true, true, false};
539543
}
540544
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch, time_upscale, is_wide));
541545
if (!decode_only) {
@@ -545,24 +549,123 @@ class TAEHV : public GGMLBlock {
545549

546550
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
547551
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
548-
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
552+
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
549553
// (W, H, C, T) -> (W, H, T, C)
550554
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
551555
}
552556
auto result = decoder->forward(ctx, z);
553-
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
557+
558+
if (sd_version_is_minimax_h3(version)) {
559+
int64_t num_frames = result->ne[3];
560+
int64_t chunk_frames = 5 * decoder->t_upscale;
561+
int64_t pad = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames;
562+
563+
result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, pad, false, false);
564+
565+
int64_t num_chunks = (num_frames + pad) / chunk_frames;
566+
auto to_trim = decoder->t_upscale - 1;
567+
std::vector<ggml_tensor*> to_concat = {};
568+
for (int i = 0; i < num_chunks; i++) {
569+
auto chunk = ggml_view_4d(ctx->ggml_ctx, result,
570+
result->ne[0], result->ne[1], result->ne[2], chunk_frames - to_trim,
571+
result->nb[1], result->nb[2], result->nb[3],
572+
i * chunk_frames * result->nb[3]);
573+
to_concat.push_back(chunk);
574+
}
575+
result = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
576+
result = ggml_view_4d(ctx->ggml_ctx, result,
577+
result->ne[0], result->ne[1], result->ne[2],
578+
result->ne[3] - decoder->t_upscale * 3,
579+
result->nb[1], result->nb[2], result->nb[3], 0);
580+
}
581+
582+
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
554583
// (W, H, T, C) -> (W, H, C, T)
555584
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
556585
}
557586
return result;
558587
}
559588

560-
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
589+
ggml_tensor* encode_h3(GGMLRunnerContext* ctx, ggml_tensor* x) {
561590
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
562-
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
591+
592+
int64_t num_frames = x->ne[3];
593+
int64_t pad = (17 - (num_frames % 17)) % 17;
594+
595+
if (pad > 0) {
596+
auto last_frame = ggml_view_4d(ctx->ggml_ctx, x,
597+
x->ne[0], x->ne[1], x->ne[2], 1,
598+
x->nb[1], x->nb[2], x->nb[3],
599+
(num_frames - 1) * x->nb[3]);
600+
for (int i = 0; i < pad; i++) {
601+
x = ggml_concat(ctx->ggml_ctx, x, last_frame, 3);
602+
}
603+
}
604+
605+
int64_t T_padded = x->ne[3];
606+
int64_t num_chunks = T_padded / 17;
607+
608+
auto zero_frame = ggml_view_4d(ctx->ggml_ctx, x,
609+
x->ne[0], x->ne[1], x->ne[2], 1,
610+
x->nb[1], x->nb[2], x->nb[3], 0);
611+
auto zeros_1 = ggml_scale(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, zero_frame), 0.0f);
612+
auto zeros_3 = zeros_1;
613+
for (int i = 1; i < 3; i++) {
614+
zeros_3 = ggml_concat(ctx->ggml_ctx, zeros_3, zeros_1, 3);
615+
}
616+
ggml_tensor* out = nullptr;
617+
if (false) {
618+
std::vector<ggml_tensor*> to_concat = {};
619+
for (int i = 0; i < num_chunks; i++) {
620+
auto chunk = ggml_view_4d(ctx->ggml_ctx, x,
621+
x->ne[0], x->ne[1], x->ne[2], 17,
622+
x->nb[1], x->nb[2], x->nb[3],
623+
i * 17 * x->nb[3]);
624+
625+
auto chunk_padded = ggml_concat(ctx->ggml_ctx, zeros_3, chunk, 3);
626+
627+
to_concat.push_back(chunk_padded);
628+
}
629+
ggml_tensor* x_in = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
630+
out = encoder->forward(ctx, x_in);
631+
} else {
632+
std::vector<ggml_tensor*> to_concat = {};
633+
for (int i = 0; i < num_chunks; i++) {
634+
auto chunk = ggml_view_4d(ctx->ggml_ctx, x,
635+
x->ne[0], x->ne[1], x->ne[2], 17,
636+
x->nb[1], x->nb[2], x->nb[3],
637+
i * 17 * x->nb[3]);
638+
639+
auto chunk_padded = ggml_concat(ctx->ggml_ctx, zeros_3, chunk, 3);
640+
641+
auto chunk_out = encoder->forward(ctx, chunk_padded);
642+
// auto chunk_out = encoder->forward_seq(ctx, chunk_padded); // ~same vram usage, and straight-up slower. it's already sequential enough
643+
644+
to_concat.push_back(chunk_out);
645+
}
646+
out = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
647+
}
648+
649+
// Return x[:, :-3] - drop the last 3 elements in the T dimension
650+
int64_t out_T = out->ne[3];
651+
out = ggml_view_4d(ctx->ggml_ctx, out,
652+
out->ne[0], out->ne[1], out->ne[2], out_T - 3,
653+
out->nb[1], out->nb[2], out->nb[3], 0);
654+
655+
return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 1, 3, 2));
656+
}
657+
658+
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
659+
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || (sd_version_is_minimax_h3(version) && x->ne[3] > 1)) {
563660
// (W, H, T, C) -> (W, H, C, T)
564661
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
565662
}
663+
if (sd_version_is_minimax_h3(version)) {
664+
return encode_h3(ctx, x);
665+
}
666+
667+
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
668+
566669
int64_t num_frames = x->ne[3];
567670
if (num_frames % encoder->t_downscale) {
568671
// pad to multiple of encoder->t_downscale at the end

‎src/stable-diffusion.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1023,11 +1023,11 @@ class StableDiffusionGGML {
10231023
tae_preview_only = false;
10241024
use_tae = true;
10251025
}
1026-
if (sd_version_is_minimax_h3(version) && use_tae) {
1027-
LOG_WARN("MiniMax-H3 does not have a compatible TAE; ignoring --taesd");
1028-
tae_preview_only = false;
1029-
use_tae = false;
1030-
}
1026+
// if (sd_version_is_minimax_h3(version) && use_tae) {
1027+
// LOG_WARN("MiniMax-H3 does not have a compatible TAE; ignoring --taesd");
1028+
// tae_preview_only = false;
1029+
// use_tae = false;
1030+
// }
10311031

10321032
auto& tensor_storage_map = model_loader.get_tensor_storage_map();
10331033

@@ -1402,7 +1402,7 @@ class StableDiffusionGGML {
14021402
}
14031403

14041404
auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
1405-
if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
1405+
if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
14061406
return std::make_shared<TinyVideoAutoEncoder>(backend_for(SDBackendModule::VAE),
14071407
tensor_storage_map,
14081408
"decoder",

0 commit comments

Comments
 (0)