Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-10 23:09:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

sd: cherry-pick changes up to master-841-6b3edaa (#2426)

Browse Source 
fix: match exact weights in LLM config detection
feat: add LTX-2.5 support
fix: correct MiniMax H3 reference audio encoding
fix: correct MiniMax H3 audio Euler steps
feat: additional `--preview-interval` values
feat: support numbering for preview images
fix: use carrier sampling for MiniMax H3 audio
feat: generalize temporal tiling across video VAEs
This commit is contained in:
Wagner Bruna
2026-09-03 05:32:59 -03:00
committed by GitHub
parent 2e04e6509d
commit b83cdc6f42
22 changed files with 1010 additions and 338 deletions
Download Patch File Download Diff File Expand all files Collapse all files
otherarch/sdcpp/examples/cli/main.cpp
+39 -25
View File
@@ -36,6 +36,7 @@ struct SDCliParams {
SDMode mode = IMG_GEN;
std::string output_path = "output.png";
int output_begin_idx = -1;
int compression_quality = 90;
std::string image_path;
std::string metadata_format = "text";
@@ -80,7 +81,7 @@ struct SDCliParams {
&metadata_format},
{"",
"--preview-path",
"path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
"path to write preview image to (default: ./preview.png). For image generation, the filename can have %03d placeholder for sequential numbering. Multi-frame previews support .avi, .webm, and animated .webp",
0,
&preview_path},
{"",
@@ -93,12 +94,16 @@ struct SDCliParams {
options.int_options = {
{"",
"--preview-interval",
"interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at every step)",
"preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1",
&preview_interval},
{"",
"--output-begin-idx",
"starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)",
&output_begin_idx},
{"",
"--compression-quality",
"compression quality of video and JPEG / WebP images (90 by default)",
&compression_quality},
};
options.bool_options = {
@@ -372,27 +377,6 @@ bool load_images_from_dir(const std::string dir,
return true;
}
void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) {
(void)step;
(void)is_noisy;
SDCliParams* cli_params = (SDCliParams*)data;
// is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
// unused in this app, it will either be always noisy or always denoised here
if (frame_count == 1) {
if (!write_image_to_file(cli_params->preview_path,
image->data,
image->width,
image->height,
image->channel)) {
LOG_ERROR("save preview image to '%s' failed", cli_params->preview_path.c_str());
}
} else {
if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps) != 0) {
LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
}
}
}
std::string format_frame_idx(std::string pattern, int frame_idx) {
std::smatch match;
std::string result = pattern;
@@ -412,6 +396,36 @@ std::string format_frame_idx(std::string pattern, int frame_idx) {
return result;
}
int continuous_preview_counter = 0;
void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) {
(void)step;
(void)is_noisy;
SDCliParams* cli_params = (SDCliParams*)data;
// is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
// unused in this app, it will either be always noisy or always denoised here
if (frame_count == 1) {
fs::path path = cli_params->preview_path;
if (encoded_image_format_from_path(path.string()) == EncodedImageFormat::UNKNOWN)
path += ".png";
if (std::regex_search(path.string(), format_specifier_regex))
path = fs::path(format_frame_idx(path.string(), continuous_preview_counter++));
if (!write_image_to_file(path.string(),
image->data,
image->width,
image->height,
image->channel,
"",
cli_params->compression_quality)) {
LOG_ERROR("save preview image to '%s' failed", path.string().c_str());
}
} else {
if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps, cli_params->compression_quality) != 0) {
LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
}
}
}
static fs::path get_video_audio_sidecar_path(const SDCliParams& cli_params) {
fs::path out_path = cli_params.output_path;
fs::path base_path = out_path;
@@ -486,7 +500,7 @@ bool save_results(const SDCliParams& cli_params,
std::string params = gen_params.embed_image_metadata
? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
: "";
const bool ok = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, 90);
const bool ok = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, cli_params.compression_quality);
LOG_INFO("save result image %d to '%s' (%s)", idx, path.string().c_str(), ok ? "success" : "failure");
return ok;
};
@@ -532,7 +546,7 @@ bool save_results(const SDCliParams& cli_params,
std::string final_ext_lower = ext.string();
std::transform(final_ext_lower.begin(), final_ext_lower.end(), final_ext_lower.begin(), ::tolower);
const bool mux_audio = generated_audio != nullptr && (final_ext_lower == ".avi" || final_ext_lower == ".webm");
if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps, 90, mux_audio ? generated_audio : nullptr) == 0) {
if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr) == 0) {
LOG_INFO("save result video to '%s'", video_path.string().c_str());
if (generated_audio != nullptr && !mux_audio) {
fs::path wav_path = video_path;
otherarch/sdcpp/examples/common/common.cpp
+2 -2
View File
@@ -1013,7 +1013,7 @@ ArgOptions SDGenerationParams::get_options() {
&extra_sample_args},
{"",
"--extra-tiling-args",
"extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
"extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (default: 4), temporal_tile_overlap (default: 1)",
(int)',',
&extra_tiling_args},
{"",
@@ -1230,7 +1230,7 @@ ArgOptions SDGenerationParams::get_options() {
&vae_tiling_params.enabled},
{"",
"--temporal-tiling",
"enable temporal tiling for LTX video VAE decode",
"enable temporal tiling for supported video VAE decode",
true,
&vae_tiling_params.temporal_tiling},
{"",
otherarch/sdcpp/examples/common/media_io.cpp
+3
View File
@@ -849,6 +849,9 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
const uint32_t audio_byte_rate = has_audio ? static_cast<uint32_t>(audio->sample_rate * audio_block_align) : 0;
const uint32_t audio_data_size = has_audio ? static_cast<uint32_t>(audio_pcm.size()) : 0;
if (mjpg_quality != quality)
LOG_DEBUG("create_mjpg_avi...(): compression quality was limited from %i to %i", quality, mjpg_quality);
std::vector<uint8_t> avi_data;
avi_data.reserve(static_cast<size_t>(num_images) * 1024);
otherarch/sdcpp/include/stable-diffusion.h
+3
View File
@@ -444,6 +444,9 @@ typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void*
SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
// In each sampling pass, a positive interval previews every Nth denoiser step, while a
// negative interval previews only completed logical step -interval. Zero previews the final
// completed step of the first sampling pass (base-resolution or high-noise).
SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
SD_API int32_t sd_get_num_physical_cores();
otherarch/sdcpp/src/conditioning/conditioner.hpp
+24 -3
View File
@@ -2978,15 +2978,36 @@ struct LTXAVEmbedder : public Conditioner {
std::shared_ptr<GemmaTokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<LTXAVTextProjectionRunner> projector;
std::string projector_prefix;
bool dual_projection = false;
// Gemma 4 keeps a per-layer output scalar that no Gemma 3 checkpoint has, and widens its
// full-attention heads to 512 so their q_proj is twice a sliding layer's.
static LLM::LLMArch detect_gemma_arch(const String2TensorStorage& tensor_storage_map,
const std::string& llm_prefix) {
if (tensor_storage_map.find(llm_prefix + ".model.layers.0.layer_scalar") != tensor_storage_map.end()) {
return LLM::LLMArch::GEMMA4_12B;
}
auto global_q = tensor_storage_map.find(llm_prefix + ".model.layers.5.self_attn.q_proj.weight");
auto sliding_q = tensor_storage_map.find(llm_prefix + ".model.layers.0.self_attn.q_proj.weight");
if (global_q != tensor_storage_map.end() &&
sliding_q != tensor_storage_map.end() &&
global_q->second.ne[1] == sliding_q->second.ne[1] * 2) {
return LLM::LLMArch::GEMMA4_12B;
}
return LLM::LLMArch::GEMMA3_12B;
}
LTXAVEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& llm_prefix = "text_encoders.llm",
const std::string& projector_prefix = "text_embedding_projection",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: projector_prefix(projector_prefix) {
LLM::LLMArch arch = detect_gemma_arch(tensor_storage_map, llm_prefix);
LOG_INFO("ltxav text encoder: %s", arch == LLM::LLMArch::GEMMA4_12B ? "gemma 4" : "gemma 3");
tokenizer = std::make_shared<GemmaTokenizer>();
llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::GEMMA3_12B,
llm = std::make_shared<LLM::LLMRunner>(arch,
backend,
tensor_storage_map,
llm_prefix,
@@ -3001,7 +3022,7 @@ struct LTXAVEmbedder : public Conditioner {
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
projector->get_param_tensors(tensors, "text_embedding_projection");
projector->get_param_tensors(tensors, projector_prefix);
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
otherarch/sdcpp/src/core/backend_fit.cpp
+8 -9
View File
@@ -364,15 +364,11 @@ namespace sd::backend_fit {
}
bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
if (prefer_temporal_tiling) {
if (tiling_params.temporal_tiling) {
return false;
}
const char* retry_mode = nullptr;
if (prefer_temporal_tiling && !tiling_params.temporal_tiling) {
tiling_params.temporal_tiling = true;
} else {
if (tiling_params.enabled) {
return false;
}
retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal";
} else if (!tiling_params.enabled) {
tiling_params.enabled = true;
if (tiling_params.tile_size_x <= 0) {
tiling_params.tile_size_x = 256;
@@ -380,10 +376,13 @@ namespace sd::backend_fit {
if (tiling_params.tile_size_y <= 0) {
tiling_params.tile_size_y = 256;
}
retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial";
} else {
return false;
}
LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling",
tiling_params.temporal_tiling ? "temporal" : "spatial");
retry_mode);
return true;
}
otherarch/sdcpp/src/model/common/block.hpp
+4 -3
View File
@@ -268,10 +268,11 @@ public:
int64_t dim_out,
int64_t mult = 4,
Activation activation = Activation::GEGLU,
bool precision_fix = false) {
bool precision_fix = false,
bool bias = true) {
int64_t inner_dim = dim * mult;
if (activation == Activation::GELU) {
blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GELU(dim, inner_dim));
blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GELU(dim, inner_dim, bias));
} else {
blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GEGLU(dim, inner_dim));
}
@@ -285,7 +286,7 @@ public:
// The purpose of the scale here is to prevent NaN issues in certain situations.
// For example, when using Vulkan without enabling force_prec_f32,
// or when using CUDA but the weights are k-quants.
blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale));
blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, bias, false, force_prec_f32, scale));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
otherarch/sdcpp/src/model/diffusion/ltxv.hpp
+65 -6
View File
@@ -129,6 +129,10 @@ namespace LTXV {
bool self_attention_gated = false;
bool cross_attention_gated = false;
bool ff_bias = true;
bool audio_ff_bias = true;
bool use_keyframes_abs_pos_embedding = false;
static std::pair<int64_t, int64_t> infer_attention_layout(int64_t hidden_size,
int64_t preferred_heads = -1) {
if (preferred_heads > 0 && hidden_size % preferred_heads == 0) {
@@ -207,6 +211,19 @@ namespace LTXV {
tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) {
config.cross_attention_gated = true;
}
// LTX 2.5 sets ff_bias=false but leaves audio_ff_bias at its default, so the two
// branches must be detected separately; older checkpoints ship both sets of biases.
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.ff.net.0.proj.bias") == tensor_storage_map.end() &&
tensor_storage_map.find(prefix + ".transformer_blocks.0.ff.net.2.bias") == tensor_storage_map.end()) {
config.ff_bias = false;
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_ff.net.0.proj.bias") == tensor_storage_map.end() &&
tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_ff.net.2.bias") == tensor_storage_map.end()) {
config.audio_ff_bias = false;
}
if (tensor_storage_map.find(prefix + ".keyframes_abs_pos_embedding") != tensor_storage_map.end()) {
config.use_keyframes_abs_pos_embedding = true;
}
if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() &&
tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) {
config.use_caption_projection = false;
@@ -874,8 +891,7 @@ namespace LTXV {
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
if (num_learnable_registers > 0) {
ggml_type wtype = get_type(prefix + "learnable_registers", tensor_storage_map, GGML_TYPE_F32);
params["learnable_registers"] = ggml_new_tensor_2d(ctx, wtype, hidden_size, num_learnable_registers);
params["learnable_registers"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_learnable_registers);
}
}
@@ -1130,7 +1146,9 @@ namespace LTXV {
int64_t a_context_dim,
bool apply_gated_attention,
bool cross_attention_adaln,
bool video_rope_interleaved)
bool video_rope_interleaved,
bool ff_bias = true,
bool audio_ff_bias = true)
: v_dim(v_dim),
a_dim(a_dim),
cross_attention_adaln(cross_attention_adaln) {
@@ -1140,8 +1158,8 @@ namespace LTXV {
blocks["audio_attn2"] = std::make_shared<CrossAttention>(a_dim, a_context_dim, a_heads, ad_head, apply_gated_attention, false);
blocks["audio_to_video_attn"] = std::make_shared<CrossAttention>(v_dim, a_dim, a_heads, ad_head, apply_gated_attention, false);
blocks["video_to_audio_attn"] = std::make_shared<CrossAttention>(a_dim, v_dim, a_heads, ad_head, apply_gated_attention, false);
blocks["ff"] = std::make_shared<FeedForward>(v_dim, v_dim, 4, FeedForward::Activation::GELU);
blocks["audio_ff"] = std::make_shared<FeedForward>(a_dim, a_dim, 4, FeedForward::Activation::GELU);
blocks["ff"] = std::make_shared<FeedForward>(v_dim, v_dim, 4, FeedForward::Activation::GELU, false, ff_bias);
blocks["audio_ff"] = std::make_shared<FeedForward>(a_dim, a_dim, 4, FeedForward::Activation::GELU, false, audio_ff_bias);
}
std::vector<ggml_tensor*> get_ada_values(GGMLRunnerContext* ctx,
@@ -1320,6 +1338,12 @@ namespace LTXV {
get_type(prefix + "audio_scale_shift_table", tensor_storage_map, GGML_TYPE_F32),
config.audio_hidden_size,
2);
if (config.use_keyframes_abs_pos_embedding) {
params["keyframes_abs_pos_embedding"] = ggml_new_tensor_2d(ctx,
get_type(prefix + "keyframes_abs_pos_embedding", tensor_storage_map, GGML_TYPE_F32),
config.hidden_size,
1);
}
}
LTXAVModelBlock(const LTXAVConfig& config)
@@ -1386,7 +1410,9 @@ namespace LTXV {
config.audio_cross_attention_dim,
config.self_attention_gated || config.cross_attention_gated,
config.cross_attention_adaln,
config.video_rope_interleaved);
config.video_rope_interleaved,
config.ff_bias,
config.audio_ff_bias);
}
blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
@@ -1534,6 +1560,38 @@ namespace LTXV {
return {v_context, a_context};
}
// The video encoder is causal, so the first latent frame covers a single pixel frame while
// every later one covers temporal_scale_factor. LTX 2.5 marks that token class with a
// learned embedding added right after patchify_proj.
ggml_tensor* apply_keyframes_abs_pos_embedding(GGMLRunnerContext* ctx,
ggml_tensor* vx,
int64_t tokens_per_latent_frame) {
if (!config.use_keyframes_abs_pos_embedding || params.count("keyframes_abs_pos_embedding") == 0) {
return vx;
}
int64_t tokens = vx->ne[1];
if (tokens_per_latent_frame <= 0 || tokens_per_latent_frame > tokens) {
return vx;
}
auto embedding = params["keyframes_abs_pos_embedding"];
auto first = ggml_cont(ctx->ggml_ctx,
ggml_view_3d(ctx->ggml_ctx, vx, vx->ne[0], tokens_per_latent_frame, vx->ne[2], vx->nb[1], vx->nb[2], 0));
first = ggml_add(ctx->ggml_ctx, first, embedding);
if (tokens_per_latent_frame == tokens) {
return first;
}
auto rest = ggml_cont(ctx->ggml_ctx,
ggml_view_3d(ctx->ggml_ctx,
vx,
vx->ne[0],
tokens - tokens_per_latent_frame,
vx->ne[2],
vx->nb[1],
vx->nb[2],
tokens_per_latent_frame * vx->nb[1]));
return ggml_concat(ctx->ggml_ctx, first, rest, 1);
}
std::vector<ggml_tensor*> get_output_scale_shift(GGMLRunnerContext* ctx,
ggml_tensor* table,
ggml_tensor* embedded_timestep,
@@ -1575,6 +1633,7 @@ namespace LTXV {
vx = patchify_video(ctx, vx, n);
vx = patchify_proj->forward(ctx, vx);
vx = apply_keyframes_abs_pos_embedding(ctx, vx, width * height);
if (ax != nullptr && ggml_nelements(ax) > 0 && audio_time > 0) {
ax = patchify_audio(ctx, ax);
ax = audio_patchify_proj->forward(ctx, ax);
otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp
+34 -32
View File
@@ -123,13 +123,6 @@ namespace MiniMaxH3 {
return to_shift * base / (1.f + (to_shift - 1.f) * base);
}
static float time_shift_slope(float sigma, float from_shift, float to_shift) {
float base = sigma / (from_shift + sigma * (1.f - from_shift));
float a = 1.f + (from_shift - 1.f) * base;
float b = 1.f + (to_shift - 1.f) * base;
return to_shift * a * a / (from_shift * b * b);
}
struct TimeEmbedder : public GGMLBlock {
TimeEmbedder(int64_t input_dim, int64_t hidden_dim, int64_t output_dim) {
blocks["proj_in"] = std::make_shared<Linear>(input_dim, hidden_dim, true, true);
@@ -594,8 +587,7 @@ namespace MiniMaxH3 {
const std::vector<TokenModulationSpan>& segments,
const std::vector<SequenceSegment>& sequence_segments,
const TokenModulationSpan& video_segment,
const TokenModulationSpan& audio_segment,
float audio_slope) {
const TokenModulationSpan& audio_segment) {
auto video_proj = std::dynamic_pointer_cast<Linear>(blocks["video_patch_proj"]);
auto audio_proj = std::dynamic_pointer_cast<Linear>(blocks["audio_patch_proj"]);
@@ -715,7 +707,7 @@ namespace MiniMaxH3 {
audio->ne[2]);
audio_out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3));
video_out = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f);
audio_out = ggml_ext_scale(ctx->ggml_ctx, audio_out, -audio_slope);
audio_out = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f);
return {video_out, audio_out};
}
};
@@ -1040,9 +1032,9 @@ namespace MiniMaxH3 {
GGML_ASSERT(!audio_input_cache.empty());
GGML_ASSERT(!context_tensor.empty());
auto video = make_input(video_input_cache);
auto audio = make_input(audio_input_cache);
auto context = make_input(context_tensor);
auto video = make_input(video_input_cache);
auto audio_carrier = make_input(audio_input_cache);
auto context = make_input(context_tensor);
std::vector<ggml_tensor*> condition_inputs;
condition_inputs.reserve(condition_videos.size());
for (const auto& condition : condition_videos) {
@@ -1054,21 +1046,26 @@ namespace MiniMaxH3 {
audio_condition_inputs.push_back(make_input(condition));
}
float sigma_v = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
float t_v = 1.f - sigma_v;
float t_a = 1.f - time_shift_sigma(sigma_v, video_shift, audio_shift);
auto layout = build_layout(context_tensor.shape()[1],
video_input_cache.shape()[2],
video_input_cache.shape()[1],
video_input_cache.shape()[0],
audio_length,
condition_videos,
condition_audios,
keyframe_indices,
reference_blocks,
text_tags,
t_v,
t_a);
float sigma_v = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
float sigma_a = time_shift_sigma(sigma_v, video_shift, audio_shift);
float audio_scale = video_shift / audio_shift;
float t_v = 1.f - sigma_v;
float t_a = 1.f - sigma_a;
// The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed
// latent follows one sigma schedule. Restore x_a for the H3 network.
auto audio = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
auto layout = build_layout(context_tensor.shape()[1],
video_input_cache.shape()[2],
video_input_cache.shape()[1],
video_input_cache.shape()[0],
audio_length,
condition_videos,
condition_audios,
keyframe_indices,
reference_blocks,
text_tags,
t_v,
t_a);
position_input_cache = sd::Tensor<float>(
{3, static_cast<int64_t>(layout.positions.size() / 3)},
@@ -1129,10 +1126,15 @@ namespace MiniMaxH3 {
layout.segments,
layout.sequence_segments,
layout.video_segment,
layout.audio_segment,
time_shift_slope(sigma_v, video_shift, audio_shift));
auto merged = merge_av_latents(compute_ctx, output.first, output.second);
auto graph = new_graph_custom(H3_GRAPH_SIZE);
layout.audio_segment);
// Convert the model's audio velocity to d(c_a) / d(sigma_v).
output.second = ggml_add(compute_ctx,
ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale),
ggml_ext_scale(compute_ctx,
output.second,
1.f + (audio_scale - 1.f) * sigma_a));
auto merged = merge_av_latents(compute_ctx, output.first, output.second);
auto graph = new_graph_custom(H3_GRAPH_SIZE);
ggml_build_forward_expand(graph, merged);
return graph;
}
otherarch/sdcpp/src/model/te/llm.hpp
+150 -27
View File
@@ -40,6 +40,7 @@ namespace LLM {
MINISTRAL_3_3B,
GEMMA3_12B,
GEMMA2_2B,
GEMMA4_12B,
GPT_OSS_20B,
ARCH_COUNT,
};
@@ -52,6 +53,7 @@ namespace LLM {
"ministral3.3b",
"gemma3_12b",
"gemma2_2b",
"gemma4_12b",
"gpt_oss_20b",
};
@@ -120,6 +122,15 @@ namespace LLM {
bool have_vision_weight = false;
bool llama_cpp_style = false;
// gemma4 config
int global_head_dim = 0;
int num_global_kv_heads = 0;
float global_partial_rotary = 1.f;
bool global_k_eq_v = false;
bool v_norm = false;
bool layer_scalar = false;
bool unscaled_attention = false;
static LLMConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
LLMArch arch) {
@@ -157,6 +168,27 @@ namespace LLM {
config.rope_thetas = {1000000.f, 10000.f};
config.rope_scales = {8.f, 1.f};
config.sliding_attention = {1024, 1024, 1024, 1024, 1024, 0};
} else if (arch == LLMArch::GEMMA4_12B) {
config.head_dim = 256;
config.num_heads = 16;
config.num_kv_heads = 8;
config.global_head_dim = 512;
config.num_global_kv_heads = 1;
config.global_partial_rotary = 0.25f;
config.global_k_eq_v = true;
config.v_norm = true;
config.layer_scalar = true;
config.unscaled_attention = true;
config.qkv_bias = false;
config.qk_norm = true;
config.rms_norm_eps = 1e-6f;
config.rms_norm_add = false;
config.normalize_input = true;
config.max_position_embeddings = 262144;
config.mlp_activation = MLPActivation::GELU_TANH;
config.rope_thetas = {1000000.f, 10000.f};
config.rope_scales = {1.f, 1.f};
config.sliding_attention = {1024, 1024, 1024, 1024, 1024, 0};
} else if (arch == LLMArch::GEMMA2_2B) {
config.head_dim = 256;
config.num_heads = 8;
@@ -232,12 +264,12 @@ namespace LLM {
}
}
}
if (contains(name, "visual.blocks.0.mlp.linear_fc1.weight") ||
contains(name, "visual.blocks.0.mlp.gate_proj.weight")) {
if (ends_with(name, "visual.blocks.0.mlp.linear_fc1.weight") ||
ends_with(name, "visual.blocks.0.mlp.gate_proj.weight")) {
config.vision.intermediate_size = tensor_storage.ne[1];
}
if (contains(name, "visual.merger.linear_fc2.weight") ||
contains(name, "visual.merger.mlp.2.weight")) {
if (ends_with(name, "visual.merger.linear_fc2.weight") ||
ends_with(name, "visual.merger.mlp.2.weight")) {
config.vision.out_hidden_size = tensor_storage.ne[1];
}
continue;
@@ -256,22 +288,26 @@ namespace LLM {
config.hidden_size = tensor_storage.ne[0];
config.vocab_size = tensor_storage.ne[1];
}
if (contains(name, "layers.0.mlp.gate_proj.weight")) {
if (ends_with(name, "layers.0.mlp.gate_proj.weight")) {
config.intermediate_size = tensor_storage.ne[1];
}
if (contains(name, "layers.0.mlp.experts.gate_up_proj.weight")) {
if (ends_with(name, "layers.0.mlp.experts.gate_up_proj.weight")) {
config.intermediate_size = tensor_storage.ne[1] / 2;
}
if (contains(name, "layers.0.mlp.experts.gate_proj.weight")) {
if (ends_with(name, "layers.0.mlp.experts.gate_proj.weight")) {
config.intermediate_size = tensor_storage.ne[1];
}
}
if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
config.num_heads = 16;
}
if (arch == LLMArch::QWEN3_VL && config.num_layers == 50 && config.hidden_size == 5120) {
config.num_heads = 64;
config.final_norm = false;
if (arch == LLMArch::QWEN3_VL &&
(config.num_layers == 50 || config.num_layers == 64) &&
config.hidden_size == 5120) {
config.num_heads = 64;
if (config.num_layers == 50) {
config.final_norm = false;
}
}
if (detected_vision_layers > 0) {
config.vision.num_layers = detected_vision_layers;
@@ -1059,6 +1095,11 @@ namespace LLM {
std::vector<float> rope_thetas;
std::vector<float> rope_scales;
bool has_attention_sinks;
bool k_eq_v;
bool v_norm;
bool unscaled_attention;
float rms_norm_eps;
int rope_pairs;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
@@ -1069,24 +1110,48 @@ namespace LLM {
}
public:
Attention(const LLMConfig& config)
Attention(const LLMConfig& config, bool global_layer = false)
: arch(config.arch),
num_heads(config.num_heads),
num_kv_heads(config.num_kv_heads),
head_dim(config.head_dim),
num_kv_heads(global_layer && config.num_global_kv_heads > 0 ? config.num_global_kv_heads : config.num_kv_heads),
head_dim(global_layer && config.global_head_dim > 0 ? config.global_head_dim : config.head_dim),
qk_norm(config.qk_norm),
max_position_embeddings(config.max_position_embeddings),
rope_thetas(config.rope_thetas),
rope_scales(config.rope_scales),
has_attention_sinks(config.arch == LLMArch::GPT_OSS_20B) {
has_attention_sinks(config.arch == LLMArch::GPT_OSS_20B),
k_eq_v(global_layer && config.global_k_eq_v),
v_norm(config.v_norm),
unscaled_attention(config.unscaled_attention),
rms_norm_eps(config.rms_norm_eps),
rope_pairs(0) {
blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
if (!k_eq_v) {
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
}
blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, config.hidden_size, config.attention_out_bias);
if (config.qk_norm) {
blocks["q_norm"] = std::make_shared<LLMRMSNorm>(head_dim, config.rms_norm_eps, config.rms_norm_add);
blocks["k_norm"] = std::make_shared<LLMRMSNorm>(head_dim, config.rms_norm_eps, config.rms_norm_add);
}
// Proportional RoPE rotates only the leading `rope_pairs` dimension pairs of the head;
// the rest are left unrotated through freq_factors (see rope_freq_factors()).
float partial = global_layer ? config.global_partial_rotary : 1.f;
rope_pairs = static_cast<int>(partial * head_dim / 2.f);
}
// ggml applies theta_i / freq_factors[i], so a huge factor collapses the angle to zero and
// leaves that pair unrotated. This reproduces transformers' "proportional" RoPE, whose
// inv_freq is zero-padded past `rope_pairs`, without reordering the head.
ggml_tensor* rope_freq_factors(ggml_context* ctx) const {
int pairs = head_dim / 2;
if (rope_pairs >= pairs) {
return nullptr;
}
auto rotated = ggml_ext_ones(ctx, rope_pairs, 1, 1, 1);
auto unrotated = ggml_ext_full(ctx, 1e30f, pairs - rope_pairs, 1, 1, 1);
return ggml_concat(ctx, rotated, unrotated, 0);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -1099,12 +1164,12 @@ namespace LLM {
int64_t N = x->ne[2];
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o_proj"]);
auto q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim]
auto k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
auto v = v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
auto q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim]
auto k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
auto v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, num_heads, head_dim]
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N); // [N, n_token, num_kv_heads, head_dim]
@@ -1117,6 +1182,10 @@ namespace LLM {
q = q_norm->forward(ctx, q);
k = k_norm->forward(ctx, k);
}
if (v_norm) {
// Gemma 4 normalizes V with a weightless RMS norm, and never rotates it.
v = ggml_rms_norm(ctx->ggml_ctx, v, rms_norm_eps);
}
if (arch == LLMArch::MISTRAL_SMALL_3_2) {
q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@@ -1187,6 +1256,35 @@ namespace LLM {
1.f,
32.f,
1.f);
} else if (arch == LLMArch::GEMMA4_12B) {
float rope_theta = (rope_index == 1 ? 10000.0f : 1000000.0f);
auto freq_factors = rope_freq_factors(ctx->ggml_ctx);
q = ggml_rope_ext(ctx->ggml_ctx,
q,
input_pos,
freq_factors,
head_dim,
GGML_ROPE_TYPE_NEOX,
static_cast<int>(max_position_embeddings),
rope_theta,
1.f,
0.f,
1.f,
32.f,
1.f);
k = ggml_rope_ext(ctx->ggml_ctx,
k,
input_pos,
freq_factors,
head_dim,
GGML_ROPE_TYPE_NEOX,
static_cast<int>(max_position_embeddings),
rope_theta,
1.f,
0.f,
1.f,
32.f,
1.f);
} else if (arch == LLMArch::GEMMA2_2B) {
q = ggml_rope_ext(ctx->ggml_ctx,
q,
@@ -1224,6 +1322,11 @@ namespace LLM {
k = ggml_rope_multi(ctx->ggml_ctx, k, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_MROPE, 128000, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
}
if (unscaled_attention) {
// Gemma 4 attends with scaling=1.0; undo the helper's own 1/sqrt(head_dim).
q = ggml_ext_scale(ctx->ggml_ctx, q, std::sqrt(static_cast<float>(head_dim)));
}
q = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)); // [N, num_heads, n_token, head_dim]
q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]); // [N*num_heads, n_token, head_dim]
@@ -1262,15 +1365,30 @@ namespace LLM {
protected:
LLMArch arch;
int sliding_attention;
bool has_layer_scalar;
std::string post_attention_norm_name;
std::string pre_ffw_norm_name;
std::string post_ffw_norm_name;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
if (has_layer_scalar) {
params["layer_scalar"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
}
}
public:
TransformerBlock(const LLMConfig& config, int layer_index)
: arch(config.arch),
sliding_attention(0) {
if (config.arch == LLMArch::GEMMA3_12B) {
sliding_attention(0),
has_layer_scalar(config.layer_scalar) {
if (config.arch == LLMArch::GEMMA4_12B) {
post_attention_norm_name = "post_attention_layernorm";
pre_ffw_norm_name = "pre_feedforward_layernorm";
post_ffw_norm_name = "post_feedforward_layernorm";
} else if (config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GEMMA4_12B) {
post_attention_norm_name = "post_attention_norm"; // attn_post_norm
pre_ffw_norm_name = "post_attention_layernorm"; // ffn_norm
post_ffw_norm_name = "post_ffw_norm"; // ffn_post_norm
@@ -1284,7 +1402,10 @@ namespace LLM {
pre_ffw_norm_name = "post_attention_layernorm"; // ffn_norm
}
blocks["self_attn"] = std::make_shared<Attention>(config);
if (!config.sliding_attention.empty()) {
sliding_attention = config.sliding_attention[layer_index % config.sliding_attention.size()];
}
blocks["self_attn"] = std::make_shared<Attention>(config, sliding_attention == 0);
if (config.arch == LLMArch::GPT_OSS_20B) {
blocks["mlp"] = std::make_shared<GPTOSSMLP>(config);
} else {
@@ -1301,9 +1422,6 @@ namespace LLM {
if (!post_ffw_norm_name.empty()) {
blocks[post_ffw_norm_name] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
}
if (!config.sliding_attention.empty()) {
sliding_attention = config.sliding_attention[layer_index % config.sliding_attention.size()];
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -1325,7 +1443,7 @@ namespace LLM {
}
ggml_tensor* block_attention_mask = attention_mask;
int rope_index = 0;
if ((arch == LLMArch::GEMMA3_12B || arch == LLMArch::GPT_OSS_20B) && sliding_attention > 0) {
if ((arch == LLMArch::GEMMA3_12B || arch == LLMArch::GEMMA4_12B || arch == LLMArch::GPT_OSS_20B) && sliding_attention > 0) {
block_attention_mask = sliding_attention_mask;
rope_index = 1;
}
@@ -1352,6 +1470,10 @@ namespace LLM {
}
x = ggml_add_inplace(ctx->ggml_ctx, x, residual);
if (has_layer_scalar) {
x = ggml_mul(ctx->ggml_ctx, x, params["layer_scalar"]);
}
return x;
}
};
@@ -1846,6 +1968,7 @@ namespace LLM {
config.arch == LLMArch::MINISTRAL_3_3B ||
config.arch == LLMArch::QWEN3 ||
config.arch == LLMArch::GEMMA3_12B ||
config.arch == LLMArch::GEMMA4_12B ||
config.arch == LLMArch::GEMMA2_2B ||
config.arch == LLMArch::GPT_OSS_20B) {
input_pos_vec.resize(n_tokens);
@@ -1910,7 +2033,7 @@ namespace LLM {
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
}
if (config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GPT_OSS_20B) {
if (config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GEMMA4_12B || config.arch == LLMArch::GPT_OSS_20B) {
int sliding_window = 0;
for (int window : config.sliding_attention) {
sliding_window = std::max(sliding_window, window);
otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp
+9
View File
@@ -758,6 +758,15 @@ namespace Hunyuan {
return "hunyuan_video_vae";
}
bool supports_temporal_tiling(VAETemporalDirection direction) const override {
return direction == VAETemporalDirection::DECODE;
}
int get_temporal_tile_output_scale(VAETemporalDirection direction) const override {
SD_UNUSED(direction);
return 4;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (!decode_only) {
encoder.get_param_tensors(tensors, weight_prefix + ".encoder");
otherarch/sdcpp/src/model/vae/ltx_vae.hpp
+36 -81
View File
@@ -1213,9 +1213,6 @@ struct LTXVideoVAE : public VAE {
static constexpr int DEFAULT_TEMPORAL_TILE_OVERLAP = 1;
bool decode_only;
bool temporal_tiling_enabled = false;
int temporal_tile_frames = DEFAULT_TEMPORAL_TILE_FRAMES;
int temporal_tile_overlap = DEFAULT_TEMPORAL_TILE_OVERLAP;
int ltx_vae_version;
bool timestep_conditioning;
int patch_size;
@@ -1248,64 +1245,24 @@ struct LTXVideoVAE : public VAE {
return "ltx_video_vae";
}
void set_temporal_tiling_enabled(bool enabled) override {
temporal_tiling_enabled = enabled;
bool supports_temporal_tiling(VAETemporalDirection direction) const override {
return direction == VAETemporalDirection::DECODE;
}
void set_tiling_params(const sd_tiling_params_t& params) override {
temporal_tiling_enabled = params.temporal_tiling;
temporal_tile_frames = DEFAULT_TEMPORAL_TILE_FRAMES;
temporal_tile_overlap = DEFAULT_TEMPORAL_TILE_OVERLAP;
int get_default_temporal_tile_frames(VAETemporalDirection direction) const override {
SD_UNUSED(direction);
return DEFAULT_TEMPORAL_TILE_FRAMES;
}
for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "LTX VAE extra tiling arg")) {
int parsed = 0;
if (!parse_strict_int(value, parsed)) {
LOG_WARN("ignoring invalid LTX VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str());
} else if (key == "temporal_tile_frames") {
temporal_tile_frames = std::max(1, parsed);
} else if (key == "temporal_tile_overlap") {
temporal_tile_overlap = std::max(0, parsed);
} else {
LOG_WARN("ignoring unknown LTX VAE extra tiling arg '%s'", key.c_str());
}
}
int get_default_temporal_tile_overlap(VAETemporalDirection direction) const override {
SD_UNUSED(direction);
return DEFAULT_TEMPORAL_TILE_OVERLAP;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
vae.get_param_tensors(tensors, weight_prefix);
}
struct TemporalTilePlan {
int frames = 1;
int overlap = 0;
int stride = 1;
int num_tiles = 1;
};
TemporalTilePlan resolve_temporal_tile_plan(int64_t total_frames) const {
TemporalTilePlan plan;
plan.frames = std::max(1, temporal_tile_frames);
plan.overlap = std::max(0, temporal_tile_overlap);
if (plan.overlap >= plan.frames) {
LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows",
plan.overlap,
plan.frames);
plan.overlap = plan.frames - 1;
}
if (total_frames > 1 && plan.overlap >= total_frames) {
LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to total latent frames (%lld), adjusting values to decode at least one tile",
plan.overlap,
(long long)total_frames);
plan.overlap = static_cast<int>(total_frames - 1);
}
plan.stride = std::max(1, plan.frames - plan.overlap);
int64_t tiled_frames = std::max<int64_t>(1, total_frames - plan.overlap);
plan.num_tiles = total_frames > 0 ? static_cast<int>((tiled_frames + plan.stride - 1) / plan.stride) : 0;
return plan;
}
std::string temporal_feat_cache_name(size_t feat_idx) const {
return "ltx_vae_temporal_feat:" + std::to_string(feat_idx);
}
@@ -1365,52 +1322,53 @@ struct LTXVideoVAE : public VAE {
sd::Tensor<float> decode_temporal_tiled_streaming(const int n_threads,
const sd::Tensor<float>& input,
size_t expected_dim) {
size_t expected_dim,
const VAETemporalTilingConfig& config) {
const int64_t total_frames = input.shape()[2];
TemporalTilePlan plan = resolve_temporal_tile_plan(total_frames);
auto plan = make_vae_temporal_tile_plan(total_frames, config);
LOG_DEBUG("Using streaming temporal tiling: temporal_tile_frames=%d, temporal_tile_overlap=%d, total latent frames=%lld, resulting in %d tiles",
plan.frames,
plan.tile_frames,
plan.overlap,
(long long)total_frames,
plan.num_tiles);
(int)plan.tiles.size());
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
sd::Tensor<float> output;
for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) {
const int64_t end = std::min<int64_t>(total_frames, start + plan.frames);
const int chunk_overlap = end < total_frames ? plan.overlap : 0;
auto z_chunk = sd::ops::slice(input, 2, start, end);
auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& z_chunk, const VAETemporalTile& tile) {
LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d",
(long long)(start / plan.stride + 1),
plan.num_tiles,
(long long)start,
(long long)end,
chunk_overlap);
(long long)tile.index + 1,
(int)plan.tiles.size(),
(long long)tile.start,
(long long)tile.end,
tile.overlap);
auto get_graph = [&]() -> ggml_cgraph* {
return build_temporal_tile_graph(z_chunk,
static_cast<int>(start),
chunk_overlap);
static_cast<int>(tile.start),
tile.overlap);
};
auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
expected_dim);
if (chunk.empty()) {
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
return {};
}
output = output.empty() ? std::move(chunk) : sd::ops::concat(output, chunk, 2);
}
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
expected_dim);
});
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
return output;
}
sd::Tensor<float> _compute_temporal_tiled(const int n_threads,
const sd::Tensor<float>& input,
VAETemporalDirection direction,
const VAETemporalTilingConfig& config) override {
GGML_ASSERT(direction == VAETemporalDirection::DECODE);
return decode_temporal_tiled_streaming(n_threads,
input,
static_cast<size_t>(input.dim()),
config);
}
ggml_cgraph* build_latent_statistics_graph(const sd::Tensor<float>& z_tensor, bool normalize) {
ggml_cgraph* gf = new_graph_custom(1024);
ggml_tensor* z = make_input(z_tensor);
@@ -1446,9 +1404,6 @@ struct LTXVideoVAE : public VAE {
input = sd::ops::slice(input, 2, 0, cropped_t);
}
}
if (decode_graph && temporal_tiling_enabled && input.dim() == 5 && input.shape()[2] > 1) {
return decode_temporal_tiled_streaming(n_threads, input, expected_dim);
}
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp
+31 -14
View File
@@ -154,8 +154,9 @@ namespace MiniMaxH3 {
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
params["q_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
params["v_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
params["q_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
params["zero_k_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
params["v_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
@@ -166,7 +167,7 @@ namespace MiniMaxH3 {
return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1);
};
auto q = ggml_add(ctx->ggml_ctx, qkv[0], bias_shape(params["q_bias"]));
auto k = qkv[1];
auto k = ggml_add(ctx->ggml_ctx, qkv[1], bias_shape(params["zero_k_bias"]));
auto v = ggml_add(ctx->ggml_ctx, qkv[2], bias_shape(params["v_bias"]));
int64_t sequence = x->ne[1];
@@ -358,22 +359,38 @@ namespace MiniMaxH3 {
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* waveform) {
GGML_ASSERT(waveform->ne[1] == 2);
GGML_ASSERT(waveform->ne[1] * waveform->ne[2] * waveform->ne[3] == 2);
auto encoder = std::dynamic_pointer_cast<AudioEncoder>(blocks["encoder"]);
auto pre = std::dynamic_pointer_cast<AudioAttentionProjection>(blocks["pre_block"]);
auto mean_proj = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["mean_proj"]);
waveform = ggml_reshape_3d(ctx->ggml_ctx, waveform, waveform->ne[0], 1, waveform->ne[1]);
auto x = encoder->forward(ctx, waveform); // [B*S, 2048, T]
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = pre->forward(ctx, x);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
auto z = mean_proj->forward(ctx, x);
// GGML's batched conv1d storage interleaves the stream dimension
// with output channels. Subsequent layers then read stereo samples
// as adjacent feature channels. Run each mono stream independently,
// matching PyTorch's reshape(B*S, 1, samples), and concatenate only
// the completed normalized latents.
const int64_t streams = waveform->ne[2] * waveform->ne[3];
waveform = ggml_reshape_3d(ctx->ggml_ctx,
waveform,
waveform->ne[0],
1,
streams);
ggml_tensor* stereo_z = nullptr;
for (int64_t stream = 0; stream < streams; ++stream) {
auto mono = ggml_ext_slice(ctx->ggml_ctx, waveform, 2, stream, stream + 1);
auto x = encoder->forward(ctx, mono);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = pre->forward(ctx, x);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
auto z = mean_proj->forward(ctx, x);
auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
auto std = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
z = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
auto std = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
z = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
stereo_z = stereo_z == nullptr ? z : ggml_concat(ctx->ggml_ctx, stereo_z, z, 1);
}
return stereo_z;
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent) {
otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp
+31 -31
View File
@@ -558,10 +558,11 @@ namespace MiniMaxH3VAE {
}
static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) {
params.enabled = true;
params.tile_size_x = 16;
params.tile_size_y = 16;
params.target_overlap = 0.25f;
params.enabled = true;
params.temporal_tiling = false;
params.tile_size_x = 16;
params.tile_size_y = 16;
params.target_overlap = 0.25f;
return params;
}
@@ -624,15 +625,13 @@ namespace MiniMaxH3VAE {
if (pad > 0) {
input = repeat_last_frame(input, pad);
}
sd::Tensor<float> result;
for (int64_t start = 0; start < input.shape()[2]; start += 17) {
auto chunk = sd::ops::slice(input, 2, start, start + 17);
auto encoded = VAE::encode(n_threads, chunk, tiling, circular_x, circular_y);
if (encoded.empty()) {
return {};
}
result = result.empty() ? std::move(encoded)
: sd::ops::concat(result, encoded, 2);
auto plan = make_vae_temporal_tile_plan(input.shape()[2], {17, 0});
auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) {
SD_UNUSED(tile);
return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y);
});
if (result.empty()) {
return {};
}
if (result.shape()[2] > 3) {
result = sd::ops::slice(result, 2, 0, result.shape()[2] - 3);
@@ -685,22 +684,21 @@ namespace MiniMaxH3VAE {
input = repeat_last_frame(input, pad_tokens);
}
sd::Tensor<float> result;
sd::Tensor<float> overlap;
for (int64_t i = 0; i < num_chunks; ++i) {
int64_t start = i * tokens_per_chunk;
int64_t end = std::min(start + tokens_per_chunk + token_overlap,
input.shape()[2]);
auto chunk = sd::ops::slice(input, 2, start, end);
auto decoded = VAE::decode(n_threads,
chunk,
tiling,
true,
circular_x,
circular_y,
silent);
auto plan = make_vae_temporal_tile_plan(
input.shape()[2],
{static_cast<int>(tokens_per_chunk + token_overlap), static_cast<int>(token_overlap)});
GGML_ASSERT(plan.tiles.size() == static_cast<size_t>(num_chunks));
auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) {
auto decoded = VAE::decode(n_threads,
chunk,
tiling,
true,
circular_x,
circular_y,
silent);
if (decoded.empty()) {
return {};
return sd::Tensor<float>();
}
int64_t first_end = std::min<int64_t>(frames_per_chunk, decoded.shape()[2]);
@@ -712,8 +710,6 @@ namespace MiniMaxH3VAE {
first = blend_temporal(overlap, first, frame_overlap);
overlap = {};
}
result = result.empty() ? std::move(first)
: sd::ops::concat(result, first, 2);
if (decoded.shape()[2] > frames_per_chunk + frame_pre_padding) {
overlap = sd::ops::slice(decoded,
@@ -721,10 +717,14 @@ namespace MiniMaxH3VAE {
frames_per_chunk + frame_pre_padding,
decoded.shape()[2]);
}
if (i == num_chunks - 1 && !overlap.empty()) {
result = sd::ops::concat(result, overlap, 2);
if (tile.last && !overlap.empty()) {
first = sd::ops::concat(first, overlap, 2);
overlap = {};
}
return first;
});
if (result.empty()) {
return {};
}
int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5;
otherarch/sdcpp/src/model/vae/tae.hpp
+15
View File
@@ -819,6 +819,21 @@ struct TinyVideoAutoEncoder : public VAE {
return "taehv";
}
bool supports_temporal_tiling(VAETemporalDirection direction) const override {
return direction == VAETemporalDirection::DECODE && !sd_version_is_minimax_h3(version);
}
int get_temporal_tile_output_scale(VAETemporalDirection direction) const override {
SD_UNUSED(direction);
int scale = 1;
for (bool upscale : taehv.time_upscale) {
if (upscale) {
scale *= 2;
}
}
return scale;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
taehv.get_param_tensors(tensors, weight_prefix);
}
otherarch/sdcpp/src/model/vae/vae.hpp
+102 -8
View File
@@ -3,6 +3,7 @@
#include "core/tensor_ggml.hpp"
#include "model/common/block.hpp"
#include "model/vae/vae_tiling.hpp"
#include "model_manager.h"
struct VAE : public GGMLRunner {
@@ -14,6 +15,87 @@ protected:
const sd::Tensor<float>& z,
bool decode_graph) = 0;
virtual bool supports_temporal_tiling(VAETemporalDirection direction) const {
SD_UNUSED(direction);
return false;
}
virtual int get_default_temporal_tile_frames(VAETemporalDirection direction) const {
SD_UNUSED(direction);
return 4;
}
virtual int get_default_temporal_tile_overlap(VAETemporalDirection direction) const {
SD_UNUSED(direction);
return 1;
}
virtual int get_temporal_tile_output_scale(VAETemporalDirection direction) const {
SD_UNUSED(direction);
return 1;
}
virtual sd::Tensor<float> _compute_temporal_tiled(const int n_threads,
const sd::Tensor<float>& input,
VAETemporalDirection direction,
const VAETemporalTilingConfig& config) {
if (direction != VAETemporalDirection::DECODE) {
return _compute(n_threads, input, false);
}
VAETemporalTilingConfig resolved_config = config;
const int output_scale = get_temporal_tile_output_scale(direction);
if (output_scale > 1 &&
resolved_config.overlap == 0 &&
input.shape()[2] > resolved_config.tile_frames) {
LOG_WARN("%s temporal decode requires at least one overlapping latent frame; using overlap=1",
get_desc().c_str());
resolved_config.overlap = 1;
}
auto plan = make_vae_temporal_tile_plan(input.shape()[2], resolved_config);
LOG_DEBUG("%s temporal tiling: tile_frames=%d, overlap=%d, total_frames=%lld, tiles=%d",
get_desc().c_str(),
plan.tile_frames,
plan.overlap,
(long long)input.shape()[2],
(int)plan.tiles.size());
return process_vae_temporal_tiles_blended(
input,
plan,
output_scale,
[&](const sd::Tensor<float>& input_tile, const VAETemporalTile& tile) {
LOG_DEBUG("%s temporal tile %d/%d: input frames [%lld, %lld)",
get_desc().c_str(),
tile.index + 1,
(int)plan.tiles.size(),
(long long)tile.start,
(long long)tile.end);
return _compute(n_threads, input_tile, true);
});
}
sd::Tensor<float> compute_with_temporal_tiling(const int n_threads,
const sd::Tensor<float>& input,
VAETemporalDirection direction,
const sd_tiling_params_t& tiling_params) {
if (!tiling_params.temporal_tiling || input.dim() != 5 || input.shape()[2] <= 1) {
return _compute(n_threads, input, direction == VAETemporalDirection::DECODE);
}
if (!supports_temporal_tiling(direction)) {
LOG_WARN("%s does not support temporal tiling for %s; processing the full temporal dimension",
get_desc().c_str(),
direction == VAETemporalDirection::DECODE ? "decode" : "encode");
return _compute(n_threads, input, direction == VAETemporalDirection::DECODE);
}
auto config = resolve_vae_temporal_tiling_config(
tiling_params,
get_default_temporal_tile_frames(direction),
get_default_temporal_tile_overlap(direction));
return _compute_temporal_tiled(n_threads, input, direction, config);
}
static inline void scale_tensor_to_minus1_1(sd::Tensor<float>* tensor) {
GGML_ASSERT(tensor != nullptr);
for (int64_t i = 0; i < tensor->numel(); ++i) {
@@ -40,10 +122,15 @@ protected:
bool circular_x,
bool circular_y,
bool decode_graph,
const sd_tiling_params_t& tiling_params,
const char* error_message,
bool silent = false) {
auto on_processing = [&](const sd::Tensor<float>& input_tile) {
auto output_tile = _compute(n_threads, input_tile, decode_graph);
auto output_tile = compute_with_temporal_tiling(
n_threads,
input_tile,
decode_graph ? VAETemporalDirection::DECODE : VAETemporalDirection::ENCODE,
tiling_params);
if (output_tile.empty()) {
LOG_ERROR("%s", error_message);
return sd::Tensor<float>();
@@ -86,6 +173,10 @@ public:
virtual int get_encoder_output_channels(int input_channels) = 0;
bool can_temporal_tile_decode() const {
return supports_temporal_tiling(VAETemporalDirection::DECODE);
}
void get_tile_sizes(int& tile_size_x,
int& tile_size_y,
float& tile_overlap,
@@ -151,9 +242,13 @@ public:
circular_x,
circular_y,
false,
tiling_params,
"vae encode compute failed while processing a tile");
} else {
output = _compute(n_threads, input, false);
output = compute_with_temporal_tiling(n_threads,
input,
VAETemporalDirection::ENCODE,
tiling_params);
}
runner_done();
@@ -177,7 +272,6 @@ public:
int64_t t0 = ggml_time_ms();
sd::Tensor<float> input = x;
sd::Tensor<float> output;
set_tiling_params(tiling_params);
if (tiling_params.enabled) {
const int scale_factor = get_scale_factor();
@@ -201,10 +295,14 @@ public:
circular_x,
circular_y,
true,
tiling_params,
"vae decode compute failed while processing a tile",
silent);
} else {
output = _compute(n_threads, input, true);
output = compute_with_temporal_tiling(n_threads,
input,
VAETemporalDirection::DECODE,
tiling_params);
}
runner_done();
@@ -226,10 +324,6 @@ public:
virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
virtual void set_tiling_params(const sd_tiling_params_t& params) {
set_temporal_tiling_enabled(params.temporal_tiling);
};
};
struct FakeVAE : public VAE {
otherarch/sdcpp/src/model/vae/vae_tiling.hpp
+213
View File
@@ -0,0 +1,213 @@
#ifndef __SD_MODEL_VAE_VAE_TILING_HPP__
#define __SD_MODEL_VAE_VAE_TILING_HPP__
#include <algorithm>
#include <cstdint>
#include <utility>
#include <vector>
#include "core/tensor.hpp"
#include "core/util.h"
enum class VAETemporalDirection {
ENCODE,
DECODE,
};
struct VAETemporalTilingConfig {
int tile_frames = 1;
int overlap = 0;
};
struct VAETemporalTile {
int index = 0;
int64_t start = 0;
int64_t end = 0;
int overlap = 0;
bool first = false;
bool last = false;
};
struct VAETemporalTilePlan {
int tile_frames = 1;
int overlap = 0;
int stride = 1;
std::vector<VAETemporalTile> tiles;
};
inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tiling_params_t& params,
int default_tile_frames,
int default_overlap) {
VAETemporalTilingConfig config;
config.tile_frames = std::max(1, default_tile_frames);
config.overlap = std::max(0, default_overlap);
for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "VAE extra tiling arg")) {
if (key != "temporal_tile_frames" && key != "temporal_tile_size" && key != "temporal_tile_overlap") {
continue;
}
int parsed = 0;
if (!parse_strict_int(value, parsed)) {
LOG_WARN("ignoring invalid VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str());
} else if (key == "temporal_tile_overlap") {
config.overlap = std::max(0, parsed);
} else {
config.tile_frames = std::max(1, parsed);
}
}
return config;
}
inline VAETemporalTilePlan make_vae_temporal_tile_plan(int64_t total_frames,
const VAETemporalTilingConfig& config) {
VAETemporalTilePlan plan;
plan.tile_frames = std::max(1, config.tile_frames);
plan.overlap = std::max(0, config.overlap);
if (total_frames <= 1) {
plan.overlap = 0;
}
if (plan.overlap >= plan.tile_frames) {
LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows",
plan.overlap,
plan.tile_frames);
plan.overlap = plan.tile_frames - 1;
}
if (total_frames > 1 && plan.overlap >= total_frames) {
LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to total frames (%lld), adjusting values to process at least one tile",
plan.overlap,
(long long)total_frames);
plan.overlap = static_cast<int>(total_frames - 1);
}
plan.stride = std::max(1, plan.tile_frames - plan.overlap);
for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) {
VAETemporalTile tile;
tile.index = static_cast<int>(plan.tiles.size());
tile.start = start;
tile.end = std::min<int64_t>(total_frames, start + plan.tile_frames);
tile.overlap = tile.end < total_frames ? plan.overlap : 0;
tile.first = start == 0;
tile.last = tile.end == total_frames;
plan.tiles.push_back(tile);
}
return plan;
}
template <typename Fn>
inline sd::Tensor<float> process_vae_temporal_tiles(const sd::Tensor<float>& input,
const VAETemporalTilePlan& plan,
Fn&& on_processing) {
sd::Tensor<float> output;
for (const auto& tile : plan.tiles) {
auto input_tile = sd::ops::slice(input, 2, tile.start, tile.end);
auto output_tile = on_processing(input_tile, tile);
if (output_tile.empty()) {
return {};
}
output = output.empty() ? std::move(output_tile)
: sd::ops::concat(output, output_tile, 2);
}
return output;
}
template <typename Fn>
inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<float>& input,
const VAETemporalTilePlan& plan,
int output_scale,
Fn&& on_processing) {
GGML_ASSERT(output_scale >= 1);
const int64_t output_frames = 1 + (input.shape()[2] - 1) * output_scale;
const int overlap_frames = plan.overlap > 0 ? 1 + (plan.overlap - 1) * output_scale : 0;
std::vector<float> weights(static_cast<size_t>(output_frames), 0.f);
sd::Tensor<float> output;
auto smootherstep = [](float value) {
return value * value * value * (value * (value * 6.f - 15.f) + 10.f);
};
for (const auto& tile : plan.tiles) {
auto input_tile = sd::ops::slice(input, 2, tile.start, tile.end);
auto output_tile = on_processing(input_tile, tile);
if (output_tile.empty()) {
return {};
}
const int64_t expected_tile_frames = 1 + (input_tile.shape()[2] - 1) * output_scale;
if (output_tile.dim() < 3 || output_tile.shape()[2] != expected_tile_frames) {
LOG_ERROR("unexpected temporal tile output shape: expected %lld frames, got %lld",
(long long)expected_tile_frames,
output_tile.dim() < 3 ? -1LL : (long long)output_tile.shape()[2]);
return {};
}
if (output.empty()) {
auto output_shape = output_tile.shape();
output_shape[2] = output_frames;
output = sd::Tensor<float>::zeros(std::move(output_shape));
} else {
if (output.dim() != output_tile.dim()) {
LOG_ERROR("temporal tile output rank mismatch: expected %lld, got %lld",
(long long)output.dim(),
(long long)output_tile.dim());
return {};
}
for (size_t dim = 0; dim < static_cast<size_t>(output.dim()); ++dim) {
if (dim != 2 && output.shape()[dim] != output_tile.shape()[dim]) {
LOG_ERROR("temporal tile output shape mismatch at dimension %zu", dim);
return {};
}
}
}
const int64_t output_start = tile.start * output_scale;
const int64_t inner = output.shape()[0] * output.shape()[1];
const int64_t outer = output.numel() / (inner * output.shape()[2]);
const int64_t tile_frames = output_tile.shape()[2];
for (int64_t frame = 0; frame < tile_frames; ++frame) {
float weight = 1.f;
if (!tile.first && overlap_frames > 0 && frame < overlap_frames) {
weight *= smootherstep(static_cast<float>(frame + 1) /
static_cast<float>(overlap_frames + 1));
}
if (!tile.last && overlap_frames > 0 && frame >= tile_frames - overlap_frames) {
weight *= smootherstep(static_cast<float>(tile_frames - frame) /
static_cast<float>(overlap_frames + 1));
}
const int64_t output_frame = output_start + frame;
GGML_ASSERT(output_frame >= 0 && output_frame < output_frames);
weights[static_cast<size_t>(output_frame)] += weight;
for (int64_t outer_index = 0; outer_index < outer; ++outer_index) {
const int64_t src_offset = (outer_index * tile_frames + frame) * inner;
const int64_t dst_offset = (outer_index * output_frames + output_frame) * inner;
for (int64_t inner_index = 0; inner_index < inner; ++inner_index) {
output[dst_offset + inner_index] += output_tile[src_offset + inner_index] * weight;
}
}
}
}
if (output.empty()) {
return {};
}
const int64_t inner = output.shape()[0] * output.shape()[1];
const int64_t outer = output.numel() / (inner * output.shape()[2]);
for (int64_t frame = 0; frame < output_frames; ++frame) {
const float weight = weights[static_cast<size_t>(frame)];
if (weight <= 0.f) {
LOG_ERROR("temporal tiling left output frame %lld uncovered", (long long)frame);
return {};
}
for (int64_t outer_index = 0; outer_index < outer; ++outer_index) {
const int64_t offset = (outer_index * output_frames + frame) * inner;
for (int64_t inner_index = 0; inner_index < inner; ++inner_index) {
output[offset + inner_index] /= weight;
}
}
}
return output;
}
#endif // __SD_MODEL_VAE_VAE_TILING_HPP__
otherarch/sdcpp/src/model/vae/wan_vae.hpp
+90 -63
View File
@@ -1219,24 +1219,40 @@ namespace WAN {
return out;
}
ggml_tensor* decode_partial(GGMLRunnerContext* ctx,
ggml_tensor* z,
int i,
int64_t b = 1) {
ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx,
ggml_tensor* z,
int chunk_idx,
int64_t b = 1) {
// z: [b*c, t, h, w]
GGML_ASSERT(b == 1);
auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
auto x = conv2->forward(ctx, z);
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode_partial.prelude", "x");
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w]
_conv_idx = 0;
auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
return out;
ggml_tensor* x;
if (is_2D) {
auto conv2_2d = std::dynamic_pointer_cast<Conv2dBut3d>(blocks["conv2"]);
x = conv2_2d->forward(ctx, z);
} else {
x = conv2->forward(ctx, z);
}
ggml_tensor* out = nullptr;
for (int64_t frame = 0; frame < x->ne[2]; ++frame) {
const int global_frame = chunk_idx + static_cast<int>(frame);
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1);
_conv_idx = 0;
auto out_frame = decoder->forward(ctx, in, b, _feat_map, _conv_idx, global_frame);
if (is_2D && global_frame > 0) {
auto repeated = out_frame;
for (int repeat = 1; repeat < 4; ++repeat) {
repeated = ggml_concat(ctx->ggml_ctx, repeated, out_frame, 2);
}
out_frame = repeated;
}
out = out == nullptr ? out_frame : ggml_concat(ctx->ggml_ctx, out, out_frame, 2);
}
return unpatchify(ctx->ggml_ctx, out, patch_size, b);
}
};
@@ -1272,6 +1288,15 @@ namespace WAN {
return "wan_vae";
}
bool supports_temporal_tiling(VAETemporalDirection direction) const override {
return direction == VAETemporalDirection::DECODE;
}
int get_temporal_tile_output_scale(VAETemporalDirection direction) const override {
SD_UNUSED(direction);
return 4;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
ae.get_param_tensors(tensors, weight_prefix);
}
@@ -1346,8 +1371,8 @@ namespace WAN {
return gf;
}
ggml_cgraph* build_graph_partial(const sd::Tensor<float>& z_tensor, bool decode_graph, int i) {
ggml_cgraph* gf = new_graph_custom(20480);
ggml_cgraph* build_temporal_tile_graph(const sd::Tensor<float>& z_tensor, int chunk_idx) {
ggml_cgraph* gf = new_graph_custom(std::max<size_t>(20480, 10240 * z_tensor.shape()[2]));
ae.clear_cache();
@@ -1360,7 +1385,7 @@ namespace WAN {
auto runner_ctx = get_context();
ggml_tensor* out = decode_graph ? ae.decode_partial(&runner_ctx, z, i) : ae.encode(&runner_ctx, z);
ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx);
for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
ggml_tensor* feat_cache = ae._feat_map[feat_idx];
@@ -1375,58 +1400,60 @@ namespace WAN {
return gf;
}
sd::Tensor<float> _compute_temporal_tiled(const int n_threads,
const sd::Tensor<float>& input,
VAETemporalDirection direction,
const VAETemporalTilingConfig& config) override {
GGML_ASSERT(direction == VAETemporalDirection::DECODE);
VAETemporalTilingConfig stateful_config = config;
stateful_config.overlap = 0;
auto plan = make_vae_temporal_tile_plan(input.shape()[2], stateful_config);
LOG_DEBUG("Wan VAE stateful temporal tiling: tile_frames=%d, total latent frames=%lld, tiles=%d",
plan.tile_frames,
(long long)input.shape()[2],
(int)plan.tiles.size());
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
ae.clear_cache();
auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& input_tile, const VAETemporalTile& tile) {
LOG_DEBUG("Wan VAE temporal tile %d/%d: latent frames [%lld, %lld)",
tile.index + 1,
(int)plan.tiles.size(),
(long long)tile.start,
(long long)tile.end);
auto get_graph = [&]() -> ggml_cgraph* {
return build_temporal_tile_graph(input_tile, static_cast<int>(tile.start));
};
return restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
static_cast<size_t>(input.dim()));
});
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
ae.clear_cache();
return output;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z,
bool decode_graph) override {
if (true) {
sd::Tensor<float> input;
if (z.dim() == 4) {
input = z.unsqueeze(2);
}
auto get_graph = [&]() -> ggml_cgraph* {
if (input.empty()) {
return build_graph(z, decode_graph);
} else {
return build_graph(input, decode_graph);
}
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
input.empty() ? z.dim() : input.dim());
if (!result.empty() && z.dim() == 4) {
result.squeeze_(2);
}
return result;
} else { // chunk 1 result is weird
ae.clear_cache();
int64_t t = z.shape()[2];
int i = 0;
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph_partial(z, decode_graph, i);
};
auto out_opt = GGMLRunner::compute<float>(get_graph, n_threads, true, true, true);
if (!out_opt.has_value()) {
return {};
}
sd::Tensor<float> out = std::move(*out_opt);
ae.clear_cache();
if (t == 1) {
return out;
}
sd::Tensor<float> output = std::move(out);
for (i = 1; i < t; i++) {
auto chunk_opt = GGMLRunner::compute<float>(get_graph, n_threads, true, true, true);
if (!chunk_opt.has_value()) {
return {};
}
out = std::move(*chunk_opt);
ae.clear_cache();
output = sd::ops::concat(output, out, 2);
}
free_cache_ctx_and_buffer();
return output;
sd::Tensor<float> input;
if (z.dim() == 4) {
input = z.unsqueeze(2);
}
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input.empty() ? z : input, decode_graph);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
input.empty() ? z.dim() : input.dim());
if (!result.empty() && z.dim() == 4) {
result.squeeze_(2);
}
return result;
}
void test() {
otherarch/sdcpp/src/name_conversion.cpp
+2
View File
@@ -149,6 +149,7 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
{"ffn_up.", "mlp.up_proj."},
{"ffn_post_norm.", "post_ffw_norm."},
{"ffn_norm.", "post_attention_layernorm."},
{"layer_output_scale.weight", "layer_scalar"},
{"output_norm.", "model.norm."},
};
@@ -1459,6 +1460,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
{"unet.", "model.diffusion_model."},
{"transformer.", "model.diffusion_model."}, // dit
{"vae.", "first_stage_model."},
{"text_encoders.llm.text_embedding_projection.", "text_embedding_projection."},
{"text_encoder.", "cond_stage_model.transformer."},
{"te.", "cond_stage_model.transformer."},
{"text_encoder.2.", "cond_stage_model.1.transformer."},
otherarch/sdcpp/src/runtime/denoiser.hpp
+44
View File
@@ -1043,6 +1043,16 @@ struct Denoiser {
const sd::Tensor<float>& latent) = 0;
virtual float noise_level_to_sigma(float noise_level) = 0;
virtual sd::Tensor<float> process_latent_in(const sd::Tensor<float>& latent) {
// An empty result means the original latent can be used unchanged.
SD_UNUSED(latent);
return {};
}
virtual sd::Tensor<float> process_latent_out(sd::Tensor<float> latent) {
return latent;
}
virtual std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) {
auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1);
std::shared_ptr<SigmaScheduler> scheduler;
@@ -1286,6 +1296,40 @@ struct DiscreteFlowDenoiser : public Denoiser {
}
};
struct H3AVFlowDenoiser : public DiscreteFlowDenoiser {
int64_t video_channels;
float audio_shift;
H3AVFlowDenoiser(float shift, float audio_shift, int64_t video_channels)
: DiscreteFlowDenoiser(shift),
video_channels(video_channels),
audio_shift(audio_shift) {
GGML_ASSERT(shift > 0.f && audio_shift > 0.f && video_channels > 0);
}
sd::Tensor<float> process_latent_in(const sd::Tensor<float>& latent) override {
return scale_audio(latent, shift / audio_shift);
}
sd::Tensor<float> process_latent_out(sd::Tensor<float> latent) override {
auto transformed = scale_audio(latent, audio_shift / shift);
if (transformed.empty()) {
return latent;
}
return transformed;
}
private:
sd::Tensor<float> scale_audio(const sd::Tensor<float>& latent, float scale) const {
if (scale == 1.f || latent.dim() < 4 || latent.shape()[3] <= video_channels) {
return {};
}
auto video = sd::ops::slice(latent, 3, 0, video_channels);
auto audio = sd::ops::slice(latent, 3, video_channels, latent.shape()[3]) * scale;
return sd::ops::concat(video, audio, 3);
}
};
struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
FluxFlowDenoiser() = default;
otherarch/sdcpp/src/runtime/preview_interval.h
+45
View File
@@ -0,0 +1,45 @@
#ifndef __SD_RUNTIME_PREVIEW_INTERVAL_H__
#define __SD_RUNTIME_PREVIEW_INTERVAL_H__
#include <cstddef>
#include <cstdint>
#include <limits>
namespace sd::preview {
constexpr std::uint64_t logical_sample_step(int step) {
return step < 0 ? static_cast<std::uint64_t>(-static_cast<std::int64_t>(step))
: static_cast<std::uint64_t>(step);
}
constexpr bool sample_step_is_complete(int step,
std::size_t total_steps,
bool terminal_sigma_is_zero) {
return step > 0 ||
(terminal_sigma_is_zero &&
step < 0 &&
logical_sample_step(step) == static_cast<std::uint64_t>(total_steps));
}
constexpr bool should_preview_sample_step(int step,
std::size_t total_steps,
bool terminal_sigma_is_zero,
int interval,
bool preview_final_step) {
if (interval > 0) {
return step % interval == 0;
}
if (!sample_step_is_complete(step, total_steps, terminal_sigma_is_zero)) {
return false;
}
std::uint64_t logical_step = logical_sample_step(step);
if (interval < 0) {
std::uint64_t requested_step = static_cast<std::uint64_t>(-static_cast<std::int64_t>(interval));
return logical_step == requested_step;
}
return preview_final_step && logical_step == static_cast<std::uint64_t>(total_steps);
}
} // namespace sd::preview
#endif // __SD_RUNTIME_PREVIEW_INTERVAL_H__
otherarch/sdcpp/src/stable-diffusion.cpp
+60 -34
View File
@@ -63,6 +63,7 @@
#include "model/vae/wan_vae.hpp"
#include "runtime/denoiser.hpp"
#include "runtime/guidance.h"
#include "runtime/preview_interval.h"
#include "runtime/sample-cache.h"
#include "upscaler.h"
@@ -2102,6 +2103,9 @@ public:
if (sd_version_is_ltxav(version)) {
LOG_INFO("running in LTXAV FLOW mode");
denoiser = std::make_shared<FluxFlowDenoiser>();
} else if (sd_version_is_minimax_h3(version)) {
LOG_INFO("running in MiniMax H3 AV FLOW mode");
denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, 3.f, get_latent_channel());
} else {
LOG_INFO("running in FLOW mode");
denoiser = std::make_shared<DiscreteFlowDenoiser>();
@@ -2635,11 +2639,9 @@ public:
sd::Tensor<float> vae_latents;
sd::Tensor<float> decoded;
if (preview_vae) {
preview_vae->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
vae_latents = preview_vae->diffusion_to_vae_latents(_latents);
decoded = preview_vae->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true);
} else {
first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
vae_latents = first_stage_model->diffusion_to_vae_latents(_latents);
decoded = first_stage_model->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true);
}
@@ -2728,8 +2730,11 @@ public:
sd_get_preview_mode()};
}
void report_sample_progress(int step, size_t total_steps, int64_t* last_progress_us) {
if (step > 0 || step == -(int)total_steps) {
void report_sample_progress(int step,
size_t total_steps,
bool terminal_sigma_is_zero,
int64_t* last_progress_us) {
if (sd::preview::sample_step_is_complete(step, total_steps, terminal_sigma_is_zero)) {
int64_t now = ggml_time_us();
int showstep = std::abs(step);
float step_seconds = last_progress_us != nullptr && *last_progress_us > 0
@@ -2791,6 +2796,7 @@ public:
int audio_length,
float frame_rate,
const sd_cache_params_t* cache_params,
bool preview_final_step,
const sd::Tensor<float>& video_positions = {}) {
struct RunnerDoneOnExit {
GGMLRunner* runner = nullptr;
@@ -2850,8 +2856,9 @@ public:
}
}
size_t steps = sigmas.size() - 1;
bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
size_t steps = sigmas.size() - 1;
bool terminal_sigma_is_zero = sigmas.back() == 0.f;
bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
if (has_skiplayer && !sd_version_is_dit(version)) {
has_skiplayer = false;
LOG_WARN("SLG is incompatible with this model type");
@@ -2878,10 +2885,14 @@ public:
int64_t last_progress_us = ggml_time_us();
SamplePreviewContext preview = prepare_sample_preview_context();
sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, init_latent)
: init_latent;
sd::Tensor<float> denoised = x_t;
sd::Tensor<float> processed_init_latent = denoiser->process_latent_in(init_latent);
const sd::Tensor<float>& sampling_init_latent = processed_init_latent.empty()
? init_latent
: processed_init_latent;
sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, sampling_init_latent)
: sampling_init_latent;
sd::Tensor<float> denoised = x_t;
auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput {
if (get_cancel_flag() == SD_CANCEL_ALL) {
@@ -2900,14 +2911,21 @@ public:
float c_out = scaling[1];
float c_in = scaling[2];
bool preview_needed = preview.callback != nullptr &&
sd::preview::should_preview_sample_step(step,
steps,
terminal_sigma_is_zero,
sd_get_preview_interval(),
preview_final_step);
std::vector<float> base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep);
std::vector<float> timesteps_vec = base_timesteps_vec;
sd::Tensor<float> audio_timesteps_tensor;
if (sd_version_is_ltxav(version) && !denoise_mask.empty()) {
timesteps_vec = process_ltxav_video_timesteps(base_timesteps_vec, init_latent, denoise_mask);
timesteps_vec = process_ltxav_video_timesteps(base_timesteps_vec, sampling_init_latent, denoise_mask);
audio_timesteps_tensor = sd::Tensor<float>({static_cast<int64_t>(base_timesteps_vec.size())}, base_timesteps_vec);
} else {
timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask, step);
timesteps_vec = process_timesteps(timesteps_vec, sampling_init_latent, denoise_mask, step);
}
const std::vector<float>& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty())
? base_timesteps_vec
@@ -2922,24 +2940,24 @@ public:
}
sd::Tensor<float> noised_input = x * c_in;
if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask);
noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
}
if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
cache_runtime.spectrum.predict(&denoised);
if (!denoise_mask.empty()) {
denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
}
if (sd_should_preview_denoised() && preview.callback != nullptr) {
if (preview_needed && sd_should_preview_denoised()) {
preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
}
report_sample_progress(step, steps, &last_progress_us);
report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us);
sd::guidance::GuiderOutput output;
output.pred = denoised;
return output;
}
if (sd_should_preview_noisy() && preview.callback != nullptr) {
if (preview_needed && sd_should_preview_noisy()) {
preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true);
}
@@ -3147,12 +3165,12 @@ public:
cache_runtime.spectrum.update(denoised);
}
if (!denoise_mask.empty()) {
denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
}
if (sd_should_preview_denoised() && preview.callback != nullptr) {
if (preview_needed && sd_should_preview_denoised()) {
preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
}
report_sample_progress(step, steps, &last_progress_us);
report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us);
output.pred = denoised;
return output;
};
@@ -3175,6 +3193,7 @@ public:
if (inverse_noise_scaling) {
x0 = denoiser->inverse_noise_scaling(sigmas[sigmas.size() - 1], x0);
}
x0 = denoiser->process_latent_out(std::move(x0));
if (control_net) {
control_net->free_control_ctx();
@@ -3324,16 +3343,14 @@ public:
if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) {
return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
}
auto latents = first_stage_model->diffusion_to_vae_latents(x);
first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
if (decoded.empty() && auto_fit_enabled) {
bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast<LTXVideoVAE>(first_stage_model) != nullptr;
if (sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) {
first_stage_model->free_compute_buffer();
first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
}
auto latents = first_stage_model->diffusion_to_vae_latents(x);
auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode();
while (decoded.empty() &&
auto_fit_enabled &&
sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) {
first_stage_model->free_compute_buffer();
decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
}
return decoded;
}
@@ -5071,7 +5088,11 @@ static sd::Tensor<float> prepare_minimax_h3_reference_waveform(const sd_audio_t&
static_cast<long double>(audio.sample_count) * target_sample_rate / audio.sample_rate));
output_samples = std::max<uint64_t>(1, output_samples);
uint64_t padded_samples = (output_samples + 799) / 800 * 800;
sd::Tensor<float> waveform({static_cast<int64_t>(padded_samples), 2, 1, 1});
// Keep stereo streams planar for the mono-per-stream audio encoder:
// [samples, 1, stereo, batch]. This avoids flattening interleaved L/R
// storage into alternating samples when the encoder folds streams into
// its batch dimension.
sd::Tensor<float> waveform({static_cast<int64_t>(padded_samples), 1, 2, 1});
for (uint64_t i = 0; i < output_samples; ++i) {
long double source_pos = static_cast<long double>(i) * audio.sample_rate / target_sample_rate;
@@ -5082,7 +5103,7 @@ static sd::Tensor<float> prepare_minimax_h3_reference_waveform(const sd_audio_t&
uint32_t source_channel = audio.channels == 1 ? 0 : std::min<uint32_t>(channel, audio.channels - 1);
float a = audio.data[source0 * audio.channels + source_channel];
float b = audio.data[source1 * audio.channels + source_channel];
waveform.index(static_cast<int64_t>(i), channel, 0, 0) =
waveform.index(static_cast<int64_t>(i), 0, channel, 0) =
std::clamp(a + (b - a) * fraction, -1.f, 1.f);
}
}
@@ -6039,7 +6060,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
1.f,
0,
static_cast<float>(request.fps),
request.cache_params);
request.cache_params,
true);
int64_t sampling_end = ggml_time_ms();
if (!x_0.empty()) {
LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000);
@@ -6160,7 +6182,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
1.f,
0,
static_cast<float>(request.fps),
request.cache_params);
request.cache_params,
false);
int64_t hires_sample_end = ggml_time_ms();
if (!x_0.empty()) {
LOG_INFO("hires sampling %d/%d completed, taking %.2fs",
@@ -7349,6 +7372,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
latents.audio_length,
static_cast<float>(request.fps),
request.cache_params,
true,
latents.video_positions);
int64_t sampling_end = ggml_time_ms();
if (x_t_sampled.empty()) {
@@ -7391,6 +7415,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
latents.audio_length,
static_cast<float>(request.fps),
request.cache_params,
plan.high_noise_sample_steps <= 0,
latents.video_positions);
int64_t sampling_end = ggml_time_ms();
@@ -7540,6 +7565,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
latents.audio_length,
static_cast<float>(hires_request.fps),
hires_request.cache_params,
false,
hires_video_positions);
sampling_end = ggml_time_ms();
if (final_latent.empty()) {