diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index 710729fd7..e14c362ea 100644 --- a/src/model/vae/tae.hpp +++ b/src/model/vae/tae.hpp @@ -426,10 +426,10 @@ class TinyVideoDecoder : public UnaryBlock { static const int num_layers = 3; int channels[num_layers + 1] = {256, 128, 64, 64}; int patch_size = 1; - int t_upscale = 1; bool is_wide = false; - + public: + int t_upscale = 1; TinyVideoDecoder(int z_channels = 4, int patch_size = 1, std::vector time_upscale = {false, true, true}, bool is_wide = false) : z_channels(z_channels), patch_size(patch_size), is_wide(is_wide) { t_upscale = 1; @@ -536,6 +536,10 @@ class TAEHV : public GGMLBlock { patch = 4; time_downscale = {true, true, true}; time_upscale = {true, true, true}; + } else if (sd_version_is_minimax_h3(version)) { + z_channels = 24; + patch = 2; + time_downscale = {true, true, false}; } blocks["decoder"] = std::shared_ptr(new TinyVideoDecoder(z_channels, patch, time_upscale, is_wide)); if (!decode_only) { @@ -545,19 +549,119 @@ class TAEHV : public GGMLBlock { ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) { auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); - if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { + if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) { // (W, H, C, T) -> (W, H, T, C) z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2)); } auto result = decoder->forward(ctx, z); - if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { + + if (sd_version_is_minimax_h3(version)) { + int64_t num_frames = result->ne[3]; + int64_t chunk_frames = 5 * decoder->t_upscale; + int64_t pad = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames; + + result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, pad, false, false); + + int64_t num_chunks = (num_frames + pad) / chunk_frames; + auto to_trim = decoder->t_upscale - 1; + std::vector to_concat = {}; + for (int i = 0; i < num_chunks; i++) { + auto chunk = ggml_view_4d(ctx->ggml_ctx, result, + result->ne[0], result->ne[1], result->ne[2], chunk_frames - to_trim, + result->nb[1], result->nb[2], result->nb[3], + i * chunk_frames * result->nb[3]); + to_concat.push_back(chunk); + } + result = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3); + result = ggml_view_4d(ctx->ggml_ctx, result, + result->ne[0], result->ne[1], result->ne[2], + result->ne[3] - decoder->t_upscale * 3, + result->nb[1], result->nb[2], result->nb[3], 0); + } + + if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) { // (W, H, T, C) -> (W, H, C, T) result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2)); } return result; } +ggml_tensor* encode_h3(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); + + int64_t num_frames = x->ne[3]; + int64_t pad = (17 - (num_frames % 17)) % 17; + + if (pad > 0) { + auto last_frame = ggml_view_4d(ctx->ggml_ctx, x, + x->ne[0], x->ne[1], x->ne[2], 1, + x->nb[1], x->nb[2], x->nb[3], + (num_frames - 1) * x->nb[3]); + for (int i = 0; i < pad; i++) { + x = ggml_concat(ctx->ggml_ctx, x, last_frame, 3); + } + } + + int64_t T_padded = x->ne[3]; + int64_t num_chunks = T_padded / 17; + + auto zero_frame = ggml_view_4d(ctx->ggml_ctx, x, + x->ne[0], x->ne[1], x->ne[2], 1, + x->nb[1], x->nb[2], x->nb[3], 0); + auto zeros_1 = ggml_scale(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, zero_frame), 0.0f); + auto zeros_3 = zeros_1; + for (int i = 1; i < 3; i++) { + zeros_3 = ggml_concat(ctx->ggml_ctx, zeros_3, zeros_1, 3); + } + +#if SEQUENTIAL_TAEH3_ENCODE + //not sure wich one is better, I don't have enough vram to test video encode easily + std::vector to_concat = {}; + // Process sequentially (equivalent to the non-parallel unbind iteration in python) + for (int i = 0; i < num_chunks; i++) { + auto chunk = ggml_view_4d(ctx->ggml_ctx, x, + x->ne[0], x->ne[1], x->ne[2], 17, + x->nb[1], x->nb[2], x->nb[3], + i * 17 * x->nb[3]); + + // Pad 3 zeros at the beginning of the T dimension -> (W, H, C, 20) + auto chunk_padded = ggml_concat(ctx->ggml_ctx, zeros_3, chunk, 3); + + // Encode the 20 frame chunk + auto chunk_out = encoder->forward(ctx, chunk_padded); + + to_concat.push_back(chunk_out); + } + ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3); +#else + std::vector to_concat = {}; + for (int i = 0; i < num_chunks; i++) { + auto chunk = ggml_view_4d(ctx->ggml_ctx, x, + x->ne[0], x->ne[1], x->ne[2], 17, + x->nb[1], x->nb[2], x->nb[3], + i * 17 * x->nb[3]); + + auto chunk_padded = ggml_concat(ctx->ggml_ctx, zeros_3, chunk, 3); + + to_concat.push_back(chunk_padded); + } + ggml_tensor* x_in = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3); + auto out = encoder->forward(ctx, x_in); +#endif + + // Return x[:, :-3] - drop the last 3 elements in the T dimension + int64_t out_T = out->ne[3]; + out = ggml_view_4d(ctx->ggml_ctx, out, + out->ne[0], out->ne[1], out->ne[2], out_T - 3, + out->nb[1], out->nb[2], out->nb[3], 0); + + return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 1, 3, 2)); + } + ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (sd_version_is_minimax_h3(version)) { + return encode_h3(ctx, x); + } auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { // (W, H, T, C) -> (W, H, C, T) diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index b74d981aa..321f77cd9 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -1023,11 +1023,11 @@ class StableDiffusionGGML { tae_preview_only = false; use_tae = true; } - if (sd_version_is_minimax_h3(version) && use_tae) { - LOG_WARN("MiniMax-H3 does not have a compatible TAE; ignoring --taesd"); - tae_preview_only = false; - use_tae = false; - } + // if (sd_version_is_minimax_h3(version) && use_tae) { + // LOG_WARN("MiniMax-H3 does not have a compatible TAE; ignoring --taesd"); + // tae_preview_only = false; + // use_tae = false; + // } auto& tensor_storage_map = model_loader.get_tensor_storage_map(); @@ -1402,7 +1402,7 @@ class StableDiffusionGGML { } auto create_tae = [&](bool decode_only) -> std::shared_ptr { - if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { + if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) { return std::make_shared(backend_for(SDBackendModule::VAE), tensor_storage_map, "decoder",