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

docs/minimax_h3.md

@@ -99,6 +99,24 @@ The C API exposes the same inputs through `ref_images`, `ref_videos`, and
 `ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own
 frame rate and optional soundtrack; non-24-fps inputs are resampled internally.
 
+## Sol-H3-VAE decoding
+
+Pass a [Sol-H3-VAE checkpoint](https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/vae)
+with `--vae` to use the distilled video decoder. The ComfyUI FP16 and INT8
+ConvRot checkpoints include the MiniMax-H3 encoder for image and video inputs.
+Supply `--audio-vae` as usual to decode the soundtrack.
+
+Sol decodes the whole clip by default. Use `--vae-tiling` for spatial tiling or
+`--temporal-tiling` for temporal tiling. Either can reduce memory use but may
+affect output quality because the decoder has less context at tile boundaries.
+
+Temporal tiles default to 10 latent frames with 5 latent frames of overlap.
+Set `temporal_tile_frames` and `temporal_tile_overlap` in
+`--extra-tiling-args` to override them. Both must be multiples of 5 latent
+frames, and overlap must be smaller than the tile. The last tile may be shorter.
+Circular decoding requires spatial tiling on each circular axis, with a tile
+size smaller than the input on that axis.
+
 ## Latent hires second pass (two-stage sampling)
 
 MiniMax-H3 supports `--hires` with two interchangeable second-stage upscalers:

docs/performance.md

@@ -92,10 +92,12 @@ dimensions down to the minimum size. Each spatial retry must reduce the effectiv
 tile size. These runtime adjustments do not change the caller's parameters.
 Execution failures are not retried, and encoding has no automatic OOM retry.
 
-`--temporal-tiling` remains independent of spatial tiling. MiniMax H3 always uses
-spatial tiling (256x256 pixels and 25% overlap by default) and its own temporal
-windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap`; explicit
-spatial sizes are honored.
+`--temporal-tiling` remains independent of spatial tiling. The original MiniMax-H3
+video VAE uses spatial tiling (256x256 pixels and 25% overlap by default) and its own
+temporal windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap` and
+explicit spatial sizes are honored. [Sol-H3-VAE](minimax_h3.md#sol-h3-vae-decoding)
+decodes the whole clip by default and supports opt-in spatial and block-aligned
+temporal tiling. Tiling may affect output quality.
 
 ## Offload weights to the CPU to save VRAM without reducing generation speed.
 

src/model/vae/minimax_h3_sol_attention.hpp

@@ -0,0 +1,204 @@
+#ifndef __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__
+#define __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__
+
+#include <algorithm>
+#include <array>
+#include <map>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/tensor.hpp"
+
+namespace MiniMaxH3VAE {
+
+    constexpr int64_t sol_window_size    = 16;
+    constexpr int64_t sol_block_size     = 5;
+    constexpr int64_t sol_global_rows    = 5;
+    constexpr int64_t sol_spatial_shift  = sol_window_size / 2;
+    constexpr int64_t sol_temporal_shift = 2;
+
+    struct SolAttentionPlan {
+        template <typename Tensor>
+        struct Batch {
+            Tensor output_rows;
+            Tensor queries;
+            Tensor keys;
+            Tensor previous_queries;
+            Tensor previous_keys;
+        };
+
+        template <typename Tensor>
+        struct Block {
+            Tensor rows;
+            std::vector<Batch<Tensor>> batches;
+        };
+
+        std::vector<Block<sd::Tensor<int32_t>>> blocks;
+
+        SolAttentionPlan(int64_t width, int64_t height, int64_t frames, bool shifted) {
+            const int64_t frame_size = width * height;
+            const int64_t pixels     = frames * frame_size;
+            const int64_t block_rows = sol_block_size * frame_size;
+            const int64_t count      = (frames + sol_block_size - 1) / sol_block_size;
+            const int64_t shift      = shifted ? sol_spatial_shift : 0;
+            const int64_t time_shift = shifted ? sol_temporal_shift : 0;
+
+            auto pack = [](std::vector<int32_t> rows) {
+                const auto size = static_cast<int64_t>(rows.size());
+                return sd::Tensor<int32_t>({size}, std::move(rows));
+            };
+            for (int64_t block = 0; block < count; ++block) {
+                std::vector<int32_t> rows;
+                const int64_t start = block * block_rows;
+                const int64_t end   = std::min(pixels, start + block_rows);
+                for (int64_t i = start; i < end; ++i) {
+                    rows.push_back(static_cast<int32_t>(i));
+                }
+                std::vector<int32_t> globals;
+                std::vector<int32_t> queries;
+                for (int64_t i = 0; i < sol_global_rows; ++i) {
+                    globals.push_back(static_cast<int32_t>(pixels + block * sol_global_rows + i));
+                    queries.push_back(static_cast<int32_t>(rows.size()));
+                    rows.push_back(globals.back());
+                }
+                blocks.push_back({pack(std::move(rows)), {}});
+                blocks.back().batches.push_back({pack(std::move(globals)), pack(std::move(queries)), {}, {}, {}});
+            }
+
+            struct Window {
+                std::vector<int32_t> output_rows;
+                std::vector<int32_t> queries;
+                std::vector<int32_t> keys;
+                std::vector<int32_t> previous_queries;
+                std::vector<int32_t> previous_keys;
+            };
+            for (int64_t t = -time_shift; t < frames; t += sol_block_size) {
+                const int64_t first       = std::max<int64_t>(0, t);
+                const int64_t end         = std::min(frames, t + sol_block_size);
+                const int64_t first_block = first / sol_block_size;
+                const int64_t last_block  = (end - 1) / sol_block_size;
+                std::map<int64_t, std::vector<Window>> groups;
+                for (int64_t y = -shift; y < height; y += sol_window_size) {
+                    for (int64_t x = -shift; x < width; x += sol_window_size) {
+                        Window window;
+                        for (int64_t i = first; i < end; ++i) {
+                            for (int64_t j = std::max<int64_t>(0, y); j < std::min(height, y + sol_window_size); ++j) {
+                                for (int64_t k = std::max<int64_t>(0, x); k < std::min(width, x + sol_window_size); ++k) {
+                                    const int64_t row = (i * height + j) * width + k;
+                                    window.output_rows.push_back(static_cast<int32_t>(row));
+                                    const int64_t block = i / sol_block_size;
+                                    auto& queries       = block == last_block ? window.queries : window.previous_queries;
+                                    queries.push_back(static_cast<int32_t>(row - block * block_rows));
+                                }
+                            }
+                        }
+                        window.keys          = window.queries;
+                        window.previous_keys = window.previous_queries;
+                        for (int64_t block = first_block; block <= last_block; ++block) {
+                            const int64_t start = std::min(block_rows, pixels - block * block_rows);
+                            auto& keys          = block == last_block ? window.keys : window.previous_keys;
+                            for (int64_t i = 0; i < sol_global_rows; ++i) {
+                                keys.push_back(static_cast<int32_t>(start + i));
+                            }
+                        }
+                        groups[window.output_rows.size()].push_back(std::move(window));
+                    }
+                }
+                // Bound the workspace of backends that fall back to unfused attention.
+                constexpr size_t max_windows = 4;
+                for (const auto& [query_count, windows] : groups) {
+                    for (size_t first_window = 0; first_window < windows.size(); first_window += max_windows) {
+                        const auto n      = static_cast<int64_t>(std::min(max_windows, windows.size() - first_window));
+                        auto pack_windows = [&](auto member) -> sd::Tensor<int32_t> {
+                            const auto rows = static_cast<int64_t>((windows[first_window].*member).size());
+                            if (rows == 0) {
+                                return {};
+                            }
+                            sd::Tensor<int32_t> indices({rows, n});
+                            for (int64_t i = 0; i < n; ++i) {
+                                const auto& source = windows[first_window + i].*member;
+                                std::copy(source.begin(), source.end(), indices.data() + i * rows);
+                            }
+                            return indices;
+                        };
+                        blocks[last_block].batches.push_back({pack_windows(&Window::output_rows),
+                                                              pack_windows(&Window::queries),
+                                                              pack_windows(&Window::keys),
+                                                              pack_windows(&Window::previous_queries),
+                                                              pack_windows(&Window::previous_keys)});
+                    }
+                }
+            }
+        }
+
+        struct Inputs {
+            std::vector<Block<ggml_tensor*>> blocks;
+            ggml_tensor* row_ids = nullptr;
+
+            template <typename Project>
+            ggml_tensor* forward(GGMLRunnerContext* runner_ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* pe,
+                                 int64_t heads,
+                                 Project project) const {
+                auto ctx          = runner_ctx->ggml_ctx;
+                const int64_t dim = x->ne[0];
+                auto output       = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, x->ne[1]);
+                auto positions    = ggml_reshape_2d(ctx, pe, pe->ne[0] * pe->ne[1] * pe->ne[2], pe->ne[3]);
+                auto gather       = [&](ggml_tensor* rows, ggml_tensor* indices) {
+                    auto flat = ggml_reshape_1d(ctx, indices, ggml_nelements(indices));
+                    return ggml_reshape_3d(ctx, ggml_get_rows(ctx, rows, flat), dim, indices->ne[0], indices->ne[1]);
+                };
+                std::array<ggml_tensor*, 3> previous{};
+                for (const auto& block : blocks) {
+                    auto local_pe = ggml_reshape_4d(ctx, ggml_get_rows(ctx, positions, block.rows),
+                                                    pe->ne[0], pe->ne[1], pe->ne[2], block.rows->ne[0]);
+                    auto current  = project(ggml_get_rows(ctx, x, block.rows), local_pe);
+                    for (const auto& batch : block.batches) {
+                        // Gather first so a crossing window never copies either full block.
+                        auto window = [&](size_t i, ggml_tensor* now, ggml_tensor* before) {
+                            // Global queries read the whole native block in its existing order.
+                            auto rows = now == nullptr ? current[i] : gather(current[i], now);
+                            if (before != nullptr) {
+                                rows = ggml_concat(ctx, gather(previous[i], before), rows, 1);
+                            }
+                            return rows;
+                        };
+                        auto attended = ggml_ext_attention_ext(runner_ctx,
+                                                               window(0, batch.queries, batch.previous_queries),
+                                                               window(1, batch.keys, batch.previous_keys),
+                                                               window(2, batch.keys, batch.previous_keys),
+                                                               heads, nullptr, false, runner_ctx->flash_attn_enabled);
+                        auto indices  = ggml_reshape_1d(ctx, batch.output_rows, ggml_nelements(batch.output_rows));
+                        output        = ggml_set_rows(ctx, output, ggml_reshape_2d(ctx, attended, dim, indices->ne[0]), indices);
+                    }
+                    // Later windows can only read this block and its immediate successor.
+                    previous = current;
+                }
+                return output;
+            }
+        };
+
+        template <typename BindInput>
+        Inputs bind(BindInput bind_input) const {
+            Inputs inputs;
+            for (const auto& block : blocks) {
+                Block<ggml_tensor*> bound{bind_input(block.rows), {}};
+                for (const auto& batch : block.batches) {
+                    auto bind_optional = [&](const auto& indices) {
+                        return indices.empty() ? nullptr : bind_input(indices);
+                    };
+                    bound.batches.push_back({bind_input(batch.output_rows), bind_input(batch.queries),
+                                             bind_optional(batch.keys), bind_optional(batch.previous_queries),
+                                             bind_optional(batch.previous_keys)});
+                }
+                inputs.blocks.push_back(std::move(bound));
+            }
+            return inputs;
+        }
+    };
+
+}  // namespace MiniMaxH3VAE
+
+#endif  // __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__

src/model/vae/minimax_h3_vae.hpp

@@ -13,26 +13,34 @@
 #include "core/rng_mt19937.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
+#include "model/vae/minimax_h3_sol_attention.hpp"
 #include "model/vae/vae.hpp"
 
 namespace MiniMaxH3VAE {
 
     constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144;
 
+    static int64_t sol_output_frames(int64_t latent_frames, int patch_size_t) {
+        const int64_t blocks = (latent_frames + sol_block_size - 1) / sol_block_size;
+        return latent_frames * patch_size_t - blocks * (patch_size_t - 1);
+    }
+
     struct MiniMaxH3VAEConfig {
         int num_layers   = 36;
         int patch_size_t = 4;
+        bool sol         = false;
 
         static MiniMaxH3VAEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) {
             MiniMaxH3VAEConfig config;
-            auto proj = tensors.find(prefix + ".decoder.proj_out.weight");
+            config.sol = tensors.count(prefix + ".decoder.blocks.0.scale1") != 0;
+            auto proj  = tensors.find(prefix + ".decoder.proj_out.weight");
             if (proj != tensors.end()) {
                 config.patch_size_t = static_cast<int>(proj->second.ne[1] / (3 * 16 * 16));
                 if (config.patch_size_t != 1 && config.patch_size_t != 4) {
                     throw std::runtime_error("unsupported MiniMax-H3 VAE temporal patch size");
                 }
                 int layers              = 0;
-                const auto layer_prefix = prefix + ".decoder.transformer_blocks.";
+                const auto layer_prefix = prefix + (config.sol ? ".decoder.blocks." : ".decoder.transformer_blocks.");
                 for (const auto& [name, _] : tensors) {
                     if (starts_with(name, layer_prefix)) {
                         layers = std::max(layers, std::atoi(name.c_str() + layer_prefix.size()) + 1);
@@ -41,7 +49,12 @@ namespace MiniMaxH3VAE {
                 if (layers > 0) {
                     config.num_layers = layers;
                 }
-                LOG_VERBOSE("minimax_h3_vae: layers=%d, patch_size_t=%d", config.num_layers, config.patch_size_t);
+            }
+            if (config.sol || proj != tensors.end()) {
+                LOG_VERBOSE("minimax_h3_vae: decoder=%s, layers=%d, patch_size_t=%d",
+                            config.sol ? "sol" : "original",
+                            config.num_layers,
+                            config.patch_size_t);
             }
             return config;
         }
@@ -283,11 +296,10 @@ namespace MiniMaxH3VAE {
             blocks["to_out"] = std::make_shared<Linear>(dim, dim, true);
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* pe) {
+        std::array<ggml_tensor*, 3> project(GGMLRunnerContext* ctx,
+                                            ggml_tensor* x,
+                                            ggml_tensor* pe) {
             auto to_qkv         = std::dynamic_pointer_cast<Linear>(blocks["to_qkv"]);
-            auto to_out         = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
             auto qkv_projection = to_qkv->forward(ctx, x);
             int64_t sequence    = x->ne[1];
             int64_t batch_size  = x->ne[2] * x->ne[3];
@@ -320,6 +332,29 @@ namespace MiniMaxH3VAE {
             k                   = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
             q                   = apply_partial_rope(ctx, q, pe);
             k                   = apply_partial_rope(ctx, k, pe);
+            return {q, k, v};
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             RMSNorm& norm,
+                             const SolAttentionPlan::Inputs* sol) {
+            auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
+            if (sol != nullptr) {
+                auto project_rows = [&](ggml_tensor* rows, ggml_tensor* local_pe) {
+                    auto [q, k, v] = project(ctx, norm.forward(ctx, rows), local_pe);
+                    auto as_rows   = [&](ggml_tensor* tensor) {
+                        return ggml_reshape_2d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, tensor), dim, rows->ne[1]);
+                    };
+                    return std::array<ggml_tensor*, 3>{
+                        as_rows(ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)),
+                        as_rows(ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)),
+                        as_rows(v)};
+                };
+                return to_out->forward(ctx, sol->forward(ctx, x, pe, num_head, project_rows));
+            }
+            auto [q, k, v]      = project(ctx, norm.forward(ctx, x), pe);
             auto out            = ggml_ext_attention_ext(ctx,
                                                          q,
                                                          k,
@@ -370,21 +405,33 @@ namespace MiniMaxH3VAE {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* x,
-                             ggml_tensor* pe) {
+                             ggml_tensor* pe,
+                             const SolAttentionPlan::Inputs* sol,
+                             int64_t rows_per_chunk) {
             auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
             auto attn  = std::dynamic_pointer_cast<DecoderAttention>(blocks["attn"]);
             auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
             auto ff    = std::dynamic_pointer_cast<DecoderFeedForward>(blocks["ff"]);
-            x          = ggml_add(ctx->ggml_ctx,
-                                  x,
-                                  ggml_mul(ctx->ggml_ctx,
-                                           attn->forward(ctx, norm1->forward(ctx, x), pe),
-                                           params["scale1"]));
-            return ggml_add(ctx->ggml_ctx,
-                            x,
-                            ggml_mul(ctx->ggml_ctx,
-                                     ff->forward(ctx, norm2->forward(ctx, x)),
-                                     params["scale2"]));
+            x           = ggml_add(ctx->ggml_ctx,
+                                   x,
+                                   ggml_mul(ctx->ggml_ctx,
+                                            attn->forward(ctx, x, pe, *norm1, sol),
+                                            params["scale1"]));
+            auto finish = [&](ggml_tensor* rows) {
+                auto update = ff->forward(ctx, norm2->forward(ctx, rows));
+                return ggml_add(ctx->ggml_ctx, rows, ggml_mul(ctx->ggml_ctx, update, params["scale2"]));
+            };
+            if (rows_per_chunk == 0 || x->ne[1] <= rows_per_chunk) {
+                return finish(x);
+            }
+            // Normalize and update each independent chunk before writing it to the layer output.
+            auto output = ggml_new_tensor_2d(ctx->ggml_ctx, x->type, x->ne[0], x->ne[1]);
+            for (int64_t row = 0; row < x->ne[1]; row += rows_per_chunk) {
+                auto input   = ggml_ext_slice(ctx, x, 1, row, std::min(row + rows_per_chunk, x->ne[1]), false);
+                auto indices = ggml_ext_slice(ctx, sol->row_ids, 0, row, std::min(row + rows_per_chunk, x->ne[1]), false);
+                output       = ggml_set_rows(ctx->ggml_ctx, output, finish(input), indices);
+            }
+            return output;
         }
     };
 
@@ -394,12 +441,13 @@ namespace MiniMaxH3VAE {
         static constexpr int num_register_tokens = 4;
         static constexpr int patch_size          = 16;
         int patch_size_t;
+        bool sol;
 
         explicit Decoder(const MiniMaxH3VAEConfig& config)
-            : num_layers(config.num_layers), patch_size_t(config.patch_size_t) {
+            : num_layers(config.num_layers), patch_size_t(config.patch_size_t), sol(config.sol) {
             blocks["x_embedder"] = std::make_shared<Linear>(24, dim, true);
             for (int i = 0; i < num_layers; ++i) {
-                blocks["transformer_blocks." + std::to_string(i)] =
+                blocks[(sol ? "blocks." : "transformer_blocks.") + std::to_string(i)] =
                     std::make_shared<DecoderBlock>();
             }
             blocks["norm_out"] = std::make_shared<LayerNorm>(dim, 1e-5f, true, true);
@@ -418,14 +466,18 @@ namespace MiniMaxH3VAE {
                                                            GGML_TYPE_F32,
                                                            dim,
                                                            num_register_tokens);
-            if (patch_size_t != 1) {
+            if (patch_size_t != 1 && !sol) {
                 params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
             }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* z,
-                             ggml_tensor* pe) {
+                             ggml_tensor* pe,
+                             ggml_tensor* shifted_pe,
+                             const std::array<SolAttentionPlan::Inputs, 2>& sol_inputs,
+                             ggml_tensor* pixel_mean,
+                             ggml_tensor* pixel_std) {
             int64_t width      = z->ne[0];
             int64_t height     = z->ne[1];
             int64_t num_frames = z->ne[2];
@@ -442,34 +494,68 @@ namespace MiniMaxH3VAE {
             auto x_embedder     = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
             auto h              = x_embedder->forward(ctx, z);
             int64_t num_patches = h->ne[1];
-            h                   = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1);
             auto zero           = ggml_ext_scale(ctx,
                                                  ggml_ext_slice(ctx, h, 1, 0, 1),
                                                  0.f);
-            h                   = ggml_concat(ctx->ggml_ctx, h, zero, 1);
+            auto registers      = ggml_concat(ctx->ggml_ctx, params["register_tokens"], zero, 1);
+            if (sol) {
+                const int64_t num_blocks = (num_frames + sol_block_size - 1) / sol_block_size;
+                registers                = ggml_repeat(ctx->ggml_ctx, registers,
+                                                       ggml_new_tensor_2d(ctx->ggml_ctx, registers->type, dim, num_blocks * sol_global_rows));
+            }
+            h = ggml_concat(ctx->ggml_ctx, h, registers, 1);
 
+            const int64_t rows_per_chunk = sol ? width * height + sol_global_rows : 0;
             for (int i = 0; i < num_layers; ++i) {
                 auto block = std::dynamic_pointer_cast<DecoderBlock>(
-                    blocks["transformer_blocks." + std::to_string(i)]);
-                h = block->forward(ctx, h, pe);
+                    blocks[(sol ? "blocks." : "transformer_blocks.") + std::to_string(i)]);
+                h = block->forward(ctx, h, sol && i % 2 != 0 ? shifted_pe : pe, sol ? &sol_inputs[i % 2] : nullptr, rows_per_chunk);
                 sd::ggml_graph_cut::mark_graph_cut(h,
                                                    "minimax_h3_vae.decoder.blocks." + std::to_string(i),
                                                    "hidden_states");
             }
 
-            auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
-            auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
-            h             = proj_out->forward(ctx, norm_out->forward(ctx, h));
-            h             = ggml_ext_slice(ctx, h, 1, 0, num_patches);
-            return DiT::unpatchify_3d(ctx,
-                                      h,
-                                      num_frames,
-                                      height,
-                                      width,
-                                      patch_size_t,
-                                      patch_size,
-                                      patch_size,
-                                      true);
+            auto norm_out              = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
+            auto proj_out              = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
+            auto denormalize_and_clamp = [&](ggml_tensor* pixels) {
+                auto scaled = ggml_mul(ctx->ggml_ctx, pixels, pixel_std);
+                return ggml_clamp(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, scaled, pixel_mean), 0.f, 1.f);
+            };
+            if (sol) {
+                const int64_t frame_rows    = width * height;
+                const int64_t output_frames = sol_output_frames(num_frames, patch_size_t);
+                const int64_t frame_size    = width * height * patch_size * patch_size;
+                const int64_t channels      = pixel_mean->ne[3];
+                auto output                 = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, frame_size, channels * output_frames);
+                int64_t frame_offset        = 0;
+                for (int64_t t = 0; t < num_frames; t += sol_block_size) {
+                    const int64_t end = std::min(t + sol_block_size, num_frames);
+                    auto rows         = ggml_ext_slice(ctx, h, 1, t * frame_rows, end * frame_rows, false);
+                    auto patches      = proj_out->forward(ctx, norm_out->forward(ctx, rows));
+                    auto pixels       = DiT::unpatchify_3d(ctx, patches, end - t, height, width,
+                                                           patch_size_t, patch_size, patch_size, true);
+                    // Keep only the last frame of the first latent patch in each native block.
+                    pixels              = ggml_ext_slice(ctx, pixels, 2, patch_size_t - 1, pixels->ne[2], false);
+                    pixels              = denormalize_and_clamp(pixels);
+                    const int64_t count = pixels->ne[2];
+                    for (int64_t c = 0; c < channels; ++c) {
+                        auto plane          = ggml_ext_slice(ctx, pixels, 3, c, c + 1, false);
+                        auto frames         = ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, plane), frame_size, count);
+                        const int64_t start = c * output_frames + frame_offset;
+                        auto indices        = ggml_ext_slice(ctx, sol_inputs[0].row_ids, 0, start, start + count, false);
+                        output              = ggml_set_rows(ctx->ggml_ctx, output, frames, indices);
+                    }
+                    // Write each block before expanding the next block's pixel graph.
+                    ctx->expand_graph(output);
+                    frame_offset += count;
+                }
+                return ggml_reshape_4d(ctx->ggml_ctx, output, width * patch_size, height * patch_size, output_frames, channels);
+            }
+            h           = proj_out->forward(ctx, norm_out->forward(ctx, h));
+            h           = ggml_ext_slice(ctx, h, 1, 0, num_patches);
+            auto pixels = DiT::unpatchify_3d(ctx, h, num_frames, height, width,
+                                             patch_size_t, patch_size, patch_size, true);
+            return denormalize_and_clamp(pixels);
         }
     };
 
@@ -479,9 +565,7 @@ namespace MiniMaxH3VAE {
         explicit MiniMaxH3VideoVAE(const MiniMaxH3VAEConfig& config)
             : single_frame(config.patch_size_t == 1) {
             blocks["encoder"]         = std::make_shared<Encoder>();
-            blocks["quant_conv"]      = std::make_shared<Conv3d>(48,
-                                                                 48,
-                                                                 std::tuple{1, 1, 1});
+            blocks["quant_conv"]      = std::make_shared<Conv3d>(48, 48, std::tuple{1, 1, 1});
             blocks["post_quant_conv"] = std::make_shared<Conv3d>(24,
                                                                  24,
                                                                  std::tuple{1, 1, 1});
@@ -504,15 +588,13 @@ namespace MiniMaxH3VAE {
         ggml_tensor* decode(GGMLRunnerContext* ctx,
                             ggml_tensor* latent,
                             ggml_tensor* pe,
+                            ggml_tensor* shifted_pe,
+                            const std::array<SolAttentionPlan::Inputs, 2>& sol_inputs,
                             ggml_tensor* pixel_mean,
                             ggml_tensor* pixel_std) {
             auto post_quant = std::dynamic_pointer_cast<Conv3d>(blocks["post_quant_conv"]);
             auto decoder    = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
-            auto pixels     = decoder->forward(ctx, post_quant->forward(ctx, latent), pe);
-            pixels          = ggml_add(ctx->ggml_ctx,
-                                       ggml_mul(ctx->ggml_ctx, pixels, pixel_std),
-                                       pixel_mean);
-            return ggml_clamp(ctx->ggml_ctx, pixels, 0.f, 1.f);
+            return decoder->forward(ctx, post_quant->forward(ctx, latent), pe, shifted_pe, sol_inputs, pixel_mean, pixel_std);
         }
     };
 
@@ -553,6 +635,9 @@ namespace MiniMaxH3VAE {
         }
 
         std::string get_desc() override {
+            if (config.sol) {
+                return "sol_h3_vae";
+            }
             return config.patch_size_t == 1 ? "p_h3_image_vae" : "minimax_h3_video_vae";
         }
 
@@ -601,6 +686,10 @@ namespace MiniMaxH3VAE {
         }
 
         sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const override {
+            return config.sol ? params : resolve_original_tiling_params(params);
+        }
+
+        static sd_tiling_params_t resolve_original_tiling_params(sd_tiling_params_t params) {
             if (!params.enabled) {
                 params.target_overlap = 0.25f;
             }
@@ -655,7 +744,7 @@ namespace MiniMaxH3VAE {
                                  bool circular_x = false,
                                  bool circular_y = false) override {
             auto input  = ensure_video_shape(x);
-            auto tiling = resolve_tiling_params(tiling_params);
+            auto tiling = resolve_original_tiling_params(tiling_params);
             if (config.patch_size_t == 1 && input.shape()[2] != 1) {
                 LOG_ERROR("P-H3-Image VAE encodes single-frame images only");
                 return {};
@@ -695,6 +784,46 @@ namespace MiniMaxH3VAE {
             return result;
         }
 
+        bool supports_temporal_tiling(VAETemporalDirection direction) const override {
+            return config.sol && direction == VAETemporalDirection::DECODE;
+        }
+
+        int get_default_temporal_tile_frames(VAETemporalDirection direction) const override {
+            SD_UNUSED(direction);
+            return 2 * sol_block_size;
+        }
+
+        int get_default_temporal_tile_overlap(VAETemporalDirection direction) const override {
+            SD_UNUSED(direction);
+            return sol_block_size;
+        }
+
+        sd::Tensor<float> _compute_temporal_tiled(int n_threads,
+                                                  const sd::Tensor<float>& input,
+                                                  VAETemporalDirection direction,
+                                                  const VAETemporalTilingConfig& tiling) override {
+            GGML_ASSERT(supports_temporal_tiling(direction));
+            if (tiling.tile_frames % sol_block_size || tiling.overlap % sol_block_size ||
+                tiling.overlap >= tiling.tile_frames) {
+                LOG_ERROR("Sol-H3-VAE temporal tile size and overlap must be multiples of %d latent frames, with overlap smaller than the tile",
+                          static_cast<int>(sol_block_size));
+                return {};
+            }
+            if (input.shape()[2] <= tiling.tile_frames) {
+                return _compute(n_threads, input, true);
+            }
+            auto plan = make_vae_temporal_tile_plan(input.shape()[2], tiling);
+            // With block-aligned starts, the same mapping gives tile lengths and output offsets.
+            auto frame_count = [&](int64_t frames) { return sol_output_frames(frames, config.patch_size_t); };
+            return process_vae_temporal_tiles_blended(input, plan, frame_count, frame_count,
+                                                      [&](const sd::Tensor<float>& tile, const VAETemporalTile& window) {
+                                                          LOG_VERBOSE("Sol-H3-VAE temporal tile %d/%d: latent frames [%lld, %lld)",
+                                                                      window.index + 1, static_cast<int>(plan.tiles.size()),
+                                                                      static_cast<long long>(window.start), static_cast<long long>(window.end));
+                                                          return _compute(n_threads, tile, true);
+                                                      });
+        }
+
         sd::Tensor<float> decode(int n_threads,
                                  const sd::Tensor<float>& x,
                                  sd_tiling_params_t tiling_params,
@@ -702,7 +831,25 @@ namespace MiniMaxH3VAE {
                                  bool circular_x   = false,
                                  bool circular_y   = false,
                                  bool silent       = false) override {
-            auto input  = ensure_video_shape(x);
+            auto input = ensure_video_shape(x);
+            if (config.sol) {
+                circular_x = circular_x || circular_x_enabled;
+                circular_y = circular_y || circular_y_enabled;
+                if (circular_x || circular_y) {
+                    int tile_w = static_cast<int>(input.shape()[0]);
+                    int tile_h = static_cast<int>(input.shape()[1]);
+                    float overlap;
+                    if (tiling_params.enabled &&
+                        !get_tile_sizes(tile_w, tile_h, overlap, tiling_params, input.shape()[0], input.shape()[1])) {
+                        return {};
+                    }
+                    if ((circular_x && tile_w >= input.shape()[0]) || (circular_y && tile_h >= input.shape()[1])) {
+                        LOG_ERROR("Sol-H3-VAE circular decoding requires spatial tiles smaller than the input on every circular axis");
+                        return {};
+                    }
+                }
+                return VAE::decode(n_threads, input, tiling_params, decode_video, circular_x, circular_y, silent);
+            }
             auto tiling = resolve_tiling_params(tiling_params);
             if (config.patch_size_t == 1 && input.shape()[2] != 1) {
                 LOG_ERROR("P-H3-Image VAE decodes single-frame latents only");
@@ -801,23 +948,29 @@ namespace MiniMaxH3VAE {
 
         sd::Tensor<float> build_rope(int64_t width,
                                      int64_t height,
-                                     int64_t num_frames) {
+                                     int64_t num_frames,
+                                     bool shifted = false) {
             std::vector<std::vector<float>> ids;
-            ids.reserve(static_cast<size_t>(width * height * num_frames + 5));
+            const int64_t global_rows = config.sol ? ((num_frames + sol_block_size - 1) / sol_block_size) * sol_global_rows : 5;
+            ids.reserve(static_cast<size_t>(width * height * num_frames + global_rows));
             constexpr float two_pi = 6.28318530717958647692f;
+            const int64_t spatial_shift = shifted ? sol_spatial_shift : 0;
+            const int64_t time_shift    = shifted ? sol_temporal_shift : 0;
             for (int64_t t = 0; t < num_frames; ++t) {
-                float pt = (config.patch_size_t == 1 ? 2.f * (0.5f / 7.f) - 1.f
-                                                     : 2.f * ((t + 0.5f) / num_frames) - 1.f) *
-                           two_pi;
+                const float time = config.sol ? ((t + time_shift) % sol_block_size + 0.5f) / sol_block_size
+                                              : (config.patch_size_t == 1 ? 0.5f / 7.f : (t + 0.5f) / num_frames);
+                float pt         = (2.f * time - 1.f) * two_pi;
                 for (int64_t h = 0; h < height; ++h) {
-                    float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi;
+                    const float y = config.sol ? ((h + spatial_shift) % sol_window_size + 0.5f) / sol_window_size : (h + 0.5f) / height;
+                    float ph      = (2.f * y - 1.f) * two_pi;
                     for (int64_t w = 0; w < width; ++w) {
-                        float pw = (2.f * ((w + 0.5f) / width) - 1.f) * two_pi;
+                        const float x = config.sol ? ((w + spatial_shift) % sol_window_size + 0.5f) / sol_window_size : (w + 0.5f) / width;
+                        float pw      = (2.f * x - 1.f) * two_pi;
                         ids.push_back({pt, ph, pw});
                     }
                 }
             }
-            for (int i = 0; i < 5; ++i) {
+            for (int64_t i = 0; i < global_rows; ++i) {
                 ids.push_back({0.f, 0.f, 0.f});
             }
             auto values = Rope::embed_nd(ids,
@@ -835,24 +988,49 @@ namespace MiniMaxH3VAE {
                                    const sd::Tensor<float>& z,
                                    bool decode_graph) override {
             auto input = ensure_video_shape(z);
+            std::unique_ptr<SolAttentionPlan> sol_plan;
+            std::unique_ptr<SolAttentionPlan> shifted_plan;
+            sd::Tensor<float> shifted_rope;
+            sd::Tensor<int32_t> sol_row_ids;
             if (decode_graph) {
                 rope_cache = build_rope(input.shape()[0],
                                         input.shape()[1],
                                         input.shape()[2]);
+                if (config.sol) {
+                    const auto& shape  = input.shape();
+                    sol_plan           = std::make_unique<SolAttentionPlan>(shape[0], shape[1], shape[2], false);
+                    shifted_plan       = std::make_unique<SolAttentionPlan>(shape[0], shape[1], shape[2], true);
+                    shifted_rope       = build_rope(shape[0], shape[1], shape[2], true);
+                    const int64_t rows = std::max(shifted_rope.shape()[3], pixel_mean.shape()[3] *
+                                                                               sol_output_frames(shape[2], config.patch_size_t));
+                    sol_row_ids        = sd::Tensor<int32_t>({rows});
+                    for (int64_t i = 0; i < rows; ++i) {
+                        sol_row_ids.data()[i] = static_cast<int32_t>(i);
+                    }
+                }
             }
             auto get_graph = [&]() -> ggml_cgraph* {
                 auto value       = make_input(input);
                 auto mean        = make_input(pixel_mean);
                 auto std         = make_input(pixel_std);
-                auto runner_ctx  = get_context();
+                auto graph       = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE);
+                auto runner_ctx  = get_context(graph);
                 ggml_tensor* out = nullptr;
                 if (decode_graph) {
-                    auto pe = make_input(rope_cache);
-                    out     = model.decode(&runner_ctx, value, pe, mean, std);
+                    auto pe         = make_input(rope_cache);
+                    auto shifted_pe = config.sol ? make_input(shifted_rope) : nullptr;
+                    std::array<SolAttentionPlan::Inputs, 2> sol_inputs;
+                    if (sol_plan) {
+                        auto bind_input       = [&](const auto& tensor) { return make_input(tensor); };
+                        sol_inputs            = {sol_plan->bind(bind_input), shifted_plan->bind(bind_input)};
+                        auto row_ids          = make_input(sol_row_ids);
+                        sol_inputs[0].row_ids = row_ids;
+                        sol_inputs[1].row_ids = row_ids;
+                    }
+                    out = model.decode(&runner_ctx, value, pe, shifted_pe, sol_inputs, mean, std);
                 } else {
                     out = model.encode(&runner_ctx, value, mean, std);
                 }
-                auto graph = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE);
                 ggml_build_forward_expand(graph, out);
                 return graph;
             };

src/model/vae/vae_tiling.hpp

@@ -116,14 +116,14 @@ inline sd::Tensor<float> process_vae_temporal_tiles(const sd::Tensor<float>& inp
     return output;
 }
 
-template <typename Fn>
+template <typename FrameCount, typename FrameOffset, typename Fn>
 inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<float>& input,
                                                             const VAETemporalTilePlan& plan,
-                                                            int output_scale,
+                                                            FrameCount frame_count,
+                                                            FrameOffset frame_offset,
                                                             Fn&& on_processing) {
-    GGML_ASSERT(output_scale >= 1);
-    const int64_t output_frames = 1 + (input.shape()[2] - 1) * output_scale;
-    const int overlap_frames    = plan.overlap > 0 ? 1 + (plan.overlap - 1) * output_scale : 0;
+    const int64_t output_frames  = frame_count(input.shape()[2]);
+    const int64_t overlap_frames = frame_count(plan.overlap);
     std::vector<float> weights(static_cast<size_t>(output_frames), 0.f);
     sd::Tensor<float> output;
 
@@ -138,7 +138,7 @@ inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<flo
             return {};
         }
 
-        const int64_t expected_tile_frames = 1 + (input_tile.shape()[2] - 1) * output_scale;
+        const int64_t expected_tile_frames = frame_count(input_tile.shape()[2]);
         if (output_tile.dim() < 3 || output_tile.shape()[2] != expected_tile_frames) {
             LOG_ERROR("unexpected temporal tile output shape: expected %lld frames, got %lld",
                       (long long)expected_tile_frames,
@@ -165,7 +165,7 @@ inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<flo
             }
         }
 
-        const int64_t output_start = tile.start * output_scale;
+        const int64_t output_start = frame_offset(tile.start);
         const int64_t inner        = output.shape()[0] * output.shape()[1];
         const int64_t outer        = output.numel() / (inner * output.shape()[2]);
         const int64_t tile_frames  = output_tile.shape()[2];
@@ -214,4 +214,15 @@ inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<flo
     return output;
 }
 
+template <typename Fn>
+inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<float>& input,
+                                                            const VAETemporalTilePlan& plan,
+                                                            int output_scale,
+                                                            Fn&& on_processing) {
+    GGML_ASSERT(output_scale >= 1);
+    auto frame_count  = [&](int64_t frames) { return frames > 0 ? 1 + (frames - 1) * output_scale : 0; };
+    auto frame_offset = [&](int64_t start) { return start * output_scale; };
+    return process_vae_temporal_tiles_blended(input, plan, frame_count, frame_offset, std::forward<Fn>(on_processing));
+}
+
 #endif  // __SD_MODEL_VAE_VAE_TILING_HPP__