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

+12 -6
View File
@@ -778,7 +778,7 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
if (strlen(SAFE_STR(sd_ctx_params->embeddings_connectors_path)) > 0) {
LOG_INFO("loading embeddings connectors from '%s'", sd_ctx_params->embeddings_connectors_path);
if (!load_model("embeddings_connectors", sd_ctx_params->embeddings_connectors_path)) {
if (!load_model("embeddings_connectors", sd_ctx_params->embeddings_connectors_path, "embeddings_connectors.")) {
LOG_WARN("loading embeddings connectors from '%s' failed", sd_ctx_params->embeddings_connectors_path);
}
}
@@ -1491,7 +1491,7 @@ bool StableDiffusionGGML::build_denoiser() {
}
case H3_AV_FLOW_PRED: {
LOG_INFO("running in MiniMax H3 AV FLOW mode");
denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, 3.f, get_latent_channel());
denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, version == VERSION_P_H3_IMAGE ? 1.f : 3.f, get_latent_channel());
break;
}
case MING_IMAGE_FLOW_PRED: {
@@ -2594,9 +2594,9 @@ 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,
audio_length,
version == VERSION_P_H3_IMAGE ? 2 : audio_length,
std::isfinite(active_flow_shift) ? active_flow_shift : 12.f,
3.f};
version == VERSION_P_H3_IMAGE ? 1.f : 3.f};
} else if (sd_version_is_ltxav(version)) {
diffusion_params.extra = LTXAVDiffusionExtra{
nullptr,
@@ -3007,7 +3007,10 @@ sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float
if (sd_version_is_pid(version) || sd_version_is_minit2i(version) || version == VERSION_IRIS) {
return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
}
auto latents = first_stage_model->diffusion_to_vae_latents(x);
auto image_latents = version == VERSION_P_H3_IMAGE && x.dim() == 5
? sd::ops::slice(x, 3, 0, get_latent_channel())
: x;
auto latents = first_stage_model->diffusion_to_vae_latents(image_latents);
auto tiling_params = first_stage_model->resolve_tiling_params(vae_tiling_params);
const bool prefer_temporal_tiling = decode_video && latents.dim() == 5 && latents.shape()[2] > 1 &&
first_stage_model->can_temporal_tile_decode();
@@ -3021,6 +3024,9 @@ 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) {
decoded = decoded.squeeze(2);
}
if (!decoded.empty() ||
!sd::backend_fit::prepare_vae_decode_retry_tiling(tiling_params, prefer_temporal_tiling,
first_stage_model->last_compute_status(),
@@ -3220,7 +3226,7 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima
if (params.force_ref_timestep_zero && !sd_version_is_krea2(version)) {
LOG_WARN("force_ref_timestep_zero is only supported by Krea2 architecture for now");
}
if (sd_version_is_minimax_h3(version)) {
if (version == VERSION_MINIMAX_H3) {
// H3 consumes references as VAE latents in its own packed sequence rather than
// through the generic conditioner/DiT reference path, so ref_image_size is the
// only arg it acts on. Everything else is reported as ineffective rather than
+55 -2
View File
@@ -144,6 +144,34 @@ namespace sd::pipeline {
return padded;
}
static sd::Tensor<float> resize_ph3_reference(const sd::Tensor<float>& image, int64_t pixels) {
double scale = std::min(1.0, std::sqrt(static_cast<double>(pixels) / (image.shape()[0] * image.shape()[1])));
int64_t width = std::max<int64_t>(32, static_cast<int64_t>(std::nearbyint(image.shape()[0] * scale / 32)) * 32);
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 sd::Tensor<float> pack_ph3_image_latent(const sd::Tensor<float>& image,
const sd::Tensor<float>& previous = {}) {
auto video = image.dim() == 4 ? image.unsqueeze(2) : image;
GGML_ASSERT(video.dim() == 5 && video.shape()[2] == 1 && video.shape()[3] == 24 && video.shape()[4] == 1);
const int64_t spatial = video.shape()[0] * video.shape()[1];
// The image DiT still attends to two stereo audio latents (32 channels).
auto shape = video.shape();
shape[3] += (128 + spatial - 1) / spatial;
auto packed = sd::zeros<float>(shape);
std::copy_n(video.data(), video.numel(), packed.data());
if (!previous.empty()) {
const int64_t offset = previous.shape()[0] * previous.shape()[1] * 24;
GGML_ASSERT(previous.numel() >= offset + 128);
std::copy_n(previous.data() + offset, 128, packed.data() + video.numel());
}
return packed;
}
static std::optional<ImageGenerationLatents> prepare_image_generation_latents(StableDiffusionGGML* sd,
const sd_img_gen_params_t* sd_img_gen_params,
GenerationRequest* request,
@@ -323,7 +351,10 @@ namespace sd::pipeline {
continue;
}
sd::Tensor<float> ref_latent;
if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd->version)) {
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)) {
LOG_VERBOSE("auto resize ref images");
double vae_width;
double vae_height;
@@ -450,6 +481,16 @@ namespace sd::pipeline {
}
ImageGenerationLatents latents;
if (sd->version == VERSION_P_H3_IMAGE) {
init_latent = pack_ph3_image_latent(init_latent);
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);
}
}
latents.init_latent = std::move(init_latent);
latents.concat_latent = std::move(concat_latent);
latents.img_uncond_concat_latent = std::move(img_uncond_concat_latent);
@@ -502,6 +543,9 @@ 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) {
cond.c_reference_blocks = latents->minimax_reference_blocks;
}
bool use_ref_latent_img_cfg = request->use_img_uncond &&
!latents->ref_images.empty() &&
@@ -546,6 +590,9 @@ 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) {
uncond.c_reference_blocks = latents->minimax_reference_blocks;
}
}
SDCondition img_uncond;
@@ -1047,13 +1094,19 @@ namespace sd::pipeline {
sd->rng->manual_seed(cur_seed);
sd->sampler_rng->manual_seed(cur_seed);
auto hires_source = sd->version == VERSION_P_H3_IMAGE
? sd::ops::slice(final_latents[b], 3, 0, 24)
: final_latents[b];
sd::Tensor<float> upscaled = upscale_hires_latent(sd,
final_latents[b],
hires_source,
request,
hires_upscaler.get());
if (upscaled.empty()) {
return false;
}
if (sd->version == VERSION_P_H3_IMAGE) {
upscaled = pack_ph3_image_latent(upscaled, final_latents[b]);
}
sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd->rng);
+11 -1
View File
@@ -110,7 +110,7 @@ namespace sd::model_builders {
return false;
}
bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) {
bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) try {
const auto* sd_ctx_params = &ctx.params;
const auto& tensor_storage_map = ctx.tensor_storage_map;
const auto version = ctx.version;
@@ -159,6 +159,13 @@ namespace sd::model_builders {
version,
weight_manager);
} else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MINIMAX_H3) {
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");
if (proj == tensor_storage_map.end() || proj->second.ne[1] != 768) {
throw std::runtime_error("P-H3-Image requires its single-frame VAE; pass --vae p-h3-image_vae_bf16.safetensors");
}
}
return std::make_shared<MiniMaxH3VAE::MiniMaxH3VideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
@@ -240,6 +247,9 @@ namespace sd::model_builders {
}
runners = std::move(result);
return true;
} catch (const std::exception& error) {
LOG_ERROR("failed to build VAE runners: %s", error.what());
return false;
}
bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner) {