feat: add MiniMax-H3-Image support (#2324)

src/pipeline/diffusion_engine.cpp

@@ -41,6 +41,7 @@
 #include "model/audio/wav2vec2.hpp"
 #include "model/diffusion/animatediff.hpp"
 #include "model/diffusion/control.hpp"
+#include "model/diffusion/minimax_h3.hpp"
 #include "model/diffusion/model.hpp"
 #include "model/diffusion/z_image.hpp"
 #include "model/vae/audio_vae.hpp"
@@ -1492,6 +1493,11 @@ bool StableDiffusionGGML::build_denoiser() {
             break;
         }
         case H3_AV_FLOW_PRED: {
+            if (version == VERSION_MINIMAX_H3_IMAGE) {
+                LOG_INFO("running in MiniMax H3 Image FLOW mode");
+                denoiser = std::make_shared<MiniMaxH3ImageDenoiser>(MiniMaxH3::pdd_sigmas);
+                break;
+            }
             LOG_INFO("running in MiniMax H3 AV FLOW mode");
             denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, version == VERSION_P_H3_IMAGE ? 1.f : 3.f, get_latent_channel());
             break;
@@ -2615,7 +2621,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
                     condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids,
                     condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios,
                     condition.c_reference_blocks.empty() ? nullptr : &condition.c_reference_blocks,
-                    version == VERSION_P_H3_IMAGE ? 2 : audio_length,
+                    version == VERSION_MINIMAX_H3_IMAGE ? 0 : (version == VERSION_P_H3_IMAGE ? 2 : audio_length),
                     std::isfinite(active_flow_shift) ? active_flow_shift : 12.f,
                     version == VERSION_P_H3_IMAGE ? 1.f : 3.f,
                     target_index > 0 && target_index < sigmas.size() ? sigmas[target_index] : NAN};
@@ -3054,7 +3060,7 @@ sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float
             return {};
         }
         auto decoded = first_stage_model->decode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y);
