From b83cdc6f42f54b32086842936a013e6690bfebad Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Thu, 3 Sep 2026 05:32:59 -0300 Subject: [PATCH] sd: cherry-pick changes up to master-841-6b3edaa (#2426) 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 --- otherarch/sdcpp/examples/cli/main.cpp | 64 ++++-- otherarch/sdcpp/examples/common/common.cpp | 4 +- otherarch/sdcpp/examples/common/media_io.cpp | 3 + otherarch/sdcpp/include/stable-diffusion.h | 3 + .../sdcpp/src/conditioning/conditioner.hpp | 27 ++- otherarch/sdcpp/src/core/backend_fit.cpp | 17 +- otherarch/sdcpp/src/model/common/block.hpp | 7 +- otherarch/sdcpp/src/model/diffusion/ltxv.hpp | 71 +++++- .../sdcpp/src/model/diffusion/minimax_h3.hpp | 66 +++--- otherarch/sdcpp/src/model/te/llm.hpp | 177 ++++++++++++--- otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp | 9 + otherarch/sdcpp/src/model/vae/ltx_vae.hpp | 117 +++------- .../src/model/vae/minimax_h3_audio_vae.hpp | 45 ++-- .../sdcpp/src/model/vae/minimax_h3_vae.hpp | 62 ++--- otherarch/sdcpp/src/model/vae/tae.hpp | 15 ++ otherarch/sdcpp/src/model/vae/vae.hpp | 110 ++++++++- otherarch/sdcpp/src/model/vae/vae_tiling.hpp | 213 ++++++++++++++++++ otherarch/sdcpp/src/model/vae/wan_vae.hpp | 153 +++++++------ otherarch/sdcpp/src/name_conversion.cpp | 2 + otherarch/sdcpp/src/runtime/denoiser.hpp | 44 ++++ .../sdcpp/src/runtime/preview_interval.h | 45 ++++ otherarch/sdcpp/src/stable-diffusion.cpp | 94 +++++--- 22 files changed, 1010 insertions(+), 338 deletions(-) create mode 100644 otherarch/sdcpp/src/model/vae/vae_tiling.hpp create mode 100644 otherarch/sdcpp/src/runtime/preview_interval.h diff --git a/otherarch/sdcpp/examples/cli/main.cpp b/otherarch/sdcpp/examples/cli/main.cpp index 1cc7a7af4..d77356b12 100644 --- a/otherarch/sdcpp/examples/cli/main.cpp +++ b/otherarch/sdcpp/examples/cli/main.cpp @@ -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; diff --git a/otherarch/sdcpp/examples/common/common.cpp b/otherarch/sdcpp/examples/common/common.cpp index 35812157c..d46be91a4 100644 --- a/otherarch/sdcpp/examples/common/common.cpp +++ b/otherarch/sdcpp/examples/common/common.cpp @@ -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}, {"", diff --git a/otherarch/sdcpp/examples/common/media_io.cpp b/otherarch/sdcpp/examples/common/media_io.cpp index aadec6f0f..4fb26e980 100644 --- a/otherarch/sdcpp/examples/common/media_io.cpp +++ b/otherarch/sdcpp/examples/common/media_io.cpp @@ -849,6 +849,9 @@ std::vector create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images const uint32_t audio_byte_rate = has_audio ? static_cast(audio->sample_rate * audio_block_align) : 0; const uint32_t audio_data_size = has_audio ? static_cast(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 avi_data; avi_data.reserve(static_cast(num_images) * 1024); diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 857680ca4..b9fc4205e 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -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(); diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index 3d7ff0397..8968676b3 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -2978,15 +2978,36 @@ struct LTXAVEmbedder : public Conditioner { std::shared_ptr tokenizer; std::shared_ptr llm; std::shared_ptr 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 weight_manager = nullptr) { + std::shared_ptr 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(); - llm = std::make_shared(LLM::LLMArch::GEMMA3_12B, + llm = std::make_shared(arch, backend, tensor_storage_map, llm_prefix, @@ -3001,7 +3022,7 @@ struct LTXAVEmbedder : public Conditioner { void get_param_tensors(std::map& 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& tensor_ops) override { diff --git a/otherarch/sdcpp/src/core/backend_fit.cpp b/otherarch/sdcpp/src/core/backend_fit.cpp index 460194740..0ba2df14f 100644 --- a/otherarch/sdcpp/src/core/backend_fit.cpp +++ b/otherarch/sdcpp/src/core/backend_fit.cpp @@ -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; } diff --git a/otherarch/sdcpp/src/model/common/block.hpp b/otherarch/sdcpp/src/model/common/block.hpp index be76dd713..6eb387d9f 100644 --- a/otherarch/sdcpp/src/model/common/block.hpp +++ b/otherarch/sdcpp/src/model/common/block.hpp @@ -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(new GELU(dim, inner_dim)); + blocks["net.0"] = std::shared_ptr(new GELU(dim, inner_dim, bias)); } else { blocks["net.0"] = std::shared_ptr(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(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale)); + blocks["net.2"] = std::shared_ptr(new Linear(inner_dim, dim_out, bias, false, force_prec_f32, scale)); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { diff --git a/otherarch/sdcpp/src/model/diffusion/ltxv.hpp b/otherarch/sdcpp/src/model/diffusion/ltxv.hpp index 39c633747..b75c9f7fa 100644 --- a/otherarch/sdcpp/src/model/diffusion/ltxv.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ltxv.hpp @@ -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 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(a_dim, a_context_dim, a_heads, ad_head, apply_gated_attention, false); blocks["audio_to_video_attn"] = std::make_shared(v_dim, a_dim, a_heads, ad_head, apply_gated_attention, false); blocks["video_to_audio_attn"] = std::make_shared(a_dim, v_dim, a_heads, ad_head, apply_gated_attention, false); - blocks["ff"] = std::make_shared(v_dim, v_dim, 4, FeedForward::Activation::GELU); - blocks["audio_ff"] = std::make_shared(a_dim, a_dim, 4, FeedForward::Activation::GELU); + blocks["ff"] = std::make_shared(v_dim, v_dim, 4, FeedForward::Activation::GELU, false, ff_bias); + blocks["audio_ff"] = std::make_shared(a_dim, a_dim, 4, FeedForward::Activation::GELU, false, audio_ff_bias); } std::vector 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(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 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); diff --git a/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp b/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp index d0683166d..a0a6c7aeb 100644 --- a/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp +++ b/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp @@ -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(input_dim, hidden_dim, true, true); @@ -594,8 +587,7 @@ namespace MiniMaxH3 { const std::vector& segments, const std::vector& sequence_segments, const TokenModulationSpan& video_segment, - const TokenModulationSpan& audio_segment, - float audio_slope) { + const TokenModulationSpan& audio_segment) { auto video_proj = std::dynamic_pointer_cast(blocks["video_patch_proj"]); auto audio_proj = std::dynamic_pointer_cast(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 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( {3, static_cast(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; } diff --git a/otherarch/sdcpp/src/model/te/llm.hpp b/otherarch/sdcpp/src/model/te/llm.hpp index f4dfa9f76..f1a057d7e 100644 --- a/otherarch/sdcpp/src/model/te/llm.hpp +++ b/otherarch/sdcpp/src/model/te/llm.hpp @@ -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 rope_thetas; std::vector 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(config.hidden_size, num_heads * head_dim, config.qkv_bias); blocks["k_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); - blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + if (!k_eq_v) { + blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + } blocks["o_proj"] = std::make_shared(num_heads * head_dim, config.hidden_size, config.attention_out_bias); if (config.qk_norm) { blocks["q_norm"] = std::make_shared(head_dim, config.rms_norm_eps, config.rms_norm_add); blocks["k_norm"] = std::make_shared(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(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(blocks["q_proj"]); auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); - auto v_proj = std::dynamic_pointer_cast(blocks["v_proj"]); + auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast(blocks["v_proj"]); auto out_proj = std::dynamic_pointer_cast(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(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(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(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(config); + if (!config.sliding_attention.empty()) { + sliding_attention = config.sliding_attention[layer_index % config.sliding_attention.size()]; + } + blocks["self_attn"] = std::make_shared(config, sliding_attention == 0); if (config.arch == LLMArch::GPT_OSS_20B) { blocks["mlp"] = std::make_shared(config); } else { @@ -1301,9 +1422,6 @@ namespace LLM { if (!post_ffw_norm_name.empty()) { blocks[post_ffw_norm_name] = std::make_shared(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); diff --git a/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp b/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp index b16126181..938f5ffe8 100644 --- a/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp @@ -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& tensors) override { if (!decode_only) { encoder.get_param_tensors(tensors, weight_prefix + ".encoder"); diff --git a/otherarch/sdcpp/src/model/vae/ltx_vae.hpp b/otherarch/sdcpp/src/model/vae/ltx_vae.hpp index a19ce820d..5629b939b 100644 --- a/otherarch/sdcpp/src/model/vae/ltx_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/ltx_vae.hpp @@ -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& 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(total_frames - 1); - } - - plan.stride = std::max(1, plan.frames - plan.overlap); - int64_t tiled_frames = std::max(1, total_frames - plan.overlap); - plan.num_tiles = total_frames > 0 ? static_cast((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 decode_temporal_tiled_streaming(const int n_threads, const sd::Tensor& 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 output; - for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) { - const int64_t end = std::min(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& 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(start), - chunk_overlap); + static_cast(tile.start), + tile.overlap); }; - auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute(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(get_graph, n_threads, true, true, true), + expected_dim); + }); free_cache_ctx_and_buffer(); cache_tensor_map.clear(); return output; } + sd::Tensor _compute_temporal_tiled(const int n_threads, + const sd::Tensor& input, + VAETemporalDirection direction, + const VAETemporalTilingConfig& config) override { + GGML_ASSERT(direction == VAETemporalDirection::DECODE); + return decode_temporal_tiled_streaming(n_threads, + input, + static_cast(input.dim()), + config); + } + ggml_cgraph* build_latent_statistics_graph(const sd::Tensor& 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); }; diff --git a/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp b/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp index 5af751e36..21f6ce167 100644 --- a/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp @@ -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(blocks["encoder"]); auto pre = std::dynamic_pointer_cast(blocks["pre_block"]); auto mean_proj = std::dynamic_pointer_cast(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) { diff --git a/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp b/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp index 163a9447a..a8aa6b5f1 100644 --- a/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp @@ -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 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& 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 result; sd::Tensor 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(tokens_per_chunk + token_overlap), static_cast(token_overlap)}); + GGML_ASSERT(plan.tiles.size() == static_cast(num_chunks)); + auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& 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(); } int64_t first_end = std::min(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; diff --git a/otherarch/sdcpp/src/model/vae/tae.hpp b/otherarch/sdcpp/src/model/vae/tae.hpp index a03e5d21a..d291bb785 100644 --- a/otherarch/sdcpp/src/model/vae/tae.hpp +++ b/otherarch/sdcpp/src/model/vae/tae.hpp @@ -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& tensors) override { taehv.get_param_tensors(tensors, weight_prefix); } diff --git a/otherarch/sdcpp/src/model/vae/vae.hpp b/otherarch/sdcpp/src/model/vae/vae.hpp index 62a6dd066..f3adb0ccc 100644 --- a/otherarch/sdcpp/src/model/vae/vae.hpp +++ b/otherarch/sdcpp/src/model/vae/vae.hpp @@ -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& 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 _compute_temporal_tiled(const int n_threads, + const sd::Tensor& 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& 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 compute_with_temporal_tiling(const int n_threads, + const sd::Tensor& 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* 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& 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(); @@ -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 input = x; sd::Tensor 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 vae_to_diffusion_latents(const sd::Tensor& latents) = 0; virtual void get_param_tensors(std::map& 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 { diff --git a/otherarch/sdcpp/src/model/vae/vae_tiling.hpp b/otherarch/sdcpp/src/model/vae/vae_tiling.hpp new file mode 100644 index 000000000..50cf08355 --- /dev/null +++ b/otherarch/sdcpp/src/model/vae/vae_tiling.hpp @@ -0,0 +1,213 @@ +#ifndef __SD_MODEL_VAE_VAE_TILING_HPP__ +#define __SD_MODEL_VAE_VAE_TILING_HPP__ + +#include +#include +#include +#include + +#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 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(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(plan.tiles.size()); + tile.start = start; + tile.end = std::min(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 +inline sd::Tensor process_vae_temporal_tiles(const sd::Tensor& input, + const VAETemporalTilePlan& plan, + Fn&& on_processing) { + sd::Tensor 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 +inline sd::Tensor process_vae_temporal_tiles_blended(const sd::Tensor& 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 weights(static_cast(output_frames), 0.f); + sd::Tensor 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::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(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(frame + 1) / + static_cast(overlap_frames + 1)); + } + if (!tile.last && overlap_frames > 0 && frame >= tile_frames - overlap_frames) { + weight *= smootherstep(static_cast(tile_frames - frame) / + static_cast(overlap_frames + 1)); + } + + const int64_t output_frame = output_start + frame; + GGML_ASSERT(output_frame >= 0 && output_frame < output_frames); + weights[static_cast(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(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__ diff --git a/otherarch/sdcpp/src/model/vae/wan_vae.hpp b/otherarch/sdcpp/src/model/vae/wan_vae.hpp index 949562096..422a57823 100644 --- a/otherarch/sdcpp/src/model/vae/wan_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/wan_vae.hpp @@ -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(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(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(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(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& 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& z_tensor, bool decode_graph, int i) { - ggml_cgraph* gf = new_graph_custom(20480); + ggml_cgraph* build_temporal_tile_graph(const sd::Tensor& z_tensor, int chunk_idx) { + ggml_cgraph* gf = new_graph_custom(std::max(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 _compute_temporal_tiled(const int n_threads, + const sd::Tensor& 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& 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(tile.start)); + }; + return restore_trailing_singleton_dims( + GGMLRunner::compute(get_graph, n_threads, true, true, true), + static_cast(input.dim())); + }); + + free_cache_ctx_and_buffer(); + cache_tensor_map.clear(); + ae.clear_cache(); + return output; + } + sd::Tensor _compute(const int n_threads, const sd::Tensor& z, bool decode_graph) override { - if (true) { - sd::Tensor 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(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(get_graph, n_threads, true, true, true); - if (!out_opt.has_value()) { - return {}; - } - sd::Tensor out = std::move(*out_opt); - ae.clear_cache(); - if (t == 1) { - return out; - } - - sd::Tensor output = std::move(out); - - for (i = 1; i < t; i++) { - auto chunk_opt = GGMLRunner::compute(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 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(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() { diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index 126a4ddbc..4815dcf61 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -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."}, diff --git a/otherarch/sdcpp/src/runtime/denoiser.hpp b/otherarch/sdcpp/src/runtime/denoiser.hpp index e4d9af020..b6f1843ec 100644 --- a/otherarch/sdcpp/src/runtime/denoiser.hpp +++ b/otherarch/sdcpp/src/runtime/denoiser.hpp @@ -1043,6 +1043,16 @@ struct Denoiser { const sd::Tensor& latent) = 0; virtual float noise_level_to_sigma(float noise_level) = 0; + virtual sd::Tensor process_latent_in(const sd::Tensor& latent) { + // An empty result means the original latent can be used unchanged. + SD_UNUSED(latent); + return {}; + } + + virtual sd::Tensor process_latent_out(sd::Tensor latent) { + return latent; + } + virtual std::vector 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 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 process_latent_in(const sd::Tensor& latent) override { + return scale_audio(latent, shift / audio_shift); + } + + sd::Tensor process_latent_out(sd::Tensor latent) override { + auto transformed = scale_audio(latent, audio_shift / shift); + if (transformed.empty()) { + return latent; + } + return transformed; + } + +private: + sd::Tensor scale_audio(const sd::Tensor& 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; diff --git a/otherarch/sdcpp/src/runtime/preview_interval.h b/otherarch/sdcpp/src/runtime/preview_interval.h new file mode 100644 index 000000000..aab997682 --- /dev/null +++ b/otherarch/sdcpp/src/runtime/preview_interval.h @@ -0,0 +1,45 @@ +#ifndef __SD_RUNTIME_PREVIEW_INTERVAL_H__ +#define __SD_RUNTIME_PREVIEW_INTERVAL_H__ + +#include +#include +#include + +namespace sd::preview { + + constexpr std::uint64_t logical_sample_step(int step) { + return step < 0 ? static_cast(-static_cast(step)) + : static_cast(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(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(-static_cast(interval)); + return logical_step == requested_step; + } + return preview_final_step && logical_step == static_cast(total_steps); + } +} // namespace sd::preview + +#endif // __SD_RUNTIME_PREVIEW_INTERVAL_H__ diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index 28701fd54..f5e0c0e0f 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -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(); + } else if (sd_version_is_minimax_h3(version)) { + LOG_INFO("running in MiniMax H3 AV FLOW mode"); + denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); } else { LOG_INFO("running in FLOW mode"); denoiser = std::make_shared(); @@ -2635,11 +2639,9 @@ public: sd::Tensor vae_latents; sd::Tensor 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& 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 x_t = !noise.empty() - ? denoiser->noise_scaling(sigmas[0], noise, init_latent) - : init_latent; - sd::Tensor denoised = x_t; + sd::Tensor processed_init_latent = denoiser->process_latent_in(init_latent); + const sd::Tensor& sampling_init_latent = processed_init_latent.empty() + ? init_latent + : processed_init_latent; + sd::Tensor x_t = !noise.empty() + ? denoiser->noise_scaling(sigmas[0], noise, sampling_init_latent) + : sampling_init_latent; + sd::Tensor denoised = x_t; auto denoise = [&](const sd::Tensor& 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 base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep); std::vector timesteps_vec = base_timesteps_vec; sd::Tensor 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({static_cast(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& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty()) ? base_timesteps_vec @@ -2922,24 +2940,24 @@ public: } sd::Tensor 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(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 prepare_minimax_h3_reference_waveform(const sd_audio_t& static_cast(audio.sample_count) * target_sample_rate / audio.sample_rate)); output_samples = std::max(1, output_samples); uint64_t padded_samples = (output_samples + 799) / 800 * 800; - sd::Tensor waveform({static_cast(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 waveform({static_cast(padded_samples), 1, 2, 1}); for (uint64_t i = 0; i < output_samples; ++i) { long double source_pos = static_cast(i) * audio.sample_rate / target_sample_rate; @@ -5082,7 +5103,7 @@ static sd::Tensor prepare_minimax_h3_reference_waveform(const sd_audio_t& uint32_t source_channel = audio.channels == 1 ? 0 : std::min(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(i), channel, 0, 0) = + waveform.index(static_cast(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(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(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(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(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(hires_request.fps), hires_request.cache_params, + false, hires_video_positions); sampling_end = ggml_time_ms(); if (final_latent.empty()) {