#ifndef __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__
#define __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__

#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <set>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "core/ggml_tensor_utils.h"

#include "core/ggml_graph_cut.h"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/model.hpp"

namespace MiniMaxH3 {

    constexpr int H3_GRAPH_SIZE          = 131072;
    constexpr float FRAME_RESCALE        = 5.f / 3.f;
    constexpr float VISUAL_COND_TIMESTEP = 0.999f;

    struct Config {
        int64_t hidden_size              = 5376;
        int64_t num_layers               = 50;
        int64_t token_refiner_num_layers = 2;
        int64_t num_attention_heads      = 56;
        int64_t attention_head_dim       = 128;
        int64_t ffn_hidden_size          = 14336;
        int64_t video_latent_channels    = 24;
        int64_t audio_latent_channels    = 32;
        int64_t text_dim                 = 5120;
        int64_t timestep_input_dim       = 256;
        int64_t time_embed_hidden_size   = 5376;
        int64_t time_embed_dim           = 2688;
        int64_t rope_inv_freq_len        = 16;
        int64_t adaln_curve_grid         = 0;
        int patch_t                      = 1;
        int patch_h                      = 2;
        int patch_w                      = 2;
        float norm_eps                   = 1e-5f;
        float qk_norm_eps                = 1e-5f;
        float final_norm_eps             = 1e-5f;

        bool uses_adaln_curves() const {
            return adaln_curve_grid > 0;
        }

        static int64_t count_blocks(const String2TensorStorage& tensors,
                                    const std::string& prefix) {
            std::set<int> indices;
            for (const auto& [name, _] : tensors) {
                if (!starts_with(name, prefix)) {
                    continue;
                }
                size_t begin = prefix.size();
                size_t end   = name.find('.', begin);
                if (end != std::string::npos) {
                    indices.insert(std::atoi(name.substr(begin, end - begin).c_str()));
                }
            }
            return static_cast<int64_t>(indices.size());
        }

        static Config detect_from_weights(const String2TensorStorage& tensors,
                                          const std::string& prefix) {
            Config config;
            auto find = [&](const std::string& suffix) -> const TensorStorage* {
                auto it = tensors.find(prefix + "." + suffix);
                return it == tensors.end() ? nullptr : &it->second;
            };

            if (const auto* weight = find("video_patch_proj.weight")) {
                config.video_latent_channels = weight->ne[0] / 4;
                config.hidden_size           = weight->ne[1];
            }
            if (const auto* weight = find("audio_patch_proj.weight")) {
                config.audio_latent_channels = weight->ne[0];
            }
            config.num_layers               = count_blocks(tensors, prefix + ".blocks.");
            config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks.");
            if (const auto* weight = find("blocks.0.attn.q_norm.weight")) {
                config.attention_head_dim = weight->ne[0];
            }
            if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) {
                config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim);
            }
            if (const auto* weight = find("blocks.0.mlp.fc1.weight")) {
                config.ffn_hidden_size = weight->ne[1] / 2;
            }
            if (const auto* weight = find("condition_proj.weight")) {
                config.text_dim = weight->ne[0];
            }
            if (const auto* table = find("adaln_t_table")) {
                config.time_embed_dim   = table->ne[0];
                config.adaln_curve_grid = table->ne[1];
            } else {
                if (const auto* weight = find("time_embedder.proj_in.weight")) {
                    config.timestep_input_dim     = weight->ne[0];
                    config.time_embed_hidden_size = weight->ne[1];
                }
                if (const auto* weight = find("time_embedder.proj_out.weight")) {
                    config.time_embed_dim = weight->ne[1];
                }
            }
            if (const auto* inv_freq = find("rope.inv_freq")) {
                config.rope_inv_freq_len = inv_freq->ne[0];
            }