-        if (version == VERSION_P_H3_IMAGE && !decoded.empty() && decoded.dim() == 5) {
+        if (sd_version_is_minimax_h3_image(version) && !decoded.empty() && decoded.dim() == 5) {
             decoded = decoded.squeeze(2);
         }
         if (!decoded.empty() ||

src/pipeline/image.cpp

@@ -172,6 +172,17 @@ namespace sd::pipeline {
         return packed;
     }
 
+    static sd::Tensor<float> resize_h3_image_reference(const sd::Tensor<float>& image, int64_t max_short_edge) {
+        const int64_t short_edge = std::min(image.shape()[0], image.shape()[1]);
+        const double scale       = static_cast<double>(std::clamp<int64_t>(short_edge, 1024, max_short_edge)) / short_edge;
+        const int64_t width      = std::max<int64_t>(32, static_cast<int64_t>(std::nearbyint(image.shape()[0] * scale / 32)) * 32);
+        const int64_t height     = std::max<int64_t>(32, static_cast<int64_t>(std::nearbyint(image.shape()[1] * scale / 32)) * 32);
+        if (width == image.shape()[0] && height == image.shape()[1]) {
+            return image;
+        }
+        return sd::ops::interpolate(image, {width, height, image.shape()[2], image.shape()[3]}, sd::ops::InterpolateMode::Lanczos);
+    }
+
     static std::optional<ImageGenerationLatents> prepare_image_generation_latents(StableDiffusionGGML* sd,
                                                                                   const sd_img_gen_params_t* sd_img_gen_params,
                                                                                   GenerationRequest* request,
@@ -354,7 +365,20 @@ namespace sd::pipeline {
                 continue;
             }
             sd::Tensor<float> ref_latent;
-            if (sd->version == VERSION_P_H3_IMAGE) {
+            if (sd->version == VERSION_MINIMAX_H3_IMAGE) {
+                ref_latent = sd->encode_first_stage(resize_h3_image_reference(ref_images[i], 2048));
+                if (!ref_latent.empty()) {
+                    const int64_t width    = ref_latent.shape()[0];
+                    const int64_t height   = ref_latent.shape()[1];
+                    const int64_t channels = ref_latent.shape()[3];
+                    auto condition_rng     = std::make_shared<MT19937RNG>(static_cast<uint64_t>(request->seed));
+                    // Upstream augments patch rows, so the RNG follows token/channel/patch order.
+                    auto noise = sd::randn<float>({2, 2, channels, width / 2, height / 2}, condition_rng)
+                                     .permute({0, 3, 1, 4, 2})
+                                     .reshape(ref_latent.shape());
+                    ref_latent = ref_latent * 0.999f + noise * (1.f - 0.999f);
+                }
+            } else if (sd->version == VERSION_P_H3_IMAGE) {
                 auto resized = resize_ph3_reference(ref_images[i], int64_t(request->width) * request->height);
                 ref_latent   = sd->encode_first_stage(resized);
             } else if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd->version)) {
@@ -484,14 +508,18 @@ namespace sd::pipeline {
         }
 
         ImageGenerationLatents latents;
-        if (sd->version == VERSION_P_H3_IMAGE) {
-            init_latent = pack_ph3_image_latent(init_latent);
+        if (sd_version_is_minimax_h3_image(sd->version)) {
+            if (sd->version == VERSION_P_H3_IMAGE) {
+                init_latent = pack_ph3_image_latent(init_latent);
+            } else if (init_latent.dim() == 4) {
+                init_latent = init_latent.unsqueeze(2);
+            }
             latent_mask = latent_mask.unsqueeze(2);
             for (size_t i = 0; i < ref_latents.size(); ++i) {
                 latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::IMAGE, static_cast<int32_t>(i), -1});
             }
             for (auto& image : ref_images) {
-                image = resize_ph3_reference(image, 1000000);
+                image = sd->version == VERSION_P_H3_IMAGE ? resize_ph3_reference(image, 1000000) : resize_h3_image_reference(image, 1024);
             }
         }
         latents.init_latent              = std::move(init_latent);
@@ -548,7 +576,7 @@ namespace sd::pipeline {
         if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
             cond.c_concat = latents->concat_latent;  // TODO: optimize
         }
-        if (sd->version == VERSION_P_H3_IMAGE && ref_image_params.pass_to_dit) {
+        if (sd_version_is_minimax_h3_image(sd->version) && ref_image_params.pass_to_dit) {
             cond.c_reference_blocks = latents->minimax_reference_blocks;
         }
 
@@ -596,7 +624,7 @@ namespace sd::pipeline {
             if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
                 uncond.c_concat = latents->concat_latent;  // TODO: optimize
             }
-            if (sd->version == VERSION_P_H3_IMAGE && ref_image_params.pass_to_dit) {
+            if (sd_version_is_minimax_h3_image(sd->version) && ref_image_params.pass_to_dit) {
                 uncond.c_reference_blocks = latents->minimax_reference_blocks;
             }
         }

src/pipeline/model_builders.cpp

@@ -159,6 +159,14 @@ namespace sd::model_builders {
                                                      version,
                                                      weight_manager);
             } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MINIMAX_H3) {
+                if (version == VERSION_MINIMAX_H3_IMAGE) {
+                    auto proj     = tensor_storage_map.find("first_stage_model.decoder.proj_out.weight");
+                    auto embedder = tensor_storage_map.find("first_stage_model.decoder.x_embedder.weight");
+                    if (proj == tensor_storage_map.end() || proj->second.ne[1] != 768 ||
+                        embedder == tensor_storage_map.end() || embedder->second.ne[1] != 1024) {
+                        throw std::runtime_error("MiniMax-H3-Image requires its 2D image VAE; pass --vae minimax_h3_image_vae_fp16.safetensors");
+                    }
+                }
                 if (version == VERSION_P_H3_IMAGE) {
                     // The video VAE also loads successfully, but uses different temporal decoding and posterior sampling.
                     auto proj = tensor_storage_map.find("first_stage_model.decoder.proj_out.weight");