feat: MiniMax-H3 latent hires second pass (#2249)

This commit is contained in:
farawayso authored and GitHub committed 2026-11-05 02:59:04 +08:00
1 parent 3e92bf990a
commit d1358945dc
3 files changed
+930 -21

No files matched your search

+51
View File
@@ -99,6 +99,57 @@ The C API exposes the same inputs through `ref_images`, `ref_videos`, and
`ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own `ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own
frame rate and optional soundtrack; non-24-fps inputs are resampled internally. frame rate and optional soundtrack; non-24-fps inputs are resampled internally.
## Latent hires second pass (two-stage sampling)
MiniMax-H3 supports `--hires` with two interchangeable second-stage upscalers:
- an interpolating latent upscaler (`Latent`, `Latent (nearest)`,
`Latent (bicubic)`, ...): after the first pass finishes, the packed AV latent
is resampled in latent space to `--hires-width` x `--hires-height` and
refined by a second sampling pass;
- a model-backed upscaler (`--hires-upscaler <name>` + `--hires-upscalers-dir`):
a H3-shaped learned latent upscaler checkpoint runs between the passes,
like the LTX spatial upsampler. This is the sd.cpp equivalent of the
two-sampler workflows used for H3 in node-based frontends (e.g. the
community 3D-conv latent upscaler from
[LBH-123-AI/Minimax_h3_latent_Upscaler](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler)).
Draft at a low megapixel count (e.g. 0.4 MP) and refine at a higher one
(e.g. 0.8 MP) instead of paying for the full resolution in every step:
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\minimax_h3_fl2va-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "a lovely cat" --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --video-frames 56 --hires --hires-upscaler "Latent (bicubic)" --hires-width 1216 --hires-height 672 --hires-steps 4 --hires-denoising-strength 0.4
```
or with the learned upscaler:
```sh
.\bin\Release\sd-cli.exe -M vid_gen ... --hires --hires-upscalers-dir ..\models\latent_upscale_models --hires-upscaler h3-latent-upscaler-3d --hires-width 1216 --hires-height 672 --hires-steps 4 --hires-denoising-strength 0.4
```
Notes for this mode:
- The learned upscaler expects a second normalization of the diffusion latent,
matching its training convention. The runner applies and reverses this extra
transform internally; callers should pass the usual H3 diffusion latents.
- The frame count is never changed by the second stage, so the packed audio
latent length stays valid and the generated soundtrack is preserved.
- Width and height still have to be multiples of 32; the hires target is
aligned like any other resolution. The learned upscaler additionally wants
an even latent size on both axes, which 32-pixel targets provide.
- `--hires-sigmas` can replace `--hires-denoising-strength` when the second
stage should resume from an explicit sigma schedule, which is how most
two-stage H3 workflows split the sampler.
- Keyframe conditioning (`--init-img` / `--end-img`) is not available together
with hires yet: those latents are encoded at the first stage resolution and
would have to be re-encoded for the hires canvas. Use T2VA or Ref2VA.
- The second stage runs at the larger latent, so it is the memory-hungry half;
combine it with quantization, `--vae-tiling` and `--offload-to-cpu` on
small cards.
- The learned upscaler is scale-conditional: the target-to-source ratio is
fed into the network as an embedding, 1.2-2.0x is the trained domain, and
downscaling is rejected.
## Shape and runtime notes ## Shape and runtime notes
- Width and height are aligned upward to a multiple of 32. - Width and height are aligned upward to a multiple of 32.
+584
View File
@@ -0,0 +1,584 @@
#ifndef __SD_MODEL_UPSCALER_H3_LATENT_UPSCALER_HPP__
#define __SD_MODEL_UPSCALER_H3_LATENT_UPSCALER_HPP__
#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>
#include "core/ggml_extend.h"
#include "core/ggml_graph_cut.h"
#include "core/ggml_runner.h"
#include "core/tensor.hpp"
#include "core/util.h"
#include "model/common/ggml_block.hpp"
#include "model/upscaler/ltx_latent_upscaler.hpp"
#include "model_loader.h"
// Ported from the community ComfyUI node LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler
// (MIT) with checkpoint LBH-123-AI/Minimax_h3_latent_Upscaler (Apache-2.0).
namespace H3VUpsampler {
constexpr int H3_UPSAMPLER_GRAPH_SIZE = 8192;
struct H3LatentUpscalerConfig {
int64_t in_channels = 24;
int64_t mid_channels = 512;
int embed_dim = 64;
int temporal_kernel = 5;
// Residual and temporal blocks share interleaved checkpoint indices.
std::vector<int> in_res, in_temporal, out_res, out_temporal;
bool valid = false;
static H3LatentUpscalerConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "") {
H3LatentUpscalerConfig config;
std::string base = prefix;
if (!base.empty()) {
base += ".";
}
auto find_weight = [&](const std::string& suffix) -> const TensorStorage* {
auto iter = tensor_storage_map.find(base + suffix);
if (iter == tensor_storage_map.end()) {
return nullptr;
}
return &iter->second;
};
const TensorStorage* conv_in = find_weight("conv_in.weight");
if (conv_in == nullptr) {
LOG_ERROR("h3 latent upsampler: 'conv_in.weight' not found under prefix '%s'", prefix.c_str());
return config;
}
// Bias shapes preserve channel counts when the loader folds 5D conv weights.
const TensorStorage* conv_in_bias = find_weight("conv_in.bias");
if (conv_in_bias != nullptr) {
config.mid_channels = conv_in_bias->ne[0];
} else if (conv_in->n_dims == 5) {
config.mid_channels = conv_in->ne[0];
}
const TensorStorage* conv_out_bias = find_weight("conv_out.bias");
if (conv_out_bias != nullptr) {
config.in_channels = conv_out_bias->ne[0];
} else if (conv_in->n_dims == 5) {
config.in_channels = conv_in->ne[1];
}
if (config.mid_channels <= 0 || config.in_channels <= 0) {
LOG_ERROR("h3 latent upsampler: could not infer channel counts from the checkpoint");
return config;
}
const TensorStorage* embed_mid = find_weight("embed.2.weight");
if (embed_mid != nullptr && embed_mid->n_dims >= 1) {
config.embed_dim = static_cast<int>(embed_mid->ne[0]);
}
std::set<int> in_res_idx, in_temporal_idx, out_res_idx, out_temporal_idx;
for (const auto& [name, _] : tensor_storage_map) {
if (!starts_with(name, base)) {
continue;
}
const std::string rel = name.substr(base.size());
const bool is_in = starts_with(rel, "in_blocks.");
const bool is_out = starts_with(rel, "out_blocks.");
if (!is_in && !is_out) {
continue;
}
const std::string rest = rel.substr(is_in ? 10 : 11);
const size_t dot = rest.find('.');
if (dot == std::string::npos) {
continue;
}
int idx = 0;
try {
idx = std::stoi(rest.substr(0, dot));
} catch (...) {
continue;
}
const std::string part = rest.substr(dot + 1);
if (starts_with(part, "dwconv.") || starts_with(part, "pwconv.")) {
(is_in ? in_temporal_idx : out_temporal_idx).insert(idx);
} else if (starts_with(part, "in_layers.") || starts_with(part, "out_norm.") ||
starts_with(part, "out_layers.") || starts_with(part, "emb_layers.")) {
(is_in ? in_res_idx : out_res_idx).insert(idx);
} else if (starts_with(part, "attn.")) {
LOG_ERROR("h3 latent upsampler: attention blocks are not supported by the v1 runner");
return config;
}
}
config.in_res = std::vector<int>(in_res_idx.begin(), in_res_idx.end());
config.in_temporal = std::vector<int>(in_temporal_idx.begin(), in_temporal_idx.end());
config.out_res = std::vector<int>(out_res_idx.begin(), out_res_idx.end());
config.out_temporal = std::vector<int>(out_temporal_idx.begin(), out_temporal_idx.end());
if (config.in_res.empty() || config.out_res.empty()) {
LOG_ERROR("h3 latent upsampler: no res blocks found in the checkpoint");
return config;
}
auto temporal_kernel_of = [&](std::vector<int> idxs, const char* stage) -> int {
for (int idx : idxs) {
const TensorStorage* dw = find_weight(std::string(stage) + "." + std::to_string(idx) + ".dwconv.weight");
if (dw != nullptr) {
// kD stays at ne[2] in both [OC, IC, kD, kH, kW] and [kW, kH, kD, OC*IC].
return static_cast<int>(dw->ne[2]);
}
}
return 0;
};
const int k_in = temporal_kernel_of(config.in_temporal, "in_blocks");
const int k_out = k_in > 0 ? k_in : temporal_kernel_of(config.out_temporal, "out_blocks");
if (k_out > 0) {
config.temporal_kernel = k_out;
} else {
config.temporal_kernel = 1;
}
config.valid = true;
LOG_VERBOSE("h3 latent upsampler: in_channels = %" PRId64 ", mid_channels = %" PRId64
", in_res = %zu, in_temporal = %zu, out_res = %zu, out_temporal = %zu, temporal_kernel = %d, embed_dim = %d",
config.in_channels,
config.mid_channels,
config.in_res.size(),
config.in_temporal.size(),
config.out_res.size(),
config.out_temporal.size(),
config.temporal_kernel,
config.embed_dim);
return config;
}
};
class H3ResBlock : public GGMLBlock {
protected:
int64_t channels;
int64_t embed_dim;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map,
const std::string prefix) override {
auto make_conv = [&](const std::string& name) {
blocks[name] = std::shared_ptr<GGMLBlock>(new Conv3d(channels,
channels,
{3, 3, 3},
{1, 1, 1},
{1, 1, 1}));
};
make_conv("in_layers.2");
make_conv("out_layers.2");
blocks["in_layers.0"] = std::shared_ptr<GGMLBlock>(new LTXVUpsampler::VideoGroupNorm(32, channels));
blocks["out_norm"] = std::shared_ptr<GGMLBlock>(new LTXVUpsampler::VideoGroupNorm(32, channels));
blocks["emb_layers.1"] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, 2 * channels, true));
}
public:
H3ResBlock(int64_t channels, int64_t embed_dim)
: channels(channels),
embed_dim(embed_dim) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb) {
auto in_layers = std::dynamic_pointer_cast<Conv3d>(blocks["in_layers.2"]);
auto in_norm = std::dynamic_pointer_cast<LTXVUpsampler::VideoGroupNorm>(blocks["in_layers.0"]);
auto out_layers = std::dynamic_pointer_cast<Conv3d>(blocks["out_layers.2"]);
auto out_norm = std::dynamic_pointer_cast<LTXVUpsampler::VideoGroupNorm>(blocks["out_norm"]);
auto emb_layers = std::dynamic_pointer_cast<Linear>(blocks["emb_layers.1"]);
ggml_tensor* h = in_norm->forward(ctx, x);
h = ggml_silu_inplace(ctx->ggml_ctx, h);
h = in_layers->forward(ctx, h);
// Keep emb intact: every residual block uses the same embedding.
ggml_tensor* emb_silu = ggml_silu(ctx->ggml_ctx, emb);
ggml_tensor* emb_out = emb_layers->forward(ctx, emb_silu);
ggml_tensor* emb_1d = ggml_reshape_1d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, emb_out), 2 * channels);
ggml_tensor* scale = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, 0, channels), 1, 1, 1, channels);
ggml_tensor* shift = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, channels, 2 * channels), 1, 1, 1, channels);
ggml_tensor* gain = ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f);
h = out_norm->forward(ctx, h);
h = ggml_mul_inplace(ctx->ggml_ctx, h, gain);
h = ggml_add_inplace(ctx->ggml_ctx, h, shift);
h = ggml_silu_inplace(ctx->ggml_ctx, h);
h = out_layers->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, h, x);
}
};
// Shifted multiply-adds preserve depthwise weights [1, 1, k, C] without
// expanding them into the dense kernels required by Conv3d.
class H3TemporalDWConv3d : public UnaryBlock {
protected:
int64_t channels;
int kernel;
bool has_bias;
std::string prefix;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map,
const std::string prefix) override {
this->prefix = prefix;
params["weight"] = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, kernel, channels);
if (has_bias) {
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
}
}
public:
H3TemporalDWConv3d(int64_t channels, int kernel, bool has_bias = true)
: channels(channels),
kernel(kernel),
has_bias(has_bias) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
if (ctx->weight_adapter) {
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
}
ggml_tensor* w32 = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F32);
ggml_tensor* xa = x;
if (xa->type != GGML_TYPE_F32) {
xa = ggml_cast(ctx->ggml_ctx, xa, GGML_TYPE_F32);
}
const int64_t T = xa->ne[2];
const int pad = kernel / 2;
// ggml_pad only pads the end; concatenate zeros for symmetric padding.
ggml_tensor* xp = xa;
if (pad > 0) {
ggml_tensor* zero = ggml_scale(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, xa, 2, 0, 1), 0.0f);
ggml_tensor* lead = zero;
for (int i = 1; i < pad; ++i) {
lead = ggml_concat(ctx->ggml_ctx, lead, zero, 2);
}
xp = ggml_concat(ctx->ggml_ctx, lead, xa, 2);
xp = ggml_concat(ctx->ggml_ctx, xp, lead, 2);
}
ggml_tensor* acc = nullptr;
for (int j = 0; j < kernel; ++j) {
ggml_tensor* xs = ggml_ext_slice(ctx->ggml_ctx, xp, 2, j, j + T);
ggml_tensor* wj = ggml_ext_slice(ctx->ggml_ctx, w32, 2, j, j + 1); // [1, 1, 1, C]
ggml_tensor* term = ggml_mul(ctx->ggml_ctx, xs, wj);
acc = acc == nullptr ? term : ggml_add(ctx->ggml_ctx, acc, term);
}
if (has_bias) {
ggml_tensor* b = params["bias"];
if (ctx->weight_adapter) {
b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
}
acc = ggml_add(ctx->ggml_ctx, acc, ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, 1, channels));
}
return acc;
}
};
class H3TemporalConv : public GGMLBlock {
protected:
int64_t channels;
int temporal_kernel;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map,
const std::string prefix) override {
blocks["dwconv"] = std::shared_ptr<GGMLBlock>(new H3TemporalDWConv3d(channels, temporal_kernel));
blocks["pwconv"] = std::shared_ptr<GGMLBlock>(new Conv3d(channels,
channels,
{1, 1, 1},
{1, 1, 1},
{0, 0, 0}));
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LTXVUpsampler::VideoGroupNorm(32, channels));
}
public:
H3TemporalConv(int64_t channels, int temporal_kernel)
: channels(channels),
temporal_kernel(temporal_kernel) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto dwconv = std::dynamic_pointer_cast<H3TemporalDWConv3d>(blocks["dwconv"]);
auto pwconv = std::dynamic_pointer_cast<Conv3d>(blocks["pwconv"]);
auto norm = std::dynamic_pointer_cast<LTXVUpsampler::VideoGroupNorm>(blocks["norm"]);
ggml_tensor* h = norm->forward(ctx, x);
h = ggml_silu_inplace(ctx->ggml_ctx, h);
h = dwconv->forward(ctx, h);
h = pwconv->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, h, x);
}
};
class H3Embed : public GGMLBlock {
protected:
int embed_dim;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map,
const std::string prefix) override {
blocks["0"] = std::shared_ptr<GGMLBlock>(new Linear(1, embed_dim, true));
blocks["2"] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, true));
}
public:
explicit H3Embed(int embed_dim)
: embed_dim(embed_dim) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* scale_input) {
auto l0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
auto l2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
ggml_tensor* h = l0->forward(ctx, ggml_reshape_2d(ctx->ggml_ctx, scale_input, 1, 1));
h = ggml_silu_inplace(ctx->ggml_ctx, h);
return l2->forward(ctx, h);
}
};
class H3LatentUpscaler : public GGMLBlock {
public:
H3LatentUpscalerConfig config;
explicit H3LatentUpscaler(H3LatentUpscalerConfig config_)
: config(std::move(config_)) {
for (int idx : config.in_res) {
blocks["in_blocks." + std::to_string(idx)] = std::shared_ptr<GGMLBlock>(new H3ResBlock(config.mid_channels,
config.embed_dim));
}
for (int idx : config.in_temporal) {
blocks["in_blocks." + std::to_string(idx)] = std::shared_ptr<GGMLBlock>(new H3TemporalConv(config.mid_channels,
config.temporal_kernel));
}
for (int idx : config.out_res) {
blocks["out_blocks." + std::to_string(idx)] = std::shared_ptr<GGMLBlock>(new H3ResBlock(config.mid_channels,
config.embed_dim));
}
for (int idx : config.out_temporal) {
blocks["out_blocks." + std::to_string(idx)] = std::shared_ptr<GGMLBlock>(new H3TemporalConv(config.mid_channels,
config.temporal_kernel));
}
blocks["conv_in"] = std::shared_ptr<GGMLBlock>(new Conv3d(config.in_channels,
config.mid_channels,
{3, 3, 3},
{1, 1, 1},
{1, 1, 1}));
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new LTXVUpsampler::VideoGroupNorm(32, config.mid_channels));
blocks["conv_out"] = std::shared_ptr<GGMLBlock>(new Conv3d(config.mid_channels,
config.in_channels,
{3, 3, 3},
{1, 1, 1},
{1, 1, 1}));
blocks["embed"] = std::shared_ptr<GGMLBlock>(new H3Embed(config.embed_dim));
}
ggml_tensor* forward_stage_a(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale_input) {
auto conv_in = std::dynamic_pointer_cast<Conv3d>(blocks["conv_in"]);
auto embed = std::dynamic_pointer_cast<H3Embed>(blocks["embed"]);
ggml_tensor* h = conv_in->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(h, "h3_latent_upscaler.conv_in", "x");
ggml_tensor* emb = embed->forward(ctx, scale_input);
std::set<int> temporal_set(config.in_temporal.begin(), config.in_temporal.end());
int max_idx = 0;
for (int idx : config.in_res) {
max_idx = std::max(max_idx, idx);
}
for (int idx : config.in_temporal) {
max_idx = std::max(max_idx, idx);
}
for (int i = 0; i <= max_idx; ++i) {
if (temporal_set.count(i)) {
auto block = std::dynamic_pointer_cast<H3TemporalConv>(blocks["in_blocks." + std::to_string(i)]);
h = block->forward(ctx, h);
} else {
auto block = std::dynamic_pointer_cast<H3ResBlock>(blocks["in_blocks." + std::to_string(i)]);
h = block->forward(ctx, h, emb);
}
sd::ggml_graph_cut::mark_graph_cut(h, "h3_latent_upscaler.in_blocks." + std::to_string(i), "x");
}
return h;
}
ggml_tensor* forward_stage_b(GGMLRunnerContext* ctx, ggml_tensor* m, ggml_tensor* scale_input) {
auto conv_out = std::dynamic_pointer_cast<Conv3d>(blocks["conv_out"]);
auto norm_out = std::dynamic_pointer_cast<LTXVUpsampler::VideoGroupNorm>(blocks["norm_out"]);
auto embed = std::dynamic_pointer_cast<H3Embed>(blocks["embed"]);
ggml_tensor* g = m;
ggml_tensor* emb = embed->forward(ctx, scale_input);
std::set<int> temporal_set(config.out_temporal.begin(), config.out_temporal.end());
int max_idx = 0;
for (int idx : config.out_res) {
max_idx = std::max(max_idx, idx);
}
for (int idx : config.out_temporal) {
max_idx = std::max(max_idx, idx);
}
for (int i = 0; i <= max_idx; ++i) {
if (temporal_set.count(i)) {
auto block = std::dynamic_pointer_cast<H3TemporalConv>(blocks["out_blocks." + std::to_string(i)]);
g = block->forward(ctx, g);
} else {
auto block = std::dynamic_pointer_cast<H3ResBlock>(blocks["out_blocks." + std::to_string(i)]);
g = block->forward(ctx, g, emb);
}
sd::ggml_graph_cut::mark_graph_cut(g, "h3_latent_upscaler.out_blocks." + std::to_string(i), "x");
}
g = norm_out->forward(ctx, g);
g = ggml_silu_inplace(ctx->ggml_ctx, g);
g = conv_out->forward(ctx, g);
sd::ggml_graph_cut::mark_graph_cut(g, "h3_latent_upscaler.conv_out", "x");
return g;
}
};
struct H3LatentUpscalerRunner : public GGMLRunner {
H3LatentUpscalerConfig config;
std::unique_ptr<H3LatentUpscaler> model;
std::string weight_prefix;
H3LatentUpscalerRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
config(H3LatentUpscalerConfig::detect_from_weights(tensor_storage_map, prefix)),
weight_prefix(prefix) {
if (!config.valid || config.embed_dim <= 0 || config.mid_channels <= 0) {
LOG_ERROR("unsupported H3 latent upsampler config: in_channels=%" PRId64 " mid_channels=%" PRId64 " embed_dim=%d temporal_kernel=%d",
config.in_channels,
config.mid_channels,
config.embed_dim,
config.temporal_kernel);
return;
}
model = std::make_unique<H3LatentUpscaler>(config);
model->init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "h3_latent_upscaler";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
if (model) {
model->get_param_tensors(tensors, weight_prefix);
}
}
// x: H3 diffusion latent [W, H, T, C] (or 5D batch 1). The output stays
// in the same normalized domain for the next sampling pass.
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<float>& x,
int64_t target_w,
int64_t target_h) {
if (model == nullptr) {
LOG_ERROR("H3 latent upsampler is not loaded");
return {};
}
if (x.dim() != 4 && x.dim() != 5) {
LOG_ERROR("H3 latent upsampler expects a 4D or 5D video latent, got dim=%lld", (long long)x.dim());
return {};
}
if (x.dim() == 5 && x.shape()[4] != 1) {
LOG_ERROR("H3 latent upsampler currently supports batch size 1, got batch=%lld", (long long)x.shape()[4]);
return {};
}
if (config.in_channels != 24 || x.shape()[3] != config.in_channels) {
LOG_ERROR("H3 latent upsampler requires 24 channels, got model=%" PRId64 " input=%lld", config.in_channels, (long long)x.shape()[3]);
return {};
}
const int64_t in_w = x.shape()[0];
const int64_t in_h = x.shape()[1];
const int64_t T = x.shape()[2];
if (target_w % 2 != 0 || target_h % 2 != 0) {
LOG_ERROR("H3 latent upsampler target latent size must be even (pixel size a multiple of 32), got %" PRId64 "x%" PRId64,
target_w,
target_h);
return {};
}
const float scale = (static_cast<float>(target_w) / static_cast<float>(in_w) +
static_cast<float>(target_h) / static_cast<float>(in_h)) /
2.0f;
if (scale < 1.0f - 1e-3f) {
LOG_ERROR("H3 latent upsampler only supports upscaling (scale >= 1.0), got %.3f", scale);
return {};
}
const size_t expected_dim = static_cast<size_t>(x.dim());
const sd::Tensor<float> mean({1, 1, 1, 24},
{0.858090341091156f, -0.9606591463088989f, 1.0661640167236328f, -0.5090325474739075f,
-0.2727581858634949f, -1.3675414323806763f, -0.2553254961967468f, -0.26907554268836975f,
-0.5376840829849243f, -0.0464097298681736f, 0.6657370328903198f, 0.19690127670764923f,
-0.5460608005523682f, -0.4035342037677765f, -0.23683024942874908f, 0.25928452610969543f,
-0.30133944749832153f, 0.211341992020607f, -1.1206848621368408f, 0.3581933379173279f,
-0.04225143790245056f, 0.2604829967021942f, 0.22864092886447906f, 0.7056031823158264f});
const sd::Tensor<float> stddev({1, 1, 1, 24},
{1.2223774194717407f, 1.2767263650894165f, 1.6831774711608887f, 1.7549455165863037f,
1.5636216402053833f, 2.194143533706665f, 0.9653137922286987f, 1.0569885969161987f,
0.841948926448822f, 0.7729952931404114f, 1.8955937623977661f, 0.946841835975647f,
0.7996809482574463f, 0.44988900423049927f, 0.7197399735450745f, 0.6936293244361877f,
2.961095094680786f, 2.7694199085235596f, 3.0496184825897217f, 2.1088054180145264f,
3.276226282119751f, 3.1627357006073f, 2.2816812992095947f, 2.6127843856811523f});
// Training applies the VAE normalization a second time to diffusion latents.
// Undo only this extra transform after the network.
const sd::Tensor<float> normalized = (x - mean) / stddev;
const sd::Tensor<float> scale_tensor = sd::Tensor<float>::from_vector({scale - 1.0f});
auto get_graph_a = [this, &normalized, &scale_tensor]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(H3_UPSAMPLER_GRAPH_SIZE);
ggml_tensor* x_t = make_input(normalized);
ggml_tensor* scale_input = make_input(scale_tensor);
auto runner_ctx = get_context(gf);
ggml_tensor* out = model->forward_stage_a(&runner_ctx, x_t, scale_input);
ggml_build_forward_expand(gf, out);
return gf;
};
std::optional<sd::Tensor<float>> mid_opt = GGMLRunner::compute(get_graph_a, n_threads, false, false);
runner_end();
if (!mid_opt.has_value() || mid_opt->empty()) {
LOG_ERROR("H3 latent upsampler stage A (in-block chain) failed");
return {};
}
sd::Tensor<float> mid = std::move(*mid_opt);
std::vector<int64_t> mid_shape = mid.shape();
mid_shape[0] = target_w;
mid_shape[1] = target_h;
sd::Tensor<float> mid_up = sd::ops::interpolate(mid, mid_shape, sd::ops::InterpolateMode::Bilinear, false, false);
if (mid_up.empty()) {
LOG_ERROR("H3 latent upsampler middle spatial interpolate failed: %" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64,
in_w, in_h, target_w, target_h);
return {};
}
auto get_graph_b = [this, &mid_up, &scale_tensor]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(H3_UPSAMPLER_GRAPH_SIZE);
ggml_tensor* m_t = make_input(mid_up);
ggml_tensor* scale_input = make_input(scale_tensor);
auto runner_ctx = get_context(gf);
ggml_tensor* out = model->forward_stage_b(&runner_ctx, m_t, scale_input);
ggml_build_forward_expand(gf, out);
return gf;
};
std::optional<sd::Tensor<float>> out_opt = GGMLRunner::compute(get_graph_b, n_threads, false, false);
runner_end();
if (!out_opt.has_value() || out_opt->empty()) {
LOG_ERROR("H3 latent upsampler stage B (out-block chain) failed");
return {};
}
sd::Tensor<float> out = std::move(*out_opt);
if (out.shape()[0] != target_w || out.shape()[1] != target_h || out.shape()[2] != T) {
LOG_ERROR("H3 latent upsampler output shape mismatch: got %" PRId64 "x%" PRId64 "x%" PRId64 " expected %" PRId64 "x%" PRId64 "x%" PRId64,
out.shape()[0], out.shape()[1], out.shape()[2], target_w, target_h, T);
return {};
}
return restore_trailing_singleton_dims(out * stddev + mean, expected_dim);
}
};
} // namespace H3VUpsampler
#endif // __SD_MODEL_UPSCALER_H3_LATENT_UPSCALER_HPP__
+295 -21
View File
@@ -12,6 +12,7 @@
#include "core/rng_philox.hpp" #include "core/rng_philox.hpp"
#include "diffusion_engine.h" #include "diffusion_engine.h"
#include "model/diffusion/minimax_h3.hpp" #include "model/diffusion/minimax_h3.hpp"
#include "model/upscaler/h3_latent_upscaler.hpp"
#include "model/upscaler/ltx_latent_upscaler.hpp" #include "model/upscaler/ltx_latent_upscaler.hpp"
#include "model/vae/audio_vae.hpp" #include "model/vae/audio_vae.hpp"
#include "model/vae/vae.hpp" #include "model/vae/vae.hpp"
@@ -1317,6 +1318,230 @@ namespace sd::pipeline {
return result_images; return result_images;
} }
static bool hires_upscaler_is_latent_interpolation(enum sd_hires_upscaler_t upscaler) {
switch (upscaler) {
case SD_HIRES_UPSCALER_LATENT:
case SD_HIRES_UPSCALER_LATENT_NEAREST:
case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT:
case SD_HIRES_UPSCALER_LATENT_ANTIALIASED:
case SD_HIRES_UPSCALER_LATENT_BICUBIC:
case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED:
return true;
default:
return false;
}
}
static sd::ops::InterpolateMode hires_latent_interpolate_mode(enum sd_hires_upscaler_t upscaler,
bool* antialias) {
if (antialias != nullptr) {
*antialias = false;
}
switch (upscaler) {
case SD_HIRES_UPSCALER_LATENT_NEAREST:
return sd::ops::InterpolateMode::Nearest;
case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT:
return sd::ops::InterpolateMode::NearestExact;
case SD_HIRES_UPSCALER_LATENT_BICUBIC:
return sd::ops::InterpolateMode::Bicubic;
case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED:
if (antialias != nullptr) {
*antialias = true;
}
return sd::ops::InterpolateMode::Bicubic;
case SD_HIRES_UPSCALER_LATENT_ANTIALIASED:
if (antialias != nullptr) {
*antialias = true;
}
return sd::ops::InterpolateMode::Bilinear;
default:
return sd::ops::InterpolateMode::Bilinear;
}
}
static sd::Tensor<float> upscale_spatial_video_latent(StableDiffusionGGML* sd,
const GenerationRequest& request,
const sd::Tensor<float>& packed_latent,
int audio_length) {
if (sd == nullptr || packed_latent.empty()) {
return {};
}
if (packed_latent.dim() != 4 && packed_latent.dim() != 5) {
LOG_ERROR("video latent spatial upscale expects a 4D or 5D latent");
return {};
}
int64_t target_width = request.hires.target_width / request.vae_scale_factor;
int64_t target_height = request.hires.target_height / request.vae_scale_factor;
if (target_width <= 0 || target_height <= 0) {
LOG_ERROR("video latent spatial upscale target must be positive: %dx%d",
request.hires.target_width,
request.hires.target_height);
return {};
}
if (packed_latent.shape()[0] == target_width && packed_latent.shape()[1] == target_height) {
LOG_DEBUG("video latent spatial upscale: target equals the draft latent size, skipping");
return packed_latent;
}
int latent_channels = sd->get_latent_channel();
sd::Tensor<float> video_latent = packed_latent;
sd::Tensor<float> audio_latent;
if (packed_latent.shape()[3] > latent_channels) {
video_latent = sd::ops::slice(packed_latent, 3, 0, latent_channels);
audio_latent = sd_version_is_minimax_h3(sd->version)
? unpack_minimax_h3_audio_latent(packed_latent, audio_length, latent_channels)
: unpack_ltxav_audio_latent(packed_latent, audio_length, latent_channels);
if (audio_latent.empty()) {
LOG_ERROR("failed to unpack audio latent before spatial latent upscale");
return {};
}
}
bool antialias = false;
auto mode = hires_latent_interpolate_mode(request.hires.upscaler, &antialias);
std::vector<int64_t> target_shape = video_latent.shape();
target_shape[0] = target_width;
target_shape[1] = target_height;
LOG_INFO("video latent %s spatial upscale: %" PRId64 "x%" PRId64 "x%" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64 "x%" PRId64 "x%" PRId64,
sd_hires_upscaler_name(request.hires.upscaler),
(int64_t)video_latent.shape()[0],
(int64_t)video_latent.shape()[1],
(int64_t)video_latent.shape()[2],
(int64_t)video_latent.shape()[3],
target_shape[0],
target_shape[1],
target_shape[2],
target_shape[3]);
sd::Tensor<float> upscaled = sd::ops::interpolate(video_latent, target_shape, mode, false, antialias);
if (upscaled.empty()) {
LOG_ERROR("video latent spatial upscale failed");
return {};
}
if (audio_latent.empty()) {
return upscaled;
}
return pack_ltxav_audio_and_video_latents(upscaled, audio_latent);
}
sd::Tensor<float> upscale_h3_learned_video_latent(StableDiffusionGGML* sd,
const char* model_path,
const sd::Tensor<float>& packed_latent,
int audio_length,
int64_t target_width,
int64_t target_height) {
if (sd == nullptr || sd->model_manager == nullptr || packed_latent.empty()) {
return {};
}
if (strlen(SAFE_STR(model_path)) == 0) {
LOG_ERROR("H3 learned latent upscale requires a hires model path");
return {};
}
// The learned network is not an identity mapping at scale 1.
if (packed_latent.shape()[0] == target_width && packed_latent.shape()[1] == target_height) {
LOG_DEBUG("H3 learned latent upscale: target equals the draft latent size, skipping");
return packed_latent;
}
if (!sd->ensure_backend_pair(SDBackendModule::UPSCALER)) {
return {};
}
int latent_channels = sd->get_latent_channel();
sd::Tensor<float> video_latent = packed_latent;
sd::Tensor<float> audio_latent;
if (packed_latent.shape()[3] > latent_channels) {
video_latent = sd::ops::slice(packed_latent, 3, 0, latent_channels);
audio_latent = unpack_minimax_h3_audio_latent(packed_latent, audio_length, latent_channels);
if (audio_latent.empty()) {
LOG_ERROR("failed to unpack H3 audio latent before learned latent upscale");
return {};
}
}
LOG_INFO("H3 learned latent spatial upscale: latent %dx%dx%dx%d -> %" PRId64 "x%" PRId64 " model output",
(int)video_latent.shape()[0],
(int)video_latent.shape()[1],
(int)video_latent.shape()[2],
(int)video_latent.shape()[3],
target_width,
target_height);
if (sd->first_stage_model) {
sd->first_stage_model->runner_end();
}
auto model_manager = sd->model_manager;
struct UpsamplerScope {
ModelManager& manager;
ModelLoader::FileId owned_source = 0;
std::unique_ptr<H3VUpsampler::H3LatentUpscalerRunner> runner;
std::vector<ggml_tensor*> params;
~UpsamplerScope() {
if (runner) {
runner->runner_end();
}
GGML_ASSERT(manager.unregister_param_tensors(params));
if (owned_source != 0) {
GGML_ASSERT(manager.del_file(owned_source));
}
}
} scope{*model_manager};
const std::string prefix = "h3_latent_upscaler";
ModelLoader candidate = model_manager->loader();
ModelLoader::FileId source_file = 0;
if (!candidate.add_file(model_path, prefix + ".", &source_file)) {
LOG_ERROR("init H3 latent upsampler model loader from file failed: '%s'", model_path);
return {};
}
const bool owns_source = model_manager->loader().file_revision(source_file) == 0;
if (!model_manager->set_loader(std::move(candidate))) {
return {};
}
scope.owned_source = owns_source ? source_file : 0;
auto& upsampler = scope.runner;
upsampler = std::make_unique<H3VUpsampler::H3LatentUpscalerRunner>(sd->backend_for(SDBackendModule::UPSCALER),
model_manager->loader().get_tensor_storage_map(),
prefix,
model_manager);
const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER);
upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes);
if (upsampler->model == nullptr) {
LOG_ERROR("init H3 latent upsampler from metadata failed");
return {};
}
std::map<std::string, ggml_tensor*> tensors;
upsampler->get_param_tensors(tensors);
for (const auto& entry : tensors) {
scope.params.push_back(entry.second);
}
if (!model_manager->register_param_tensors(ModelComponent::LatentUpsampler,
std::move(tensors),
ModelManager::ResidencyMode::ParamBackend,
sd->backend_for(SDBackendModule::UPSCALER),
sd->params_backend_for(SDBackendModule::UPSCALER)) ||
!model_manager->validate_registered_tensors()) {
LOG_ERROR("register H3 latent upsampler tensors with model manager failed");
return {};
}
sd::Tensor<float> upscaled = upsampler->compute(sd->n_threads, video_latent, target_width, target_height);
upsampler->runner_end();
if (upscaled.empty()) {
LOG_ERROR("H3 learned latent spatial upscale failed");
return {};
}
if (!audio_latent.empty()) {
upscaled = pack_ltxav_audio_and_video_latents(upscaled, audio_latent);
}
return upscaled;
}
sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd, sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd,
const char* model_path, const char* model_path,
const sd::Tensor<float>& packed_latent, const sd::Tensor<float>& packed_latent,
@@ -1630,17 +1855,41 @@ namespace sd::pipeline {
bool latent_upscale_enabled = request.hires.enabled; bool latent_upscale_enabled = request.hires.enabled;
GenerationRequest hires_request = request; GenerationRequest hires_request = request;
if (latent_upscale_enabled) { if (latent_upscale_enabled) {
if (!sd_version_is_ltxav(sd->version)) { const bool is_ltxav = sd_version_is_ltxav(sd->version);
LOG_ERROR("LTX latent spatial upscale is only supported for LTX video models"); const bool is_h3 = sd_version_is_minimax_h3(sd->version);
if (!is_ltxav && !is_h3) {
LOG_ERROR("video latent spatial upscale is only supported for LTX and MiniMax-H3 video models");
return false; return false;
} }
if (request.hires.upscaler != SD_HIRES_UPSCALER_MODEL) { if (is_ltxav) {
LOG_ERROR("LTX latent spatial upscale currently requires hires upscaler MODEL"); if (request.hires.upscaler != SD_HIRES_UPSCALER_MODEL) {
return false; LOG_ERROR("LTX latent spatial upscale currently requires hires upscaler MODEL");
} return false;
if (strlen(SAFE_STR(request.hires.model_path)) == 0) { }
LOG_ERROR("LTX latent spatial upscale is enabled but hires model path was not provided"); if (strlen(SAFE_STR(request.hires.model_path)) == 0) {
return false; LOG_ERROR("LTX latent spatial upscale is enabled but hires model path was not provided");
return false;
}
} else {
const bool model_upscaler = request.hires.upscaler == SD_HIRES_UPSCALER_MODEL;
if (model_upscaler) {
if (strlen(SAFE_STR(request.hires.model_path)) == 0) {
LOG_ERROR("MiniMax-H3 learned hires upscaler is enabled but hires model path was not provided");
return false;
}
} else if (!hires_upscaler_is_latent_interpolation(request.hires.upscaler)) {
LOG_ERROR("MiniMax-H3 hires requires an interpolating latent upscaler (e.g. \"Latent\" or \"Latent (bicubic)\") or a MODEL upscaler checkpoint");
return false;
}
if (sd_vid_gen_params->init_image.data != nullptr ||
sd_vid_gen_params->end_image.data != nullptr) {
LOG_ERROR("MiniMax-H3 hires does not support --init-img/--end-img keyframes yet; use T2VA or Ref2VA");
return false;
}
if (request.hires.target_width <= 0 || request.hires.target_height <= 0) {
LOG_ERROR("MiniMax-H3 hires requires a positive hires target size");
return false;
}
} }
} }
@@ -1668,7 +1917,7 @@ namespace sd::pipeline {
} }
ImageGenerationEmbeds embeds = std::move(*embeds_opt); ImageGenerationEmbeds embeds = std::move(*embeds_opt);
if (latent_upscale_enabled) { if (latent_upscale_enabled) {
LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale", LOG_INFO("generate_video %dx%dx%d -> video latent spatial upscale",
request.width, request.width,
request.height, request.height,
request.frames); request.frames);
@@ -1780,16 +2029,31 @@ namespace sd::pipeline {
LOG_ERROR("cancelling generation before latent upscale"); LOG_ERROR("cancelling generation before latent upscale");
return false; return false;
} }
int64_t upscale_start = ggml_time_ms(); int64_t upscale_start = ggml_time_ms();
sd::Tensor<float> upscaled_latent = upscale_ltx_spatial_video_latent(sd, sd::Tensor<float> upscaled_latent;
request.hires.model_path, if (sd_version_is_ltxav(sd->version)) {
final_latent, upscaled_latent = upscale_ltx_spatial_video_latent(sd,
latents.audio_length); request.hires.model_path,
int64_t upscale_end = ggml_time_ms(); final_latent,
latents.audio_length);
} else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
upscaled_latent = upscale_h3_learned_video_latent(sd,
request.hires.model_path,
final_latent,
latents.audio_length,
request.hires.target_width / request.vae_scale_factor,
request.hires.target_height / request.vae_scale_factor);
} else {
upscaled_latent = upscale_spatial_video_latent(sd,
request,
final_latent,
latents.audio_length);
}
int64_t upscale_end = ggml_time_ms();
if (upscaled_latent.empty()) { if (upscaled_latent.empty()) {
return false; return false;
} }
LOG_INFO("LTX latent spatial upscale completed, taking %.2fs", LOG_INFO("video latent spatial upscale completed, taking %.2fs",
(upscale_end - upscale_start) * 1.0f / 1000); (upscale_end - upscale_start) * 1.0f / 1000);
x_t = std::move(upscaled_latent); x_t = std::move(upscaled_latent);
@@ -1798,7 +2062,7 @@ namespace sd::pipeline {
int upscaled_latent_frames = static_cast<int>(x_t.shape()[2]); int upscaled_latent_frames = static_cast<int>(x_t.shape()[2]);
int upscaled_frames = sd->latent_frames_to_video_frames(upscaled_latent_frames); int upscaled_frames = sd->latent_frames_to_video_frames(upscaled_latent_frames);
if (upscaled_frames != hires_request.frames) { if (upscaled_frames != hires_request.frames) {
LOG_INFO("LTX latent upsampler output latent frames %d, frames %d -> %d", LOG_INFO("video latent upsampler output latent frames %d, frames %d -> %d",
upscaled_latent_frames, upscaled_latent_frames,
hires_request.frames, hires_request.frames,
upscaled_frames); upscaled_frames);
@@ -1829,7 +2093,8 @@ namespace sd::pipeline {
latents.audio_length = target_audio_length; latents.audio_length = target_audio_length;
} }
} }
if ((request.hires.target_width > 0 || request.hires.target_height > 0) && if (sd_version_is_ltxav(sd->version) &&
(request.hires.target_width > 0 || request.hires.target_height > 0) &&
(request.hires.target_width != hires_request.width || request.hires.target_height != hires_request.height)) { (request.hires.target_width != hires_request.width || request.hires.target_height != hires_request.height)) {
LOG_WARN("LTX latent spatial upsampler output is %dx%d; ignoring hires target %dx%d", LOG_WARN("LTX latent spatial upsampler output is %dx%d; ignoring hires target %dx%d",
hires_request.width, hires_request.width,
@@ -1843,7 +2108,8 @@ namespace sd::pipeline {
LOG_ERROR("cancelling generation before latent upscale refine"); LOG_ERROR("cancelling generation before latent upscale refine");
return false; return false;
} }
if (!apply_ltxv_refine_image_conditioning(sd, if (sd_version_is_ltxav(sd->version) &&
!apply_ltxv_refine_image_conditioning(sd,
sd_vid_gen_params, sd_vid_gen_params,
hires_request, hires_request,
latents, latents,
@@ -1852,6 +2118,14 @@ namespace sd::pipeline {
&hires_video_positions)) { &hires_video_positions)) {
return false; return false;
} }
// H3 keyframe conditioning must be re-encoded at the hires resolution.
if (sd_version_is_minimax_h3(sd->version) &&
!latents.denoise_mask.empty() &&
(latents.denoise_mask.shape()[0] != x_t.shape()[0] ||
latents.denoise_mask.shape()[1] != x_t.shape()[1])) {
LOG_ERROR("MiniMax-H3 hires cannot rescale an existing video denoise mask");
return false;
}
noise = sd::Tensor<float>::randn_like(x_t, sd->rng); noise = sd::Tensor<float>::randn_like(x_t, sd->rng);
W = hires_request.width / hires_request.vae_scale_factor; W = hires_request.width / hires_request.vae_scale_factor;
@@ -1872,7 +2146,7 @@ namespace sd::pipeline {
hires_sample_method); hires_sample_method);
LOG_VERBOSE("sample(latent upscale) %dx%dx%d", W, H, T); LOG_VERBOSE("sample(latent upscale) %dx%dx%d", W, H, T);
LOG_INFO("LTX latent spatial upscale refine: scheduler_steps=%d, denoising_strength=%.2f, sampler=%s, sigma_sched_size=%zu%s", LOG_INFO("video latent spatial upscale refine: scheduler_steps=%d, denoising_strength=%.2f, sampler=%s, sigma_sched_size=%zu%s",
hires_scheduler_steps, hires_scheduler_steps,
request.hires.denoising_strength, request.hires.denoising_strength,
sampling_methods_str[hires_sample_method], sampling_methods_str[hires_sample_method],