mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 03:49:09 +00:00
feat: add P-H3-Image generation and editing support (#2272)
This commit is contained in:
1 parent
984eb0686c
commit
3e92bf990a
18 files changed
+421
-41
No files matched your search
@@ -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;
|
||||
|
||||
@@ -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__
|
||||
@@ -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) {
|
||||
|
||||
Reference in new issue
Block a user