feat: add P-H3-Image generation and editing support (#2272)

This commit is contained in:
leejet authored and GitHub committed 2026-11-04 23:24:51 +08:00
1 parent 984eb0686c
commit 3e92bf990a
18 files changed
+421 -41

No files matched your search

+9
View File
@@ -40,6 +40,15 @@ const std::vector<ModelDefinition>& minimax_h3_definitions() {
definition.build = build_minimax_h3;
result.push_back(std::move(definition));
}
{
ModelDefinition definition = result.front();
definition.version = VERSION_P_H3_IMAGE;
definition.name = "P-H3-Image";
definition.traits.supports_video_generation = false;
definition.sampling.default_flow_shift = 1.f;
definition.ref_image_preset = "qwen";
result.push_back(std::move(definition));
}
return result;
}();
return definitions;
+70
View File
@@ -0,0 +1,70 @@
#ifndef __SD_MODEL_TE_PH3_IMAGE_CONNECTOR_HPP__
#define __SD_MODEL_TE_PH3_IMAGE_CONNECTOR_HPP__
#include "core/ggml_runner.h"
#include "model/common/ggml_block.hpp"
struct PH3ImageTextConnector : public UnaryBlock {
static constexpr int64_t input_dim = 5 * 2560;
static constexpr int64_t output_dim = 5120;
PH3ImageTextConnector() {
blocks["proj"] = std::make_shared<Linear>(input_dim, output_dim, false, true, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
for (const auto* name : {"input_mean", "input_std"}) {
params[name] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, input_dim);
}
for (const auto* name : {"output_mean", "output_std", "sink_embedding"}) {
params[name] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_dim);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
x = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, x, params["input_mean"]), params["input_std"]);
x = proj->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, params["output_std"]), params["output_mean"]);
auto sink = ggml_repeat(ctx->ggml_ctx, params["sink_embedding"], ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, 1));
if (x->ne[1] == 1) {
return sink;
}
return ggml_concat(ctx->ggml_ctx, sink, ggml_ext_slice(ctx->ggml_ctx, x, 1, 1, x->ne[1]), 1);
}
};
struct PH3ImageTextConnectorRunner : public GGMLRunner {
PH3ImageTextConnector model;
static constexpr const char* prefix = "text_encoders.ph3_image_connector";
PH3ImageTextConnectorRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "ph3_image_text_connector";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
model.get_param_tensors(tensors, prefix);
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& hidden_states) {
auto get_graph = [&]() {
auto gf = ggml_new_graph(compute_ctx);
auto x = make_input(hidden_states);
auto runner_ctx = get_context();
ggml_build_forward_expand(gf, model.forward(&runner_ctx, x));
return gf;
};
return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false));
}
};
#endif // __SD_MODEL_TE_PH3_IMAGE_CONNECTOR_HPP__
+68 -11
View File
@@ -10,6 +10,7 @@
#include <utility>
#include <vector>
#include "core/rng_mt19937.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/vae/vae.hpp"
@@ -18,6 +19,34 @@ namespace MiniMaxH3VAE {
constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144;
struct MiniMaxH3VAEConfig {
int num_layers = 36;
int patch_size_t = 4;
static MiniMaxH3VAEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) {
MiniMaxH3VAEConfig config;
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.";
for (const auto& [name, _] : tensors) {
if (starts_with(name, layer_prefix)) {
layers = std::max(layers, std::atoi(name.c_str() + layer_prefix.size()) + 1);
}
}
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);
}
return config;
}
};
struct CausalConv3d : public Conv3d {
std::tuple<int, int, int> temporal_padding;
@@ -361,13 +390,14 @@ namespace MiniMaxH3VAE {
};
struct Decoder : public GGMLBlock {
static constexpr int dim = 2048;
static constexpr int num_layers = 36;
static constexpr int dim = 2048;
int num_layers;
static constexpr int num_register_tokens = 4;
static constexpr int patch_size = 16;
static constexpr int patch_size_t = 4;
int patch_size_t;
Decoder() {
explicit Decoder(const MiniMaxH3VAEConfig& config)
: num_layers(config.num_layers), patch_size_t(config.patch_size_t) {
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)] =
@@ -389,7 +419,9 @@ namespace MiniMaxH3VAE {
GGML_TYPE_F32,
dim,
num_register_tokens);
params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
if (patch_size_t != 1) {
params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -443,7 +475,10 @@ namespace MiniMaxH3VAE {
};
struct MiniMaxH3VideoVAE : public GGMLBlock {
MiniMaxH3VideoVAE() {
bool single_frame;
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,
@@ -451,7 +486,7 @@ namespace MiniMaxH3VAE {
blocks["post_quant_conv"] = std::make_shared<Conv3d>(24,
24,
std::tuple{1, 1, 1});
blocks["decoder"] = std::make_shared<Decoder>();
blocks["decoder"] = std::make_shared<Decoder>(config);
}
ggml_tensor* encode(GGMLRunnerContext* ctx,
@@ -464,7 +499,7 @@ namespace MiniMaxH3VAE {
auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
auto quant = std::dynamic_pointer_cast<Conv3d>(blocks["quant_conv"]);
auto moments = quant->forward(ctx, encoder->forward(ctx, pixels));
return ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24);
return single_frame ? moments : ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24);
}
ggml_tensor* decode(GGMLRunnerContext* ctx,
@@ -483,6 +518,7 @@ namespace MiniMaxH3VAE {
};
struct MiniMaxH3VideoVAERunner : public VAE {
MiniMaxH3VAEConfig config;
MiniMaxH3VideoVAE model;
sd::Tensor<float> pixel_mean;
sd::Tensor<float> pixel_std;
@@ -495,6 +531,8 @@ namespace MiniMaxH3VAE {
const std::string& prefix = "first_stage_model",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(VERSION_MINIMAX_H3, backend, prefix, weight_manager),
config(MiniMaxH3VAEConfig::detect_from_weights(tensor_storage_map, prefix)),
model(config),
pixel_mean({1, 1, 1, 3}, {0.485f, 0.456f, 0.406f}),
pixel_std({1, 1, 1, 3}, {0.229f, 0.224f, 0.225f}),
latents_mean({1, 1, 1, 24},
@@ -516,12 +554,12 @@ namespace MiniMaxH3VAE {
}
std::string get_desc() override {
return "minimax_h3_video_vae";
return config.patch_size_t == 1 ? "p_h3_image_vae" : "minimax_h3_video_vae";
}
int get_encoder_output_channels(int input_channels) override {
SD_UNUSED(input_channels);
return 24;
return config.patch_size_t == 1 ? 48 : 24;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
@@ -531,6 +569,15 @@ namespace MiniMaxH3VAE {
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output,
std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
if (config.patch_size_t == 1) {
auto mean = sd::ops::slice(vae_output, 3, 0, 24);
auto logvar = sd::ops::slice(vae_output, 3, 24, 48);
auto noise = MT19937RNG(42).randn(static_cast<uint32_t>(mean.numel()));
for (int64_t i = 0; i < mean.numel(); ++i) {
mean[i] += std::exp(0.5f * std::clamp(logvar[i], -30.f, 20.f)) * noise[i];
}
return mean;
}
return vae_output;
}
@@ -610,6 +657,10 @@ namespace MiniMaxH3VAE {
bool circular_y = false) override {
auto input = ensure_video_shape(x);
auto tiling = resolve_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 {};
}
if (input.shape()[2] == 1) {
auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y);
if (!encoded.empty() && encoded.shape()[2] > 1) {
@@ -654,6 +705,10 @@ namespace MiniMaxH3VAE {
bool silent = false) override {
auto input = ensure_video_shape(x);
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");
return {};
}
if (input.shape()[2] == 1) {
auto decoded = VAE::decode(n_threads,
input,
@@ -752,7 +807,9 @@ namespace MiniMaxH3VAE {
ids.reserve(static_cast<size_t>(width * height * num_frames + 5));
constexpr float two_pi = 6.28318530717958647692f;
for (int64_t t = 0; t < num_frames; ++t) {
float pt = (2.f * ((t + 0.5f) / num_frames) - 1.f) * two_pi;
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;
for (int64_t h = 0; h < height; ++h) {
float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi;
for (int64_t w = 0; w < width; ++w) {