            LOG_VERBOSE("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64
                        ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64,
                        config.num_layers,
                        config.hidden_size,
                        config.num_attention_heads,
                        config.attention_head_dim,
                        config.ffn_hidden_size,
                        config.adaln_curve_grid);
            return config;
        }
    };

    static float time_shift_sigma(float sigma, float from_shift, float to_shift) {
        float base = sigma / (from_shift + sigma * (1.f - from_shift));
        return to_shift * base / (1.f + (to_shift - 1.f) * base);
    }

    struct TimeEmbedder : public GGMLBlock {
        TimeEmbedder(int64_t input_dim, int64_t hidden_dim, int64_t output_dim) {
            blocks["proj_in"]  = std::make_shared<Linear>(input_dim, hidden_dim, true, true);
            blocks["proj_out"] = std::make_shared<Linear>(hidden_dim, output_dim, true, true);
        }

        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
            auto proj_in  = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
            auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
            return proj_out->forward(ctx, ggml_silu(ctx->ggml_ctx, proj_in->forward(ctx, x)));
        }
    };

    struct MLP : public UnaryBlock {
        MLP(int64_t hidden_size,
            int64_t ffn_hidden_size) {
            blocks["fc1"] = std::make_shared<Linear>(hidden_size, ffn_hidden_size * 2, false, false, true, 1.f / 128.f);
            blocks["fc2"] = std::make_shared<Linear>(ffn_hidden_size, hidden_size, false, false, true, 1.f / 128.f);
        }

        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
            auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
            auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
            auto uv  = ggml_ext_chunk(ctx->ggml_ctx, fc1->forward(ctx, x), 2, 0);
            return fc2->forward(ctx, ggml_mul(ctx->ggml_ctx,
                                              ggml_silu(ctx->ggml_ctx, uv[0]),
                                              uv[1]));
        }
    };

    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;
        GGML_ASSERT(rot_dim <= x->ne[0]);
        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 Attention : public GGMLBlock {
        int64_t heads;
        int64_t head_dim;

        Attention(int64_t hidden_size,
                  int64_t heads,
                  int64_t head_dim,
                  float eps)
            : heads(heads), head_dim(head_dim) {
            int64_t inner      = heads * head_dim;
            blocks["qkv_proj"] = std::make_shared<Linear>(hidden_size, inner * 3, false);
            blocks["q_norm"]   = std::make_shared<RMSNorm>(head_dim, eps);
            blocks["k_norm"]   = std::make_shared<RMSNorm>(head_dim, eps);
            blocks["out_proj"] = std::make_shared<Linear>(inner, hidden_size, false);
        }

        ggml_tensor* forward(GGMLRunnerContext* ctx,
                             ggml_tensor* x,
                             ggml_tensor* pe = nullptr) {
            auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_proj"]);
            auto q_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
            auto k_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
            auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);

            int64_t sequence = x->ne[1];
            int64_t batch    = x->ne[2] * x->ne[3];
            auto qkv         = ggml_ext_chunk(ctx->ggml_ctx, qkv_proj->forward(ctx, x), 3, 0);
            auto q           = ggml_reshape_4d(ctx->ggml_ctx, qkv[0], head_dim, heads, sequence, batch);
            auto k           = ggml_reshape_4d(ctx->ggml_ctx, qkv[1], head_dim, heads, sequence, batch);
            auto v           = ggml_reshape_4d(ctx->ggml_ctx, qkv[2], head_dim, heads, sequence, batch);
            q                = q_norm->forward(ctx, q);
            k                = k_norm->forward(ctx, k);
            if (pe != nullptr) {
                q = apply_partial_rope(ctx->ggml_ctx, q, pe);
                k = apply_partial_rope(ctx->ggml_ctx, k, pe);
            } else {
                q = attention_layout(ctx->ggml_ctx, q);
                k = attention_layout(ctx->ggml_ctx, k);
            }
            auto out = ggml_ext_attention_ext(ctx,
                                              q,
                                              k,
                                              v,
                                              static_cast<int>(heads),
                                              nullptr,
                                              true,
                                              ctx->flash_attn_enabled,
                                              1.f / 128.f);
            return out_proj->forward(ctx, out);
        }
    };

    struct TokenRefinerBlock : public GGMLBlock {
        TokenRefinerBlock(const Config& config) {
            blocks["norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
            blocks["norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
            blocks["attn"]  = std::make_shared<Attention>(config.hidden_size,
                                                         config.num_attention_heads,
                                                         config.attention_head_dim,
                                                         config.qk_norm_eps);
            blocks["mlp"]   = std::make_shared<MLP>(config.hidden_size, config.ffn_hidden_size);
        }

        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
            auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
            auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
            auto attn  = std::dynamic_pointer_cast<Attention>(blocks["attn"]);
            auto mlp   = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
            x          = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, norm1->forward(ctx, x)));
            return ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, norm2->forward(ctx, x)));
        }
    };

    struct TokenRefiner : public GGMLBlock {
        int64_t num_layers;

        explicit TokenRefiner(const Config& config)
            : num_layers(config.token_refiner_num_layers) {
            for (int64_t i = 0; i < num_layers; ++i) {
                blocks["blocks." + std::to_string(i)] = std::make_shared<TokenRefinerBlock>(config);
            }
            blocks["final_norm"] = std::make_shared<RMSNorm>(config.hidden_size, config.final_norm_eps);
        }

        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
            for (int64_t i = 0; i < num_layers; ++i) {
                auto block = std::dynamic_pointer_cast<TokenRefinerBlock>(blocks["blocks." + std::to_string(i)]);
                x          = block->forward(ctx, x);
            }
            return std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm"])->forward(ctx, x);
        }
    };

    struct AdaLayerNormModulation : public GGMLBlock {
        int64_t hidden_size;
        int expand;
        int modalities;
        bool apply_silu;

        AdaLayerNormModulation(int64_t time_dim,
                               int64_t hidden_size,
                               int expand,
                               int modalities,
                               bool apply_silu,
                               bool force_f32)
            : hidden_size(hidden_size),
              expand(expand),
              modalities(modalities),
              apply_silu(apply_silu) {
            blocks["linear"] = std::make_shared<Linear>(time_dim,
                                                        hidden_size * expand * modalities,
                                                        true,
                                                        force_f32);
        }

        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t_emb) {
            if (apply_silu) {
                t_emb = ggml_silu(ctx->ggml_ctx, t_emb);
            }
            return std::dynamic_pointer_cast<Linear>(blocks["linear"])->forward(ctx, t_emb);
        }
    };

    struct TokenModulationSpan {
        int64_t start;
        int64_t end;
        int64_t modulation_row;
    };

    enum class SequenceKind {
        TEXT,
        CONDITION_VIDEO,
        CONDITION_AUDIO,
        TARGET_AUDIO,
        TARGET_VIDEO,
    };

    struct SequenceSegment {
        int64_t start;
        int64_t end;
        SequenceKind kind;
        int32_t source_index = -1;
    };

    static std::vector<ggml_tensor*> modulation_row(ggml_context* ctx,
                                                    ggml_tensor* projection,
                                                    int64_t hidden_size,
                                                    int expand,
                                                    int modalities,
                                                    int64_t row) {
        int64_t timestep_rows = projection->ne[1];
        auto reshaped         = ggml_reshape_2d(ctx,
                                                projection,
                                                hidden_size * expand,
                                                timestep_rows * modalities);
        auto selected         = ggml_ext_slice(ctx, reshaped, 1, row, row + 1);
        return ggml_ext_chunk(ctx, selected, expand, 0);
    }

    static ggml_tensor* modulate_segments(ggml_context* ctx,
                                          ggml_tensor* x,
                                          ggml_tensor* projection,
                                          const std::vector<TokenModulationSpan>& segments,
                                          int64_t hidden_size,
                                          int expand,
                                          int modalities,
                                          int shift_index,
                                          int scale_index) {
        ggml_tensor* out = nullptr;
        for (const auto& segment : segments) {
            auto mods = modulation_row(ctx,
                                       projection,
                                       hidden_size,
                                       expand,
                                       modalities,
                                       segment.modulation_row);
            auto part = ggml_ext_slice(ctx, x, 1, segment.start, segment.end);
            part      = ggml_add(ctx,
                                 ggml_add(ctx, part, ggml_mul(ctx, part, mods[scale_index])),
                                 mods[shift_index]);
            out       = out == nullptr ? part : ggml_concat(ctx, out, part, 1);
        }
        return out;
    }

    static ggml_tensor* gated_residual_segments(ggml_context* ctx,
                                                ggml_tensor* x,
                                                ggml_tensor* update,
                                                ggml_tensor* projection,
                                                const std::vector<TokenModulationSpan>& segments,
                                                int64_t hidden_size,
                                                int gate_index) {
        ggml_tensor* out = nullptr;
        for (const auto& segment : segments) {
            auto mods = modulation_row(ctx, projection, hidden_size, 6, 3, segment.modulation_row);
            auto base = ggml_ext_slice(ctx, x, 1, segment.start, segment.end);
            auto add  = ggml_ext_slice(ctx, update, 1, segment.start, segment.end);
            auto part = ggml_add(ctx, base, ggml_mul(ctx, add, mods[gate_index]));
            out       = out == nullptr ? part : ggml_concat(ctx, out, part, 1);
        }
        return out;
    }

    struct TransformerBlock : public GGMLBlock {
        Config config;

        explicit TransformerBlock(const Config& config)
            : config(config) {
            blocks["norm1"]      = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
            blocks["norm2"]      = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
            blocks["attn"]       = std::make_shared<Attention>(config.hidden_size,
                                                         config.num_attention_heads,
                                                         config.attention_head_dim,
                                                         config.qk_norm_eps);
            blocks["mlp"]        = std::make_shared<MLP>(config.hidden_size,
                                                  config.ffn_hidden_size);
            blocks["adaln_proj"] = std::make_shared<AdaLayerNormModulation>(config.time_embed_dim,
                                                                            config.hidden_size,
                                                                            6,
                                                                            3,
                                                                            !config.uses_adaln_curves(),
                                                                            config.uses_adaln_curves());
        }

        ggml_tensor* forward(GGMLRunnerContext* ctx,
                             ggml_tensor* x,
                             ggml_tensor* t_emb,
                             const std::vector<TokenModulationSpan>& segments,
                             ggml_tensor* pe) {
            auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
            auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
            auto attn  = std::dynamic_pointer_cast<Attention>(blocks["attn"]);
            auto mlp   = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
            auto adaln = std::dynamic_pointer_cast<AdaLayerNormModulation>(blocks["adaln_proj"]);
            auto mods  = adaln->forward(ctx, t_emb);

            auto h = modulate_segments(ctx->ggml_ctx,
                                       norm1->forward(ctx, x),
                                       mods,
                                       segments,
                                       config.hidden_size,
                                       6,
                                       3,
                                       0,
                                       1);
            x      = gated_residual_segments(ctx->ggml_ctx,
                                             x,
                                             attn->forward(ctx, h, pe),
                                             mods,
                                             segments,
                                             config.hidden_size,
                                             2);
            h      = modulate_segments(ctx->ggml_ctx,
                                       norm2->forward(ctx, x),
                                       mods,
                                       segments,
                                       config.hidden_size,
                                       6,
                                       3,
                                       3,
                                       4);
            return gated_residual_segments(ctx->ggml_ctx,
                                           x,
                                           mlp->forward(ctx, h),
                                           mods,
                                           segments,
                                           config.hidden_size,
                                           5);
        }
    };

    struct FinalLayer : public GGMLBlock {
        Config config;

        explicit FinalLayer(const Config& config)
            : config(config) {
            int64_t video_dim    = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w;
            blocks["norm"]       = std::make_shared<RMSNorm>(config.hidden_size, config.final_norm_eps);
            blocks["adaln_proj"] = std::make_shared<AdaLayerNormModulation>(config.time_embed_dim,
                                                                            config.hidden_size,
                                                                            2,
                                                                            1,
                                                                            !config.uses_adaln_curves(),
                                                                            config.uses_adaln_curves());
            blocks["video_out"]  = std::make_shared<Linear>(config.hidden_size, video_dim, true, true);
            blocks["audio_out"]  = std::make_shared<Linear>(config.hidden_size, config.audio_latent_channels, true, true);
        }

        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
                                                      ggml_tensor* x,
                                                      ggml_tensor* t_emb,
                                                      const TokenModulationSpan& video,
                                                      const TokenModulationSpan& audio) {
            auto norm      = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
            auto adaln     = std::dynamic_pointer_cast<AdaLayerNormModulation>(blocks["adaln_proj"]);
            auto video_out = std::dynamic_pointer_cast<Linear>(blocks["video_out"]);
            auto audio_out = std::dynamic_pointer_cast<Linear>(blocks["audio_out"]);
            auto mods      = adaln->forward(ctx, t_emb);
            auto apply     = [&](const TokenModulationSpan& segment) {
                auto row   = modulation_row(ctx->ggml_ctx, mods, config.hidden_size, 2, 1, segment.modulation_row);
                auto value = norm->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end));
                return ggml_add(ctx->ggml_ctx,
                                    ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])),
                                    row[0]);
            };
            return {video_out->forward(ctx, apply(video)),
                    audio_out->forward(ctx, apply(audio))};
        }
    };

    struct MiniMaxH3Transformer3DModel : public GGMLBlock {
        Config config;

        explicit MiniMaxH3Transformer3DModel(const Config& config)
            : config(config) {
            int64_t video_dim          = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w;
            blocks["video_patch_proj"] = std::make_shared<Linear>(video_dim, config.hidden_size, true, true);
            blocks["audio_patch_proj"] = std::make_shared<Linear>(config.audio_latent_channels, config.hidden_size, true, true);
            blocks["condition_proj"]   = std::make_shared<Linear>(config.text_dim, config.hidden_size, true);
            if (!config.uses_adaln_curves()) {
                blocks["time_embedder"] = std::make_shared<TimeEmbedder>(config.timestep_input_dim,
                                                                         config.time_embed_hidden_size,
                                                                         config.time_embed_dim);
            }
            blocks["token_refiner"] = std::make_shared<TokenRefiner>(config);
            for (int64_t i = 0; i < config.num_layers; ++i) {
                blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(config);
            }
            blocks["final_layer"] = std::make_shared<FinalLayer>(config);
        }

        void init_params(ggml_context* ctx,
                         const String2TensorStorage& tensors = {},
                         const std::string prefix            = "") override {
            GGMLBlock::init_params(ctx, tensors, prefix);
            params["rope.inv_freq"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.rope_inv_freq_len);
            if (config.uses_adaln_curves()) {
                params["adaln_t_table"] = ggml_new_tensor_2d(ctx,
                                                             GGML_TYPE_F32,
                                                             config.time_embed_dim,
                                                             config.adaln_curve_grid);
            }
        }

        ggml_tensor* refine_context(GGMLRunnerContext* ctx, ggml_tensor* context) {
            if (context->ne[0] == config.hidden_size) {
                return context;
            }
            GGML_ASSERT(context->ne[0] == config.text_dim);
            auto condition_proj = std::dynamic_pointer_cast<Linear>(blocks["condition_proj"]);
            auto token_refiner  = std::dynamic_pointer_cast<TokenRefiner>(blocks["token_refiner"]);
            return token_refiner->forward(ctx, condition_proj->forward(ctx, context));
        }

        ggml_tensor* time_embedding(GGMLRunnerContext* ctx,
                                    ggml_tensor* timestep_features,
                                    ggml_tensor* curve_indices,
                                    ggml_tensor* curve_upper_indices,
                                    ggml_tensor* curve_fractions) {
            if (!config.uses_adaln_curves()) {
                return std::dynamic_pointer_cast<TimeEmbedder>(blocks["time_embedder"])->forward(ctx, timestep_features);
            }
            auto lower = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_indices);
            auto upper = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_upper_indices);
            return ggml_add(ctx->ggml_ctx,
                            lower,
                            ggml_mul(ctx->ggml_ctx,
                                     ggml_sub(ctx->ggml_ctx, upper, lower),
                                     curve_fractions));
        }

        ggml_tensor* build_rope(GGMLRunnerContext* ctx,
                                ggml_tensor* position_ids) {
            auto inv            = ggml_reshape_2d(ctx->ggml_ctx,
                                                  params["rope.inv_freq"],
                                                  config.rope_inv_freq_len,
                                                  1);
            ggml_tensor* angles = nullptr;
            for (int axis = 0; axis < 3; ++axis) {
                auto pos          = ggml_ext_slice(ctx->ggml_ctx, position_ids, 0, axis, axis + 1);
                auto expanded_inv = ggml_repeat_4d(ctx->ggml_ctx,
                                                   inv,
                                                   inv->ne[0],
                                                   pos->ne[1],
                                                   1,
                                                   1);
                auto a            = ggml_mul(ctx->ggml_ctx, expanded_inv, pos);
                angles            = angles == nullptr ? a : ggml_concat(ctx->ggml_ctx, angles, a, 0);
            }
            auto c  = ggml_reshape_4d(ctx->ggml_ctx, ggml_cos(ctx->ggml_ctx, angles), 1, angles->ne[0], angles->ne[1], 1);
            auto s  = ggml_reshape_4d(ctx->ggml_ctx, ggml_sin(ctx->ggml_ctx, angles), 1, angles->ne[0], angles->ne[1], 1);
            auto ns = ggml_neg(ctx->ggml_ctx, s);
            auto pe = ggml_concat(ctx->ggml_ctx, c, ns, 0);
            pe      = ggml_concat(ctx->ggml_ctx, pe, s, 0);
            pe      = ggml_concat(ctx->ggml_ctx, pe, c, 0);
            return ggml_reshape_4d(ctx->ggml_ctx, pe, 2, 2, angles->ne[0], angles->ne[1]);
        }

        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
                                                      ggml_tensor* video,
                                                      ggml_tensor* audio,
                                                      ggml_tensor* context,
                                                      const std::vector<ggml_tensor*>& condition_videos,
                                                      const std::vector<ggml_tensor*>& condition_audios,
                                                      ggml_tensor* position_ids,
                                                      ggml_tensor* timestep_features,
                                                      ggml_tensor* curve_indices,
                                                      ggml_tensor* curve_upper_indices,
                                                      ggml_tensor* curve_fractions,
                                                      const std::vector<TokenModulationSpan>& segments,
                                                      const std::vector<SequenceSegment>& sequence_segments,
                                                      const TokenModulationSpan& video_segment,
                                                      const TokenModulationSpan& audio_segment) {
            auto video_proj = std::dynamic_pointer_cast<Linear>(blocks["video_patch_proj"]);
            auto audio_proj = std::dynamic_pointer_cast<Linear>(blocks["audio_patch_proj"]);

            std::vector<std::pair<int64_t, int64_t>> condition_video_ranges;
            ggml_tensor* video_rows = nullptr;
            int64_t video_offset    = 0;
            for (auto condition : condition_videos) {
                auto rows = DiT::patchify_3d(ctx->ggml_ctx,
                                             condition,
                                             config.patch_t,
                                             config.patch_h,
                                             config.patch_w,
                                             1,
                                             true);
                condition_video_ranges.push_back({video_offset, video_offset + rows->ne[1]});
                video_offset += rows->ne[1];
                video_rows = video_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, video_rows, rows, 1);
            }
            auto target_video_rows                         = DiT::patchify_3d(ctx->ggml_ctx,
                                                                              video,
                                                                              config.patch_t,
                                                                              config.patch_h,
                                                                              config.patch_w,
                                                                              1,
                                                                              true);
            std::pair<int64_t, int64_t> target_video_range = {video_offset, video_offset + target_video_rows->ne[1]};
            video_rows                                     = video_rows == nullptr ? target_video_rows
                                                                                   : ggml_concat(ctx->ggml_ctx, video_rows, target_video_rows, 1);
            auto video_embeds                              = video_proj->forward(ctx, video_rows);

            auto pack_audio_rows = [&](ggml_tensor* value) {
                value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, value, 2, 0, 1, 3));
                return ggml_reshape_3d(ctx->ggml_ctx,
                                       value,
                                       value->ne[0],
                                       value->ne[1] * value->ne[2],
                                       value->ne[3]);
            };
            std::vector<std::pair<int64_t, int64_t>> condition_audio_ranges;
            ggml_tensor* audio_rows = nullptr;
            int64_t audio_offset    = 0;
            for (auto condition : condition_audios) {
                auto rows = pack_audio_rows(condition);
                condition_audio_ranges.push_back({audio_offset, audio_offset + rows->ne[1]});
                audio_offset += rows->ne[1];
                audio_rows = audio_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, audio_rows, rows, 1);
            }
            audio                                          = pack_audio_rows(audio);
            std::pair<int64_t, int64_t> target_audio_range = {audio_offset, audio_offset + audio->ne[1]};
            audio_rows                                     = audio_rows == nullptr ? audio : ggml_concat(ctx->ggml_ctx, audio_rows, audio, 1);
            auto audio_embeds                              = audio_proj->forward(ctx, audio_rows);
            context                                        = refine_context(ctx, context);

            ggml_tensor* h = nullptr;
            auto append    = [&](ggml_tensor* value) {
                h = h == nullptr ? value : ggml_concat(ctx->ggml_ctx, h, value, 1);
            };
            for (const auto& sequence : sequence_segments) {
                if (sequence.kind == SequenceKind::TEXT) {
                    append(context);
                } else if (sequence.kind == SequenceKind::CONDITION_VIDEO) {
                    GGML_ASSERT(sequence.source_index >= 0 &&
                                sequence.source_index < static_cast<int32_t>(condition_video_ranges.size()));
                    auto range = condition_video_ranges[static_cast<size_t>(sequence.source_index)];
                    append(ggml_ext_slice(ctx->ggml_ctx, video_embeds, 1, range.first, range.second));
                } else if (sequence.kind == SequenceKind::CONDITION_AUDIO) {
                    GGML_ASSERT(sequence.source_index >= 0 &&
                                sequence.source_index < static_cast<int32_t>(condition_audio_ranges.size()));
                    auto range = condition_audio_ranges[static_cast<size_t>(sequence.source_index)];
                    append(ggml_ext_slice(ctx->ggml_ctx, audio_embeds, 1, range.first, range.second));
                } else if (sequence.kind == SequenceKind::TARGET_AUDIO) {
                    append(ggml_ext_slice(ctx->ggml_ctx,
                                          audio_embeds,
                                          1,
                                          target_audio_range.first,
                                          target_audio_range.second));
                } else {
                    append(ggml_ext_slice(ctx->ggml_ctx,
                                          video_embeds,
                                          1,
                                          target_video_range.first,
                                          target_video_range.second));
                }
            }
            GGML_ASSERT(h != nullptr);

            auto t_emb = time_embedding(ctx,
                                        timestep_features,
                                        curve_indices,
                                        curve_upper_indices,
                                        curve_fractions);
            auto pe    = build_rope(ctx, position_ids);
            for (int64_t i = 0; i < config.num_layers; ++i) {
                auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
                h          = block->forward(ctx, h, t_emb, segments, pe);
                sd::ggml_graph_cut::mark_graph_cut(h,
                                                   "minimax_h3.blocks." + std::to_string(i),
                                                   "hidden_states");
            }

            auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
            auto output      = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment);
            auto video_out   = DiT::unpatchify_3d(ctx->ggml_ctx,
                                                  output.first,
                                                  video->ne[2] / config.patch_t,
                                                  video->ne[1] / config.patch_h,
                                                  video->ne[0] / config.patch_w,
                                                  config.patch_t,
                                                  config.patch_h,
                                                  config.patch_w,
                                                  true);
            auto audio_out   = ggml_reshape_4d(ctx->ggml_ctx,
                                               output.second,
                                               config.audio_latent_channels,
                                               audio->ne[1] / 2,
                                               2,
                                               audio->ne[2]);
            audio_out        = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3));
            video_out        = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f);
            audio_out        = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f);
            return {video_out, audio_out};
        }
    };

    struct PackedSequenceLayout {
        std::vector<float> positions;
        std::vector<TokenModulationSpan> segments;
        std::vector<SequenceSegment> sequence_segments;
        TokenModulationSpan video_segment{};
        TokenModulationSpan audio_segment{};
        std::vector<float> timesteps;
    };

    static float video_span(int64_t frame) {
        static const int spans[5] = {1, 4, 4, 4, 4};
        return FRAME_RESCALE * spans[frame % 5];
    }

    static std::vector<float> spatial_axis(int64_t dim,
                                           float sqrt_area) {
        int64_t count = dim / 2;
        float ratio   = static_cast<float>(dim) / sqrt_area;
        std::vector<float> result(static_cast<size_t>(count));
        for (int64_t i = 0; i < count; ++i) {
            result[static_cast<size_t>(i)] =
                (static_cast<float>(i) * (ratio / count) + (1.f - ratio) * 0.5f) * 32.f;
        }
        return result;
    }

    static int find_or_add_timestep(std::vector<float>* values, float value) {
        auto it = std::find(values->begin(), values->end(), value);
        if (it != values->end()) {
            return static_cast<int>(it - values->begin());
        }
        values->push_back(value);
        return static_cast<int>(values->size() - 1);
    }

    static PackedSequenceLayout build_layout(int64_t text_len,
                                             int64_t latent_t,
                                             int64_t latent_h,
                                             int64_t latent_w,
                                             int64_t audio_t,
                                             const std::vector<sd::Tensor<float>>& condition_videos,
                                             const std::vector<sd::Tensor<float>>& condition_audios,
                                             const sd::Tensor<int32_t>& keyframe_indices,
                                             const std::vector<MiniMaxH3ReferenceBlock>& reference_blocks,
                                             const sd::Tensor<int32_t>& text_tags,
                                             float video_t,
                                             float audio_timestep) {
        PackedSequenceLayout layout;
        float sqrt_area    = std::sqrt(static_cast<float>(latent_h * latent_w));
        auto h_axis        = spatial_axis(latent_h, sqrt_area);
        auto w_axis        = spatial_axis(latent_w, sqrt_area);
        int64_t frame_rows = static_cast<int64_t>(h_axis.size() * w_axis.size());
        int64_t row        = 0;

        auto append_position = [&](float t, float h, float w) {
            layout.positions.push_back(t);
            layout.positions.push_back(h);
            layout.positions.push_back(w);
        };
        for (int64_t i = 0; i < text_len; ++i) {
            append_position(static_cast<float>(i), 0.f, 0.f);
        }
        layout.sequence_segments.push_back({0, text_len, SequenceKind::TEXT});

        int video_time_row           = find_or_add_timestep(&layout.timesteps, video_t);
        int audio_time_row           = find_or_add_timestep(&layout.timesteps, audio_timestep);
        int condition_time_row       = find_or_add_timestep(&layout.timesteps,
                                                            std::max(video_t, VISUAL_COND_TIMESTEP));
        int audio_condition_time_row = find_or_add_timestep(&layout.timesteps,
                                                            std::max(audio_timestep, 1.f));

        int64_t run_start = 0;
        int current_tag   = text_tags.empty() ? 1 : text_tags[0];
        for (int64_t i = 1; i <= text_len; ++i) {
            int tag = i < text_len && !text_tags.empty() ? text_tags[i] : -1;
            if (i == text_len || tag != current_tag) {
                layout.segments.push_back({run_start,
                                           i,
                                           video_time_row * 3 + current_tag});
                run_start   = i;
                current_tag = tag;
            }
        }
        row = text_len;

        auto condition_spatial_axes = [&](const sd::Tensor<float>& condition) {
            float area = std::sqrt(static_cast<float>(condition.shape()[0] * condition.shape()[1]));
            return std::make_pair(spatial_axis(condition.shape()[1], area),
                                  spatial_axis(condition.shape()[0], area));
        };
        auto append_video_positions = [&](const sd::Tensor<float>& condition,
                                          float cursor) {
            auto axes = condition_spatial_axes(condition);
            for (int64_t t = 0; t < condition.shape()[2]; ++t) {
                for (float h : axes.first) {
                    for (float w : axes.second) {
                        append_position(cursor, h, w);
                    }
                }
                cursor += video_span(t);
            }
            return cursor;
        };
        auto append_audio_positions = [&](int64_t length,
                                          float cursor,
                                          float w_low,
                                          float w_high) {
            for (int channel = 0; channel < 2; ++channel) {
                float w = channel == 0 ? w_low : w_high;
                for (int64_t t = 0; t < length; ++t) {
                    append_position(cursor + static_cast<float>(t), 0.f, w);
                }
            }
        };

        float cursor = static_cast<float>(text_len);
        if (reference_blocks.empty()) {
            float video_duration = 0.f;
            for (int64_t t = 0; t < latent_t; ++t) {
                video_duration += video_span(t);
            }
            for (size_t index = 0; index < condition_videos.size(); ++index) {
                const auto& condition = condition_videos[index];
                auto axes             = condition_spatial_axes(condition);
                int64_t count         = condition.shape()[2] *
                                static_cast<int64_t>(axes.first.size() * axes.second.size());
                bool is_first    = keyframe_indices.empty() || keyframe_indices[static_cast<int64_t>(index)] == 0;
                float keyframe_t = is_first ? static_cast<float>(text_len)
                                            : static_cast<float>(text_len) + video_duration - FRAME_RESCALE;
                for (int64_t t = 0; t < condition.shape()[2]; ++t) {
                    for (float h : axes.first) {
                        for (float w : axes.second) {
                            append_position(keyframe_t, h, w);
                        }
                    }
                }
                layout.sequence_segments.push_back({row,
                                                    row + count,
                                                    SequenceKind::CONDITION_VIDEO,
                                                    static_cast<int32_t>(index)});
                layout.segments.push_back({row, row + count, condition_time_row * 3});
                row += count;
            }
        } else {
            for (const auto& block : reference_blocks) {
                const sd::Tensor<float>* ref_video = nullptr;
                const sd::Tensor<float>* ref_audio = nullptr;
                if (block.video_index >= 0) {
                    GGML_ASSERT(block.video_index < static_cast<int32_t>(condition_videos.size()));
                    ref_video = &condition_videos[static_cast<size_t>(block.video_index)];
                }
                if (block.audio_index >= 0) {
                    GGML_ASSERT(block.audio_index < static_cast<int32_t>(condition_audios.size()));
                    ref_audio = &condition_audios[static_cast<size_t>(block.audio_index)];
                }

                float block_end = cursor;
                if (block.kind == MiniMaxH3ReferenceKind::AUDIO ||
                    block.kind == MiniMaxH3ReferenceKind::VIDEO_AUDIO) {
                    GGML_ASSERT(ref_audio != nullptr);
                    float w_low  = w_axis.front();
                    float w_high = w_axis.back();
                    if (ref_video != nullptr) {
                        auto axes = condition_spatial_axes(*ref_video);
                        w_low     = axes.second.front();
                        w_high    = axes.second.back();
                    }
                    int64_t count = ref_audio->shape()[0] * 2;
                    append_audio_positions(ref_audio->shape()[0], cursor, w_low, w_high);
                    layout.sequence_segments.push_back({row,
                                                        row + count,
                                                        SequenceKind::CONDITION_AUDIO,
                                                        block.audio_index});
                    layout.segments.push_back({row,
                                               row + count,
                                               audio_condition_time_row * 3 + 2});
                    row += count;
                    block_end = std::max(block_end, cursor + static_cast<float>(ref_audio->shape()[0]));
                }

                if (block.kind != MiniMaxH3ReferenceKind::AUDIO) {
                    GGML_ASSERT(ref_video != nullptr);
                    auto axes     = condition_spatial_axes(*ref_video);
                    int64_t count = ref_video->shape()[2] *
                                    static_cast<int64_t>(axes.first.size() * axes.second.size());
                    float video_end = append_video_positions(*ref_video, cursor);
                    layout.sequence_segments.push_back({row,
                                                        row + count,
                                                        SequenceKind::CONDITION_VIDEO,
                                                        block.video_index});
                    layout.segments.push_back({row, row + count, condition_time_row * 3});
                    row += count;
                    block_end = block.kind == MiniMaxH3ReferenceKind::IMAGE
                                    ? std::max(block_end, cursor + 1.f)
                                    : std::max(block_end, video_end);
                }
                cursor = block_end;
            }
        }

        int64_t audio_start = row;
        append_audio_positions(audio_t, cursor, w_axis.front(), w_axis.back());
        layout.audio_segment = {audio_start, row + audio_t * 2, audio_time_row};
        layout.sequence_segments.push_back({audio_start,
                                            row + audio_t * 2,
                                            SequenceKind::TARGET_AUDIO});
        layout.segments.push_back({audio_start,
                                   row + audio_t * 2,
                                   audio_time_row * 3 + 2});
        row += audio_t * 2;

        int64_t video_start = row;
        for (int64_t t = 0; t < latent_t; ++t) {
            for (float h : h_axis) {
                for (float w : w_axis) {
                    append_position(cursor, h, w);
                }
            }
            cursor += video_span(t);
        }
        int64_t video_rows   = latent_t * frame_rows;
        layout.video_segment = {video_start, video_start + video_rows, video_time_row};
        layout.sequence_segments.push_back({video_start,
                                            video_start + video_rows,
                                            SequenceKind::TARGET_VIDEO});
        layout.segments.push_back({video_start,
                                   video_start + video_rows,
                                   video_time_row * 3});
        return layout;
    }

    struct MiniMaxH3Runner : public DiffusionModelRunner {
        Config config;
        MiniMaxH3Transformer3DModel model;
        sd::Tensor<float> video_input_cache;
        sd::Tensor<float> audio_input_cache;
        sd::Tensor<float> position_input_cache;
        sd::Tensor<float> timestep_feature_input_cache;
        sd::Tensor<int32_t> curve_index_input_cache;
        sd::Tensor<int32_t> curve_upper_index_input_cache;
        sd::Tensor<float> curve_fraction_input_cache;

        MiniMaxH3Runner(ggml_backend_t backend,
                        const String2TensorStorage& tensors,
                        const std::string& prefix                           = "model.diffusion_model",
                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
            : DiffusionModelRunner(backend, prefix, weight_manager),
              config(Config::detect_from_weights(tensors, prefix)),
              model(config) {
            model.init(params_ctx, tensors, prefix);
        }

        std::string get_desc() override {
            return "minimax_h3";
        }

        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors,
                               const std::string& prefix) override {
            model.get_param_tensors(tensors, prefix);
        }

        std::pair<sd::Tensor<float>, sd::Tensor<float>> split_av_latents(const sd::Tensor<float>& packed,
                                                                         int audio_length) const {
            GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5);
            int64_t spatial      = packed.shape()[0] * packed.shape()[1] * packed.shape()[2];
            int64_t video_values = spatial * config.video_latent_channels;
            sd::Tensor<float> video({packed.shape()[0],
                                     packed.shape()[1],
                                     packed.shape()[2],
                                     config.video_latent_channels,
                                     1});
            std::copy_n(packed.data(), static_cast<size_t>(video_values), video.data());
            if (audio_length <= 0) {
                return {video, {}};
            }
            int64_t audio_values = audio_length * 2 * config.audio_latent_channels;
            GGML_ASSERT(packed.numel() >= video_values + audio_values);
            sd::Tensor<float> audio({audio_length, 2, config.audio_latent_channels, 1});
            std::copy_n(packed.data() + video_values,
                        static_cast<size_t>(audio_values),
                        audio.data());
            return {video, audio};
        }

        ggml_tensor* merge_av_latents(ggml_context* ctx,
                                      ggml_tensor* video,
                                      ggml_tensor* audio) const {
            int64_t divisor = video->ne[0] * video->ne[1] * video->ne[2];
            int64_t values  = ggml_nelements(audio);
            int64_t padding = (divisor - values % divisor) % divisor;
            audio           = ggml_reshape_4d(ctx, ggml_cont(ctx, audio), values, 1, 1, 1);
            if (padding > 0) {
                audio = ggml_ext_pad(ctx, audio, static_cast<int>(padding), 0, 0, 0);
            }
            audio = ggml_reshape_4d(ctx,
                                    audio,
                                    video->ne[0],
                                    video->ne[1],
                                    video->ne[2],
                                    (values + padding) / divisor);
            return ggml_concat(ctx, video, audio, 3);
        }

        ggml_cgraph* build_graph(const sd::Tensor<float>& packed,
                                 const sd::Tensor<float>& timestep,
                                 const sd::Tensor<float>& context_tensor,
                                 const std::vector<sd::Tensor<float>>& condition_videos,
                                 const std::vector<sd::Tensor<float>>& condition_audios,
                                 const sd::Tensor<int32_t>& text_tags,
                                 const sd::Tensor<int32_t>& keyframe_indices,
                                 const std::vector<MiniMaxH3ReferenceBlock>& reference_blocks,
                                 int audio_length,
                                 float video_shift,
                                 float audio_shift) {
            auto split        = split_av_latents(packed, audio_length);
            video_input_cache = std::move(split.first);
            audio_input_cache = std::move(split.second);
            GGML_ASSERT(!audio_input_cache.empty());
            GGML_ASSERT(!context_tensor.empty());

            auto video         = make_input(video_input_cache);
            auto audio_carrier = make_input(audio_input_cache);
            auto context       = make_input(context_tensor);
            std::vector<ggml_tensor*> condition_inputs;
            condition_inputs.reserve(condition_videos.size());
            for (const auto& condition : condition_videos) {
                condition_inputs.push_back(make_input(condition));
            }
            std::vector<ggml_tensor*> audio_condition_inputs;
            audio_condition_inputs.reserve(condition_audios.size());
            for (const auto& condition : condition_audios) {
                audio_condition_inputs.push_back(make_input(condition));
            }

            float sigma_v     = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
            float sigma_a     = time_shift_sigma(sigma_v, video_shift, audio_shift);
            float audio_scale = video_shift / audio_shift;
            float t_v         = 1.f - sigma_v;
            float t_a         = 1.f - sigma_a;
            // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed
            // latent follows one sigma schedule. Restore x_a for the H3 network.
            auto audio  = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
            auto layout = build_layout(context_tensor.shape()[1],
                                       video_input_cache.shape()[2],
                                       video_input_cache.shape()[1],
                                       video_input_cache.shape()[0],
                                       audio_length,
                                       condition_videos,
                                       condition_audios,
                                       keyframe_indices,
                                       reference_blocks,
                                       text_tags,
                                       t_v,
                                       t_a);

            position_input_cache = sd::Tensor<float>(
                {3, static_cast<int64_t>(layout.positions.size() / 3)},
                layout.positions);
            auto positions = make_input(position_input_cache);

            ggml_tensor* timestep_features   = nullptr;
            ggml_tensor* curve_indices       = nullptr;
            ggml_tensor* curve_upper_indices = nullptr;
            ggml_tensor* curve_fractions     = nullptr;
            if (config.uses_adaln_curves()) {
                std::vector<int32_t> indices(layout.timesteps.size());
                std::vector<int32_t> upper_indices(layout.timesteps.size());
                std::vector<float> fractions(layout.timesteps.size());
                for (size_t i = 0; i < layout.timesteps.size(); ++i) {
                    float position   = std::clamp(layout.timesteps[i], 0.f, 1.f) * (config.adaln_curve_grid - 1);
                    int index        = std::min(static_cast<int>(std::floor(position)),
                                                static_cast<int>(config.adaln_curve_grid - 2));
                    indices[i]       = index;
                    upper_indices[i] = index + 1;
                    fractions[i]     = position - index;
                }
                curve_index_input_cache = sd::Tensor<int32_t>(
                    {static_cast<int64_t>(indices.size())},
                    indices);
                curve_upper_index_input_cache = sd::Tensor<int32_t>(
                    {static_cast<int64_t>(upper_indices.size())},
                    upper_indices);
                curve_fraction_input_cache = sd::Tensor<float>(
                    {1, static_cast<int64_t>(fractions.size())},
                    fractions);
                curve_indices       = make_input(curve_index_input_cache);
                curve_upper_indices = make_input(curve_upper_index_input_cache);
                curve_fractions     = make_input(curve_fraction_input_cache);
            } else {
                timestep_feature_input_cache = sd::Tensor<float>(
                    {config.timestep_input_dim, static_cast<int64_t>(layout.timesteps.size())},
                    timestep_embedding(layout.timesteps,
                                       static_cast<int>(config.timestep_input_dim),
                                       10000,
                                       true,
                                       1.f));
                timestep_features = make_input(timestep_feature_input_cache);
            }

            auto runner_ctx = get_context();
            auto output     = model.forward(&runner_ctx,
                                            video,
                                            audio,
                                            context,
                                            condition_inputs,
                                            audio_condition_inputs,
                                            positions,
                                            timestep_features,
                                            curve_indices,
                                            curve_upper_indices,
                                            curve_fractions,
                                            layout.segments,
                                            layout.sequence_segments,
                                            layout.video_segment,
                                            layout.audio_segment);
            // Convert the model's audio velocity to d(c_a) / d(sigma_v).
            output.second = ggml_add(compute_ctx,
                                     ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale),
                                     ggml_ext_scale(compute_ctx,
                                                    output.second,
                                                    1.f + (audio_scale - 1.f) * sigma_a));
            auto merged   = merge_av_latents(compute_ctx, output.first, output.second);
            auto graph    = new_graph_custom(H3_GRAPH_SIZE);
            ggml_build_forward_expand(graph, merged);
            return graph;
        }

        sd::Tensor<float> compute(int n_threads,
                                  const DiffusionParams& params) override {
            GGML_ASSERT(params.x != nullptr && params.timesteps != nullptr && params.context != nullptr);
            const auto* extra = diffusion_extra_as<MiniMaxH3DiffusionExtra>(params);
            static const std::vector<sd::Tensor<float>> empty_conditions;
            static const std::vector<MiniMaxH3ReferenceBlock> empty_reference_blocks;
            const auto& conditions       = params.ref_latents == nullptr ? empty_conditions : *params.ref_latents;
            const auto& audio_conditions = extra->reference_audio_latents == nullptr
                                               ? empty_conditions
                                               : *extra->reference_audio_latents;
            const auto& reference_blocks = extra->reference_blocks == nullptr
                                               ? empty_reference_blocks
                                               : *extra->reference_blocks;
            const sd::Tensor<int32_t> empty_int;
            auto get_graph = [&]() {
                return build_graph(*params.x,
                                   *params.timesteps,
                                   *params.context,
                                   conditions,
                                   audio_conditions,
                                   extra->text_token_tags == nullptr ? empty_int : *extra->text_token_tags,
                                   extra->keyframe_indices == nullptr ? empty_int : *extra->keyframe_indices,
                                   reference_blocks,
                                   extra->audio_length,
                                   extra->video_sigma_shift,
                                   extra->audio_sigma_shift);
            };
            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph,
                                                                       n_threads,
                                                                       false),
                                                   params.x->dim());
        }
    };

}  // namespace MiniMaxH3

#endif  // __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__