diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index 5af751e36..7fffcaa61 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/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(blocks["encoder"]); auto pre = std::dynamic_pointer_cast(blocks["pre_block"]); auto mean_proj = std::dynamic_pointer_cast(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) { diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index b74d981aa..01f4704ce 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -4728,7 +4728,11 @@ static sd::Tensor prepare_minimax_h3_reference_waveform(const sd_audio_t& static_cast(audio.sample_count) * target_sample_rate / audio.sample_rate)); output_samples = std::max(1, output_samples); uint64_t padded_samples = (output_samples + 799) / 800 * 800; - sd::Tensor waveform({static_cast(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 waveform({static_cast(padded_samples), 1, 2, 1}); for (uint64_t i = 0; i < output_samples; ++i) { long double source_pos = static_cast(i) * audio.sample_rate / target_sample_rate; @@ -4739,7 +4743,7 @@ static sd::Tensor prepare_minimax_h3_reference_waveform(const sd_audio_t& uint32_t source_channel = audio.channels == 1 ? 0 : std::min(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(i), channel, 0, 0) = + waveform.index(static_cast(i), 0, channel, 0) = std::clamp(a + (b - a) * fraction, -1.f, 1.f); } }