fix: correct MiniMax H3 reference audio encoding

Author: jk212h20

src/model/vae/minimax_h3_audio_vae.hpp

@@ -154,8 +154,9 @@ namespace MiniMaxH3 {
                          const String2TensorStorage& tensor_storage_map = {},
                          const std::string prefix                       = "") override {
             GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
-            params["q_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
-            params["v_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+            params["q_bias"]      = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+            params["zero_k_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
+            params["v_bias"]      = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
@@ -166,7 +167,7 @@ namespace MiniMaxH3 {
                 return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1);
             };
             auto q = ggml_add(ctx->ggml_ctx, qkv[0], bias_shape(params["q_bias"]));
-            auto k = qkv[1];
+            auto k = ggml_add(ctx->ggml_ctx, qkv[1], bias_shape(params["zero_k_bias"]));
             auto v = ggml_add(ctx->ggml_ctx, qkv[2], bias_shape(params["v_bias"]));
 
             int64_t sequence = x->ne[1];
@@ -358,22 +359,38 @@ namespace MiniMaxH3 {
         }
 
         ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* waveform) {
-            GGML_ASSERT(waveform->ne[1] == 2);
+            GGML_ASSERT(waveform->ne[1] * waveform->ne[2] * waveform->ne[3] == 2);
             auto encoder   = std::dynamic_pointer_cast<AudioEncoder>(blocks["encoder"]);
             auto pre       = std::dynamic_pointer_cast<AudioAttentionProjection>(blocks["pre_block"]);
             auto mean_proj = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["mean_proj"]);
 
-            waveform = ggml_reshape_3d(ctx->ggml_ctx, waveform, waveform->ne[0], 1, waveform->ne[1]);
-            auto x   = encoder->forward(ctx, waveform);  // [B*S, 2048, T]
-            x        = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
-            x        = pre->forward(ctx, x);
-            x        = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
-            auto z   = mean_proj->forward(ctx, x);
-
-            auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
-            auto std  = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
-            z         = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
-            return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
+            // GGML's batched conv1d storage interleaves the stream dimension
+            // with output channels. Subsequent layers then read stereo samples
+            // as adjacent feature channels. Run each mono stream independently,
+            // matching PyTorch's reshape(B*S, 1, samples), and concatenate only
+            // the completed normalized latents.
+            const int64_t streams = waveform->ne[2] * waveform->ne[3];
+            waveform = ggml_reshape_3d(ctx->ggml_ctx,
+                                       waveform,
+                                       waveform->ne[0],
+                                       1,
+                                       streams);
+            ggml_tensor* stereo_z = nullptr;
+            for (int64_t stream = 0; stream < streams; ++stream) {
+                auto mono = ggml_ext_slice(ctx->ggml_ctx, waveform, 2, stream, stream + 1);
+                auto x    = encoder->forward(ctx, mono);
+                x         = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+                x         = pre->forward(ctx, x);
+                x         = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+                auto z    = mean_proj->forward(ctx, x);
+
+                auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
+                auto std  = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
+                z         = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
+                z         = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
+                stereo_z  = stereo_z == nullptr ? z : ggml_concat(ctx->ggml_ctx, stereo_z, z, 1);
+            }
+            return stereo_z;
         }
 
         ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent) {

src/stable-diffusion.cpp

@@ -4728,7 +4728,11 @@ static sd::Tensor<float> prepare_minimax_h3_reference_waveform(const sd_audio_t&
         static_cast<long double>(audio.sample_count) * target_sample_rate / audio.sample_rate));
     output_samples          = std::max<uint64_t>(1, output_samples);
     uint64_t padded_samples = (output_samples + 799) / 800 * 800;
-    sd::Tensor<float> waveform({static_cast<int64_t>(padded_samples), 2, 1, 1});
+    // Keep stereo streams planar for the mono-per-stream audio encoder:
+    // [samples, 1, stereo, batch]. This avoids flattening interleaved L/R
+    // storage into alternating samples when the encoder folds streams into
+    // its batch dimension.
+    sd::Tensor<float> waveform({static_cast<int64_t>(padded_samples), 1, 2, 1});
 
     for (uint64_t i = 0; i < output_samples; ++i) {
         long double source_pos = static_cast<long double>(i) * audio.sample_rate / target_sample_rate;
@@ -4739,7 +4743,7 @@ static sd::Tensor<float> prepare_minimax_h3_reference_waveform(const sd_audio_t&
             uint32_t source_channel = audio.channels == 1 ? 0 : std::min<uint32_t>(channel, audio.channels - 1);
             float a                 = audio.data[source0 * audio.channels + source_channel];
             float b                 = audio.data[source1 * audio.channels + source_channel];
-            waveform.index(static_cast<int64_t>(i), channel, 0, 0) =
+            waveform.index(static_cast<int64_t>(i), 0, channel, 0) =
                 std::clamp(a + (b - a) * fraction, -1.f, 1.f);
         }
     }