From 04a8b9d6357989df682f0c4fd81ffb602cfe5e06 Mon Sep 17 00:00:00 2001 From: Komorebi623 Date: Thu, 6 Aug 2026 16:01:52 +0800 Subject: [PATCH 1/4] feat: add MiniMax-H3 audio video generation --- examples/cli/cli_common.hpp | 45 +- examples/cli/main.cpp | 102 +- examples/cli/sample_main.cpp | 10 +- include/edge-dit.h | 5 + src/core/runtime/model_loader.cpp | 19 + src/core/runtime/model_loader.h | 6 + .../autoencoders/minimax_h3_audio_vae.hpp | 452 +++++++ .../autoencoders/minimax_h3_vae.hpp | 802 +++++++++++ .../components/autoencoders/vae.hpp | 2 +- .../components/common/common_dit.hpp | 64 + .../components/text_encoders/llm.hpp | 31 +- src/dit_models/diffusion_model.hpp | 19 + src/dit_models/models/minimax_h3_full.hpp | 1177 +++++++++++++++++ src/dit_models/pipelines/dit_pipeline.cpp | 4 + .../pipelines/minimax_h3_pipeline.cpp | 405 ++++++ .../pipelines/minimax_h3_pipeline.hpp | 79 ++ src/edge_dit.cpp | 5 + src/utils/name_conversion.cpp | 52 + 18 files changed, 3252 insertions(+), 27 deletions(-) create mode 100644 src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp create mode 100644 src/dit_models/components/autoencoders/minimax_h3_vae.hpp create mode 100644 src/dit_models/models/minimax_h3_full.hpp create mode 100644 src/dit_models/pipelines/minimax_h3_pipeline.cpp create mode 100644 src/dit_models/pipelines/minimax_h3_pipeline.hpp diff --git a/examples/cli/cli_common.hpp b/examples/cli/cli_common.hpp index 4d7443ad..301b3c52 100644 --- a/examples/cli/cli_common.hpp +++ b/examples/cli/cli_common.hpp @@ -454,9 +454,12 @@ struct FluxCliArgs { const char* model_path = nullptr; const char* diffusion_model_path = nullptr; const char* vae_path = nullptr; + const char* audio_vae_path = nullptr; const char* clip_l_path = nullptr; const char* clip_g_path = nullptr; const char* t5xxl_path = nullptr; + const char* llm_path = nullptr; + const char* llm_vision_path = nullptr; const char* prompt = nullptr; const char* negative_prompt = nullptr; const char* image_path = nullptr; @@ -550,6 +553,12 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { if (std::strcmp(key, "--video") == 0) { args->video = true; + } else if (std::strcmp(key, "-M") == 0 || std::strcmp(key, "--mode") == 0) { + const char* v = require_value(key); + if (!v) return false; + if (std::strcmp(v, "vid_gen") == 0) { + args->video = true; + } } else if (std::strcmp(key, "--video-format") == 0) { args->video_format = require_value(key); } else if (std::strcmp(key, "--model") == 0 || std::strcmp(key, "--model_path") == 0) { @@ -558,12 +567,18 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { args->diffusion_model_path = require_value(key); } else if (std::strcmp(key, "--vae") == 0) { args->vae_path = require_value(key); + } else if (std::strcmp(key, "--audio-vae") == 0 || std::strcmp(key, "--audio_vae") == 0) { + args->audio_vae_path = require_value(key); } else if (std::strcmp(key, "--clip_l") == 0) { args->clip_l_path = require_value(key); } else if (std::strcmp(key, "--clip_g") == 0) { args->clip_g_path = require_value(key); } else if (std::strcmp(key, "--t5xxl") == 0) { args->t5xxl_path = require_value(key); + } else if (std::strcmp(key, "--llm") == 0) { + args->llm_path = require_value(key); + } else if (std::strcmp(key, "--llm_vision") == 0 || std::strcmp(key, "--llm-vision") == 0) { + args->llm_vision_path = require_value(key); } else if (std::strcmp(key, "--prompt") == 0 || std::strcmp(key, "-p") == 0) { args->prompt = require_value(key); } else if (std::strcmp(key, "--prompt_file") == 0 || std::strcmp(key, "--prompt-file") == 0) { @@ -574,9 +589,17 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { std::strcmp(key, "--negative_prompt") == 0) { args->negative_prompt = require_value(key); if (!args->negative_prompt) return false; - } else if (std::strcmp(key, "--image") == 0 || std::strcmp(key, "-i") == 0) { + } else if (std::strcmp(key, "--image") == 0 || std::strcmp(key, "--init-img") == 0 || + std::strcmp(key, "--init_img") == 0 || std::strcmp(key, "-i") == 0) { args->image_path = require_value(key); if (!args->image_path) return false; + } else if (std::strcmp(key, "--end-img") == 0 || std::strcmp(key, "--end_img") == 0 || + std::strcmp(key, "--ref-image") == 0 || std::strcmp(key, "--ref_image") == 0 || + std::strcmp(key, "--ref-video") == 0 || std::strcmp(key, "--ref_video") == 0 || + std::strcmp(key, "--ref-audio") == 0 || std::strcmp(key, "--ref_audio") == 0 || + std::strcmp(key, "-r") == 0) { + const char* ignored_path = require_value(key); + if (!ignored_path) return false; } else if (std::strcmp(key, "--output") == 0 || std::strcmp(key, "-o") == 0) { args->output_path = require_value(key); } else if (std::strcmp(key, "--width") == 0 || std::strcmp(key, "-W") == 0) { @@ -587,7 +610,8 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { const char* v = require_value(key); if (!v) return false; args->height = parse_int_value(v, args->height); - } else if (std::strcmp(key, "--frames") == 0) { + } else if (std::strcmp(key, "--frames") == 0 || std::strcmp(key, "--video-frames") == 0 || + std::strcmp(key, "--video_frames") == 0) { const char* v = require_value(key); if (!v) return false; args->frames = parse_int_value(v, args->frames); @@ -777,11 +801,18 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { if (!v) return false; args->max_vram = parse_float_value(v, 0.0f); } else if (std::strcmp(key, "--flash-attention") == 0 || - std::strcmp(key, "--flash-attn") == 0) { + std::strcmp(key, "--flash-attn") == 0 || + std::strcmp(key, "--diffusion-fa") == 0 || + std::strcmp(key, "--diffusion_fa") == 0) { args->flash_attention = true; } else if (std::strcmp(key, "--no-flash-attention") == 0 || std::strcmp(key, "--no-flash-attn") == 0) { args->flash_attention = false; + } else if (std::strcmp(key, "--rng") == 0) { + const char* ignored_rng = require_value(key); + if (!ignored_rng) return false; + } else if (std::strcmp(key, "-v") == 0 || std::strcmp(key, "--verbose") == 0) { + // Accepted for sd.cpp CLI compatibility; edge-dit logging is controlled externally. } else if (std::strcmp(key, "--cfg-parallel-size") == 0 || std::strcmp(key, "--cfg-size") == 0) { const char* v = require_value(key); @@ -820,9 +851,13 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { args->vae_path != nullptr && std::strlen(args->vae_path) > 0 && args->clip_l_path != nullptr && std::strlen(args->clip_l_path) > 0 && (args->no_t5 || (args->t5xxl_path != nullptr && std::strlen(args->t5xxl_path) > 0)); + const bool has_minimax_h3_components = + args->diffusion_model_path != nullptr && std::strlen(args->diffusion_model_path) > 0 && + args->vae_path != nullptr && std::strlen(args->vae_path) > 0 && + args->llm_path != nullptr && std::strlen(args->llm_path) > 0; - if (!has_full_model && !has_components) { - std::fprintf(stderr, "--model or the full --diffusion-model/--vae/--clip_l/(--t5xxl or --no-t5) set is required\n"); + if (!has_full_model && !has_components && !has_minimax_h3_components) { + std::fprintf(stderr, "--model, --diffusion-model/--vae/--clip_l/(--t5xxl or --no-t5), or --diffusion-model/--vae/--llm is required\n"); return false; } diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 3c9cc547..e1a089de 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -31,20 +31,23 @@ static void print_usage(const char* prog) { "Usage:\n" " %s --model --prompt [options]\n" " %s --diffusion-model --vae --clip_l [--clip_g ] (--t5xxl | --no-t5) --prompt [options]\n" + " %s -M vid_gen --diffusion-model --vae [--audio-vae ] --llm --prompt [options]\n" "Options:\n" - " --video Generate video frames instead of an image\n" + " --video, -M vid_gen Generate video frames instead of an image\n" " --video-format Video format: auto, avi, mp4, mov, mkv, webm. Default: auto\n" " -i, --image Input/reference image for image-edit models\n" " --diffusion-model Standalone DiT transformer weights\n" " --vae Standalone VAE weights\n" + " --audio-vae Standalone audio VAE weights (MiniMax-H3)\n" " --clip_l CLIP-L text encoder weights\n" " --clip_g CLIP-G text encoder weights\n" " --t5xxl T5XXL text encoder weights\n" + " --llm LLM text encoder weights (MiniMax-H3/Qwen)\n" " --negative-prompt Negative prompt text, default: empty\n" " -o, --output Output image/video path, default: output.png\n" " -W, --width Image width, default: 1024\n" " -H, --height Image height, default: 1024\n" - " --frames Video frame count, default: 1\n" + " --frames, --video-frames Video frame count, default: 1\n" " --fps Video fps, default: 16\n" " --steps Sampling steps, default: 20\n" " -s, --seed Seed, default: -1\n" @@ -117,6 +120,7 @@ static void print_usage(const char* prog) { " Print graph-cut profile from every parallel rank\n" " --help Show this help\n", prog, + prog, prog ); } @@ -442,10 +446,9 @@ static bool save_mjpg_avi(const char* path, const ed_video_t& video, int fps, in AviIndexEntry entry{}; std::memcpy(entry.fourcc, "00dc", 4); entry.flags = 0x10; - // idx1 offsets are relative to the start of the 'movi' chunk data (i.e. the - // byte after the "movi" FourCC), not absolute file offsets. movi_size_pos - // points at the LIST size field; +8 skips that size field and the "movi" tag. - entry.offset = static_cast(avi.size() - (movi_size_pos + 8)); + // idx1 offsets are relative to the "movi" FourCC. movi_size_pos points at + // the LIST size field, so the first media chunk starts four bytes later. + entry.offset = static_cast(avi.size() - (movi_size_pos + 4)); entry.size = static_cast(jpg.size()); write_fourcc(avi, "00dc"); @@ -556,6 +559,81 @@ static bool save_video(const char* path, const ed_video_t& video, int fps) { return false; } +static bool save_wav(const char* path, const ed_video_t& video) { + if (path == nullptr || video.audio == nullptr || video.audio_sample_count <= 0 || + video.audio_channels <= 0 || video.audio_sample_rate <= 0) { + return false; + } + const uint32_t channels = static_cast(video.audio_channels); + const uint32_t sample_rate = static_cast(video.audio_sample_rate); + const uint32_t sample_count = static_cast(video.audio_sample_count); + const uint32_t data_size = sample_count * channels * sizeof(int16_t); + FILE* file = std::fopen(path, "wb"); + if (file == nullptr) { + return false; + } + auto put_u16 = [&](uint16_t value) { return std::fwrite(&value, sizeof(value), 1, file) == 1; }; + auto put_u32 = [&](uint32_t value) { return std::fwrite(&value, sizeof(value), 1, file) == 1; }; + const bool header_ok = std::fwrite("RIFF", 1, 4, file) == 4 && put_u32(36 + data_size) && + std::fwrite("WAVEfmt ", 1, 8, file) == 8 && put_u32(16) && put_u16(1) && + put_u16(static_cast(channels)) && put_u32(sample_rate) && + put_u32(sample_rate * channels * sizeof(int16_t)) && + put_u16(static_cast(channels * sizeof(int16_t))) && put_u16(16) && + std::fwrite("data", 1, 4, file) == 4 && put_u32(data_size); + bool samples_ok = header_ok; + for (uint32_t sample = 0; samples_ok && sample < sample_count; ++sample) { + for (uint32_t channel = 0; channel < channels; ++channel) { + const float value = std::clamp(video.audio[sample * channels + channel], -1.0f, 1.0f); + const int16_t pcm = static_cast(std::lround(value * 32767.0f)); + samples_ok = std::fwrite(&pcm, sizeof(pcm), 1, file) == 1; + if (!samples_ok) break; + } + } + return std::fclose(file) == 0 && samples_ok; +} + +static bool save_video_with_audio(const char* path, const ed_video_t& video, int fps, bool* audio_muxed) { + if (audio_muxed != nullptr) { + *audio_muxed = false; + } + if (video.audio == nullptr || video.audio_sample_count <= 0) { + return save_video(path, video, fps); + } + const fs::path output(path); + const fs::path video_tmp = output.string() + ".video.tmp.avi"; + const fs::path wav_path = output.string() + ".wav"; + if (!save_video(video_tmp.c_str(), video, fps) || !save_wav(wav_path.c_str(), video)) { + std::error_code error; + fs::remove(video_tmp, error); + return false; + } + const std::string ext = path_extension(path); + const char* audio_codec = ext == ".avi" ? "pcm_s16le" : (ext == ".webm" ? "libopus" : "aac"); + const std::string command = shell_quote(find_ffmpeg_binary().c_str()) + + " -hide_banner -loglevel error -y -i " + shell_quote(video_tmp.c_str()) + + " -i " + shell_quote(wav_path.c_str()) + + " -c:v copy -c:a " + audio_codec + " " + shell_quote(path); + const int status = std::system(command.c_str()); + std::error_code error; + if (status != 0) { + fs::remove(output, error); + fs::rename(video_tmp, output, error); + if (error) { + std::fprintf(stderr, "ffmpeg failed while muxing audio, status=%d; video remains at %s and WAV at %s\n", + status, video_tmp.c_str(), wav_path.c_str()); + return false; + } + std::fprintf(stderr, "ffmpeg failed while muxing audio, status=%d; saved AVI and WAV sidecar at %s\n", + status, wav_path.c_str()); + return true; + } + fs::remove(video_tmp, error); + if (audio_muxed != nullptr) { + *audio_muxed = true; + } + return true; +} + int main(int argc, char** argv) { for (int i = 1; i < argc; ++i) { @@ -607,9 +685,12 @@ int main(int argc, char** argv) { ctx_params.model_path = args.model_path; ctx_params.diffusion_model_path = args.diffusion_model_path; ctx_params.vae_path = args.vae_path; + ctx_params.audio_vae_path = args.audio_vae_path; ctx_params.clip_l_path = args.clip_l_path; ctx_params.clip_g_path = args.clip_g_path; ctx_params.t5xxl_path = args.t5xxl_path; + ctx_params.llm_path = args.llm_path; + ctx_params.llm_vision_path = args.llm_vision_path; ctx_params.cfg_parallel_size = args.cfg_parallel_size; ctx_params.tp_parallel_size = args.tp_parallel_size; ctx_params.sp_parallel_size = args.sp_parallel_size; @@ -719,7 +800,8 @@ int main(int argc, char** argv) { const std::string output_path = video_output_path(args.output_path, args.video_format); auto ed_wall_save0 = ed_wall_clock::now(); - bool ed_save_ok = save_video(output_path.c_str(), output, args.fps); + bool audio_muxed = false; + bool ed_save_ok = save_video_with_audio(output_path.c_str(), output, args.fps, &audio_muxed); ed_wall_save = ed_wall_ms(ed_wall_save0, ed_wall_clock::now()); if (!ed_save_ok) { std::fprintf(stderr, "failed to save output video: %s\n", output_path.c_str()); @@ -728,7 +810,11 @@ int main(int argc, char** argv) { return 5; } - std::printf("saved video to %s\n", output_path.c_str()); + if (output.audio != nullptr && output.audio_sample_count > 0 && !audio_muxed) { + std::printf("saved video to %s with WAV sidecar %s.wav\n", output_path.c_str(), output_path.c_str()); + } else { + std::printf("saved video%s to %s\n", audio_muxed ? " with audio" : "", output_path.c_str()); + } ed_free_video(&output); } else { diff --git a/examples/cli/sample_main.cpp b/examples/cli/sample_main.cpp index 23d80ec2..afd772f0 100644 --- a/examples/cli/sample_main.cpp +++ b/examples/cli/sample_main.cpp @@ -57,7 +57,8 @@ void print_usage(const char* prog) { " %s --model --prompt_file --output_dir [options]\n\n" "Required:\n" " --model Diffusers model directory (or component flags:\n" - " --diffusion-model/--vae/--clip_l/[--clip_g]/(--t5xxl|--no-t5))\n" + " --diffusion-model/--vae/--clip_l/[--clip_g]/(--t5xxl|--no-t5),\n" + " or MiniMax-H3 --diffusion-model/--vae/[--audio-vae]/--llm)\n" " --prompt_file Text file, one prompt per line\n" " --output_dir Output run directory (created)\n\n" "Sampling:\n" @@ -71,8 +72,8 @@ void print_usage(const char* prog) { " --flow_shift Flow scheduler shift. Default model default\n" " --negative_prompt Negative prompt (used when cfg_scale != 1). Default empty\n" " --image Input/reference image for image-editing models\n" - " --video Generate video frames (calls ed_generate_video; writes .avi)\n" - " --frames Video frame count (with --video). Default 1\n" + " --video, -M vid_gen Generate video frames (calls ed_generate_video; writes .avi)\n" + " --frames, --video-frames Video frame count (with --video). Default 1\n" " --fps Video fps (with --video). Default 16\n" " --start_index First prompt index (inclusive). Default 0\n" " --end_index Last prompt index (exclusive). Default all\n" @@ -529,9 +530,12 @@ int main(int argc, char** argv) { ctx_params.model_path = args.model_path; ctx_params.diffusion_model_path = args.diffusion_model_path; ctx_params.vae_path = args.vae_path; + ctx_params.audio_vae_path = args.audio_vae_path; ctx_params.clip_l_path = args.clip_l_path; ctx_params.clip_g_path = args.clip_g_path; ctx_params.t5xxl_path = args.t5xxl_path; + ctx_params.llm_path = args.llm_path; + ctx_params.llm_vision_path = args.llm_vision_path; ctx_params.cfg_parallel_size = args.cfg_parallel_size; ctx_params.tp_parallel_size = args.tp_parallel_size; ctx_params.sp_parallel_size = args.sp_parallel_size; diff --git a/include/edge-dit.h b/include/edge-dit.h index 3101391f..d95b7d47 100644 --- a/include/edge-dit.h +++ b/include/edge-dit.h @@ -122,6 +122,10 @@ typedef struct ed_image_batch_t { typedef struct ed_video_t { ed_image_t * frames; int frame_count; + float * audio; + int audio_sample_count; + int audio_channels; + int audio_sample_rate; } ed_video_t; typedef struct ed_lora_t { @@ -152,6 +156,7 @@ typedef struct ed_context_params_t { const char * llm_path; const char * llm_vision_path; const char * vae_path; + const char * audio_vae_path; const char * taesd_path; const char * control_net_path; diff --git a/src/core/runtime/model_loader.cpp b/src/core/runtime/model_loader.cpp index e8829b86..cf63a0c6 100644 --- a/src/core/runtime/model_loader.cpp +++ b/src/core/runtime/model_loader.cpp @@ -195,6 +195,7 @@ const char* ed_version_name(SDVersion version) { case VERSION_WAN2_2_TI2V: return "wan2.2-ti2v"; case VERSION_QWEN_IMAGE: return "qwen-image"; case VERSION_QWEN_IMAGE_EDIT: return "qwen-image-edit"; + case VERSION_MINIMAX_H3: return "minimax-h3"; case VERSION_ANIMA: return "anima"; case VERSION_FLUX2: return "flux2"; case VERSION_FLUX2_KLEIN: return "flux2-klein"; @@ -831,6 +832,17 @@ bool ModelLoader::load_model_files(const ed_context_params_t& params, loaded_any = true; } + if (non_empty(params.audio_vae_path)) { + if (!load_optional_file(params.audio_vae_path, + "audio_vae.", + "audio vae", + true, + error)) { + return false; + } + loaded_any = true; + } + if (non_empty(params.taesd_path)) { const bool ok = load_optional_file(params.taesd_path, "tae.", @@ -1359,6 +1371,10 @@ SDVersion ModelLoader::get_ld_version() { if (contains(name, "model.diffusion_model.transformer_blocks.0.img_mod.1.weight")) { return VERSION_QWEN_IMAGE; } + if (contains(name, "model.diffusion_model.video_patch_proj.weight") && + tensor_storage_map_.find("model.diffusion_model.audio_patch_proj.weight") != tensor_storage_map_.end()) { + return VERSION_MINIMAX_H3; + } // Wan video DiT: blocks carry a cross_attn sub-module (text conditioning) // that no other supported architecture uses, and a 3-D patch_embedding. // Diffusers loading recognizes Wan via config.json's "Wan" class, but a @@ -1867,6 +1883,9 @@ bool ModelLoader::load_tensors(std::map& tensors, } for (const auto& item : tensors) { + if (starts_with(item.first, "__ed_")) { + continue; + } if (tensor_names_in_file.find(item.first) == tensor_names_in_file.end()) { LOG_ERROR("tensor '%s' not in model file", item.first.c_str()); return false; diff --git a/src/core/runtime/model_loader.h b/src/core/runtime/model_loader.h index c7a2d683..e78cdf60 100644 --- a/src/core/runtime/model_loader.h +++ b/src/core/runtime/model_loader.h @@ -40,6 +40,7 @@ enum SDVersion { VERSION_WAN2_2_TI2V, VERSION_QWEN_IMAGE, VERSION_QWEN_IMAGE_EDIT, + VERSION_MINIMAX_H3, VERSION_ANIMA, VERSION_FLUX2, VERSION_FLUX2_KLEIN, @@ -102,6 +103,10 @@ static inline bool ed_version_is_qwen_image_edit(SDVersion version) { return version == VERSION_QWEN_IMAGE_EDIT; } +static inline bool ed_version_is_minimax_h3(SDVersion version) { + return version == VERSION_MINIMAX_H3; +} + static inline bool ed_version_is_anima(SDVersion version) { return version == VERSION_ANIMA; } @@ -128,6 +133,7 @@ static inline bool ed_version_is_dit(SDVersion version) { return ed_version_is_flux(version) || ed_version_is_flux2(version) || ed_version_is_sd3(version) || ed_version_is_wan(version) || ed_version_is_qwen_image(version) || ed_version_is_qwen_image_edit(version) || + ed_version_is_minimax_h3(version) || ed_version_is_anima(version) || ed_version_is_z_image(version) || ed_version_is_ernie_image(version); } diff --git a/src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp b/src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp new file mode 100644 index 00000000..44d85823 --- /dev/null +++ b/src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp @@ -0,0 +1,452 @@ +#ifndef __ED_MINIMAX_H3_AUDIO_VAE_HPP__ +#define __ED_MINIMAX_H3_AUDIO_VAE_HPP__ +#include +#include +#include +#include +#include +#include "backend/ggml/ggml_extend.hpp" +namespace MiniMaxH3Audio { +namespace Ops { + static ggml_type audio_conv_weight_type(ggml_type type) { + return type == GGML_TYPE_BF16 ? GGML_TYPE_F16 : type; + } + + static ggml_tensor* repeat_with_vulkan_f32_workaround(ggml_backend_t backend, + ggml_context* ctx, + ggml_tensor* x, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + if (x->type != GGML_TYPE_F32 && + (x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && + sd_backend_is(backend, "vulkan")) { + auto x_f32 = ggml_cast(ctx, x, GGML_TYPE_F32); + auto repeated = ggml_repeat_4d(ctx, + x_f32, + ne0, + ne1, + ne2, + ne3); + return ggml_cast(ctx, repeated, x->type); + } + return ggml_repeat_4d(ctx, x, ne0, ne1, ne2, ne3); + } + + static ggml_tensor* repeat_1d_value(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t count) { + auto ctx = runner_ctx->ggml_ctx; + GGML_ASSERT(x->ne[0] == 1); + return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, x, count, x->ne[1], x->ne[2], x->ne[3]); + } + + static ggml_tensor* replicate_pad_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t left, int64_t right) { + auto ctx = runner_ctx->ggml_ctx; + if (left > 0) { + auto first = ggml_ext_slice(ctx, x, 0, 0, 1); + x = ggml_concat(ctx, repeat_1d_value(runner_ctx, first, left), x, 0); + } + if (right > 0) { + auto last = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]); + x = ggml_concat(ctx, x, repeat_1d_value(runner_ctx, last, right), 0); + } + return x; + } + + static ggml_tensor* tile_depthwise_filter_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* filter, int64_t channels) { + auto ctx = runner_ctx->ggml_ctx; + ggml_tensor* base = filter; + if (ggml_n_dims(base) == 3) { + base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1); + } else if (ggml_n_dims(base) == 1) { + base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1); + } + return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, base, base->ne[0], 1, channels, 1); + } + + static ggml_tensor* depthwise_conv1d(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* filter, + int stride, + int padding) { + auto ctx = runner_ctx->ggml_ctx; + GGML_ASSERT(x->ne[3] == 1); + auto tiled = tile_depthwise_filter_1d(runner_ctx, filter, x->ne[1]); + auto out = ggml_conv_1d_dw(ctx, tiled, x, stride, padding, 1); + return ggml_reshape_4d(ctx, out, out->ne[0], out->ne[1], 1, 1); + } + + static ggml_tensor* reverse_1d_filter(ggml_context* ctx, ggml_tensor* filter) { + GGML_ASSERT(ctx != nullptr); + GGML_ASSERT(filter != nullptr); + GGML_ASSERT(filter->ne[1] == 1); + GGML_ASSERT(filter->ne[2] == 1); + GGML_ASSERT(filter->ne[3] == 1); + + ggml_tensor* reversed = nullptr; + for (int64_t k = filter->ne[0] - 1; k >= 0; --k) { + auto slice = ggml_ext_slice(ctx, filter, 0, k, k + 1); + reversed = reversed == nullptr ? slice : ggml_concat(ctx, reversed, slice, 0); + } + return reversed; + } + + static ggml_tensor* depthwise_conv_transpose1d(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* filter, + int stride) { + GGML_ASSERT(x->ne[2] == 1 && x->ne[3] == 1); + GGML_ASSERT(filter->ne[1] == 1); + GGML_ASSERT(filter->ne[2] == 1 && filter->ne[3] == 1); + + const int64_t time = x->ne[0]; + const int64_t channels = x->ne[1]; + const int64_t kernel_size = filter->ne[0]; + const int64_t out_time = (time - 1) * stride + kernel_size; + + auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels); + if (stride > 1) { + auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f); + auto zero_tail = zero_unit; + for (int i = 1; i < stride - 1; ++i) { + zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0); + } + x_flat = ggml_concat(ctx, x_flat, zero_tail, 0); + } + x_flat = ggml_reshape_3d(ctx, x_flat, time * stride, 1, channels); + + auto reversed_filter = reverse_1d_filter(ctx, filter); + auto out = ggml_conv_1d(ctx, reversed_filter, x_flat, 1, static_cast(kernel_size - 1), 1); + if (out->ne[0] > out_time) { + out = ggml_ext_slice(ctx, out, 0, 0, out_time); + } + GGML_ASSERT(out->ne[0] == out_time); + GGML_ASSERT(out->ne[1] == 1); + GGML_ASSERT(out->ne[2] == channels); + + out = ggml_ext_scale(ctx, out, static_cast(stride)); + return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1); + } + struct Conv1D : public UnaryBlock { + int64_t in_channels; + int64_t out_channels; + int kernel_size; + int stride; + int padding; + int dilation; + bool bias; + std::string prefix; + + Conv1D(int64_t in_channels, + int64_t out_channels, + int kernel_size, + int stride = 1, + int padding = 0, + int dilation = 1, + bool bias = true) + : in_channels(in_channels), + out_channels(out_channels), + kernel_size(kernel_size), + stride(stride), + padding(padding), + dilation(dilation), + bias(bias) {} + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + this->prefix = prefix; + ggml_type wtype = audio_conv_weight_type(get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16)); + params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size, in_channels, out_channels, 1); + if (bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + x = ggml_conv_1d(ctx->ggml_ctx, params["weight"], x, stride, padding, dilation); + if (bias) { + auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, params["bias"]->ne[0], 1, 1); + x = ggml_add_inplace(ctx->ggml_ctx, x, b); + } + return x; + } + }; + + struct ConvTranspose1D : public UnaryBlock { + int64_t in_channels; + int64_t out_channels; + int kernel_size; + int stride; + int padding; + int dilation; + bool bias; + + ConvTranspose1D(int64_t in_channels, + int64_t out_channels, + int kernel_size, + int stride, + int padding, + int dilation = 1, + bool bias = true) + : in_channels(in_channels), + out_channels(out_channels), + kernel_size(kernel_size), + stride(stride), + padding(padding), + dilation(dilation), + bias(bias) {} + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ED_UNUSED(tensor_storage_map); + ED_UNUSED(prefix); + params["weight"] = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kernel_size, out_channels, in_channels, 1); + if (bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + GGML_ASSERT(dilation == 1); + x = ggml_conv_transpose_1d(ctx->ggml_ctx, params["weight"], x, stride, 0, dilation); + if (padding > 0) { + x = ggml_ext_slice(ctx->ggml_ctx, x, 0, padding, x->ne[0] - padding); + } + if (bias) { + auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, params["bias"]->ne[0], 1, 1); + x = ggml_add_inplace(ctx->ggml_ctx, x, b); + } + return x; + } + }; + + struct SnakeBeta1D : public UnaryBlock { + int64_t channels; + float eps = 1e-9f; + + explicit SnakeBeta1D(int64_t channels) + : channels(channels) {} + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ED_UNUSED(tensor_storage_map); + ED_UNUSED(prefix); + params["alpha"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + params["beta"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto alpha = ggml_exp(ctx->ggml_ctx, params["alpha"]); + auto beta = ggml_exp(ctx->ggml_ctx, params["beta"]); + alpha = ggml_reshape_4d(ctx->ggml_ctx, alpha, 1, alpha->ne[0], 1, 1); + beta = ggml_reshape_4d(ctx->ggml_ctx, beta, 1, beta->ne[0], 1, 1); + auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha)); + oscillation = ggml_mul(ctx->ggml_ctx, oscillation, oscillation); + auto eps_tensor = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), eps); + oscillation = ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, beta, eps_tensor)); + return ggml_add(ctx->ggml_ctx, x, oscillation); + } + }; + + struct Activation1D : public GGMLBlock { + int64_t channels; + int up_ratio = 2; + int down_ratio = 2; + int up_kernel_size = 12; + int down_kernel_size = 12; + + explicit Activation1D(int64_t channels) + : channels(channels) { + blocks["act"] = std::make_shared(channels); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ggml_type down_type = audio_conv_weight_type(get_type(prefix + "downsample.lowpass.filter", tensor_storage_map, GGML_TYPE_F16)); + params["downsample.lowpass.filter"] = ggml_new_tensor_3d(ctx, down_type, down_kernel_size, 1, 1); + params["upsample.filter"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, up_kernel_size, 1, 1); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto act = std::dynamic_pointer_cast(blocks["act"]); + auto up_filter = params["upsample.filter"]; + auto down_filter = params["downsample.lowpass.filter"]; + + int up_pad = up_kernel_size / up_ratio - 1; + int up_pad_left = up_pad * up_ratio + (up_kernel_size - up_ratio) / 2; + int up_pad_right = up_pad * up_ratio + (up_kernel_size - up_ratio + 1) / 2; + + x = replicate_pad_1d(ctx, x, up_pad, up_pad); + x = depthwise_conv_transpose1d(ctx->ggml_ctx, x, up_filter, up_ratio); + x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right); + + x = act->forward(ctx, x); + + int down_pad_left = down_kernel_size / 2 - (down_kernel_size % 2 == 0 ? 1 : 0); + int down_pad_right = down_kernel_size / 2; + x = replicate_pad_1d(ctx, x, down_pad_left, down_pad_right); + x = depthwise_conv1d(ctx, x, down_filter, down_ratio, 0); + return x; + } + }; +} // namespace Ops + + struct AudioAMPBlock : public GGMLBlock { + int channels; + + AudioAMPBlock(int channels, + int kernel_size, + const std::array& dilations) + : channels(channels) { + for (int i = 0; i < 3; ++i) { + blocks["activations." + std::to_string(i * 2)] = + std::make_shared(channels); + blocks["activations." + std::to_string(i * 2 + 1)] = + std::make_shared(channels); + blocks["convs1." + std::to_string(i)] = + std::make_shared(channels, + channels, + kernel_size, + 1, + (kernel_size * dilations[i] - dilations[i]) / 2, + dilations[i]); + blocks["convs2." + std::to_string(i)] = + std::make_shared(channels, + channels, + kernel_size, + 1, + kernel_size / 2); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + for (int i = 0; i < 3; ++i) { + auto act1 = std::dynamic_pointer_cast( + blocks["activations." + std::to_string(i * 2)]); + auto act2 = std::dynamic_pointer_cast( + blocks["activations." + std::to_string(i * 2 + 1)]); + auto conv1 = std::dynamic_pointer_cast( + blocks["convs1." + std::to_string(i)]); + auto conv2 = std::dynamic_pointer_cast( + blocks["convs2." + std::to_string(i)]); + + auto h = conv1->forward(ctx, act1->forward(ctx, x)); + h = conv2->forward(ctx, act2->forward(ctx, h)); + x = ggml_add(ctx->ggml_ctx, x, h); + } + return x; + } + }; + + struct BigVGAN : public GGMLBlock { + static constexpr int initial_channels = 1024; + static constexpr int num_kernels = 3; + static constexpr int num_upsamples = 7; + static constexpr std::array rates = {5, 5, 2, 2, 2, 2, 2}; + static constexpr std::array kernels = {9, 9, 4, 4, 4, 4, 4}; + static constexpr std::array res_kernels = {3, 7, 11}; + + BigVGAN() { + blocks["conv_pre"] = std::make_shared(2048, + initial_channels, + 7, + 1, + 3); + int channels = initial_channels; + for (int i = 0; i < num_upsamples; ++i) { + int next_channels = initial_channels / (1 << (i + 1)); + blocks["ups." + std::to_string(i) + ".0"] = + std::make_shared(channels, + next_channels, + kernels[i], + rates[i], + (kernels[i] - rates[i]) / 2); + for (int j = 0; j < num_kernels; ++j) { + blocks["resblocks." + std::to_string(i * num_kernels + j)] = + std::make_shared(next_channels, + res_kernels[j], + std::array{1, 3, 5}); + } + channels = next_channels; + } + blocks["activation_post"] = std::make_shared(channels); + blocks["conv_post"] = std::make_shared(channels, + 1, + 7, + 1, + 3, + 1, + false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv_pre = std::dynamic_pointer_cast(blocks["conv_pre"]); + x = conv_pre->forward(ctx, x); + for (int i = 0; i < num_upsamples; ++i) { + auto up = std::dynamic_pointer_cast( + blocks["ups." + std::to_string(i) + ".0"]); + x = up->forward(ctx, x); + + ggml_tensor* sum = nullptr; + for (int j = 0; j < num_kernels; ++j) { + auto block = std::dynamic_pointer_cast( + blocks["resblocks." + std::to_string(i * num_kernels + j)]); + auto value = block->forward(ctx, x); + sum = sum == nullptr ? value : ggml_add(ctx->ggml_ctx, sum, value); + } + x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.f / num_kernels); + } + auto activation = std::dynamic_pointer_cast(blocks["activation_post"]); + auto conv_post = std::dynamic_pointer_cast(blocks["conv_post"]); + return ggml_clamp(ctx->ggml_ctx, + conv_post->forward(ctx, activation->forward(ctx, x)), + -1.f, + 1.f); + } + }; +struct AudioDecoder : public GGMLBlock { + static constexpr int kLatentChannels = 32; + AudioDecoder() { + blocks["dec_in_proj"] = std::make_shared(kLatentChannels, 2048, 1); + blocks["decoder"] = std::make_shared(); + } + void init_params(ggml_context* ctx, const String2TensorStorage& storage = {}, const std::string prefix = "") override { + ED_UNUSED(storage); ED_UNUSED(prefix); + params["latents_mean"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels); + params["latents_std"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels); + } + ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent) { + GGML_ASSERT(latent->ne[1] == 2 && latent->ne[2] == kLatentChannels); + latent = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, latent, 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); + latent = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, latent, std), mean); + auto dec_in = std::dynamic_pointer_cast(blocks["dec_in_proj"]); + auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); + const int64_t streams = latent->ne[2] * latent->ne[3]; + latent = ggml_reshape_3d(ctx->ggml_ctx, latent, latent->ne[0], latent->ne[1], streams); + ggml_tensor* waveform = nullptr; + for (int64_t stream = 0; stream < streams; ++stream) { + auto value = decoder->forward(ctx, dec_in->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1))); + waveform = waveform == nullptr ? value : ggml_concat(ctx->ggml_ctx, waveform, value, 2); + } + return ggml_reshape_4d(ctx->ggml_ctx, waveform, waveform->ne[0], streams, 1, 1); + } +}; +struct AudioVAERunner : public GGMLRunner { + AudioDecoder model; + AudioVAERunner(ggml_backend_t backend, bool offload, const String2TensorStorage& storage, const std::string& prefix = "audio_vae") + : GGMLRunner(backend, offload) { model.init(params_ctx, storage, prefix); } + std::string get_desc() override { return "minimax_h3_audio_vae"; } + void get_param_tensors(std::map& tensors, const std::string& prefix) { model.get_param_tensors(tensors, prefix); } + sd::Tensor decode(int n_threads, const sd::Tensor& latent) { + auto get_graph = [&]() { auto input = make_input(latent); auto runner_ctx = get_context(); auto graph = new_graph_custom(655360); ggml_build_forward_expand(graph, model.decode(&runner_ctx, input)); return graph; }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4); + } +}; +} // namespace MiniMaxH3Audio +#endif diff --git a/src/dit_models/components/autoencoders/minimax_h3_vae.hpp b/src/dit_models/components/autoencoders/minimax_h3_vae.hpp new file mode 100644 index 00000000..0e5368b9 --- /dev/null +++ b/src/dit_models/components/autoencoders/minimax_h3_vae.hpp @@ -0,0 +1,802 @@ +#ifndef __ED_DIT_MODELS_COMPONENTS_AUTOENCODERS_MINIMAX_H3_VAE_HPP__ +#define __ED_DIT_MODELS_COMPONENTS_AUTOENCODERS_MINIMAX_H3_VAE_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "dit_models/components/common/rope.hpp" +#include "dit_models/components/common/common_dit.hpp" +#include "dit_models/components/autoencoders/vae.hpp" + +namespace MiniMaxH3VAE { + + constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144; + + struct CausalConv3d : public Conv3d { + std::tuple temporal_padding; + + CausalConv3d(int64_t in_channels, + int64_t out_channels, + std::tuple kernel_size, + std::tuple stride = {1, 1, 1}, + std::tuple padding = {0, 0, 0}) + : Conv3d(in_channels, + out_channels, + kernel_size, + stride, + {0, 0, 0}), + temporal_padding(padding) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto reflect_pad = [&](ggml_tensor* value, int dim, int amount) { + for (int i = 0; i < amount; ++i) { + GGML_ASSERT(value->ne[dim] > 1); + auto left = ggml_ext_slice(ctx->ggml_ctx, value, dim, 1, 2); + auto right = ggml_ext_slice(ctx->ggml_ctx, + value, + dim, + value->ne[dim] - 2, + value->ne[dim] - 1); + value = ggml_concat(ctx->ggml_ctx, left, value, dim); + value = ggml_concat(ctx->ggml_ctx, value, right, dim); + } + return value; + }; + + x = reflect_pad(x, 0, std::get<2>(temporal_padding)); + x = reflect_pad(x, 1, std::get<1>(temporal_padding)); + int temporal_pad = std::get<0>(temporal_padding) * 2; + if (temporal_pad > 0) { + x = ggml_ext_pad_ext(ctx->ggml_ctx, + x, + 0, + 0, + 0, + 0, + temporal_pad, + 0, + 0, + 0); + } + return Conv3d::forward(ctx, x); + } + }; + + struct TemporalGroupNorm : public GroupNorm { + explicit TemporalGroupNorm(int64_t channels) + : GroupNorm(32, channels, 1e-6f, true) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* result = nullptr; + for (int64_t t = 0; t < x->ne[2]; ++t) { + auto frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, t, t + 1); + GGML_ASSERT(frame->ne[3] % num_channels == 0); + int64_t batch_size = frame->ne[3] / num_channels; + frame = ggml_cont(ctx->ggml_ctx, frame); + frame = ggml_reshape_4d(ctx->ggml_ctx, + frame, + frame->ne[0], + frame->ne[1], + num_channels, + batch_size); + frame = GroupNorm::forward(ctx, frame); + frame = ggml_reshape_4d(ctx->ggml_ctx, + frame, + frame->ne[0], + frame->ne[1], + 1, + num_channels * batch_size); + result = result == nullptr ? frame : ggml_concat(ctx->ggml_ctx, result, frame, 2); + } + return result; + } + }; + + struct Downsample3D : public GGMLBlock { + int spatial_stride; + + Downsample3D(int64_t in_channels, + int64_t out_channels, + int temporal_stride, + int spatial_stride) + : spatial_stride(spatial_stride) { + blocks["conv"] = std::make_shared(in_channels, + out_channels, + std::tuple{3, 3, 3}, + std::tuple{temporal_stride, spatial_stride, spatial_stride}, + std::tuple{1, 0, 0}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (spatial_stride == 2) { + GGML_ASSERT(x->ne[0] > 1 && x->ne[1] > 1); + auto right = ggml_ext_slice(ctx->ggml_ctx, x, 0, x->ne[0] - 2, x->ne[0] - 1); + x = ggml_concat(ctx->ggml_ctx, x, right, 0); + auto bottom = ggml_ext_slice(ctx->ggml_ctx, x, 1, x->ne[1] - 2, x->ne[1] - 1); + x = ggml_concat(ctx->ggml_ctx, x, bottom, 1); + } + return std::dynamic_pointer_cast(blocks["conv"])->forward(ctx, x); + } + }; + + struct ResnetBlock3D : public GGMLBlock { + int64_t in_channels; + int64_t out_channels; + + ResnetBlock3D(int64_t in_channels, + int64_t out_channels) + : in_channels(in_channels), out_channels(out_channels) { + blocks["norm1"] = std::make_shared(in_channels); + blocks["norm2"] = std::make_shared(out_channels); + blocks["conv1"] = std::make_shared(in_channels, + out_channels, + std::tuple{3, 3, 3}, + std::tuple{1, 1, 1}, + std::tuple{1, 1, 1}); + blocks["conv2"] = std::make_shared(out_channels, + out_channels, + std::tuple{3, 3, 3}, + std::tuple{1, 1, 1}, + std::tuple{1, 1, 1}); + if (in_channels != out_channels) { + blocks["nin_shortcut"] = std::make_shared(in_channels, + out_channels, + std::tuple{1, 1, 1}); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + auto h = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, x))); + h = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, h))); + if (in_channels != out_channels) { + x = std::dynamic_pointer_cast(blocks["nin_shortcut"])->forward(ctx, x); + } + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct Encoder : public GGMLBlock { + static constexpr int levels = 6; + static constexpr std::array multipliers = {1, 2, 2, 4, 4, 8}; + static constexpr std::array spatial_down = {2, 2, 2, 2, 1, 1}; + static constexpr std::array temporal_down = {1, 2, 2, 1, 1, 1}; + + Encoder() { + constexpr int ch = 128; + blocks["conv_in"] = std::make_shared(3, + ch, + std::tuple{3, 3, 3}, + std::tuple{1, 1, 1}, + std::tuple{1, 1, 1}); + int64_t previous = ch; + for (int level = 0; level < levels; ++level) { + int64_t current = ch * multipliers[level]; + for (int block = 0; block < 2; ++block) { + blocks["down." + std::to_string(level) + ".block." + std::to_string(block)] = + std::make_shared(block == 0 ? previous : current, + current); + } + if (spatial_down[level] * temporal_down[level] > 1) { + blocks["down." + std::to_string(level) + ".downsample"] = + std::make_shared(current, + current, + temporal_down[level], + spatial_down[level]); + } + previous = current; + } + blocks["norm_out"] = std::make_shared(previous); + blocks["conv_out"] = std::make_shared(previous, + 48, + std::tuple{3, 3, 3}, + std::tuple{1, 1, 1}, + std::tuple{1, 1, 1}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + x = std::dynamic_pointer_cast(blocks["conv_in"])->forward(ctx, x); + for (int level = 0; level < levels; ++level) { + for (int block = 0; block < 2; ++block) { + x = std::dynamic_pointer_cast( + blocks["down." + std::to_string(level) + ".block." + std::to_string(block)]) + ->forward(ctx, x); + } + auto downsample = blocks.find("down." + std::to_string(level) + ".downsample"); + if (downsample != blocks.end()) { + x = std::dynamic_pointer_cast(downsample->second)->forward(ctx, x); + } + } + auto norm = std::dynamic_pointer_cast(blocks["norm_out"]); + auto conv = std::dynamic_pointer_cast(blocks["conv_out"]); + return conv->forward(ctx, ggml_silu(ctx->ggml_ctx, norm->forward(ctx, x))); + } + }; + + static ggml_tensor* attention_layout(ggml_context* ctx, ggml_tensor* x) { + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); + return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]); + } + + static ggml_tensor* apply_partial_rope(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* pe) { + int64_t rot_dim = pe->ne[2] * 2; + auto rotated = Rope::apply_rope(ctx, + ggml_ext_slice(ctx, x, 0, 0, rot_dim), + pe, + false); + if (rot_dim == x->ne[0]) { + return rotated; + } + auto tail = attention_layout(ctx, + ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0])); + return ggml_concat(ctx, rotated, tail, 0); + } + + struct DecoderAttention : public GGMLBlock { + static constexpr int num_head = 32; + static constexpr int head_dim = 64; + static constexpr int dim = num_head * head_dim; + + DecoderAttention() { + blocks["to_qkv"] = std::make_shared(dim, dim * 3, true); + blocks["to_out"] = std::make_shared(dim, dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe) { + auto to_qkv = std::dynamic_pointer_cast(blocks["to_qkv"]); + auto to_out = std::dynamic_pointer_cast(blocks["to_out"]); + auto qkv_projection = to_qkv->forward(ctx, x); + int64_t sequence = x->ne[1]; + int64_t batch_size = x->ne[2] * x->ne[3]; + qkv_projection = ggml_reshape_4d(ctx->ggml_ctx, + qkv_projection, + 3 * head_dim, + num_head, + sequence, + batch_size); + auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0); + auto q = ggml_reshape_4d(ctx->ggml_ctx, + qkv[0], + head_dim, + num_head, + sequence, + batch_size); + auto k = ggml_reshape_4d(ctx->ggml_ctx, + qkv[1], + head_dim, + num_head, + sequence, + batch_size); + auto v = ggml_reshape_4d(ctx->ggml_ctx, + qkv[2], + head_dim, + num_head, + sequence, + batch_size); + q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f); + k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f); + q = apply_partial_rope(ctx->ggml_ctx, q, pe); + k = apply_partial_rope(ctx->ggml_ctx, k, pe); + auto out = ggml_ext_attention_ext(ctx->ggml_ctx, + ctx->backend, + q, + k, + v, + num_head, + nullptr, + true, + ctx->flash_attn_enabled); + return to_out->forward(ctx, out); + } + }; + + struct DecoderFeedForward : public GGMLBlock { + static constexpr int dim = 2048; + static constexpr int kInnerDim = dim * 4; + + DecoderFeedForward() { + blocks["w1"] = std::make_shared(dim, kInnerDim * 2, true); + blocks["w2"] = std::make_shared(kInnerDim, dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto w1 = std::dynamic_pointer_cast(blocks["w1"]); + auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + auto gate = ggml_ext_chunk(ctx->ggml_ctx, w1->forward(ctx, x), 2, 0); + return w2->forward(ctx, + ggml_mul(ctx->ggml_ctx, + ggml_silu(ctx->ggml_ctx, gate[0]), + gate[1])); + } + }; + + struct DecoderBlock : public GGMLBlock { + static constexpr int dim = 2048; + + DecoderBlock() { + blocks["norm1"] = std::make_shared(dim, 1e-5f); + blocks["attn"] = std::make_shared(); + blocks["norm2"] = std::make_shared(dim, 1e-5f); + blocks["ff"] = std::make_shared(); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ED_UNUSED(tensor_storage_map); + ED_UNUSED(prefix); + params["scale1"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim); + params["scale2"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto ff = std::dynamic_pointer_cast(blocks["ff"]); + x = ggml_add(ctx->ggml_ctx, + x, + ggml_mul(ctx->ggml_ctx, + attn->forward(ctx, norm1->forward(ctx, x), pe), + params["scale1"])); + return ggml_add(ctx->ggml_ctx, + x, + ggml_mul(ctx->ggml_ctx, + ff->forward(ctx, norm2->forward(ctx, x)), + params["scale2"])); + } + }; + + struct Decoder : public GGMLBlock { + static constexpr int dim = 2048; + static constexpr int num_layers = 36; + static constexpr int num_register_tokens = 4; + static constexpr int patch_size = 16; + static constexpr int patch_size_t = 4; + + Decoder() { + blocks["x_embedder"] = std::make_shared(24, dim, true); + for (int i = 0; i < num_layers; ++i) { + blocks["transformer_blocks." + std::to_string(i)] = + std::make_shared(); + } + blocks["norm_out"] = std::make_shared(dim, 1e-5f, true, true); + blocks["proj_out"] = std::make_shared(dim, + 3 * patch_size_t * patch_size * patch_size, + true, + true); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ED_UNUSED(tensor_storage_map); + ED_UNUSED(prefix); + params["register_tokens"] = ggml_new_tensor_2d(ctx, + GGML_TYPE_F32, + dim, + num_register_tokens); + params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* z, + ggml_tensor* pe) { + int64_t width = z->ne[0]; + int64_t height = z->ne[1]; + int64_t num_frames = z->ne[2]; + int64_t batch_size = z->ne[3] / 24; + GGML_ASSERT(batch_size == 1); + + z = ggml_cont(ctx->ggml_ctx, + ggml_ext_torch_permute(ctx->ggml_ctx, z, 3, 0, 1, 2)); + z = ggml_reshape_3d(ctx->ggml_ctx, + z, + 24, + width * height * num_frames, + batch_size); + auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); + auto h = x_embedder->forward(ctx, z); + int64_t num_patches = h->ne[1]; + h = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1); + auto zero = ggml_ext_scale(ctx->ggml_ctx, + ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, 1), + 0.f); + h = ggml_concat(ctx->ggml_ctx, h, zero, 1); + + for (int i = 0; i < num_layers; ++i) { + auto block = std::dynamic_pointer_cast( + blocks["transformer_blocks." + std::to_string(i)]); + h = block->forward(ctx, h, pe); + sd::ggml_graph_cut::mark_graph_cut(h, + "minimax_h3_vae.decoder.blocks." + std::to_string(i), + "hidden_states"); + } + + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + h = proj_out->forward(ctx, norm_out->forward(ctx, h)); + h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, num_patches); + return DiT::unpatchify_3d(ctx->ggml_ctx, + h, + num_frames, + height, + width, + patch_size_t, + patch_size, + patch_size, + true); + } + }; + + struct MiniMaxH3VideoVAE : public GGMLBlock { + MiniMaxH3VideoVAE() { + blocks["encoder"] = std::make_shared(); + blocks["quant_conv"] = std::make_shared(48, + 48, + std::tuple{1, 1, 1}); + blocks["post_quant_conv"] = std::make_shared(24, + 24, + std::tuple{1, 1, 1}); + blocks["decoder"] = std::make_shared(); + } + + ggml_tensor* encode(GGMLRunnerContext* ctx, + ggml_tensor* pixels, + ggml_tensor* pixel_mean, + ggml_tensor* pixel_std) { + pixels = ggml_div(ctx->ggml_ctx, + ggml_sub(ctx->ggml_ctx, pixels, pixel_mean), + pixel_std); + auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); + auto quant = std::dynamic_pointer_cast(blocks["quant_conv"]); + auto moments = quant->forward(ctx, encoder->forward(ctx, pixels)); + return ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24); + } + + ggml_tensor* decode(GGMLRunnerContext* ctx, + ggml_tensor* latent, + ggml_tensor* pe, + ggml_tensor* pixel_mean, + ggml_tensor* pixel_std) { + auto post_quant = std::dynamic_pointer_cast(blocks["post_quant_conv"]); + auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); + auto pixels = decoder->forward(ctx, post_quant->forward(ctx, latent), pe); + pixels = ggml_add(ctx->ggml_ctx, + ggml_mul(ctx->ggml_ctx, pixels, pixel_std), + pixel_mean); + return ggml_clamp(ctx->ggml_ctx, pixels, 0.f, 1.f); + } + }; + + struct MiniMaxH3VideoVAERunner : public VAE { + MiniMaxH3VideoVAE model; + sd::Tensor pixel_mean; + sd::Tensor pixel_std; + sd::Tensor latents_mean; + sd::Tensor latents_std; + sd::Tensor rope_cache; + + MiniMaxH3VideoVAERunner(ggml_backend_t backend, + bool offload_params_to_cpu, + const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "first_stage_model") + : VAE(VERSION_MINIMAX_H3, backend, offload_params_to_cpu), + pixel_mean({1, 1, 1, 3}, {0.485f, 0.456f, 0.406f}), + pixel_std({1, 1, 1, 3}, {0.229f, 0.224f, 0.225f}), + latents_mean({1, 1, 1, 24}, + {0.858090341091156f, -0.960659146308899f, 1.066164016723633f, -0.509032547473907f, + -0.272758185863495f, -1.367541432380676f, -0.255325496196747f, -0.269075542688370f, + -0.537684082984924f, -0.046409729868174f, 0.665737032890320f, 0.196901276707649f, + -0.546060800552368f, -0.403534203767776f, -0.236830249428749f, 0.259284526109695f, + -0.301339447498322f, 0.211341992020607f, -1.120684862136841f, 0.358193337917328f, + -0.042251437902451f, 0.260482996702194f, 0.228640928864479f, 0.705603182315826f}), + latents_std({1, 1, 1, 24}, + {1.222377419471741f, 1.276726365089417f, 1.683177471160889f, 1.754945516586304f, + 1.563621640205383f, 2.194143533706665f, 0.965313792228699f, 1.056988596916199f, + 0.841948926448822f, 0.772995293140411f, 1.895593762397766f, 0.946841835975647f, + 0.799680948257446f, 0.449889004230499f, 0.719739973545075f, 0.693629324436188f, + 2.961095094680786f, 2.769419908523560f, 3.049618482589722f, 2.108805418014527f, + 3.276226282119751f, 3.162735700607300f, 2.281681299209595f, 2.612784385681153f}) { + scale_input = false; + model.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "minimax_h3_video_vae"; + } + + int get_encoder_output_channels(int input_channels) override { + ED_UNUSED(input_channels); + return 24; + } + + void get_param_tensors(std::map& tensors, const std::string prefix) override { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, + std::shared_ptr rng) override { + ED_UNUSED(rng); + return vae_output; + } + + sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) override { + return latents * latents_std + latents_mean; + } + + sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) override { + return (latents - latents_mean) / latents_std; + } + + static sd::Tensor ensure_video_shape(const sd::Tensor& tensor) { + if (tensor.dim() == 5) { + return tensor; + } + GGML_ASSERT(tensor.dim() == 4); + return tensor.reshape({tensor.shape()[0], + tensor.shape()[1], + 1, + tensor.shape()[2], + tensor.shape()[3]}); + } + + static ed_tiling_params_t h3_tiling(ed_tiling_params_t params) { + params.enabled = true; + params.tile_size_x = 16; + params.tile_size_y = 16; + params.rel_size_x = 0.0f; + params.rel_size_y = 0.0f; + params.target_overlap = 0.25f; + return params; + } + + static sd::Tensor repeat_last_frame(const sd::Tensor& input, + int64_t count) { + auto result = input; + auto last = sd::ops::slice(input, 2, input.shape()[2] - 1, input.shape()[2]); + for (int64_t i = 0; i < count; ++i) { + result = sd::ops::concat(result, last, 2); + } + return result; + } + + static sd::Tensor blend_temporal(const sd::Tensor& previous, + const sd::Tensor& current, + int64_t extent) { + auto output = current; + extent = std::min({extent, previous.shape()[2], current.shape()[2]}); + int64_t previous_start = previous.shape()[2] - extent; + for (int64_t b = 0; b < current.shape()[4]; ++b) { + for (int64_t c = 0; c < current.shape()[3]; ++c) { + for (int64_t t = 0; t < extent; ++t) { + float wb = static_cast(t) / extent; + float wa = 1.f - wb; + for (int64_t h = 0; h < current.shape()[1]; ++h) { + for (int64_t w = 0; w < current.shape()[0]; ++w) { + output.index(w, h, t, c, b) = + previous.index(w, h, previous_start + t, c, b) * wa + + current.index(w, h, t, c, b) * wb; + } + } + } + } + } + return output; + } + + sd::Tensor encode(int n_threads, + const sd::Tensor& x, + ed_tiling_params_t tiling_params, + bool circular_x = false, + bool circular_y = false) { + auto input = ensure_video_shape(x); + auto tiling = h3_tiling(tiling_params); + if (input.shape()[2] == 1) { + auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y); + if (!encoded.empty() && encoded.shape()[2] > 1) { + encoded = sd::ops::slice(encoded, + 2, + encoded.shape()[2] - 1, + encoded.shape()[2]); + } + return encoded; + } + + int64_t pad = (-input.shape()[2]) % 17; + if (pad < 0) { + pad += 17; + } + 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); + } + if (result.shape()[2] > 3) { + result = sd::ops::slice(result, 2, 0, result.shape()[2] - 3); + } + return result; + } + + sd::Tensor decode(int n_threads, + const sd::Tensor& x, + ed_tiling_params_t tiling_params, + bool decode_video = false, + bool circular_x = false, + bool circular_y = false, + bool silent = false) { + auto input = ensure_video_shape(x); + auto tiling = h3_tiling(tiling_params); + if (input.shape()[2] == 1) { + auto decoded = VAE::decode(n_threads, + input, + tiling, + decode_video, + circular_x, + circular_y, + silent); + if (!decoded.empty() && decoded.shape()[2] > 1) { + decoded = sd::ops::slice(decoded, + 2, + decoded.shape()[2] - 1, + decoded.shape()[2]); + } + return decoded; + } + + constexpr int64_t tokens_per_chunk = 5; + constexpr int64_t token_drop = 3; + constexpr int64_t token_overlap = 2; + constexpr int64_t frames_per_chunk = 20; + constexpr int64_t frame_pre_padding = 3; + constexpr int64_t frame_overlap = 5; + + int64_t pseudo_tokens = input.shape()[2] + token_drop; + int64_t pad_tokens = (tokens_per_chunk - pseudo_tokens % tokens_per_chunk) % tokens_per_chunk; + pseudo_tokens += pad_tokens; + int64_t num_chunks = pseudo_tokens / tokens_per_chunk - 1; + if (num_chunks < 1) { + pad_tokens += tokens_per_chunk; + num_chunks += 1; + } + if (pad_tokens > 0) { + 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); + if (decoded.empty()) { + return {}; + } + + int64_t first_end = std::min(frames_per_chunk, decoded.shape()[2]); + auto first = sd::ops::slice(decoded, + 2, + std::min(frame_pre_padding, first_end), + first_end); + if (!overlap.empty()) { + 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, + 2, + frames_per_chunk + frame_pre_padding, + decoded.shape()[2]); + } + if (i == num_chunks - 1 && !overlap.empty()) { + result = sd::ops::concat(result, overlap, 2); + overlap = {}; + } + } + + int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5; + expected_frames = std::max(1, expected_frames); + if (result.shape()[2] > expected_frames) { + result = sd::ops::slice(result, 2, 0, expected_frames); + } + return result; + } + + sd::Tensor build_rope(int64_t width, + int64_t height, + int64_t num_frames) { + std::vector> ids; + ids.reserve(static_cast(width * height * num_frames + 5)); + constexpr float two_pi = 6.28318530717958647692f; + for (int64_t t = 0; t < num_frames; ++t) { + float pt = (2.f * ((t + 0.5f) / num_frames) - 1.f) * two_pi; + for (int64_t h = 0; h < height; ++h) { + float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi; + for (int64_t w = 0; w < width; ++w) { + float pw = (2.f * ((w + 0.5f) / width) - 1.f) * two_pi; + ids.push_back({pt, ph, pw}); + } + } + } + for (int i = 0; i < 5; ++i) { + ids.push_back({0.f, 0.f, 0.f}); + } + auto values = Rope::embed_nd(ids, + 1, + 100.f, + std::vector{16, 16, 16}); + return sd::Tensor({2, + 2, + 24, + static_cast(ids.size())}, + std::move(values)); + } + + sd::Tensor _compute(const int n_threads, + const sd::Tensor& z, + bool decode_graph) override { + auto input = ensure_video_shape(z); + if (decode_graph) { + rope_cache = build_rope(input.shape()[0], + input.shape()[1], + input.shape()[2]); + } + auto get_graph = [&]() -> ggml_cgraph* { + auto value = make_input(input); + auto mean = make_input(pixel_mean); + auto std = make_input(pixel_std); + auto runner_ctx = get_context(); + ggml_tensor* out = nullptr; + if (decode_graph) { + auto pe = make_input(rope_cache); + out = model.decode(&runner_ctx, value, pe, mean, std); + } else { + out = model.encode(&runner_ctx, value, mean, std); + } + auto graph = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE); + ggml_build_forward_expand(graph, out); + return graph; + }; + return restore_trailing_singleton_dims( + GGMLRunner::compute(get_graph, n_threads, false), + 5); + } + }; + +} // namespace MiniMaxH3VAE + +#endif // __ED_DIT_MODELS_COMPONENTS_AUTOENCODERS_MINIMAX_H3_VAE_HPP__ diff --git a/src/dit_models/components/autoencoders/vae.hpp b/src/dit_models/components/autoencoders/vae.hpp index 12a7d5c6..bc109e70 100644 --- a/src/dit_models/components/autoencoders/vae.hpp +++ b/src/dit_models/components/autoencoders/vae.hpp @@ -67,7 +67,7 @@ struct VAE : public GGMLRunner { int get_scale_factor() { int scale_factor = 8; - if (version == VERSION_WAN2_2_TI2V) { + if (version == VERSION_MINIMAX_H3 || version == VERSION_WAN2_2_TI2V) { scale_factor = 16; } else if (ed_version_uses_flux2_vae(version)) { scale_factor = 16; diff --git a/src/dit_models/components/common/common_dit.hpp b/src/dit_models/components/common/common_dit.hpp index 2d757994..6002bda3 100644 --- a/src/dit_models/components/common/common_dit.hpp +++ b/src/dit_models/components/common/common_dit.hpp @@ -103,6 +103,70 @@ namespace DiT { x = ggml_ext_slice(ctx, x, 0, 0, W); // [N, C, H, W] return x; } + + inline ggml_tensor* patchify_3d(ggml_context* ctx, + ggml_tensor* x, + int pt, + int ph, + int pw, + int64_t N = 1, + bool patch_last = true) { + int64_t C = x->ne[3] / N; + int64_t T = x->ne[2]; + int64_t H = x->ne[1]; + int64_t W = x->ne[0]; + int64_t t_len = T / pt; + int64_t h_len = H / ph; + int64_t w_len = W / pw; + + GGML_ASSERT(C * N == x->ne[3]); + GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W); + + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); + if (patch_last) { + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + } else { + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); + } + return ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); + } + + inline ggml_tensor* unpatchify_3d(ggml_context* ctx, + ggml_tensor* x, + int64_t t_len, + int64_t h_len, + int64_t w_len, + int pt, + int ph, + int pw, + bool patch_last = true) { + int64_t N = x->ne[2]; + int64_t C = x->ne[0] / pt / ph / pw; + + GGML_ASSERT(C * pt * ph * pw == x->ne[0]); + + if (patch_last) { + x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + } else { + x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); + } + + x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); + return ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); + } } // namespace DiT #endif // __COMMON_DIT_HPP__ diff --git a/src/dit_models/components/text_encoders/llm.hpp b/src/dit_models/components/text_encoders/llm.hpp index 56871024..ea6a322b 100644 --- a/src/dit_models/components/text_encoders/llm.hpp +++ b/src/dit_models/components/text_encoders/llm.hpp @@ -265,9 +265,10 @@ namespace LLM { int head_dim = 128; bool qkv_bias = true; bool qk_norm = false; - int64_t vocab_size = 152064; - float rms_norm_eps = 1e-06f; - LLMVisionParams vision; + int64_t vocab_size = 152064; + float rms_norm_eps = 1e-06f; + bool final_norm = true; + LLMVisionParams vision; }; struct LLMImageEmbedInfo { @@ -1240,19 +1241,23 @@ namespace LLM { protected: int64_t num_layers; bool diffusers_text_dtype; + bool final_norm; public: TextModel(const LLMParams& params) : num_layers(params.num_layers), diffusers_text_dtype(params.arch == LLMArch::QWEN2_5_VL && - qwen_align_diffusers_text_dtype_enabled()) { + qwen_align_diffusers_text_dtype_enabled()), + final_norm(params.final_norm) { const bool cast_rms_output_to_input_type = params.arch == LLMArch::QWEN2_5_VL; blocks["embed_tokens"] = std::shared_ptr(new Embedding(params.vocab_size, params.hidden_size)); for (int i = 0; i < num_layers; i++) { blocks["layers." + std::to_string(i)] = std::shared_ptr(new TransformerBlock(params)); } - blocks["norm"] = std::shared_ptr( - new RMSNorm(params.hidden_size, params.rms_norm_eps, false, cast_rms_output_to_input_type)); + if (final_norm) { + blocks["norm"] = std::shared_ptr( + new RMSNorm(params.hidden_size, params.rms_norm_eps, false, cast_rms_output_to_input_type)); + } } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -1266,7 +1271,9 @@ namespace LLM { // return: [N, n_token, hidden_size] auto embed_tokens = std::dynamic_pointer_cast(blocks["embed_tokens"]); - auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto norm = final_norm + ? std::dynamic_pointer_cast(blocks["norm"]) + : nullptr; auto x = embed_tokens->forward(ctx, input_ids); if (diffusers_text_dtype && x->type != GGML_TYPE_BF16) { @@ -1366,7 +1373,7 @@ namespace LLM { for (int i = 1; i < intermediate_outputs.size(); i++) { x = ggml_concat(ctx->ggml_ctx, x, intermediate_outputs[i], 0); } - } else { + } else if (norm != nullptr) { x = norm->forward(ctx, x); if (debug_target == "final_norm") { return x; @@ -1458,7 +1465,7 @@ namespace LLM { params.rms_norm_eps = 1e-5f; } else if (arch == LLMArch::QWEN3) { params.head_dim = 128; - params.num_heads = 32; + params.num_heads = 64; params.num_kv_heads = 8; params.qkv_bias = false; params.qk_norm = true; @@ -1501,6 +1508,10 @@ namespace LLM { if (arch == LLMArch::QWEN3 && params.num_layers == 28) { // Qwen3 2B params.num_heads = 16; } + if (arch == LLMArch::QWEN3 && params.num_layers == 50 && params.hidden_size == 5120) { + params.num_heads = 64; + params.final_norm = false; + } LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64, params.num_layers, params.vocab_size, @@ -1673,7 +1684,7 @@ namespace LLM { std::set out_layers, const std::vector& image_embed_infos = {}, const std::string& debug_target = "") { - ggml_cgraph* gf = ggml_new_graph(compute_ctx); + ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE); ggml_tensor* input_ids = make_input(input_ids_tensor); std::vector> image_embeds; image_embeds.reserve(image_embeds_tensor.size()); diff --git a/src/dit_models/diffusion_model.hpp b/src/dit_models/diffusion_model.hpp index a37e8ddd..ac25ae56 100644 --- a/src/dit_models/diffusion_model.hpp +++ b/src/dit_models/diffusion_model.hpp @@ -19,6 +19,18 @@ using sd::DiffusionCacheResult; +enum class MiniMaxH3ReferenceKind : int32_t { + IMAGE, + VIDEO, + AUDIO, + VIDEO_AUDIO, +}; + +struct MiniMaxH3ReferenceBlock { + MiniMaxH3ReferenceKind kind = MiniMaxH3ReferenceKind::IMAGE; + int32_t video_index = -1; + int32_t audio_index = -1; +}; struct DiffusionParams { const sd::Tensor* x = nullptr; @@ -37,6 +49,13 @@ struct DiffusionParams { const sd::Tensor* vace_context = nullptr; float vace_strength = 1.f; const std::vector* skip_layers = nullptr; + const sd::Tensor* minimax_text_token_tags = nullptr; + const sd::Tensor* minimax_keyframe_indices = nullptr; + const std::vector>* minimax_reference_audio_latents = nullptr; + const std::vector* minimax_reference_blocks = nullptr; + int minimax_audio_length = 0; + float minimax_video_sigma_shift = 12.f; + float minimax_audio_sigma_shift = 3.f; }; template diff --git a/src/dit_models/models/minimax_h3_full.hpp b/src/dit_models/models/minimax_h3_full.hpp new file mode 100644 index 00000000..00c4f5a0 --- /dev/null +++ b/src/dit_models/models/minimax_h3_full.hpp @@ -0,0 +1,1177 @@ +#ifndef __ED_DIT_MODELS_MODELS_MINIMAX_H3_FULL_HPP__ +#define __ED_DIT_MODELS_MODELS_MINIMAX_H3_FULL_HPP__ + +#include +#include +#include +#include +#include +#include +#include + +#include "backend/ggml/ggml_graph_cut.h" +#include "dit_models/components/common/common_dit.hpp" +#include "backend/ggml/ggml_extend.hpp" +#include "dit_models/components/common/rope.hpp" + +namespace MiniMaxH3 { + + constexpr int H3_GRAPH_SIZE = 131072; + constexpr float FRAME_RESCALE = 5.f / 3.f; + constexpr float VISUAL_COND_TIMESTEP = 0.999f; + + struct Config { + int64_t hidden_size = 5376; + int64_t num_layers = 50; + int64_t token_refiner_num_layers = 2; + int64_t num_attention_heads = 56; + int64_t attention_head_dim = 128; + int64_t ffn_hidden_size = 14336; + int64_t video_latent_channels = 24; + int64_t audio_latent_channels = 32; + int64_t text_dim = 5120; + int64_t timestep_input_dim = 256; + int64_t time_embed_hidden_size = 5376; + int64_t time_embed_dim = 2688; + int64_t rope_inv_freq_len = 16; + int64_t adaln_curve_grid = 0; + int patch_t = 1; + int patch_h = 2; + int patch_w = 2; + float norm_eps = 1e-5f; + float qk_norm_eps = 1e-5f; + float final_norm_eps = 1e-5f; + + bool uses_adaln_curves() const { + return adaln_curve_grid > 0; + } + + static int64_t count_blocks(const String2TensorStorage& tensors, + const std::string& prefix) { + std::set indices; + for (const auto& [name, _] : tensors) { + if (!starts_with(name, prefix)) { + continue; + } + size_t begin = prefix.size(); + size_t end = name.find('.', begin); + if (end != std::string::npos) { + indices.insert(std::atoi(name.substr(begin, end - begin).c_str())); + } + } + return static_cast(indices.size()); + } + + static Config detect_from_weights(const String2TensorStorage& tensors, + const std::string& prefix) { + Config config; + auto find = [&](const std::string& suffix) -> const TensorStorage* { + auto it = tensors.find(prefix + "." + suffix); + return it == tensors.end() ? nullptr : &it->second; + }; + + if (const auto* weight = find("video_patch_proj.weight")) { + config.video_latent_channels = weight->ne[0] / 4; + config.hidden_size = weight->ne[1]; + } + if (const auto* weight = find("audio_patch_proj.weight")) { + config.audio_latent_channels = weight->ne[0]; + } + config.num_layers = count_blocks(tensors, prefix + ".blocks."); + config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks."); + if (const auto* weight = find("blocks.0.attn.q_norm.weight")) { + config.attention_head_dim = weight->ne[0]; + } + if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) { + config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim); + } + if (const auto* weight = find("blocks.0.mlp.fc1.weight")) { + config.ffn_hidden_size = weight->ne[1] / 2; + } + if (const auto* weight = find("condition_proj.weight")) { + config.text_dim = weight->ne[0]; + } + if (const auto* table = find("adaln_t_table")) { + config.time_embed_dim = table->ne[0]; + config.adaln_curve_grid = table->ne[1]; + } else { + if (const auto* weight = find("time_embedder.proj_in.weight")) { + config.timestep_input_dim = weight->ne[0]; + config.time_embed_hidden_size = weight->ne[1]; + } + if (const auto* weight = find("time_embedder.proj_out.weight")) { + config.time_embed_dim = weight->ne[1]; + } + } + if (const auto* inv_freq = find("rope.inv_freq")) { + config.rope_inv_freq_len = inv_freq->ne[0]; + } + + LOG_DEBUG("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 + ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64, + config.num_layers, + config.hidden_size, + config.num_attention_heads, + config.attention_head_dim, + config.ffn_hidden_size, + config.adaln_curve_grid); + return config; + } + }; + + static float time_shift_sigma(float sigma, float from_shift, float to_shift) { + float base = sigma / (from_shift + sigma * (1.f - from_shift)); + 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); + blocks["proj_out"] = std::make_shared(hidden_dim, output_dim, true, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto proj_in = std::dynamic_pointer_cast(blocks["proj_in"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + return proj_out->forward(ctx, ggml_silu(ctx->ggml_ctx, proj_in->forward(ctx, x))); + } + }; + + struct MLP : public UnaryBlock { + MLP(int64_t hidden_size, + int64_t ffn_hidden_size) { + blocks["fc1"] = std::make_shared(hidden_size, ffn_hidden_size * 2, false, false, true, 1.f / 128.f); + blocks["fc2"] = std::make_shared(ffn_hidden_size, hidden_size, false, false, true, 1.f / 128.f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + auto uv = ggml_ext_chunk(ctx->ggml_ctx, fc1->forward(ctx, x), 2, 0); + return fc2->forward(ctx, ggml_mul(ctx->ggml_ctx, + ggml_silu(ctx->ggml_ctx, uv[0]), + uv[1])); + } + }; + + static ggml_tensor* attention_layout(ggml_context* ctx, ggml_tensor* x) { + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); + return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]); + } + + static ggml_tensor* apply_partial_rope(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* pe) { + int64_t rot_dim = pe->ne[2] * 2; + GGML_ASSERT(rot_dim <= x->ne[0]); + auto rotated = Rope::apply_rope(ctx, + ggml_ext_slice(ctx, x, 0, 0, rot_dim), + pe, + false); + if (rot_dim == x->ne[0]) { + return rotated; + } + auto tail = attention_layout(ctx, ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0])); + return ggml_concat(ctx, rotated, tail, 0); + } + + struct Attention : public GGMLBlock { + int64_t heads; + int64_t head_dim; + + Attention(int64_t hidden_size, + int64_t heads, + int64_t head_dim, + float eps) + : heads(heads), head_dim(head_dim) { + int64_t inner = heads * head_dim; + blocks["qkv_proj"] = std::make_shared(hidden_size, inner * 3, false); + blocks["q_norm"] = std::make_shared(head_dim, eps); + blocks["k_norm"] = std::make_shared(head_dim, eps); + blocks["out_proj"] = std::make_shared(inner, hidden_size, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe = nullptr) { + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv_proj"]); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); + + int64_t sequence = x->ne[1]; + int64_t batch = x->ne[2] * x->ne[3]; + auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_proj->forward(ctx, x), 3, 0); + auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv[0], head_dim, heads, sequence, batch); + auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv[1], head_dim, heads, sequence, batch); + auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv[2], head_dim, heads, sequence, batch); + q = q_norm->forward(ctx, q); + k = k_norm->forward(ctx, k); + if (pe != nullptr) { + q = apply_partial_rope(ctx->ggml_ctx, q, pe); + k = apply_partial_rope(ctx->ggml_ctx, k, pe); + } else { + q = attention_layout(ctx->ggml_ctx, q); + k = attention_layout(ctx->ggml_ctx, k); + } + auto out = ggml_ext_attention_ext(ctx->ggml_ctx, + ctx->backend, + q, + k, + v, + static_cast(heads), + nullptr, + true, + ctx->flash_attn_enabled, + 1.f / 128.f); + return out_proj->forward(ctx, out); + } + }; + + struct TokenRefinerBlock : public GGMLBlock { + TokenRefinerBlock(const Config& config) { + blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["attn"] = std::make_shared(config.hidden_size, + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps); + blocks["mlp"] = std::make_shared(config.hidden_size, config.ffn_hidden_size); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); + x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, norm1->forward(ctx, x))); + return ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, norm2->forward(ctx, x))); + } + }; + + struct TokenRefiner : public GGMLBlock { + int64_t num_layers; + + explicit TokenRefiner(const Config& config) + : num_layers(config.token_refiner_num_layers) { + for (int64_t i = 0; i < num_layers; ++i) { + blocks["blocks." + std::to_string(i)] = std::make_shared(config); + } + blocks["final_norm"] = std::make_shared(config.hidden_size, config.final_norm_eps); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + for (int64_t i = 0; i < num_layers; ++i) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + x = block->forward(ctx, x); + } + return std::dynamic_pointer_cast(blocks["final_norm"])->forward(ctx, x); + } + }; + + struct AdaLayerNormModulation : public GGMLBlock { + int64_t hidden_size; + int expand; + int modalities; + bool apply_silu; + + AdaLayerNormModulation(int64_t time_dim, + int64_t hidden_size, + int expand, + int modalities, + bool apply_silu, + bool force_f32) + : hidden_size(hidden_size), + expand(expand), + modalities(modalities), + apply_silu(apply_silu) { + blocks["linear"] = std::make_shared(time_dim, + hidden_size * expand * modalities, + true, + force_f32); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t_emb) { + if (apply_silu) { + t_emb = ggml_silu(ctx->ggml_ctx, t_emb); + } + return std::dynamic_pointer_cast(blocks["linear"])->forward(ctx, t_emb); + } + }; + + struct TokenModulationSpan { + int64_t start; + int64_t end; + int64_t modulation_row; + }; + + enum class SequenceKind { + TEXT, + CONDITION_VIDEO, + CONDITION_AUDIO, + TARGET_AUDIO, + TARGET_VIDEO, + }; + + struct SequenceSegment { + int64_t start; + int64_t end; + SequenceKind kind; + int32_t source_index = -1; + }; + + static std::vector modulation_row(ggml_context* ctx, + ggml_tensor* projection, + int64_t hidden_size, + int expand, + int modalities, + int64_t row) { + int64_t timestep_rows = projection->ne[1]; + auto reshaped = ggml_reshape_2d(ctx, + projection, + hidden_size * expand, + timestep_rows * modalities); + auto selected = ggml_ext_slice(ctx, reshaped, 1, row, row + 1); + return ggml_ext_chunk(ctx, selected, expand, 0); + } + + static ggml_tensor* modulate_segments(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* projection, + const std::vector& segments, + int64_t hidden_size, + int expand, + int modalities, + int shift_index, + int scale_index) { + ggml_tensor* out = nullptr; + for (const auto& segment : segments) { + auto mods = modulation_row(ctx, + projection, + hidden_size, + expand, + modalities, + segment.modulation_row); + auto part = ggml_ext_slice(ctx, x, 1, segment.start, segment.end); + part = ggml_add(ctx, + ggml_add(ctx, part, ggml_mul(ctx, part, mods[scale_index])), + mods[shift_index]); + out = out == nullptr ? part : ggml_concat(ctx, out, part, 1); + } + return out; + } + + static ggml_tensor* gated_residual_segments(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* update, + ggml_tensor* projection, + const std::vector& segments, + int64_t hidden_size, + int gate_index) { + ggml_tensor* out = nullptr; + for (const auto& segment : segments) { + auto mods = modulation_row(ctx, projection, hidden_size, 6, 3, segment.modulation_row); + auto base = ggml_ext_slice(ctx, x, 1, segment.start, segment.end); + auto add = ggml_ext_slice(ctx, update, 1, segment.start, segment.end); + auto part = ggml_add(ctx, base, ggml_mul(ctx, add, mods[gate_index])); + out = out == nullptr ? part : ggml_concat(ctx, out, part, 1); + } + return out; + } + + struct TransformerBlock : public GGMLBlock { + Config config; + + explicit TransformerBlock(const Config& config) + : config(config) { + blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["attn"] = std::make_shared(config.hidden_size, + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps); + blocks["mlp"] = std::make_shared(config.hidden_size, + config.ffn_hidden_size); + blocks["adaln_proj"] = std::make_shared(config.time_embed_dim, + config.hidden_size, + 6, + 3, + !config.uses_adaln_curves(), + config.uses_adaln_curves()); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* t_emb, + const std::vector& segments, + ggml_tensor* pe) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); + auto adaln = std::dynamic_pointer_cast(blocks["adaln_proj"]); + auto mods = adaln->forward(ctx, t_emb); + + auto h = modulate_segments(ctx->ggml_ctx, + norm1->forward(ctx, x), + mods, + segments, + config.hidden_size, + 6, + 3, + 0, + 1); + x = gated_residual_segments(ctx->ggml_ctx, + x, + attn->forward(ctx, h, pe), + mods, + segments, + config.hidden_size, + 2); + h = modulate_segments(ctx->ggml_ctx, + norm2->forward(ctx, x), + mods, + segments, + config.hidden_size, + 6, + 3, + 3, + 4); + return gated_residual_segments(ctx->ggml_ctx, + x, + mlp->forward(ctx, h), + mods, + segments, + config.hidden_size, + 5); + } + }; + + struct FinalLayer : public GGMLBlock { + Config config; + + explicit FinalLayer(const Config& config) + : config(config) { + int64_t video_dim = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w; + blocks["norm"] = std::make_shared(config.hidden_size, config.final_norm_eps); + blocks["adaln_proj"] = std::make_shared(config.time_embed_dim, + config.hidden_size, + 2, + 1, + !config.uses_adaln_curves(), + config.uses_adaln_curves()); + blocks["video_out"] = std::make_shared(config.hidden_size, video_dim, true, true); + blocks["audio_out"] = std::make_shared(config.hidden_size, config.audio_latent_channels, true, true); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* t_emb, + const TokenModulationSpan& video, + const TokenModulationSpan& audio) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto adaln = std::dynamic_pointer_cast(blocks["adaln_proj"]); + auto video_out = std::dynamic_pointer_cast(blocks["video_out"]); + auto audio_out = std::dynamic_pointer_cast(blocks["audio_out"]); + auto mods = adaln->forward(ctx, t_emb); + auto apply = [&](const TokenModulationSpan& segment) { + auto row = modulation_row(ctx->ggml_ctx, mods, config.hidden_size, 2, 1, segment.modulation_row); + auto value = norm->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end)); + return ggml_add(ctx->ggml_ctx, + ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])), + row[0]); + }; + return {video_out->forward(ctx, apply(video)), + audio_out->forward(ctx, apply(audio))}; + } + }; + + struct MiniMaxH3Transformer3DModel : public GGMLBlock { + Config config; + + explicit MiniMaxH3Transformer3DModel(const Config& config) + : config(config) { + int64_t video_dim = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w; + blocks["video_patch_proj"] = std::make_shared(video_dim, config.hidden_size, true, true); + blocks["audio_patch_proj"] = std::make_shared(config.audio_latent_channels, config.hidden_size, true, true); + blocks["condition_proj"] = std::make_shared(config.text_dim, config.hidden_size, true); + if (!config.uses_adaln_curves()) { + blocks["time_embedder"] = std::make_shared(config.timestep_input_dim, + config.time_embed_hidden_size, + config.time_embed_dim); + } + blocks["token_refiner"] = std::make_shared(config); + for (int64_t i = 0; i < config.num_layers; ++i) { + blocks["blocks." + std::to_string(i)] = std::make_shared(config); + } + blocks["final_layer"] = std::make_shared(config); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensors = {}, + const std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensors, prefix); + params["rope.inv_freq"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.rope_inv_freq_len); + if (config.uses_adaln_curves()) { + params["adaln_t_table"] = ggml_new_tensor_2d(ctx, + GGML_TYPE_F32, + config.time_embed_dim, + config.adaln_curve_grid); + } + } + + ggml_tensor* refine_context(GGMLRunnerContext* ctx, ggml_tensor* context) { + if (context->ne[0] == config.hidden_size) { + return context; + } + GGML_ASSERT(context->ne[0] == config.text_dim); + auto condition_proj = std::dynamic_pointer_cast(blocks["condition_proj"]); + auto token_refiner = std::dynamic_pointer_cast(blocks["token_refiner"]); + return token_refiner->forward(ctx, condition_proj->forward(ctx, context)); + } + + ggml_tensor* time_embedding(GGMLRunnerContext* ctx, + ggml_tensor* timestep_features, + ggml_tensor* curve_indices, + ggml_tensor* curve_upper_indices, + ggml_tensor* curve_fractions) { + if (!config.uses_adaln_curves()) { + return std::dynamic_pointer_cast(blocks["time_embedder"])->forward(ctx, timestep_features); + } + auto lower = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_indices); + auto upper = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_upper_indices); + return ggml_add(ctx->ggml_ctx, + lower, + ggml_mul(ctx->ggml_ctx, + ggml_sub(ctx->ggml_ctx, upper, lower), + curve_fractions)); + } + + ggml_tensor* build_rope(GGMLRunnerContext* ctx, + ggml_tensor* position_ids) { + auto inv = ggml_reshape_2d(ctx->ggml_ctx, + params["rope.inv_freq"], + config.rope_inv_freq_len, + 1); + ggml_tensor* angles = nullptr; + for (int axis = 0; axis < 3; ++axis) { + auto pos = ggml_ext_slice(ctx->ggml_ctx, position_ids, 0, axis, axis + 1); + auto expanded_inv = ggml_repeat_4d(ctx->ggml_ctx, + inv, + inv->ne[0], + pos->ne[1], + 1, + 1); + auto a = ggml_mul(ctx->ggml_ctx, expanded_inv, pos); + angles = angles == nullptr ? a : ggml_concat(ctx->ggml_ctx, angles, a, 0); + } + auto c = ggml_reshape_4d(ctx->ggml_ctx, ggml_cos(ctx->ggml_ctx, angles), 1, angles->ne[0], angles->ne[1], 1); + auto s = ggml_reshape_4d(ctx->ggml_ctx, ggml_sin(ctx->ggml_ctx, angles), 1, angles->ne[0], angles->ne[1], 1); + auto ns = ggml_neg(ctx->ggml_ctx, s); + auto pe = ggml_concat(ctx->ggml_ctx, c, ns, 0); + pe = ggml_concat(ctx->ggml_ctx, pe, s, 0); + pe = ggml_concat(ctx->ggml_ctx, pe, c, 0); + return ggml_reshape_4d(ctx->ggml_ctx, pe, 2, 2, angles->ne[0], angles->ne[1]); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* video, + ggml_tensor* audio, + ggml_tensor* context, + const std::vector& condition_videos, + const std::vector& condition_audios, + ggml_tensor* position_ids, + ggml_tensor* timestep_features, + ggml_tensor* curve_indices, + ggml_tensor* curve_upper_indices, + ggml_tensor* curve_fractions, + const std::vector& segments, + const std::vector& sequence_segments, + const TokenModulationSpan& video_segment, + const TokenModulationSpan& audio_segment, + float audio_slope) { + auto video_proj = std::dynamic_pointer_cast(blocks["video_patch_proj"]); + auto audio_proj = std::dynamic_pointer_cast(blocks["audio_patch_proj"]); + + std::vector> condition_video_ranges; + ggml_tensor* video_rows = nullptr; + int64_t video_offset = 0; + for (auto condition : condition_videos) { + auto rows = DiT::patchify_3d(ctx->ggml_ctx, + condition, + config.patch_t, + config.patch_h, + config.patch_w, + 1, + true); + condition_video_ranges.push_back({video_offset, video_offset + rows->ne[1]}); + video_offset += rows->ne[1]; + video_rows = video_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, video_rows, rows, 1); + } + auto target_video_rows = DiT::patchify_3d(ctx->ggml_ctx, + video, + config.patch_t, + config.patch_h, + config.patch_w, + 1, + true); + std::pair target_video_range = {video_offset, video_offset + target_video_rows->ne[1]}; + video_rows = video_rows == nullptr ? target_video_rows + : ggml_concat(ctx->ggml_ctx, video_rows, target_video_rows, 1); + auto video_embeds = video_proj->forward(ctx, video_rows); + + auto pack_audio_rows = [&](ggml_tensor* value) { + value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, value, 2, 0, 1, 3)); + return ggml_reshape_3d(ctx->ggml_ctx, + value, + value->ne[0], + value->ne[1] * value->ne[2], + value->ne[3]); + }; + std::vector> condition_audio_ranges; + ggml_tensor* audio_rows = nullptr; + int64_t audio_offset = 0; + for (auto condition : condition_audios) { + auto rows = pack_audio_rows(condition); + condition_audio_ranges.push_back({audio_offset, audio_offset + rows->ne[1]}); + audio_offset += rows->ne[1]; + audio_rows = audio_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, audio_rows, rows, 1); + } + audio = pack_audio_rows(audio); + std::pair target_audio_range = {audio_offset, audio_offset + audio->ne[1]}; + audio_rows = audio_rows == nullptr ? audio : ggml_concat(ctx->ggml_ctx, audio_rows, audio, 1); + auto audio_embeds = audio_proj->forward(ctx, audio_rows); + context = refine_context(ctx, context); + + ggml_tensor* h = nullptr; + auto append = [&](ggml_tensor* value) { + h = h == nullptr ? value : ggml_concat(ctx->ggml_ctx, h, value, 1); + }; + for (const auto& sequence : sequence_segments) { + if (sequence.kind == SequenceKind::TEXT) { + append(context); + } else if (sequence.kind == SequenceKind::CONDITION_VIDEO) { + GGML_ASSERT(sequence.source_index >= 0 && + sequence.source_index < static_cast(condition_video_ranges.size())); + auto range = condition_video_ranges[static_cast(sequence.source_index)]; + append(ggml_ext_slice(ctx->ggml_ctx, video_embeds, 1, range.first, range.second)); + } else if (sequence.kind == SequenceKind::CONDITION_AUDIO) { + GGML_ASSERT(sequence.source_index >= 0 && + sequence.source_index < static_cast(condition_audio_ranges.size())); + auto range = condition_audio_ranges[static_cast(sequence.source_index)]; + append(ggml_ext_slice(ctx->ggml_ctx, audio_embeds, 1, range.first, range.second)); + } else if (sequence.kind == SequenceKind::TARGET_AUDIO) { + append(ggml_ext_slice(ctx->ggml_ctx, + audio_embeds, + 1, + target_audio_range.first, + target_audio_range.second)); + } else { + append(ggml_ext_slice(ctx->ggml_ctx, + video_embeds, + 1, + target_video_range.first, + target_video_range.second)); + } + } + GGML_ASSERT(h != nullptr); + + auto t_emb = time_embedding(ctx, + timestep_features, + curve_indices, + curve_upper_indices, + curve_fractions); + auto pe = build_rope(ctx, position_ids); + for (int64_t i = 0; i < config.num_layers; ++i) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + h = block->forward(ctx, h, t_emb, segments, pe); + sd::ggml_graph_cut::mark_graph_cut(h, + "minimax_h3.blocks." + std::to_string(i), + "hidden_states"); + } + + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + auto output = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment); + auto video_out = DiT::unpatchify_3d(ctx->ggml_ctx, + output.first, + video->ne[2] / config.patch_t, + video->ne[1] / config.patch_h, + video->ne[0] / config.patch_w, + config.patch_t, + config.patch_h, + config.patch_w, + true); + auto audio_out = ggml_reshape_4d(ctx->ggml_ctx, + output.second, + config.audio_latent_channels, + audio->ne[1] / 2, + 2, + 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); + return {video_out, audio_out}; + } + }; + + struct PackedSequenceLayout { + std::vector positions; + std::vector segments; + std::vector sequence_segments; + TokenModulationSpan video_segment{}; + TokenModulationSpan audio_segment{}; + std::vector timesteps; + }; + + static float video_span(int64_t frame) { + static const int spans[5] = {1, 4, 4, 4, 4}; + return FRAME_RESCALE * spans[frame % 5]; + } + + static std::vector spatial_axis(int64_t dim, + float sqrt_area) { + int64_t count = dim / 2; + float ratio = static_cast(dim) / sqrt_area; + std::vector result(static_cast(count)); + for (int64_t i = 0; i < count; ++i) { + result[static_cast(i)] = + (static_cast(i) * (ratio / count) + (1.f - ratio) * 0.5f) * 32.f; + } + return result; + } + + static int find_or_add_timestep(std::vector* values, float value) { + auto it = std::find(values->begin(), values->end(), value); + if (it != values->end()) { + return static_cast(it - values->begin()); + } + values->push_back(value); + return static_cast(values->size() - 1); + } + + static PackedSequenceLayout build_layout(int64_t text_len, + int64_t latent_t, + int64_t latent_h, + int64_t latent_w, + int64_t audio_t, + const std::vector>& condition_videos, + const std::vector>& condition_audios, + const sd::Tensor& keyframe_indices, + const std::vector& reference_blocks, + const sd::Tensor& text_tags, + float video_t, + float audio_timestep) { + PackedSequenceLayout layout; + float sqrt_area = std::sqrt(static_cast(latent_h * latent_w)); + auto h_axis = spatial_axis(latent_h, sqrt_area); + auto w_axis = spatial_axis(latent_w, sqrt_area); + int64_t frame_rows = static_cast(h_axis.size() * w_axis.size()); + int64_t row = 0; + + auto append_position = [&](float t, float h, float w) { + layout.positions.push_back(t); + layout.positions.push_back(h); + layout.positions.push_back(w); + }; + for (int64_t i = 0; i < text_len; ++i) { + append_position(static_cast(i), 0.f, 0.f); + } + layout.sequence_segments.push_back({0, text_len, SequenceKind::TEXT}); + + int video_time_row = find_or_add_timestep(&layout.timesteps, video_t); + int audio_time_row = find_or_add_timestep(&layout.timesteps, audio_timestep); + int condition_time_row = find_or_add_timestep(&layout.timesteps, + std::max(video_t, VISUAL_COND_TIMESTEP)); + int audio_condition_time_row = find_or_add_timestep(&layout.timesteps, + std::max(audio_timestep, 1.f)); + + int64_t run_start = 0; + int current_tag = text_tags.empty() ? 1 : text_tags[0]; + for (int64_t i = 1; i <= text_len; ++i) { + int tag = i < text_len && !text_tags.empty() ? text_tags[i] : -1; + if (i == text_len || tag != current_tag) { + layout.segments.push_back({run_start, + i, + video_time_row * 3 + current_tag}); + run_start = i; + current_tag = tag; + } + } + row = text_len; + + auto condition_spatial_axes = [&](const sd::Tensor& condition) { + float area = std::sqrt(static_cast(condition.shape()[0] * condition.shape()[1])); + return std::make_pair(spatial_axis(condition.shape()[1], area), + spatial_axis(condition.shape()[0], area)); + }; + auto append_video_positions = [&](const sd::Tensor& condition, + float cursor) { + auto axes = condition_spatial_axes(condition); + for (int64_t t = 0; t < condition.shape()[2]; ++t) { + for (float h : axes.first) { + for (float w : axes.second) { + append_position(cursor, h, w); + } + } + cursor += video_span(t); + } + return cursor; + }; + auto append_audio_positions = [&](int64_t length, + float cursor, + float w_low, + float w_high) { + for (int channel = 0; channel < 2; ++channel) { + float w = channel == 0 ? w_low : w_high; + for (int64_t t = 0; t < length; ++t) { + append_position(cursor + static_cast(t), 0.f, w); + } + } + }; + + float cursor = static_cast(text_len); + if (reference_blocks.empty()) { + float video_duration = 0.f; + for (int64_t t = 0; t < latent_t; ++t) { + video_duration += video_span(t); + } + for (size_t index = 0; index < condition_videos.size(); ++index) { + const auto& condition = condition_videos[index]; + auto axes = condition_spatial_axes(condition); + int64_t count = condition.shape()[2] * + static_cast(axes.first.size() * axes.second.size()); + bool is_first = keyframe_indices.empty() || keyframe_indices[static_cast(index)] == 0; + float keyframe_t = is_first ? static_cast(text_len) + : static_cast(text_len) + video_duration - FRAME_RESCALE; + for (int64_t t = 0; t < condition.shape()[2]; ++t) { + for (float h : axes.first) { + for (float w : axes.second) { + append_position(keyframe_t, h, w); + } + } + } + layout.sequence_segments.push_back({row, + row + count, + SequenceKind::CONDITION_VIDEO, + static_cast(index)}); + layout.segments.push_back({row, row + count, condition_time_row * 3}); + row += count; + } + } else { + for (const auto& block : reference_blocks) { + const sd::Tensor* ref_video = nullptr; + const sd::Tensor* ref_audio = nullptr; + if (block.video_index >= 0) { + GGML_ASSERT(block.video_index < static_cast(condition_videos.size())); + ref_video = &condition_videos[static_cast(block.video_index)]; + } + if (block.audio_index >= 0) { + GGML_ASSERT(block.audio_index < static_cast(condition_audios.size())); + ref_audio = &condition_audios[static_cast(block.audio_index)]; + } + + float block_end = cursor; + if (block.kind == MiniMaxH3ReferenceKind::AUDIO || + block.kind == MiniMaxH3ReferenceKind::VIDEO_AUDIO) { + GGML_ASSERT(ref_audio != nullptr); + float w_low = w_axis.front(); + float w_high = w_axis.back(); + if (ref_video != nullptr) { + auto axes = condition_spatial_axes(*ref_video); + w_low = axes.second.front(); + w_high = axes.second.back(); + } + int64_t count = ref_audio->shape()[0] * 2; + append_audio_positions(ref_audio->shape()[0], cursor, w_low, w_high); + layout.sequence_segments.push_back({row, + row + count, + SequenceKind::CONDITION_AUDIO, + block.audio_index}); + layout.segments.push_back({row, + row + count, + audio_condition_time_row * 3 + 2}); + row += count; + block_end = std::max(block_end, cursor + static_cast(ref_audio->shape()[0])); + } + + if (block.kind != MiniMaxH3ReferenceKind::AUDIO) { + GGML_ASSERT(ref_video != nullptr); + auto axes = condition_spatial_axes(*ref_video); + int64_t count = ref_video->shape()[2] * + static_cast(axes.first.size() * axes.second.size()); + float video_end = append_video_positions(*ref_video, cursor); + layout.sequence_segments.push_back({row, + row + count, + SequenceKind::CONDITION_VIDEO, + block.video_index}); + layout.segments.push_back({row, row + count, condition_time_row * 3}); + row += count; + block_end = block.kind == MiniMaxH3ReferenceKind::IMAGE + ? std::max(block_end, cursor + 1.f) + : std::max(block_end, video_end); + } + cursor = block_end; + } + } + + int64_t audio_start = row; + append_audio_positions(audio_t, cursor, w_axis.front(), w_axis.back()); + layout.audio_segment = {audio_start, row + audio_t * 2, audio_time_row}; + layout.sequence_segments.push_back({audio_start, + row + audio_t * 2, + SequenceKind::TARGET_AUDIO}); + layout.segments.push_back({audio_start, + row + audio_t * 2, + audio_time_row * 3 + 2}); + row += audio_t * 2; + + int64_t video_start = row; + for (int64_t t = 0; t < latent_t; ++t) { + for (float h : h_axis) { + for (float w : w_axis) { + append_position(cursor, h, w); + } + } + cursor += video_span(t); + } + int64_t video_rows = latent_t * frame_rows; + layout.video_segment = {video_start, video_start + video_rows, video_time_row}; + layout.sequence_segments.push_back({video_start, + video_start + video_rows, + SequenceKind::TARGET_VIDEO}); + layout.segments.push_back({video_start, + video_start + video_rows, + video_time_row * 3}); + return layout; + } + + struct MiniMaxH3Runner : public GGMLRunner { + Config config; + MiniMaxH3Transformer3DModel model; + sd::Tensor video_input_cache; + sd::Tensor audio_input_cache; + sd::Tensor position_input_cache; + sd::Tensor timestep_feature_input_cache; + sd::Tensor curve_index_input_cache; + sd::Tensor curve_upper_index_input_cache; + sd::Tensor curve_fraction_input_cache; + + MiniMaxH3Runner(ggml_backend_t backend, + const String2TensorStorage& tensors, + const std::string& prefix = "model.diffusion_model", + bool offload_params_to_cpu = false) + : GGMLRunner(backend, offload_params_to_cpu), + config(Config::detect_from_weights(tensors, prefix)), + model(config) { + model.init(params_ctx, tensors, prefix); + } + + std::string get_desc() override { + return "minimax_h3"; + } + + void get_param_tensors(std::map& tensors, + const std::string& prefix) { + model.get_param_tensors(tensors, prefix); + } + + std::pair, sd::Tensor> split_av_latents(const sd::Tensor& packed, + int audio_length) const { + GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5); + int64_t spatial = packed.shape()[0] * packed.shape()[1] * packed.shape()[2]; + int64_t video_values = spatial * config.video_latent_channels; + sd::Tensor video({packed.shape()[0], + packed.shape()[1], + packed.shape()[2], + config.video_latent_channels, + 1}); + std::copy_n(packed.data(), static_cast(video_values), video.data()); + if (audio_length <= 0) { + return {video, {}}; + } + int64_t audio_values = audio_length * 2 * config.audio_latent_channels; + GGML_ASSERT(packed.numel() >= video_values + audio_values); + sd::Tensor audio({audio_length, 2, config.audio_latent_channels, 1}); + std::copy_n(packed.data() + video_values, + static_cast(audio_values), + audio.data()); + return {video, audio}; + } + + ggml_tensor* merge_av_latents(ggml_context* ctx, + ggml_tensor* video, + ggml_tensor* audio) const { + int64_t divisor = video->ne[0] * video->ne[1] * video->ne[2]; + int64_t values = ggml_nelements(audio); + int64_t padding = (divisor - values % divisor) % divisor; + audio = ggml_reshape_4d(ctx, ggml_cont(ctx, audio), values, 1, 1, 1); + if (padding > 0) { + audio = ggml_ext_pad(ctx, audio, static_cast(padding), 0, 0, 0); + } + audio = ggml_reshape_4d(ctx, + audio, + video->ne[0], + video->ne[1], + video->ne[2], + (values + padding) / divisor); + return ggml_concat(ctx, video, audio, 3); + } + + ggml_cgraph* build_graph(const sd::Tensor& packed, + const sd::Tensor& timestep, + const sd::Tensor& context_tensor, + const std::vector>& condition_videos, + const std::vector>& condition_audios, + const sd::Tensor& text_tags, + const sd::Tensor& keyframe_indices, + const std::vector& reference_blocks, + int audio_length, + float video_shift, + float audio_shift) { + auto split = split_av_latents(packed, audio_length); + video_input_cache = std::move(split.first); + audio_input_cache = std::move(split.second); + 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); + std::vector condition_inputs; + condition_inputs.reserve(condition_videos.size()); + for (const auto& condition : condition_videos) { + condition_inputs.push_back(make_input(condition)); + } + std::vector audio_condition_inputs; + audio_condition_inputs.reserve(condition_audios.size()); + for (const auto& condition : condition_audios) { + 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); + + position_input_cache = sd::Tensor( + {3, static_cast(layout.positions.size() / 3)}, + layout.positions); + auto positions = make_input(position_input_cache); + + ggml_tensor* timestep_features = nullptr; + ggml_tensor* curve_indices = nullptr; + ggml_tensor* curve_upper_indices = nullptr; + ggml_tensor* curve_fractions = nullptr; + if (config.uses_adaln_curves()) { + std::vector indices(layout.timesteps.size()); + std::vector upper_indices(layout.timesteps.size()); + std::vector fractions(layout.timesteps.size()); + for (size_t i = 0; i < layout.timesteps.size(); ++i) { + float position = std::clamp(layout.timesteps[i], 0.f, 1.f) * (config.adaln_curve_grid - 1); + int index = std::min(static_cast(std::floor(position)), + static_cast(config.adaln_curve_grid - 2)); + indices[i] = index; + upper_indices[i] = index + 1; + fractions[i] = position - index; + } + curve_index_input_cache = sd::Tensor( + {static_cast(indices.size())}, + indices); + curve_upper_index_input_cache = sd::Tensor( + {static_cast(upper_indices.size())}, + upper_indices); + curve_fraction_input_cache = sd::Tensor( + {1, static_cast(fractions.size())}, + fractions); + curve_indices = make_input(curve_index_input_cache); + curve_upper_indices = make_input(curve_upper_index_input_cache); + curve_fractions = make_input(curve_fraction_input_cache); + } else { + timestep_feature_input_cache = sd::Tensor( + {config.timestep_input_dim, static_cast(layout.timesteps.size())}, + timestep_embedding(layout.timesteps, + static_cast(config.timestep_input_dim), + 10000, + true, + 1.f)); + timestep_features = make_input(timestep_feature_input_cache); + } + + auto runner_ctx = get_context(); + auto output = model.forward(&runner_ctx, + video, + audio, + context, + condition_inputs, + audio_condition_inputs, + positions, + timestep_features, + curve_indices, + curve_upper_indices, + curve_fractions, + 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); + ggml_build_forward_expand(graph, merged); + return graph; + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& params) { + GGML_ASSERT(params.x != nullptr && params.timesteps != nullptr && params.context != nullptr); + static const std::vector> empty_conditions; + static const std::vector empty_reference_blocks; + const auto& conditions = params.ref_latents == nullptr ? empty_conditions : *params.ref_latents; + const auto& audio_conditions = params.minimax_reference_audio_latents == nullptr + ? empty_conditions + : *params.minimax_reference_audio_latents; + const auto& reference_blocks = params.minimax_reference_blocks == nullptr + ? empty_reference_blocks + : *params.minimax_reference_blocks; + const sd::Tensor empty_int; + get_graph_cb_t get_graph = [&]() -> ggml_cgraph* { + return build_graph(*params.x, + *params.timesteps, + *params.context, + conditions, + audio_conditions, + params.minimax_text_token_tags == nullptr ? empty_int : *params.minimax_text_token_tags, + params.minimax_keyframe_indices == nullptr ? empty_int : *params.minimax_keyframe_indices, + reference_blocks, + params.minimax_audio_length, + params.minimax_video_sigma_shift, + params.minimax_audio_sigma_shift); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, + n_threads, + false, + false), + params.x->dim()); + } + }; + +} // namespace MiniMaxH3 + +#endif // __ED_DIT_MODELS_MODELS_MINIMAX_H3_FULL_HPP__ diff --git a/src/dit_models/pipelines/dit_pipeline.cpp b/src/dit_models/pipelines/dit_pipeline.cpp index 25e4494f..1a5b6311 100644 --- a/src/dit_models/pipelines/dit_pipeline.cpp +++ b/src/dit_models/pipelines/dit_pipeline.cpp @@ -2,6 +2,7 @@ #include "dit_models/pipelines/flux_kontext_pipeline.hpp" #include "dit_models/pipelines/flux_pipeline.hpp" +#include "dit_models/pipelines/minimax_h3_pipeline.hpp" #include "dit_models/pipelines/qwen_image_edit_pipeline.hpp" #include "dit_models/pipelines/qwen_image_pipeline.hpp" #include "dit_models/pipelines/sd3_pipeline.hpp" @@ -30,6 +31,9 @@ std::unique_ptr create_dit_pipeline(SDVersion version, if (ed_version_is_wan(version)) { return std::make_unique(version); } + if (ed_version_is_minimax_h3(version)) { + return std::make_unique(version); + } const std::string msg = "unsupported DiT pipeline version: " + std::string(ed_version_name(version)); diff --git a/src/dit_models/pipelines/minimax_h3_pipeline.cpp b/src/dit_models/pipelines/minimax_h3_pipeline.cpp new file mode 100644 index 00000000..15de1e7d --- /dev/null +++ b/src/dit_models/pipelines/minimax_h3_pipeline.cpp @@ -0,0 +1,405 @@ +#include "dit_models/pipelines/minimax_h3_pipeline.hpp" + +#include +#include +#include +#include +#include + +#include "dit_models/diffusion_model.hpp" +#include "dit_models/components/autoencoders/minimax_h3_vae.hpp" +#include "dit_models/components/autoencoders/minimax_h3_audio_vae.hpp" +#include "dit_models/components/text_encoders/llm.hpp" +#include "dit_models/models/minimax_h3_full.hpp" +#include "utils/rng_philox.hpp" +#include "utils/util.h" + +namespace edgedit { +namespace { + +bool set_minimax_error(std::string* error, const char* message) { + if (error != nullptr) { + *error = message; + } + LOG_ERROR("%s", message); + return false; +} + +uint8_t h3_to_u8(float value) { + return static_cast(std::round(std::clamp(value, 0.0f, 1.0f) * 255.0f)); +} + +int64_t h3_resolve_seed(int64_t seed) { + return seed >= 0 ? seed : static_cast(std::time(nullptr)); +} + +float h3_discrete_flow_sigma(int step, int steps, float shift) { + const float timestep = 999.0f - 999.0f * static_cast(step) / static_cast(steps); + const float unit_timestep = (timestep + 1.0f) / 1000.0f; + return shift == 1.0f ? unit_timestep + : shift * unit_timestep / (1.0f + (shift - 1.0f) * unit_timestep); +} + +struct MiniMaxH3DetectedConfig { + int64_t hidden_size = 0; + int64_t num_layers = 0; + int64_t token_refiner_num_layers = 0; + int64_t num_attention_heads = 0; + int64_t attention_head_dim = 0; + int64_t ffn_hidden_size = 0; + int64_t video_latent_channels = 0; + int64_t audio_latent_channels = 0; + int64_t text_dim = 0; + int64_t adaln_curve_grid = 0; +}; + +int64_t count_minimax_blocks(const String2TensorStorage& tensors, const std::string& prefix) { + std::set indices; + for (const auto& item : tensors) { + const std::string& name = item.first; + if (!starts_with(name, prefix)) { + continue; + } + const size_t begin = prefix.size(); + const size_t end = name.find('.', begin); + if (end != std::string::npos) { + indices.insert(std::atoi(name.substr(begin, end - begin).c_str())); + } + } + return static_cast(indices.size()); +} + +MiniMaxH3DetectedConfig detect_minimax_config(const String2TensorStorage& tensors) { + MiniMaxH3DetectedConfig config; + const std::string prefix = "model.diffusion_model"; + auto find = [&](const std::string& suffix) -> const TensorStorage* { + auto it = tensors.find(prefix + "." + suffix); + return it == tensors.end() ? nullptr : &it->second; + }; + if (const auto* weight = find("video_patch_proj.weight")) { + config.video_latent_channels = weight->ne[0] / 4; + config.hidden_size = weight->ne[1]; + } + if (const auto* weight = find("audio_patch_proj.weight")) { + config.audio_latent_channels = weight->ne[0]; + } + config.num_layers = count_minimax_blocks(tensors, prefix + ".blocks."); + config.token_refiner_num_layers = count_minimax_blocks(tensors, prefix + ".token_refiner.blocks."); + if (const auto* weight = find("blocks.0.attn.q_norm.weight")) { + config.attention_head_dim = weight->ne[0]; + } + if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) { + config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim); + } + if (const auto* weight = find("blocks.0.mlp.fc1.weight")) { + config.ffn_hidden_size = weight->ne[1] / 2; + } + if (const auto* weight = find("condition_proj.weight")) { + config.text_dim = weight->ne[0]; + } + if (const auto* table = find("adaln_t_table")) { + config.adaln_curve_grid = table->ne[1]; + } + return config; +} + +} // namespace + +MiniMaxH3Pipeline::MiniMaxH3Pipeline(SDVersion version) + : version_(version) { +} + +MiniMaxH3Pipeline::~MiniMaxH3Pipeline() { + if (sentinel_ctx_ != nullptr) { + ggml_free(sentinel_ctx_); + sentinel_ctx_ = nullptr; + sentinel_tensor_ = nullptr; + } +} + +bool MiniMaxH3Pipeline::prepare(const ed_context_params_t& params, + ModelRuntime& runtime, + const ModelLoader& loader, + PipelineTensorRegistry& registry, + std::string* error) { + (void)params; + runtime_ = &runtime; + registry.clear(); + const MiniMaxH3DetectedConfig config = detect_minimax_config(loader.get_tensor_storage_map()); + if (config.num_layers <= 0 || config.hidden_size <= 0 || + config.video_latent_channels <= 0 || config.audio_latent_channels <= 0) { + return set_minimax_error(error, "MiniMax-H3 diffusion transformer signature is incomplete"); + } + + const bool diffusion_offload = runtime.dit_offload_params_to_cpu(); + diffusion_ = std::make_unique(runtime.backend(), + loader.get_tensor_storage_map(), + "model.diffusion_model", + diffusion_offload); + diffusion_->set_max_graph_vram_bytes(runtime.max_graph_vram_bytes()); + diffusion_->set_flash_attention_enabled(runtime.flash_attention()); + if (auto process_group = runtime.graph_process_group_ref()) { + diffusion_->set_process_group(process_group); + } + + diffusion_->alloc_params_buffer(); + diffusion_->get_param_tensors(registry.tensors(), "model.diffusion_model"); + + const bool text_offload = runtime.clip_offload_params_to_cpu(); + conditioner_ = std::make_unique(LLM::LLMArch::QWEN3, + runtime.clip_backend(), + text_offload, + loader.get_tensor_storage_map(), + "text_encoders.llm", + false); + conditioner_->alloc_params_buffer(); + conditioner_->get_param_tensors(registry.tensors(), "text_encoders.llm"); + + vae_ = std::make_unique(runtime.vae_backend(), + runtime.vae_offload_params_to_cpu(), + loader.get_tensor_storage_map(), + "first_stage_model"); + vae_->set_max_graph_vram_bytes(runtime.max_graph_vram_bytes()); + vae_->set_flash_attention_enabled(runtime.flash_attention()); + vae_->alloc_params_buffer(); + vae_->get_param_tensors(registry.tensors(), "first_stage_model"); + + audio_vae_ = std::make_unique(runtime.vae_backend(), + runtime.vae_offload_params_to_cpu(), + loader.get_tensor_storage_map(), + "audio_vae"); + audio_vae_->set_max_graph_vram_bytes(runtime.max_graph_vram_bytes()); + audio_vae_->set_flash_attention_enabled(runtime.flash_attention()); + audio_vae_->alloc_params_buffer(); + audio_vae_->get_param_tensors(registry.tensors(), "audio_vae"); + + registry.ignore_prefix("first_stage_model.encoder."); + registry.ignore_prefix("text_encoders.llm.visual."); + registry.ignore_prefix("text_encoders.llm."); + + if (sentinel_ctx_ == nullptr) { + ggml_init_params init_params{}; + init_params.mem_size = ggml_tensor_overhead() + 1024; + init_params.mem_buffer = nullptr; + init_params.no_alloc = false; + sentinel_ctx_ = ggml_init(init_params); + if (sentinel_ctx_ == nullptr) { + return set_minimax_error(error, "failed to allocate MiniMax-H3 sentinel tensor context"); + } + sentinel_tensor_ = ggml_new_tensor_1d(sentinel_ctx_, GGML_TYPE_F32, 1); + } + registry.add("__ed_minimax_h3_sentinel.weight", sentinel_tensor_); + + ready_ = true; + LOG_INFO("MiniMax-H3 detected: layers=%lld token_refiner_layers=%lld hidden=%lld heads=%lld head_dim=%lld ffn=%lld video_latent=%lld audio_latent=%lld text_dim=%lld adaln_curve_grid=%lld", + (long long)config.num_layers, + (long long)config.token_refiner_num_layers, + (long long)config.hidden_size, + (long long)config.num_attention_heads, + (long long)config.attention_head_dim, + (long long)config.ffn_hidden_size, + (long long)config.video_latent_channels, + (long long)config.audio_latent_channels, + (long long)config.text_dim, + (long long)config.adaln_curve_grid); + return true; +} + +void MiniMaxH3Pipeline::mark_ready() { + ready_ = true; +} + +ed_status_t MiniMaxH3Pipeline::generate_image(const ed_image_generation_params_t*, + ed_image_batch_t*, + std::string* error) { + set_minimax_error(error, "MiniMax-H3 supports video generation only"); + return ED_STATUS_UNSUPPORTED; +} + +bool MiniMaxH3Pipeline::build_text_context(const char* prompt, + sd::Tensor* context, + sd::Tensor* token_tags, + std::string* error) { + if (context == nullptr || token_tags == nullptr || conditioner_ == nullptr || runtime_ == nullptr) { + return set_minimax_error(error, "MiniMax-H3 text conditioner is not initialized"); + } + const std::string text = "<|im_start|>user\n" + + std::string(prompt == nullptr ? "" : prompt) + + "<|im_end|>\n<|im_start|>assistant\n"; + const std::vector tokens = conditioner_->tokenizer->tokenize(text, nullptr, true, 0, 4096, false); + if (tokens.empty()) { + return set_minimax_error(error, "MiniMax-H3 prompt tokenization produced no tokens"); + } + sd::Tensor ids({static_cast(tokens.size())}, tokens); + *context = conditioner_->model.compute(runtime_->n_threads(), ids, {}, {}, {50}); + if (context->empty()) { + return set_minimax_error(error, "MiniMax-H3 text encoder compute failed"); + } + *token_tags = sd::Tensor({context->shape()[1]}, + std::vector(static_cast(context->shape()[1]), 1)); + return true; +} + +ed_status_t MiniMaxH3Pipeline::decode_video_latent(const sd::Tensor& latent, + int requested_frames, + ed_video_t* out, + std::string* error) { + sd::Tensor vae_latent = vae_->diffusion_to_vae_latents(latent); + ed_tiling_params_t tiling = runtime_->vae_tiling(); + if (tiling.force_disable) { + tiling.enabled = false; + } + sd::Tensor video = vae_->decode(runtime_->n_threads(), vae_latent, tiling, true); + if (video.empty() || video.dim() != 5) { + set_minimax_error(error, "MiniMax-H3 video VAE decode failed"); + return ED_STATUS_GENERATION_FAILED; + } + const int decoded_frames = static_cast(video.shape()[2]); + const int frames_count = std::max(decoded_frames, requested_frames); + ed_image_t* frames = static_cast(std::calloc(static_cast(frames_count), sizeof(ed_image_t))); + if (frames == nullptr) { + set_minimax_error(error, "failed to allocate MiniMax-H3 output frames"); + return ED_STATUS_OUT_OF_MEMORY; + } + const size_t width = static_cast(video.shape()[0]); + const size_t height = static_cast(video.shape()[1]); + const size_t channels = static_cast(video.shape()[3]); + const size_t pixels = width * height; + for (int frame = 0; frame < frames_count; ++frame) { + frames[frame].width = static_cast(width); + frames[frame].height = static_cast(height); + frames[frame].channels = static_cast(channels); + frames[frame].data = static_cast(std::malloc(pixels * channels)); + if (frames[frame].data == nullptr) { + for (int index = 0; index < frame; ++index) std::free(frames[index].data); + std::free(frames); + set_minimax_error(error, "failed to allocate MiniMax-H3 frame pixels"); + return ED_STATUS_OUT_OF_MEMORY; + } + for (size_t pixel = 0; pixel < pixels; ++pixel) { + for (size_t channel = 0; channel < channels; ++channel) { + frames[frame].data[pixel * channels + channel] = + h3_to_u8(video.index(pixel % width, + pixel / width, + std::min(frame, decoded_frames - 1), + channel, + 0)); + } + } + } + out->frames = frames; + out->frame_count = frames_count; + return ED_STATUS_OK; +} + +bool MiniMaxH3Pipeline::decode_audio_latent(const sd::Tensor& latent, + ed_video_t* out, + std::string* error) { + if (audio_vae_ == nullptr || latent.empty()) { + return set_minimax_error(error, "MiniMax-H3 audio VAE is not initialized"); + } + sd::Tensor waveform = audio_vae_->decode(runtime_->n_threads(), latent); + if (waveform.empty() || waveform.dim() != 4 || waveform.shape()[1] != 2) { + return set_minimax_error(error, "MiniMax-H3 audio VAE decode failed"); + } + const int64_t sample_count = waveform.shape()[0]; + const int channels = static_cast(waveform.shape()[1]); + if (sample_count <= 0 || sample_count > std::numeric_limits::max()) { + return set_minimax_error(error, "MiniMax-H3 audio sample count is invalid"); + } + float* samples = static_cast(std::malloc(static_cast(sample_count) * channels * sizeof(float))); + if (samples == nullptr) { + return set_minimax_error(error, "failed to allocate MiniMax-H3 audio output"); + } + for (int64_t sample = 0; sample < sample_count; ++sample) { + for (int channel = 0; channel < channels; ++channel) { + samples[sample * channels + channel] = std::clamp(waveform.index(sample, channel, 0, 0), -1.0f, 1.0f); + } + } + out->audio = samples; + out->audio_sample_count = static_cast(sample_count); + out->audio_channels = channels; + out->audio_sample_rate = 32000; + return true; +} + +ed_status_t MiniMaxH3Pipeline::generate_video(const ed_video_generation_params_t* params, + ed_video_t* out, + std::string* error) { + if (out == nullptr || params == nullptr) { + set_minimax_error(error, "MiniMax-H3 video parameters or output are null"); + return ED_STATUS_INVALID_ARGUMENT; + } + out->frames = nullptr; + out->frame_count = 0; + if (!ready_ || runtime_ == nullptr || !diffusion_ || !conditioner_ || !vae_) { + set_minimax_error(error, "MiniMax-H3 pipeline is not ready"); + return ED_STATUS_MODEL_LOAD_FAILED; + } + if (params->width <= 0 || params->height <= 0 || params->frames <= 0 || + params->width % 32 != 0 || params->height % 32 != 0) { + set_minimax_error(error, "MiniMax-H3 width and height must be positive multiples of 32"); + return ED_STATUS_INVALID_ARGUMENT; + } + const int frames = std::max(5, params->frames); + if (frames % 17 != 5) { + set_minimax_error(error, "MiniMax-H3 frame count must satisfy 17k + 5"); + return ED_STATUS_INVALID_ARGUMENT; + } + sd::Tensor context; + sd::Tensor token_tags; + if (!build_text_context(params->prompt, &context, &token_tags, error)) return ED_STATUS_GENERATION_FAILED; + const int latent_frames = frames <= 5 ? 2 : ((frames - 5) / 17) * 5 + 2; + const int audio_length = std::max(1, static_cast(std::lround(static_cast(frames) * 40.0 / 24.0))); + const int latent_width = params->width / 16; + const int latent_height = params->height / 16; + const int64_t video_spatial = static_cast(latent_width) * latent_height * latent_frames; + sd::Tensor video = sd::zeros({latent_width, latent_height, latent_frames, 24, 1}); + sd::Tensor audio = sd::zeros({audio_length, 2, 32, 1}); + const int64_t audio_channels = (audio.numel() + video_spatial - 1) / video_spatial; + sd::Tensor packed({latent_width, + latent_height, + latent_frames, + 24 + audio_channels, + 1}); + auto rng = std::make_shared(static_cast(h3_resolve_seed(params->seed))); + packed = sd::randn_like(packed, rng); + const int steps = params->sample.steps > 0 ? params->sample.steps : 20; + const float video_sigma_shift = params->sample.flow_shift > 0.0f ? params->sample.flow_shift : 12.0f; + for (int step = 0; step < steps; ++step) { + const float sigma = h3_discrete_flow_sigma(step, steps, video_sigma_shift); + const float sigma_next = step + 1 == steps ? 0.0f + : h3_discrete_flow_sigma(step + 1, steps, video_sigma_shift); + sd::Tensor timestep({1}, {sigma * 1000.0f}); + DiffusionParams diffusion_params{}; + diffusion_params.x = &packed; + diffusion_params.timesteps = ×tep; + diffusion_params.context = &context; + diffusion_params.minimax_text_token_tags = &token_tags; + diffusion_params.minimax_audio_length = audio_length; + diffusion_params.minimax_video_sigma_shift = video_sigma_shift; + diffusion_params.minimax_audio_sigma_shift = 3.0f; + sd::Tensor velocity = diffusion_->compute(runtime_->n_threads(), diffusion_params); + if (velocity.empty()) { + set_minimax_error(error, "MiniMax-H3 diffusion compute failed"); + return ED_STATUS_GENERATION_FAILED; + } + packed += velocity * (sigma_next - sigma); + } + auto av = diffusion_->split_av_latents(packed, audio_length); + ed_status_t status = decode_video_latent(av.first, frames, out, error); + if (status != ED_STATUS_OK) { + return status; + } + if (!decode_audio_latent(av.second, out, error)) { + ed_free_video(out); + return ED_STATUS_GENERATION_FAILED; + } + return ED_STATUS_OK; +} + +ed_scheduler_t MiniMaxH3Pipeline::default_scheduler(ed_sampler_t) const { + return ED_SCHEDULER_DISCRETE; +} + +} // namespace edgedit diff --git a/src/dit_models/pipelines/minimax_h3_pipeline.hpp b/src/dit_models/pipelines/minimax_h3_pipeline.hpp new file mode 100644 index 00000000..aafac715 --- /dev/null +++ b/src/dit_models/pipelines/minimax_h3_pipeline.hpp @@ -0,0 +1,79 @@ +#pragma once + +#include +#include + +#include "dit_models/pipelines/dit_pipeline.hpp" + +namespace MiniMaxH3 { +struct MiniMaxH3Runner; +} +namespace LLM { +struct LLMEmbedder; +} +namespace MiniMaxH3VAE { +struct MiniMaxH3VideoVAERunner; +} +namespace MiniMaxH3Audio { +struct AudioVAERunner; +} + +namespace edgedit { + +class MiniMaxH3Pipeline final : public DiTPipeline { +public: + explicit MiniMaxH3Pipeline(SDVersion version = VERSION_MINIMAX_H3); + ~MiniMaxH3Pipeline() override; + + const char* name() const override { return "minimax-h3"; } + + bool prepare(const ed_context_params_t& params, + ModelRuntime& runtime, + const ModelLoader& loader, + PipelineTensorRegistry& registry, + std::string* error) override; + + void mark_ready() override; + + ed_status_t generate_image(const ed_image_generation_params_t* params, + ed_image_batch_t* out, + std::string* error) override; + + ed_status_t generate_video(const ed_video_generation_params_t* params, + ed_video_t* out, + std::string* error) override; + + SDVersion version() const override { return version_; } + bool ready() const override { return ready_; } + + bool supports_image_generation() const override { return false; } + bool supports_video_generation() const override { return ready_; } + + ed_sampler_t default_sample_method() const override { return ED_SAMPLER_EULER; } + ed_scheduler_t default_scheduler(ed_sampler_t method) const override; + +private: + bool ready_ = false; + ModelRuntime* runtime_ = nullptr; + SDVersion version_ = VERSION_MINIMAX_H3; + std::unique_ptr diffusion_; + std::unique_ptr conditioner_; + std::unique_ptr vae_; + std::unique_ptr audio_vae_; + ggml_context* sentinel_ctx_ = nullptr; + ggml_tensor* sentinel_tensor_ = nullptr; + + bool build_text_context(const char* prompt, + sd::Tensor* context, + sd::Tensor* token_tags, + std::string* error); + ed_status_t decode_video_latent(const sd::Tensor& latent, + int requested_frames, + ed_video_t* out, + std::string* error); + bool decode_audio_latent(const sd::Tensor& latent, + ed_video_t* out, + std::string* error); +}; + +} // namespace edgedit diff --git a/src/edge_dit.cpp b/src/edge_dit.cpp index 6d94329c..d6bde5f4 100644 --- a/src/edge_dit.cpp +++ b/src/edge_dit.cpp @@ -336,6 +336,11 @@ void ed_free_video(ed_video_t* video) { std::free(video->frames); video->frames = nullptr; video->frame_count = 0; + std::free(video->audio); + video->audio = nullptr; + video->audio_sample_count = 0; + video->audio_channels = 0; + video->audio_sample_rate = 0; } const char * ed_get_last_error(const ed_context_t * ctx) { diff --git a/src/utils/name_conversion.cpp b/src/utils/name_conversion.cpp index d5c183ab..63243caf 100644 --- a/src/utils/name_conversion.cpp +++ b/src/utils/name_conversion.cpp @@ -177,6 +177,45 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix) return name; } +std::string convert_qwen3_vl_vision_name(std::string name) { + static const std::vector> qwen3_vl_deepstack_name_map{ + {"v.deepstack_merger_list.", "deepstack_merger_list."}, + {"v.deepstack.5.", "deepstack_merger_list.0."}, + {"v.deepstack.8.", "deepstack_merger_list.0."}, + {"v.deepstack.11.", "deepstack_merger_list.1."}, + {"v.deepstack.16.", "deepstack_merger_list.1."}, + {"v.deepstack.17.", "deepstack_merger_list.2."}, + {"v.deepstack.24.", "deepstack_merger_list.2."}, + {"fc1.", "linear_fc1."}, + {"fc2.", "linear_fc2."}, + {"ffn_up.", "linear_fc1."}, + {"ffn_down.", "linear_fc2."}, + {"ffn_norm.", "norm."}, + }; + static const std::vector> qwen3_vl_vision_name_map{ + {"mm.0.", "merger.linear_fc1."}, + {"mm.2.", "merger.linear_fc2."}, + {"v.post_ln.", "merger.norm."}, + {"v.position_embd.weight", "pos_embed.weight"}, + {"v.patch_embd.weight.1", "patch_embed.proj.1.weight"}, + {"v.patch_embd.weight", "patch_embed.proj.0.weight"}, + {"v.patch_embd.bias", "patch_embed.bias"}, + {"v.blk.", "blocks."}, + {"attn_qkv.", "attn.qkv."}, + {"attn_out.", "attn.proj."}, + {"ffn_up.", "mlp.linear_fc1."}, + {"ffn_down.", "mlp.linear_fc2."}, + {"ln1.", "norm1."}, + {"ln2.", "norm2."}, + }; + if (contains(name, "v.deepstack_merger_list.") || contains(name, "v.deepstack.")) { + replace_with_name_map(name, qwen3_vl_deepstack_name_map); + return name; + } + replace_with_name_map(name, qwen3_vl_vision_name_map); + return name; +} + // ref: https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py std::string convert_diffusers_unet_to_original_sd1(std::string name) { // (stable-diffusion, HF Diffusers) @@ -878,6 +917,9 @@ std::string convert_diffusers_qwen_image_vae_to_wan(std::string name) { } std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) { + if (ed_version_is_minimax_h3(version)) { + return name; + } if (ed_version_is_qwen_image(version) || ed_version_is_qwen_image_edit(version) || ed_version_is_wan(version)) { return convert_diffusers_qwen_image_vae_to_wan(name); } @@ -1201,6 +1243,16 @@ std::string convert_tensor_name(std::string name, SDVersion version) { replace_with_prefix_map(name, prefix_map); + if (ed_version_is_minimax_h3(version)) { + const std::string hf_vision_prefix = "text_encoders.llm.model.visual."; + if (starts_with(name, hf_vision_prefix)) { + name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size()); + } + if (starts_with(name, "text_encoders.llm.visual.")) { + name = convert_qwen3_vl_vision_name(std::move(name)); + } + } + // diffusion model { for (const auto& prefix : diffuison_model_prefix_vec) { From e1e96086f347c3f896616990ea7dcfc208ed3e24 Mon Sep 17 00:00:00 2001 From: Komorebi623 Date: Thu, 6 Aug 2026 17:38:17 +0800 Subject: [PATCH 2/4] feat: add MiniMax-H3 keyframe conditioning --- examples/cli/cli_common.hpp | 7 +- examples/cli/main.cpp | 21 +++ .../pipelines/minimax_h3_pipeline.cpp | 164 ++++++++++++++++-- 3 files changed, 173 insertions(+), 19 deletions(-) diff --git a/examples/cli/cli_common.hpp b/examples/cli/cli_common.hpp index 301b3c52..d0a24c83 100644 --- a/examples/cli/cli_common.hpp +++ b/examples/cli/cli_common.hpp @@ -463,6 +463,7 @@ struct FluxCliArgs { const char* prompt = nullptr; const char* negative_prompt = nullptr; const char* image_path = nullptr; + const char* end_image_path = nullptr; const char* output_path = "output.png"; const char* video_format = nullptr; const char* backend = nullptr; @@ -593,8 +594,10 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { std::strcmp(key, "--init_img") == 0 || std::strcmp(key, "-i") == 0) { args->image_path = require_value(key); if (!args->image_path) return false; - } else if (std::strcmp(key, "--end-img") == 0 || std::strcmp(key, "--end_img") == 0 || - std::strcmp(key, "--ref-image") == 0 || std::strcmp(key, "--ref_image") == 0 || + } else if (std::strcmp(key, "--end-img") == 0 || std::strcmp(key, "--end_img") == 0) { + args->end_image_path = require_value(key); + if (!args->end_image_path) return false; + } else if (std::strcmp(key, "--ref-image") == 0 || std::strcmp(key, "--ref_image") == 0 || std::strcmp(key, "--ref-video") == 0 || std::strcmp(key, "--ref_video") == 0 || std::strcmp(key, "--ref-audio") == 0 || std::strcmp(key, "--ref_audio") == 0 || std::strcmp(key, "-r") == 0) { diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index e1a089de..19d6355e 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -756,6 +756,8 @@ int main(int argc, char** argv) { if (args.video) { ed_video_generation_params_t gen_params; ed_video_generation_params_init(&gen_params); + ed_image_t init_image = {}; + ed_image_t end_image = {}; gen_params.prompt = args.prompt; gen_params.negative_prompt = args.negative_prompt; @@ -771,6 +773,21 @@ int main(int argc, char** argv) { gen_params.sample.distilled_guidance = args.guidance; gen_params.sample.flow_shift = args.flow_shift; apply_cache_args(args, &gen_params.sample); + if (args.image_path != nullptr && std::strlen(args.image_path) > 0) { + if (!load_image(args.image_path, &init_image)) { + ed_free_context(ctx); + return 6; + } + gen_params.init_image = &init_image; + } + if (args.end_image_path != nullptr && std::strlen(args.end_image_path) > 0) { + if (!load_image(args.end_image_path, &end_image)) { + ed_free_image(&init_image); + ed_free_context(ctx); + return 6; + } + gen_params.end_image = &end_image; + } ed_video_t output; auto ed_wall_gen0 = ed_wall_clock::now(); @@ -782,6 +799,8 @@ int main(int argc, char** argv) { if (err != nullptr && std::strlen(err) > 0) { std::fprintf(stderr, "last error: %s\n", err); } + ed_free_image(&init_image); + ed_free_image(&end_image); ed_free_context(ctx); return 3; } @@ -817,6 +836,8 @@ int main(int argc, char** argv) { } ed_free_video(&output); + ed_free_image(&init_image); + ed_free_image(&end_image); } else { ed_image_generation_params_t gen_params; ed_image_generation_params_init(&gen_params); diff --git a/src/dit_models/pipelines/minimax_h3_pipeline.cpp b/src/dit_models/pipelines/minimax_h3_pipeline.cpp index 15de1e7d..fa171a8d 100644 --- a/src/dit_models/pipelines/minimax_h3_pipeline.cpp +++ b/src/dit_models/pipelines/minimax_h3_pipeline.cpp @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include #include @@ -34,12 +36,96 @@ int64_t h3_resolve_seed(int64_t seed) { } float h3_discrete_flow_sigma(int step, int steps, float shift) { - const float timestep = 999.0f - 999.0f * static_cast(step) / static_cast(steps); + if (steps <= 1) { + return 1.0f; + } + const float timestep = 999.0f - 999.0f * static_cast(step) / static_cast(steps - 1); const float unit_timestep = (timestep + 1.0f) / 1000.0f; return shift == 1.0f ? unit_timestep : shift * unit_timestep / (1.0f + (shift - 1.0f) * unit_timestep); } +bool h3_trace_enabled() { + const char* value = std::getenv("ED_MINIMAX_H3_TRACE"); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + +void h3_trace_tensor(const char* name, const sd::Tensor& tensor) { + if (!h3_trace_enabled()) { + return; + } + if (tensor.empty()) { + LOG_INFO("minimax-h3 trace %s: empty", name); + return; + } + + uint64_t hash = 1469598103934665603ULL; + double sum = 0.0; + double squared_sum = 0.0; + float minimum = std::numeric_limits::infinity(); + float maximum = -std::numeric_limits::infinity(); + for (float value : tensor.values()) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + hash ^= bits; + hash *= 1099511628211ULL; + sum += value; + squared_sum += static_cast(value) * value; + minimum = std::min(minimum, value); + maximum = std::max(maximum, value); + } + const double count = static_cast(tensor.numel()); + LOG_INFO("minimax-h3 trace %s: shape=%s n=%lld hash=%016llx mean=%.8g rms=%.8g min=%.8g max=%.8g", + name, + sd::tensor_shape_to_string(tensor.shape()).c_str(), + static_cast(tensor.numel()), + static_cast(hash), + sum / count, + std::sqrt(squared_sum / count), + minimum, + maximum); +} + +sd::Tensor h3_pack_audio_and_video_latents(const sd::Tensor& video, + const sd::Tensor& audio) { + if (audio.empty()) { + return video; + } + GGML_ASSERT(video.dim() == 5 && video.shape()[4] == 1); + GGML_ASSERT(audio.dim() == 4 && audio.shape()[3] == 1); + + const int64_t spatial_size = video.shape()[0] * video.shape()[1] * video.shape()[2]; + const int64_t extra_channels = (audio.numel() + spatial_size - 1) / spatial_size; + std::vector packed_shape = video.shape(); + packed_shape[3] += extra_channels; + sd::Tensor packed = sd::zeros(packed_shape); + std::copy_n(video.data(), video.numel(), packed.data()); + std::copy_n(audio.data(), audio.numel(), packed.data() + video.numel()); + return packed; +} + +sd::Tensor h3_image_to_tensor(const ed_image_t& image, int width, int height) { + if (image.data == nullptr || image.width <= 0 || image.height <= 0 || image.channels <= 0) { + return {}; + } + const int source_width = static_cast(image.width); + const int source_height = static_cast(image.height); + const int source_channels = static_cast(image.channels); + sd::Tensor tensor({width, height, 3, 1}); + for (int y = 0; y < height; ++y) { + const int source_y = std::min(source_height - 1, static_cast((static_cast(y) * source_height) / height)); + for (int x = 0; x < width; ++x) { + const int source_x = std::min(source_width - 1, static_cast((static_cast(x) * source_width) / width)); + const uint8_t* pixel = image.data + + (static_cast(source_y) * source_width + static_cast(source_x)) * source_channels; + for (int channel = 0; channel < 3; ++channel) { + tensor.index(x, y, channel, 0) = static_cast(pixel[std::min(channel, source_channels - 1)]) / 255.0f; + } + } + } + return tensor; +} + struct MiniMaxH3DetectedConfig { int64_t hidden_size = 0; int64_t num_layers = 0; @@ -164,14 +250,21 @@ bool MiniMaxH3Pipeline::prepare(const ed_context_params_t& params, vae_->alloc_params_buffer(); vae_->get_param_tensors(registry.tensors(), "first_stage_model"); - audio_vae_ = std::make_unique(runtime.vae_backend(), - runtime.vae_offload_params_to_cpu(), - loader.get_tensor_storage_map(), - "audio_vae"); - audio_vae_->set_max_graph_vram_bytes(runtime.max_graph_vram_bytes()); - audio_vae_->set_flash_attention_enabled(runtime.flash_attention()); - audio_vae_->alloc_params_buffer(); - audio_vae_->get_param_tensors(registry.tensors(), "audio_vae"); + const bool has_audio_vae = std::any_of(loader.get_tensor_storage_map().begin(), + loader.get_tensor_storage_map().end(), + [](const auto& item) { return starts_with(item.first, "audio_vae."); }); + if (has_audio_vae) { + audio_vae_ = std::make_unique(runtime.vae_backend(), + runtime.vae_offload_params_to_cpu(), + loader.get_tensor_storage_map(), + "audio_vae"); + audio_vae_->set_max_graph_vram_bytes(runtime.max_graph_vram_bytes()); + audio_vae_->set_flash_attention_enabled(runtime.flash_attention()); + audio_vae_->alloc_params_buffer(); + audio_vae_->get_param_tensors(registry.tensors(), "audio_vae"); + } else { + LOG_INFO("MiniMax-H3 audio VAE not provided; generated video will have no decoded audio track"); + } registry.ignore_prefix("first_stage_model.encoder."); registry.ignore_prefix("text_encoders.llm.visual."); @@ -237,6 +330,7 @@ bool MiniMaxH3Pipeline::build_text_context(const char* prompt, } *token_tags = sd::Tensor({context->shape()[1]}, std::vector(static_cast(context->shape()[1]), 1)); + h3_trace_tensor("text_context", *context); return true; } @@ -353,17 +447,40 @@ ed_status_t MiniMaxH3Pipeline::generate_video(const ed_video_generation_params_t const int audio_length = std::max(1, static_cast(std::lround(static_cast(frames) * 40.0 / 24.0))); const int latent_width = params->width / 16; const int latent_height = params->height / 16; - const int64_t video_spatial = static_cast(latent_width) * latent_height * latent_frames; + std::vector> keyframe_latents; + std::vector keyframe_indices; + auto add_keyframe = [&](const ed_image_t* image, int32_t frame_index, const char* name) -> bool { + if (image == nullptr || image->data == nullptr) { + return true; + } + sd::Tensor image_tensor = h3_image_to_tensor(*image, params->width, params->height); + if (image_tensor.empty()) { + return set_minimax_error(error, "MiniMax-H3 keyframe image is invalid"); + } + sd::Tensor video_image = image_tensor.reshape({params->width, params->height, 1, 3, 1}); + sd::Tensor vae_latent = vae_->encode(runtime_->n_threads(), video_image, runtime_->vae_tiling()); + if (vae_latent.empty()) { + return set_minimax_error(error, "MiniMax-H3 keyframe VAE encode failed"); + } + sd::Tensor latent = vae_->vae_to_diffusion_latents(vae_latent); + auto condition_rng = std::make_shared(static_cast(h3_resolve_seed(params->seed))); + latent = latent * MiniMaxH3::VISUAL_COND_TIMESTEP + + sd::randn_like(latent, condition_rng) * (1.0f - MiniMaxH3::VISUAL_COND_TIMESTEP); + h3_trace_tensor((std::string(name) + "_keyframe_latent").c_str(), latent); + keyframe_latents.push_back(std::move(latent)); + keyframe_indices.push_back(frame_index); + return true; + }; + if (!add_keyframe(params->init_image, 0, "init") || + !add_keyframe(params->end_image, frames - 1, "end")) { + return ED_STATUS_GENERATION_FAILED; + } sd::Tensor video = sd::zeros({latent_width, latent_height, latent_frames, 24, 1}); sd::Tensor audio = sd::zeros({audio_length, 2, 32, 1}); - const int64_t audio_channels = (audio.numel() + video_spatial - 1) / video_spatial; - sd::Tensor packed({latent_width, - latent_height, - latent_frames, - 24 + audio_channels, - 1}); + sd::Tensor packed = h3_pack_audio_and_video_latents(video, audio); auto rng = std::make_shared(static_cast(h3_resolve_seed(params->seed))); packed = sd::randn_like(packed, rng); + h3_trace_tensor("initial_packed_noise", packed); const int steps = params->sample.steps > 0 ? params->sample.steps : 20; const float video_sigma_shift = params->sample.flow_shift > 0.0f ? params->sample.flow_shift : 12.0f; for (int step = 0; step < steps; ++step) { @@ -376,6 +493,12 @@ ed_status_t MiniMaxH3Pipeline::generate_video(const ed_video_generation_params_t diffusion_params.timesteps = ×tep; diffusion_params.context = &context; diffusion_params.minimax_text_token_tags = &token_tags; + diffusion_params.ref_latents = keyframe_latents.empty() ? nullptr : &keyframe_latents; + sd::Tensor keyframe_index_tensor; + if (!keyframe_indices.empty()) { + keyframe_index_tensor = sd::Tensor({static_cast(keyframe_indices.size())}, keyframe_indices); + diffusion_params.minimax_keyframe_indices = &keyframe_index_tensor; + } diffusion_params.minimax_audio_length = audio_length; diffusion_params.minimax_video_sigma_shift = video_sigma_shift; diffusion_params.minimax_audio_sigma_shift = 3.0f; @@ -384,14 +507,21 @@ ed_status_t MiniMaxH3Pipeline::generate_video(const ed_video_generation_params_t set_minimax_error(error, "MiniMax-H3 diffusion compute failed"); return ED_STATUS_GENERATION_FAILED; } + if (h3_trace_enabled()) { + LOG_INFO("minimax-h3 trace step=%d sigma=%.8g sigma_next=%.8g", step, sigma, sigma_next); + h3_trace_tensor(("step_" + std::to_string(step) + "_velocity").c_str(), velocity); + } packed += velocity * (sigma_next - sigma); + if (h3_trace_enabled()) { + h3_trace_tensor(("step_" + std::to_string(step) + "_packed").c_str(), packed); + } } auto av = diffusion_->split_av_latents(packed, audio_length); ed_status_t status = decode_video_latent(av.first, frames, out, error); if (status != ED_STATUS_OK) { return status; } - if (!decode_audio_latent(av.second, out, error)) { + if (audio_vae_ != nullptr && !decode_audio_latent(av.second, out, error)) { ed_free_video(out); return ED_STATUS_GENERATION_FAILED; } From b044f8facbe8c32762dafa1fd91d5abe9f738837 Mon Sep 17 00:00:00 2001 From: Komorebi623 Date: Fri, 7 Aug 2026 10:48:44 +0800 Subject: [PATCH 3/4] feat: add MiniMax-H3 Ref2VA support --- examples/cli/cli_common.hpp | 28 +- examples/cli/main.cpp | 125 +++++- include/edge-dit.h | 21 + .../autoencoders/minimax_h3_audio_vae.hpp | 172 +++++++- .../components/text_encoders/llm.hpp | 403 +++++++++++++++--- .../pipelines/minimax_h3_pipeline.cpp | 268 +++++++++++- .../pipelines/minimax_h3_pipeline.hpp | 6 + 7 files changed, 939 insertions(+), 84 deletions(-) diff --git a/examples/cli/cli_common.hpp b/examples/cli/cli_common.hpp index d0a24c83..c8638cad 100644 --- a/examples/cli/cli_common.hpp +++ b/examples/cli/cli_common.hpp @@ -464,6 +464,10 @@ struct FluxCliArgs { const char* negative_prompt = nullptr; const char* image_path = nullptr; const char* end_image_path = nullptr; + std::vector ref_image_paths; + std::vector ref_video_paths; + std::vector ref_video_audio_paths; + std::vector ref_audio_paths; const char* output_path = "output.png"; const char* video_format = nullptr; const char* backend = nullptr; @@ -598,11 +602,22 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { args->end_image_path = require_value(key); if (!args->end_image_path) return false; } else if (std::strcmp(key, "--ref-image") == 0 || std::strcmp(key, "--ref_image") == 0 || - std::strcmp(key, "--ref-video") == 0 || std::strcmp(key, "--ref_video") == 0 || - std::strcmp(key, "--ref-audio") == 0 || std::strcmp(key, "--ref_audio") == 0 || std::strcmp(key, "-r") == 0) { - const char* ignored_path = require_value(key); - if (!ignored_path) return false; + const char* path = require_value(key); + if (!path) return false; + args->ref_image_paths.emplace_back(path); + } else if (std::strcmp(key, "--ref-video") == 0 || std::strcmp(key, "--ref_video") == 0) { + const char* path = require_value(key); + if (!path) return false; + args->ref_video_paths.emplace_back(path); + } else if (std::strcmp(key, "--ref-video-audio") == 0 || std::strcmp(key, "--ref_video_audio") == 0) { + const char* path = require_value(key); + if (!path) return false; + args->ref_video_audio_paths.emplace_back(path); + } else if (std::strcmp(key, "--ref-audio") == 0 || std::strcmp(key, "--ref_audio") == 0) { + const char* path = require_value(key); + if (!path) return false; + args->ref_audio_paths.emplace_back(path); } else if (std::strcmp(key, "--output") == 0 || std::strcmp(key, "-o") == 0) { args->output_path = require_value(key); } else if (std::strcmp(key, "--width") == 0 || std::strcmp(key, "-W") == 0) { @@ -888,6 +903,11 @@ inline bool parse_args(int argc, char** argv, FluxCliArgs* args) { return false; } + if (args->ref_video_audio_paths.size() > args->ref_video_paths.size()) { + std::fprintf(stderr, "each --ref-video-audio needs a corresponding --ref-video\n"); + return false; + } + const std::string video_format = normalized_video_format(args->video_format); if (video_format != "auto" && video_format != "avi" && diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 19d6355e..726e7375 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -14,9 +14,63 @@ #include #include #include +#include namespace fs = std::filesystem; +inline bool load_image(const char* path, ed_image_t* image); + +static uint16_t read_le16(const uint8_t* data) { return static_cast(data[0]) | (static_cast(data[1]) << 8); } +static uint32_t read_le32(const uint8_t* data) { return static_cast(data[0]) | (static_cast(data[1]) << 8) | (static_cast(data[2]) << 16) | (static_cast(data[3]) << 24); } + +static bool load_wav(const std::string& path, std::vector* samples, ed_audio_t* audio) { + std::ifstream file(path, std::ios::binary); + uint8_t header[12]; + if (!file.read(reinterpret_cast(header), sizeof(header)) || std::memcmp(header, "RIFF", 4) != 0 || std::memcmp(header + 8, "WAVE", 4) != 0) return false; + uint16_t format = 0, channels = 0, bits = 0; uint32_t rate = 0, data_size = 0; std::streampos data_pos = -1; + while (file.good()) { uint8_t chunk[8]; if (!file.read(reinterpret_cast(chunk), 8)) break; uint32_t size = read_le32(chunk + 4); auto pos = file.tellg(); + if (std::memcmp(chunk, "fmt ", 4) == 0) { std::vector value(size); if (!file.read(reinterpret_cast(value.data()), size) || size < 16) return false; format = read_le16(value.data()); channels = read_le16(value.data() + 2); rate = read_le32(value.data() + 4); bits = read_le16(value.data() + 14); if (format == 0xfffe && size >= 40) format = read_le16(value.data() + 24); } + else if (std::memcmp(chunk, "data", 4) == 0) { data_pos = pos; data_size = size; file.seekg(size, std::ios::cur); } + else file.seekg(size, std::ios::cur); if (size & 1) file.seekg(1, std::ios::cur); + } + if (data_pos == std::streampos(-1) || channels == 0 || rate == 0 || + !((format == 1 && (bits == 8 || bits == 16 || bits == 24 || bits == 32)) || + (format == 3 && (bits == 32 || bits == 64)))) return false; + const size_t bytes_per = bits / 8, count = data_size / bytes_per; samples->resize(count); file.clear(); file.seekg(data_pos); + std::vector bytes(data_size); if (!file.read(reinterpret_cast(bytes.data()), data_size)) return false; + for (size_t index = 0; index < count; ++index) { + const uint8_t* sample = bytes.data() + index * bytes_per; + if (format == 3 && bits == 32) { + float value; + std::memcpy(&value, sample, sizeof(value)); + (*samples)[index] = std::clamp(value, -1.f, 1.f); + } else if (format == 3) { + double value; + std::memcpy(&value, sample, sizeof(value)); + (*samples)[index] = std::clamp(static_cast(value), -1.f, 1.f); + } else if (bits == 8) { + (*samples)[index] = (static_cast(sample[0]) - 128.f) / 128.f; + } else if (bits == 16) { + (*samples)[index] = static_cast(static_cast(read_le16(sample))) / 32768.f; + } else if (bits == 24) { + int32_t value = static_cast(sample[0]) | + (static_cast(sample[1]) << 8) | + (static_cast(sample[2]) << 16); + if (value & 0x00800000) value |= ~0x00ffffff; + (*samples)[index] = static_cast(value) / 8388608.f; + } else { + (*samples)[index] = static_cast(static_cast(read_le32(sample))) / 2147483648.f; + } + } + *audio = {rate, channels, count / channels, samples->data()}; return true; +} + +static bool load_images_from_dir(const std::string& directory, std::vector* frames) { + if (!fs::is_directory(directory)) return false; std::vector paths; + for (const auto& entry : fs::directory_iterator(directory)) if (entry.is_regular_file()) { auto ext = entry.path().extension().string(); std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower); if (ext == ".png" || ext == ".jpg" || ext == ".jpeg" || ext == ".bmp" || ext == ".webp") paths.push_back(entry.path()); } + std::sort(paths.begin(), paths.end()); frames->resize(paths.size()); for (size_t index = 0; index < paths.size(); ++index) if (!load_image(paths[index].c_str(), &(*frames)[index])) return false; return !frames->empty(); +} + #define STB_IMAGE_IMPLEMENTATION #include "ggml/examples/stb_image.h" @@ -600,14 +654,14 @@ static bool save_video_with_audio(const char* path, const ed_video_t& video, int return save_video(path, video, fps); } const fs::path output(path); - const fs::path video_tmp = output.string() + ".video.tmp.avi"; + const std::string ext = path_extension(path); + const fs::path video_tmp = output.string() + ".video.tmp" + (ext == ".avi" ? ".avi" : ext); const fs::path wav_path = output.string() + ".wav"; if (!save_video(video_tmp.c_str(), video, fps) || !save_wav(wav_path.c_str(), video)) { std::error_code error; fs::remove(video_tmp, error); return false; } - const std::string ext = path_extension(path); const char* audio_codec = ext == ".avi" ? "pcm_s16le" : (ext == ".webm" ? "libopus" : "aac"); const std::string command = shell_quote(find_ffmpeg_binary().c_str()) + " -hide_banner -loglevel error -y -i " + shell_quote(video_tmp.c_str()) + @@ -758,6 +812,19 @@ int main(int argc, char** argv) { ed_video_generation_params_init(&gen_params); ed_image_t init_image = {}; ed_image_t end_image = {}; + std::vector ref_images(args.ref_image_paths.size()); + std::vector> ref_video_frames(args.ref_video_paths.size()); + std::vector> ref_video_audio_samples(args.ref_video_paths.size()); + std::vector ref_videos(args.ref_video_paths.size()); + std::vector> ref_audio_samples(args.ref_audio_paths.size()); + std::vector ref_audios(args.ref_audio_paths.size()); + auto free_video_references = [&]() { + for (auto& frames : ref_video_frames) { + for (ed_image_t& frame : frames) { + ed_free_image(&frame); + } + } + }; gen_params.prompt = args.prompt; gen_params.negative_prompt = args.negative_prompt; @@ -788,6 +855,44 @@ int main(int argc, char** argv) { } gen_params.end_image = &end_image; } + for (size_t index = 0; index < args.ref_image_paths.size(); ++index) { + if (!load_image(args.ref_image_paths[index].c_str(), &ref_images[index])) { + for (ed_image_t& image : ref_images) ed_free_image(&image); + ed_free_image(&init_image); + ed_free_image(&end_image); + ed_free_context(ctx); + return 6; + } + } + if (!ref_images.empty()) { + gen_params.ref_images = ref_images.data(); + gen_params.ref_image_count = static_cast(ref_images.size()); + } + for (size_t index = 0; index < args.ref_audio_paths.size(); ++index) { + if (!load_wav(args.ref_audio_paths[index], &ref_audio_samples[index], &ref_audios[index])) { + std::fprintf(stderr, "failed to load reference WAV '%s'\n", args.ref_audio_paths[index].c_str()); + for (ed_image_t& image : ref_images) ed_free_image(&image); + free_video_references(); + ed_free_image(&init_image); ed_free_image(&end_image); ed_free_context(ctx); return 6; + } + } + if (!ref_audios.empty()) { gen_params.ref_audios = ref_audios.data(); gen_params.ref_audio_count = static_cast(ref_audios.size()); } + for (size_t index = 0; index < args.ref_video_paths.size(); ++index) { + if (!load_images_from_dir(args.ref_video_paths[index], &ref_video_frames[index])) { + std::fprintf(stderr, "failed to load reference video frames from '%s'\n", args.ref_video_paths[index].c_str()); + for (ed_image_t& image : ref_images) ed_free_image(&image); + free_video_references(); + ed_free_image(&init_image); ed_free_image(&end_image); ed_free_context(ctx); return 6; + } + ref_videos[index].frames = ref_video_frames[index].data(); ref_videos[index].frame_count = static_cast(ref_video_frames[index].size()); ref_videos[index].fps = 24; + if (index < args.ref_video_audio_paths.size() && !load_wav(args.ref_video_audio_paths[index], &ref_video_audio_samples[index], &ref_videos[index].audio)) { + std::fprintf(stderr, "failed to load reference video WAV '%s'\n", args.ref_video_audio_paths[index].c_str()); + for (ed_image_t& image : ref_images) ed_free_image(&image); + free_video_references(); + ed_free_image(&init_image); ed_free_image(&end_image); ed_free_context(ctx); return 6; + } + } + if (!ref_videos.empty()) { gen_params.ref_videos = ref_videos.data(); gen_params.ref_video_count = static_cast(ref_videos.size()); } ed_video_t output; auto ed_wall_gen0 = ed_wall_clock::now(); @@ -801,18 +906,28 @@ int main(int argc, char** argv) { } ed_free_image(&init_image); ed_free_image(&end_image); + for (ed_image_t& image : ref_images) ed_free_image(&image); + free_video_references(); ed_free_context(ctx); return 3; } if (!ed_context_parallel_is_root(ctx)) { ed_free_video(&output); + free_video_references(); + for (ed_image_t& image : ref_images) ed_free_image(&image); + ed_free_image(&init_image); + ed_free_image(&end_image); ed_free_context(ctx); return 0; } if (output.frame_count <= 0 || output.frames == nullptr) { std::fprintf(stderr, "generation succeeded but video output is empty\n"); + free_video_references(); + for (ed_image_t& image : ref_images) ed_free_image(&image); + ed_free_image(&init_image); + ed_free_image(&end_image); ed_free_context(ctx); return 4; } @@ -825,6 +940,10 @@ int main(int argc, char** argv) { if (!ed_save_ok) { std::fprintf(stderr, "failed to save output video: %s\n", output_path.c_str()); ed_free_video(&output); + free_video_references(); + for (ed_image_t& image : ref_images) ed_free_image(&image); + ed_free_image(&init_image); + ed_free_image(&end_image); ed_free_context(ctx); return 5; } @@ -838,6 +957,8 @@ int main(int argc, char** argv) { ed_free_video(&output); ed_free_image(&init_image); ed_free_image(&end_image); + for (ed_image_t& image : ref_images) ed_free_image(&image); + free_video_references(); } else { ed_image_generation_params_t gen_params; ed_image_generation_params_init(&gen_params); diff --git a/include/edge-dit.h b/include/edge-dit.h index d95b7d47..d93267f7 100644 --- a/include/edge-dit.h +++ b/include/edge-dit.h @@ -114,6 +114,20 @@ typedef struct ed_image_t { uint8_t * data; } ed_image_t; +typedef struct ed_audio_t { + uint32_t sample_rate; + uint32_t channels; + uint64_t sample_count; + const float * data; +} ed_audio_t; + +typedef struct ed_ref_video_t { + const ed_image_t * frames; + int frame_count; + int fps; + ed_audio_t audio; +} ed_ref_video_t; + typedef struct ed_image_batch_t { ed_image_t * images; int count; @@ -260,6 +274,13 @@ typedef struct ed_video_generation_params_t { const ed_image_t * init_image; const ed_image_t * end_image; + const ed_image_t * ref_images; + int ref_image_count; + const ed_ref_video_t * ref_videos; + int ref_video_count; + const ed_audio_t * ref_audios; + int ref_audio_count; + const ed_image_t * control_frames; int control_frame_count; diff --git a/src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp b/src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp index 44d85823..a0977c2d 100644 --- a/src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp +++ b/src/dit_models/components/autoencoders/minimax_h3_audio_vae.hpp @@ -408,6 +408,125 @@ namespace Ops { 1.f); } }; +struct AudioSnake1D : public UnaryBlock { + int64_t channels; + explicit AudioSnake1D(int64_t value) : channels(value) {} + void init_params(ggml_context* ctx, const String2TensorStorage& storage = {}, const std::string prefix = "") override { + ED_UNUSED(storage); ED_UNUSED(prefix); + params["alpha"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, channels, 1); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto alpha = params["alpha"]; + auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha)); + oscillation = ggml_mul(ctx->ggml_ctx, oscillation, oscillation); + auto eps = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), 1e-9f); + return ggml_add(ctx->ggml_ctx, x, ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, alpha, eps))); + } +}; + +struct AudioEncoderResidualUnit : public GGMLBlock { + AudioEncoderResidualUnit(int64_t channels, int dilation) { + blocks["block.0"] = std::make_shared(channels); + blocks["block.1"] = std::make_shared(channels, channels, 7, 1, 3 * dilation, dilation); + blocks["block.2"] = std::make_shared(channels); + blocks["block.3"] = std::make_shared(channels, channels, 1); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto act1 = std::dynamic_pointer_cast(blocks["block.0"]); + auto conv1 = std::dynamic_pointer_cast(blocks["block.1"]); + auto act2 = std::dynamic_pointer_cast(blocks["block.2"]); + auto conv2 = std::dynamic_pointer_cast(blocks["block.3"]); + auto h = conv2->forward(ctx, act2->forward(ctx, conv1->forward(ctx, act1->forward(ctx, x)))); + if (x->ne[0] != h->ne[0]) { + const int64_t pad = (x->ne[0] - h->ne[0]) / 2; + x = ggml_ext_slice(ctx->ggml_ctx, x, 0, pad, x->ne[0] - pad); + } + return ggml_add(ctx->ggml_ctx, x, h); + } +}; + +struct AudioEncoderBlock : public GGMLBlock { + AudioEncoderBlock(int64_t out_channels, int stride) { + const int64_t in_channels = out_channels / 2; + blocks["block.0"] = std::make_shared(in_channels, 1); + blocks["block.1"] = std::make_shared(in_channels, 3); + blocks["block.2"] = std::make_shared(in_channels, 9); + blocks["block.3"] = std::make_shared(in_channels); + blocks["block.4"] = std::make_shared(in_channels, out_channels, 2 * stride, stride, static_cast(std::ceil(stride / 2.f))); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + for (int index = 0; index < 3; ++index) x = std::dynamic_pointer_cast(blocks["block." + std::to_string(index)])->forward(ctx, x); + return std::dynamic_pointer_cast(blocks["block.4"])->forward(ctx, std::dynamic_pointer_cast(blocks["block.3"])->forward(ctx, x)); + } +}; + +struct AudioEncoder : public GGMLBlock { + static constexpr std::array strides = {2, 4, 4, 5, 5}; + AudioEncoder() { + int64_t channels = 64; + blocks["block.0"] = std::make_shared(1, channels, 7, 1, 3); + for (size_t index = 0; index < strides.size(); ++index) { channels *= 2; blocks["block." + std::to_string(index + 1)] = std::make_shared(channels, strides[index]); } + blocks["block.6"] = std::make_shared(channels); + blocks["block.7"] = std::make_shared(channels, 2048, 3, 1, 1); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + x = std::dynamic_pointer_cast(blocks["block.0"])->forward(ctx, x); + for (size_t index = 0; index < strides.size(); ++index) x = std::dynamic_pointer_cast(blocks["block." + std::to_string(index + 1)])->forward(ctx, x); + return std::dynamic_pointer_cast(blocks["block.7"])->forward(ctx, std::dynamic_pointer_cast(blocks["block.6"])->forward(ctx, x)); + } +}; + +struct AudioGeGLUMLP : public GGMLBlock { + AudioGeGLUMLP(int64_t hidden_size, int64_t intermediate_size) { + blocks["norm"] = std::make_shared(hidden_size); + blocks["w0"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["w1"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["w2"] = std::make_shared(intermediate_size, hidden_size, true); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + x = std::dynamic_pointer_cast(blocks["norm"])->forward(ctx, x); + auto gate = ggml_ext_gelu(ctx->ggml_ctx, std::dynamic_pointer_cast(blocks["w0"])->forward(ctx, x), true); + return std::dynamic_pointer_cast(blocks["w2"])->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, std::dynamic_pointer_cast(blocks["w1"])->forward(ctx, x))); + } +}; + +struct AudioCausalAttention : public GGMLBlock { + static constexpr int64_t in_channels = 2048, out_channels = 32, num_head = 8, head_dim = in_channels / num_head; + AudioCausalAttention() { blocks["qkv"] = std::make_shared(in_channels, in_channels * 3, false); blocks["proj"] = std::make_shared(out_channels, out_channels, true); } + void init_params(ggml_context* ctx, const String2TensorStorage& storage = {}, const std::string prefix = "") override { + GGMLBlock::init_params(ctx, storage, 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); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto qkv = ggml_ext_chunk(ctx->ggml_ctx, std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x), 3, 0); + auto shape_bias = [&](ggml_tensor* value) { return ggml_reshape_4d(ctx->ggml_ctx, value, value->ne[0], 1, 1, 1); }; + auto q = ggml_add(ctx->ggml_ctx, qkv[0], shape_bias(params["q_bias"])); + auto v = ggml_add(ctx->ggml_ctx, qkv[2], shape_bias(params["v_bias"])); + const int64_t sequence = x->ne[1]; + auto mask = ggml_diag_mask_inf(ctx->ggml_ctx, ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1), 0); + auto out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, qkv[1], v, num_head, mask, false, ctx->flash_attn_enabled); + const int64_t batch = out->ne[2] * out->ne[3]; + out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, ggml_reshape_4d(ctx->ggml_ctx, out, head_dim, num_head, sequence, batch), 1, 0, 2, 3)); + out = ggml_mean(ctx->ggml_ctx, out); + out = ggml_reshape_3d(ctx->ggml_ctx, out, head_dim, sequence, batch); + out = ggml_mean(ctx->ggml_ctx, ggml_reshape_4d(ctx->ggml_ctx, out, head_dim / out_channels, out_channels, sequence, batch)); + out = ggml_reshape_3d(ctx->ggml_ctx, out, out_channels, sequence, batch); + return std::dynamic_pointer_cast(blocks["proj"])->forward(ctx, out); + } +}; + +struct AudioAttentionProjection : public GGMLBlock { + AudioAttentionProjection() { + blocks["norm1"] = std::make_shared(2048); blocks["attn"] = std::make_shared(); blocks["proj"] = std::make_shared(2048, 32, true); + blocks["norm3"] = std::make_shared(2048); blocks["norm2"] = std::make_shared(32); blocks["mlp"] = std::make_shared(32, 64); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto h = ggml_add(ctx->ggml_ctx, + std::dynamic_pointer_cast(blocks["proj"])->forward(ctx, std::dynamic_pointer_cast(blocks["norm3"])->forward(ctx, x)), + std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, std::dynamic_pointer_cast(blocks["norm1"])->forward(ctx, x))); + return ggml_add(ctx->ggml_ctx, h, std::dynamic_pointer_cast(blocks["mlp"])->forward(ctx, std::dynamic_pointer_cast(blocks["norm2"])->forward(ctx, h))); + } +}; + struct AudioDecoder : public GGMLBlock { static constexpr int kLatentChannels = 32; AudioDecoder() { @@ -437,8 +556,54 @@ struct AudioDecoder : public GGMLBlock { return ggml_reshape_4d(ctx->ggml_ctx, waveform, waveform->ne[0], streams, 1, 1); } }; +struct AudioVAE : public GGMLBlock { + static constexpr int kLatentChannels = 32; + AudioVAE() { + blocks["encoder"] = std::make_shared(); + blocks["pre_block"] = std::make_shared(); + blocks["mean_proj"] = std::make_shared(kLatentChannels, kLatentChannels, 1); + blocks["logs_proj"] = std::make_shared(kLatentChannels, kLatentChannels, 1); + blocks["dec_in_proj"] = std::make_shared(kLatentChannels, 2048, 1); + blocks["decoder"] = std::make_shared(); + } + void init_params(ggml_context* ctx, const String2TensorStorage& storage = {}, const std::string prefix = "") override { + GGMLBlock::init_params(ctx, storage, prefix); + params["latents_mean"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels); + params["latents_std"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels); + } + ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* waveform) { + GGML_ASSERT(waveform->ne[1] == 2); + waveform = ggml_reshape_3d(ctx->ggml_ctx, waveform, waveform->ne[0], 1, waveform->ne[1]); + auto x = std::dynamic_pointer_cast(blocks["encoder"])->forward(ctx, waveform); + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = std::dynamic_pointer_cast(blocks["pre_block"])->forward(ctx, x); + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + auto z = std::dynamic_pointer_cast(blocks["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); + return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std), 0, 2, 1, 3)); + } + ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent) { + GGML_ASSERT(latent->ne[1] == 2 && latent->ne[2] == kLatentChannels); + latent = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, latent, 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); + latent = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, latent, std), mean); + auto decoder_input = std::dynamic_pointer_cast(blocks["dec_in_proj"]); + auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); + const int64_t streams = latent->ne[2] * latent->ne[3]; + latent = ggml_reshape_3d(ctx->ggml_ctx, latent, latent->ne[0], latent->ne[1], streams); + ggml_tensor* waveform = nullptr; + for (int64_t stream = 0; stream < streams; ++stream) { + auto value = decoder->forward(ctx, decoder_input->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1))); + waveform = waveform == nullptr ? value : ggml_concat(ctx->ggml_ctx, waveform, value, 2); + } + return ggml_reshape_4d(ctx->ggml_ctx, waveform, waveform->ne[0], streams, 1, 1); + } +}; + struct AudioVAERunner : public GGMLRunner { - AudioDecoder model; + AudioVAE model; AudioVAERunner(ggml_backend_t backend, bool offload, const String2TensorStorage& storage, const std::string& prefix = "audio_vae") : GGMLRunner(backend, offload) { model.init(params_ctx, storage, prefix); } std::string get_desc() override { return "minimax_h3_audio_vae"; } @@ -447,6 +612,11 @@ struct AudioVAERunner : public GGMLRunner { auto get_graph = [&]() { auto input = make_input(latent); auto runner_ctx = get_context(); auto graph = new_graph_custom(655360); ggml_build_forward_expand(graph, model.decode(&runner_ctx, input)); return graph; }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4); } + sd::Tensor encode(int n_threads, const sd::Tensor& waveform) { + auto get_graph = [&]() { auto input = make_input(waveform); auto runner_ctx = get_context(); auto graph = new_graph_custom(655360); ggml_build_forward_expand(graph, model.encode(&runner_ctx, input)); return graph; }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4); + } + int input_sample_rate() const { return 32000; } }; } // namespace MiniMaxH3Audio #endif diff --git a/src/dit_models/components/text_encoders/llm.hpp b/src/dit_models/components/text_encoders/llm.hpp index ea6a322b..b110f3a2 100644 --- a/src/dit_models/components/text_encoders/llm.hpp +++ b/src/dit_models/components/text_encoders/llm.hpp @@ -229,6 +229,7 @@ namespace LLM { enum class LLMArch { QWEN2_5_VL, QWEN3, + QWEN3_VL, MISTRAL_SMALL_3_2, MINISTRAL_3_3B, ARCH_COUNT, @@ -237,11 +238,18 @@ namespace LLM { static const char* llm_arch_to_str[] = { "qwen2.5vl", "qwen3", + "qwen3vl", "mistral_small3.2", "ministral3.3b", }; + enum class LLMVisionArch { + QWEN2_5_VL, + QWEN3_VL, + }; + struct LLMVisionParams { + LLMVisionArch arch = LLMVisionArch::QWEN2_5_VL; int num_layers = 32; int64_t hidden_size = 1280; int64_t intermediate_size = 3420; @@ -252,6 +260,7 @@ namespace LLM { int patch_size = 14; int spatial_merge_size = 2; int window_size = 112; + std::vector deepstack_visual_indexes; std::set fullatt_block_indexes = {7, 15, 23, 31}; }; @@ -363,6 +372,7 @@ namespace LLM { struct VisionPatchEmbed : public GGMLBlock { protected: bool llama_cpp_style; + bool bias; int patch_size; int temporal_patch_size; int64_t in_channels; @@ -370,11 +380,13 @@ namespace LLM { public: VisionPatchEmbed(bool llama_cpp_style, + LLMVisionArch arch, int patch_size = 14, int temporal_patch_size = 2, int64_t in_channels = 3, int64_t embed_dim = 1152) : llama_cpp_style(llama_cpp_style), + bias(arch == LLMVisionArch::QWEN3_VL), patch_size(patch_size), temporal_patch_size(temporal_patch_size), in_channels(in_channels), @@ -402,7 +414,7 @@ namespace LLM { kernel_size, // stride {0, 0, 0}, // padding {1, 1, 1}, // dilation - false, + bias, true)); } } @@ -492,14 +504,23 @@ namespace LLM { struct PatchMerger : public GGMLBlock { protected: + LLMVisionArch arch_; int64_t hidden_size; public: - PatchMerger(int64_t dim, + PatchMerger(LLMVisionArch arch, + int64_t dim, int64_t context_dim, - int64_t spatial_merge_size) { + int64_t spatial_merge_size) + : arch_(arch) { const bool diffusers_dtype = qwen_align_diffusers_vision_dtype_enabled(); hidden_size = context_dim * spatial_merge_size * spatial_merge_size; + if (arch_ == LLMVisionArch::QWEN3_VL) { + blocks["norm"] = std::make_shared(context_dim, 1e-6f); + blocks["linear_fc1"] = std::make_shared(hidden_size, hidden_size, true); + blocks["linear_fc2"] = std::make_shared(hidden_size, dim, true); + return; + } blocks["ln_q"] = std::shared_ptr(new RMSNorm(context_dim, 1e-6f, false, true)); blocks["mlp.0"] = std::shared_ptr(new Linear(hidden_size, hidden_size, @@ -526,6 +547,19 @@ namespace LLM { ggml_tensor* x, const std::string& debug_target = "", const std::string& debug_prefix = "") { + if (arch_ == LLMVisionArch::QWEN3_VL) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto linear_fc1 = std::dynamic_pointer_cast(blocks["linear_fc1"]); + auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); + if (x->type != GGML_TYPE_F32) { + x = ggml_cast(ctx->ggml_ctx, x, GGML_TYPE_F32); + } + x = norm->forward(ctx, x); + x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size); + x = linear_fc1->forward(ctx, x); + x = ggml_gelu_erf(ctx->ggml_ctx, x); + return linear_fc2->forward(ctx, x); + } auto ln_q = std::dynamic_pointer_cast(blocks["ln_q"]); auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); @@ -564,24 +598,86 @@ namespace LLM { } }; + struct Qwen3VLDeepStackMerger : public GGMLBlock { + protected: + int64_t merge_dim; + + public: + Qwen3VLDeepStackMerger(int64_t dim, int64_t context_dim, int64_t spatial_merge_size) + : merge_dim(context_dim * spatial_merge_size * spatial_merge_size) { + blocks["norm"] = std::make_shared(merge_dim, 1e-6f); + blocks["linear_fc1"] = std::make_shared(merge_dim, merge_dim, true); + blocks["linear_fc2"] = std::make_shared(merge_dim, dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto linear_fc1 = std::dynamic_pointer_cast(blocks["linear_fc1"]); + auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); + if (x->type != GGML_TYPE_F32) { + x = ggml_cast(ctx->ggml_ctx, x, GGML_TYPE_F32); + } + x = ggml_reshape_2d(ctx->ggml_ctx, x, merge_dim, ggml_nelements(x) / merge_dim); + x = norm->forward(ctx, x); + x = linear_fc1->forward(ctx, x); + x = ggml_gelu_erf(ctx->ggml_ctx, x); + return linear_fc2->forward(ctx, x); + } + }; + + struct VisionMLP : public GGMLBlock { + protected: + LLMVisionArch arch_; + + public: + VisionMLP(LLMVisionArch arch, int64_t hidden_size, int64_t intermediate_size) + : arch_(arch) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + blocks["linear_fc1"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["linear_fc2"] = std::make_shared(intermediate_size, hidden_size, true); + } else { + blocks["gate_proj"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["up_proj"] = std::make_shared(hidden_size, intermediate_size, true); + blocks["down_proj"] = std::make_shared(intermediate_size, hidden_size, true); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + x = std::dynamic_pointer_cast(blocks["linear_fc1"])->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x); + return std::dynamic_pointer_cast(blocks["linear_fc2"])->forward(ctx, x); + } + auto gate_proj = std::dynamic_pointer_cast(blocks["gate_proj"]); + auto up_proj = std::dynamic_pointer_cast(blocks["up_proj"]); + auto down_proj = std::dynamic_pointer_cast(blocks["down_proj"]); + auto h = ggml_silu_inplace(ctx->ggml_ctx, gate_proj->forward(ctx, x)); + h = ggml_mul_inplace(ctx->ggml_ctx, h, up_proj->forward(ctx, x)); + return down_proj->forward(ctx, h); + } + }; + struct VisionAttention : public GGMLBlock { protected: bool llama_cpp_style; + LLMVisionArch arch_; int head_dim; int num_heads; public: VisionAttention(bool llama_cpp_style, + LLMVisionArch arch, int64_t hidden_size, int num_heads) - : llama_cpp_style(llama_cpp_style), num_heads(num_heads) { + : llama_cpp_style(llama_cpp_style), arch_(arch), num_heads(num_heads) { head_dim = static_cast(hidden_size / num_heads); GGML_ASSERT(num_heads * head_dim == hidden_size); - const bool diffusers_dtype = qwen_align_diffusers_vision_dtype_enabled(); + const bool diffusers_dtype = arch_ == LLMVisionArch::QWEN2_5_VL && qwen_align_diffusers_vision_dtype_enabled(); + const bool bias = arch_ == LLMVisionArch::QWEN2_5_VL; if (llama_cpp_style) { blocks["q_proj"] = std::shared_ptr(new Linear(hidden_size, hidden_size, - true, + bias, false, diffusers_dtype, 1.f, @@ -590,7 +686,7 @@ namespace LLM { diffusers_dtype)); blocks["k_proj"] = std::shared_ptr(new Linear(hidden_size, hidden_size, - true, + bias, false, diffusers_dtype, 1.f, @@ -599,7 +695,7 @@ namespace LLM { diffusers_dtype)); blocks["v_proj"] = std::shared_ptr(new Linear(hidden_size, hidden_size, - true, + bias, false, diffusers_dtype, 1.f, @@ -609,7 +705,7 @@ namespace LLM { } else { blocks["qkv"] = std::shared_ptr(new Linear(hidden_size, hidden_size * 3, - true, + bias, false, diffusers_dtype, 1.f, @@ -619,7 +715,7 @@ namespace LLM { } blocks["proj"] = std::shared_ptr(new Linear(hidden_size, hidden_size, - true, + bias, false, diffusers_dtype, 1.f, @@ -802,14 +898,35 @@ namespace LLM { }; struct VisionBlock : public GGMLBlock { + protected: + LLMVisionArch arch_; + + ggml_tensor* forward_norm(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) { + if (arch_ == LLMVisionArch::QWEN3_VL) { + if (x->type != GGML_TYPE_F32) { + x = ggml_cast(ctx->ggml_ctx, x, GGML_TYPE_F32); + } + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); + } + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); + } + public: VisionBlock(bool llama_cpp_style, + LLMVisionArch arch, int64_t hidden_size, int64_t intermediate_size, int num_heads, - float eps = 1e-6f) { + float eps = 1e-6f) + : arch_(arch) { const bool diffusers_dtype = qwen_align_diffusers_vision_dtype_enabled(); - blocks["attn"] = std::shared_ptr(new VisionAttention(llama_cpp_style, hidden_size, num_heads)); + blocks["attn"] = std::shared_ptr(new VisionAttention(llama_cpp_style, arch_, hidden_size, num_heads)); + if (arch_ == LLMVisionArch::QWEN3_VL) { + blocks["mlp"] = std::shared_ptr(new VisionMLP(arch_, hidden_size, intermediate_size)); + blocks["norm1"] = std::shared_ptr(new LayerNorm(hidden_size, eps)); + blocks["norm2"] = std::shared_ptr(new LayerNorm(hidden_size, eps)); + return; + } blocks["mlp"] = std::shared_ptr(new MLP(hidden_size, intermediate_size, true, @@ -829,9 +946,6 @@ namespace LLM { const std::string& debug_prefix = "") { // x: [N, n_token, hidden_size] auto attn = std::dynamic_pointer_cast(blocks["attn"]); - auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); - auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); - auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); auto is_debug_target = [&](const char* suffix) { return !debug_prefix.empty() && debug_target == debug_prefix + suffix; }; @@ -840,7 +954,7 @@ namespace LLM { if (is_debug_target(".input")) { return x; } - x = norm1->forward(ctx, x); + x = forward_norm(ctx, "norm1", x); x = qwen_align_maybe_bf16_llm_roundtrip(ctx->ggml_ctx, x); if (is_debug_target(".norm1")) { return x; @@ -856,12 +970,16 @@ namespace LLM { } residual = x; - x = norm2->forward(ctx, x); + x = forward_norm(ctx, "norm2", x); x = qwen_align_maybe_bf16_llm_roundtrip(ctx->ggml_ctx, x); if (is_debug_target(".norm2")) { return x; } - x = mlp->forward(ctx, x, debug_target, debug_prefix); + if (arch_ == LLMVisionArch::QWEN3_VL) { + x = std::dynamic_pointer_cast(blocks["mlp"])->forward(ctx, x); + } else { + x = std::dynamic_pointer_cast(blocks["mlp"])->forward(ctx, x, debug_target, debug_prefix); + } if (!debug_prefix.empty() && debug_target.rfind(debug_prefix + ".mlp.", 0) == 0) { return x; } @@ -877,41 +995,34 @@ namespace LLM { struct VisionModel : public GGMLBlock { protected: + LLMVisionArch arch_; int num_layers; int spatial_merge_size; std::set fullatt_block_indexes; + std::vector deepstack_visual_indexes; public: VisionModel(bool llama_cpp_style, - int num_layers, - int64_t in_channels, - int64_t hidden_size, - int64_t out_hidden_size, - int64_t intermediate_size, - int num_heads, - int spatial_merge_size, - int patch_size, - int temporal_patch_size, - int window_size, - std::set fullatt_block_indexes = {7, 15, 23, 31}, + const LLMVisionParams& vision_params, float eps = 1e-6f) - : num_layers(num_layers), fullatt_block_indexes(std::move(fullatt_block_indexes)), spatial_merge_size(spatial_merge_size) { - blocks["patch_embed"] = std::shared_ptr(new VisionPatchEmbed(llama_cpp_style, - patch_size, - temporal_patch_size, - in_channels, - hidden_size)); + : arch_(vision_params.arch), num_layers(vision_params.num_layers), spatial_merge_size(vision_params.spatial_merge_size), + fullatt_block_indexes(vision_params.fullatt_block_indexes), deepstack_visual_indexes(vision_params.deepstack_visual_indexes) { + blocks["patch_embed"] = std::shared_ptr(new VisionPatchEmbed(llama_cpp_style, arch_, vision_params.patch_size, + vision_params.temporal_patch_size, vision_params.in_channels, vision_params.hidden_size)); for (int i = 0; i < num_layers; i++) { blocks["blocks." + std::to_string(i)] = std::shared_ptr(new VisionBlock(llama_cpp_style, - hidden_size, - intermediate_size, - num_heads, + arch_, vision_params.hidden_size, + vision_params.intermediate_size, + vision_params.num_heads, eps)); } - blocks["merger"] = std::shared_ptr(new PatchMerger(out_hidden_size, hidden_size, spatial_merge_size)); + blocks["merger"] = std::shared_ptr(new PatchMerger(arch_, vision_params.out_hidden_size, vision_params.hidden_size, spatial_merge_size)); + for (size_t i = 0; i < deepstack_visual_indexes.size(); ++i) { + blocks["deepstack_merger_list." + std::to_string(i)] = std::make_shared(vision_params.out_hidden_size, vision_params.hidden_size, spatial_merge_size); + } } - ggml_tensor* forward(GGMLRunnerContext* ctx, + std::vector forward_outputs(GGMLRunnerContext* ctx, ggml_tensor* pixel_values, ggml_tensor* pe, ggml_tensor* window_index, @@ -929,16 +1040,19 @@ namespace LLM { auto x = patch_embed->forward(ctx, pixel_values); sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.prelude", "x"); if (debug_target == "patch_embed") { - return x; + return {x}; } - x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]); - x = ggml_get_rows(ctx->ggml_ctx, x, window_index); - x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] / spatial_merge_size / spatial_merge_size, x->ne[1] * spatial_merge_size * spatial_merge_size, x->ne[2], x->ne[3]); + if (window_index != nullptr) { + x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]); + x = ggml_get_rows(ctx->ggml_ctx, x, window_index); + x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] / spatial_merge_size / spatial_merge_size, x->ne[1] * spatial_merge_size * spatial_merge_size, x->ne[2], x->ne[3]); + } if (debug_target == "windowed") { - return x; + return {x}; } + std::vector deepstack_outputs; for (int i = 0; i < num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); @@ -951,32 +1065,52 @@ namespace LLM { const std::string block_debug_prefix = "block" + std::to_string(i); x = block->forward(ctx, x, pe, mask, layer_cu_window_seqlens, debug_target, block_debug_prefix); if (debug_target.rfind(block_debug_prefix + ".", 0) == 0) { - return x; + return {x}; } if (qwen_align_bf16_vision_activations_enabled()) { x = qwen_align_bf16_roundtrip_to_f32(ctx->ggml_ctx, x); } + auto deepstack_it = std::find(deepstack_visual_indexes.begin(), deepstack_visual_indexes.end(), i); + if (deepstack_it != deepstack_visual_indexes.end()) { + size_t deepstack_index = static_cast(std::distance(deepstack_visual_indexes.begin(), deepstack_it)); + auto deepstack_merger = std::dynamic_pointer_cast(blocks["deepstack_merger_list." + std::to_string(deepstack_index)]); + deepstack_outputs.push_back(deepstack_merger->forward(ctx, x)); + } sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x"); if (debug_target == "block" + std::to_string(i)) { - return x; + return {x}; } } x = merger->forward(ctx, x, debug_target, "merger"); if (debug_target.rfind("merger.", 0) == 0) { - return x; + return {x}; } if (qwen_align_bf16_vision_activations_enabled()) { x = qwen_align_bf16_roundtrip_to_f32(ctx->ggml_ctx, x); } sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.final", "x"); if (debug_target == "merged") { - return x; + return {x}; } - x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index); + if (window_inverse_index != nullptr) { + x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index); + } + std::vector outputs = {x}; + outputs.insert(outputs.end(), deepstack_outputs.begin(), deepstack_outputs.end()); + return outputs; + } - return x; + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* pixel_values, + ggml_tensor* pe, + ggml_tensor* window_index, + ggml_tensor* window_inverse_index, + ggml_tensor* window_mask, + const std::vector* cu_window_seqlens = nullptr, + const std::string& debug_target = "") { + return forward_outputs(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask, cu_window_seqlens, debug_target)[0]; } }; @@ -1110,6 +1244,10 @@ namespace LLM { } else if (arch == LLMArch::QWEN3) { q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); + } else if (arch == LLMArch::QWEN3_VL) { + int sections[4] = {24, 20, 20, 0}; + q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); + k = ggml_rope_multi(ctx->ggml_ctx, k, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); } else { int sections[4] = {16, 24, 24, 0}; q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_MROPE, 128000, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); @@ -1265,6 +1403,7 @@ namespace LLM { ggml_tensor* input_pos, ggml_tensor* attention_mask, std::vector> image_embeds, + const std::vector>>& deepstack_image_embeds, std::set out_layers, const std::string& debug_target = "") { // input_ids: [N, n_token] @@ -1349,6 +1488,13 @@ namespace LLM { const std::string block_debug_prefix = "block" + std::to_string(i); x = block->forward(ctx, x, input_pos, attention_mask, debug_target, block_debug_prefix); + if (i < static_cast(deepstack_image_embeds.size())) { + for (const auto& [index, image_embed] : deepstack_image_embeds[static_cast(i)]) { + auto visual_embed = ggml_ext_slice(ctx->ggml_ctx, x, 1, index, index + image_embed->ne[1]); + visual_embed = ggml_add(ctx->ggml_ctx, visual_embed, image_embed); + x = ggml_set_2d(ctx->ggml_ctx, x, visual_embed, x->nb[1], index * x->nb[1]); + } + } if (debug_target.rfind(block_debug_prefix + ".", 0) == 0) { return x; } @@ -1393,18 +1539,7 @@ namespace LLM { : enable_vision(enable_vision), params(params) { blocks["model"] = std::shared_ptr(new TextModel(params)); if (enable_vision) { - blocks["visual"] = std::shared_ptr(new VisionModel(llama_cpp_style, - params.vision.num_layers, - params.vision.in_channels, - params.vision.hidden_size, - params.vision.out_hidden_size, - params.vision.intermediate_size, - params.vision.num_heads, - params.vision.spatial_merge_size, - params.vision.patch_size, - params.vision.temporal_patch_size, - params.vision.window_size, - params.vision.fullatt_block_indexes)); + blocks["visual"] = std::shared_ptr(new VisionModel(llama_cpp_style, params.vision)); } } @@ -1413,12 +1548,13 @@ namespace LLM { ggml_tensor* input_pos, ggml_tensor* attention_mask, std::vector> image_embeds, + const std::vector>>& deepstack_image_embeds, std::set out_layers, const std::string& debug_target = "") { // input_ids: [N, n_token] auto model = std::dynamic_pointer_cast(blocks["model"]); - auto x = model->forward(ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers, debug_target); + auto x = model->forward(ctx, input_ids, input_pos, attention_mask, image_embeds, deepstack_image_embeds, out_layers, debug_target); return x; } @@ -1434,6 +1570,17 @@ namespace LLM { auto vision_model = std::dynamic_pointer_cast(blocks["visual"]); return vision_model->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask, cu_window_seqlens, debug_target); } + + std::vector vision_forward_outputs(GGMLRunnerContext* ctx, + ggml_tensor* pixel_values, + ggml_tensor* pe, + ggml_tensor* window_index, + ggml_tensor* window_inverse_index, + ggml_tensor* window_mask) { + GGML_ASSERT(enable_vision); + auto vision_model = std::dynamic_pointer_cast(blocks["visual"]); + return vision_model->forward_outputs(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask); + } }; struct LLMRunner : public GGMLRunner { @@ -1463,16 +1610,20 @@ namespace LLM { params.num_kv_heads = 8; params.qkv_bias = false; params.rms_norm_eps = 1e-5f; - } else if (arch == LLMArch::QWEN3) { + } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) { params.head_dim = 128; params.num_heads = 64; params.num_kv_heads = 8; params.qkv_bias = false; params.qk_norm = true; params.rms_norm_eps = 1e-6f; + if (arch == LLMArch::QWEN3_VL) { + params.vision.arch = LLMVisionArch::QWEN3_VL; + } } bool have_vision_weight = false; bool llama_cpp_style = false; + int detected_vision_layers = 0; params.num_layers = 0; for (auto pair : tensor_storage_map) { std::string tensor_name = pair.first; @@ -1484,6 +1635,26 @@ namespace LLM { if (contains(tensor_name, "attn.q_proj")) { llama_cpp_style = true; } + if (contains(tensor_name, "visual.patch_embed.proj.weight")) { + params.vision.patch_size = static_cast(pair.second.ne[0]); + } + if (contains(tensor_name, "visual.patch_embed.proj.bias")) { + params.vision.hidden_size = pair.second.ne[0]; + } + if (contains(tensor_name, "visual.blocks.")) { + auto items = split_string(tensor_name.substr(pos), '.'); + if (items.size() > 2) { + detected_vision_layers = std::max(detected_vision_layers, std::atoi(items[2].c_str()) + 1); + } + } + if (contains(tensor_name, "visual.blocks.0.mlp.linear_fc1.weight") || + contains(tensor_name, "visual.blocks.0.mlp.gate_proj.weight")) { + params.vision.intermediate_size = pair.second.ne[1]; + } + if (contains(tensor_name, "visual.merger.linear_fc2.weight") || + contains(tensor_name, "visual.merger.mlp.2.weight")) { + params.vision.out_hidden_size = pair.second.ne[1]; + } continue; } pos = tensor_name.find("layers."); @@ -1505,13 +1676,27 @@ namespace LLM { params.intermediate_size = pair.second.ne[1]; } } - if (arch == LLMArch::QWEN3 && params.num_layers == 28) { // Qwen3 2B + if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && params.num_layers == 28) { // Qwen3 2B params.num_heads = 16; } - if (arch == LLMArch::QWEN3 && params.num_layers == 50 && params.hidden_size == 5120) { + if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && params.num_layers == 50 && params.hidden_size == 5120) { params.num_heads = 64; params.final_norm = false; } + if (detected_vision_layers > 0) { + params.vision.num_layers = detected_vision_layers; + } + if (arch == LLMArch::QWEN3_VL) { + if (params.vision.num_layers == 24) { + params.vision.deepstack_visual_indexes = {5, 11, 17}; + } else if (params.vision.num_layers == 27) { + params.vision.deepstack_visual_indexes = {8, 16, 24}; + } + } + LOG_DEBUG("llm vision: arch=%d layers=%d hidden=%" PRId64 " heads=%d patch=%d temporal=%d merge=%d", + static_cast(params.vision.arch), params.vision.num_layers, params.vision.hidden_size, + params.vision.num_heads, params.vision.patch_size, params.vision.temporal_patch_size, + params.vision.spatial_merge_size); LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64, params.num_layers, params.vocab_size, @@ -1544,9 +1729,10 @@ namespace LLM { ggml_tensor* input_pos, ggml_tensor* attention_mask, std::vector> image_embeds, + const std::vector>>& deepstack_image_embeds, std::set out_layers, const std::string& debug_target = "") { - auto hidden_states = model.forward(ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers, debug_target); // [N, n_token, hidden_size] + auto hidden_states = model.forward(ctx, input_ids, input_pos, attention_mask, image_embeds, deepstack_image_embeds, out_layers, debug_target); // [N, n_token, hidden_size] return hidden_states; } @@ -1681,6 +1867,7 @@ namespace LLM { ggml_cgraph* build_graph(const sd::Tensor& input_ids_tensor, const sd::Tensor& attention_mask_tensor, const std::vector>>& image_embeds_tensor, + const std::vector>>>& deepstack_image_embeds_tensor, std::set out_layers, const std::vector& image_embed_infos = {}, const std::string& debug_target = "") { @@ -1692,6 +1879,12 @@ namespace LLM { ggml_tensor* embed = make_input(embed_tensor); image_embeds.emplace_back(idx, embed); } + std::vector>> deepstack_image_embeds(deepstack_image_embeds_tensor.size()); + for (size_t layer = 0; layer < deepstack_image_embeds_tensor.size(); ++layer) { + for (const auto& [idx, embed_tensor] : deepstack_image_embeds_tensor[layer]) { + deepstack_image_embeds[layer].emplace_back(idx, make_input(embed_tensor)); + } + } int64_t n_tokens = input_ids->ne[0]; if (params.arch == LLMArch::MISTRAL_SMALL_3_2 || params.arch == LLMArch::MINISTRAL_3_3B || params.arch == LLMArch::QWEN3) { @@ -1737,7 +1930,7 @@ namespace LLM { auto runner_ctx = get_context(); - ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers, debug_target); + ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, attention_mask, image_embeds, deepstack_image_embeds, out_layers, debug_target); ggml_build_forward_expand(gf, hidden_states); @@ -1750,9 +1943,10 @@ namespace LLM { const std::vector>>& image_embeds, std::set out_layers, const std::vector& image_embed_infos = {}, - const std::string& debug_target = "") { + const std::string& debug_target = "", + const std::vector>>>& deepstack_image_embeds = {}) { auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(input_ids, attention_mask, image_embeds, out_layers, image_embed_infos, debug_target); + return build_graph(input_ids, attention_mask, image_embeds, deepstack_image_embeds, out_layers, image_embed_infos, debug_target); }; return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); } @@ -1872,6 +2066,35 @@ namespace LLM { GGML_ASSERT(image_height % (params.vision.patch_size * params.vision.spatial_merge_size) == 0); GGML_ASSERT(image_width % (params.vision.patch_size * params.vision.spatial_merge_size) == 0); + if (params.vision.arch == LLMVisionArch::QWEN3_VL) { + const int grid_h = static_cast(image_height) / params.vision.patch_size; + const int grid_w = static_cast(image_width) / params.vision.patch_size; + const int head_dim = static_cast(params.vision.hidden_size / params.vision.num_heads); + window_index_vec.resize(static_cast((grid_h / params.vision.spatial_merge_size) * + (grid_w / params.vision.spatial_merge_size))); + for (size_t index = 0; index < window_index_vec.size(); ++index) { + window_index_vec[index] = static_cast(index); + } + pe_vec = Rope::gen_qwen2vl_pe(grid_h, + grid_w, + params.vision.spatial_merge_size, + window_index_vec, + 10000, + {head_dim / 2, head_dim / 2}); + const int pos_len = static_cast(pe_vec.size() / head_dim / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + auto runner_ctx = get_context(); + auto outputs = model.vision_forward_outputs(&runner_ctx, + pixel_values, + pe, + nullptr, + nullptr, + nullptr); + ggml_build_forward_expand(gf, outputs[0]); + return gf; + } + int grid_t = 1; int grid_h = static_cast(image_height) / params.vision.patch_size; int grid_w = static_cast(image_width) / params.vision.patch_size; @@ -2094,6 +2317,44 @@ namespace LLM { }; return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false)); } + + std::vector> encode_image_outputs(const int n_threads, + const sd::Tensor& image) { + if (params.vision.arch != LLMVisionArch::QWEN3_VL) { + auto output = encode_image(n_threads, image); + return output.empty() ? std::vector>() : std::vector>{std::move(output)}; + } + const int64_t image_width = image.shape()[0]; + const int64_t image_height = image.shape()[1]; + const auto pixel_values = process_image_patches_host(image); + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* graph = new_graph_custom(LLM_GRAPH_SIZE); + auto pixels = make_input(pixel_values); + const int grid_h = static_cast(image_height / params.vision.patch_size); + const int grid_w = static_cast(image_width / params.vision.patch_size); + const int head_dim = static_cast(params.vision.hidden_size / params.vision.num_heads); + window_index_vec.resize(static_cast((grid_h / params.vision.spatial_merge_size) * (grid_w / params.vision.spatial_merge_size))); + for (size_t index = 0; index < window_index_vec.size(); ++index) window_index_vec[index] = static_cast(index); + pe_vec = Rope::gen_qwen2vl_pe(grid_h, grid_w, params.vision.spatial_merge_size, window_index_vec, 10000, {head_dim / 2, head_dim / 2}); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, static_cast(pe_vec.size() / head_dim / 2)); + set_backend_tensor_data(pe, pe_vec.data()); + auto runner_ctx = get_context(); + auto outputs = model.vision_forward_outputs(&runner_ctx, pixels, pe, nullptr, nullptr, nullptr); + auto combined = outputs[0]; + for (size_t index = 1; index < outputs.size(); ++index) combined = ggml_concat(compute_ctx, combined, outputs[index], 0); + ggml_build_forward_expand(graph, combined); + return graph; + }; + auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); + if (combined.empty()) return {}; + const size_t count = params.vision.deepstack_visual_indexes.size() + 1; + std::vector> outputs; + outputs.reserve(count); + for (size_t index = 0; index < count; ++index) { + outputs.push_back(sd::ops::slice(combined, 0, static_cast(index) * params.hidden_size, static_cast(index + 1) * params.hidden_size)); + } + return outputs; + } }; struct LLMEmbedder { diff --git a/src/dit_models/pipelines/minimax_h3_pipeline.cpp b/src/dit_models/pipelines/minimax_h3_pipeline.cpp index fa171a8d..eb5b68e8 100644 --- a/src/dit_models/pipelines/minimax_h3_pipeline.cpp +++ b/src/dit_models/pipelines/minimax_h3_pipeline.cpp @@ -6,7 +6,9 @@ #include #include #include +#include #include +#include #include "dit_models/diffusion_model.hpp" #include "dit_models/components/autoencoders/minimax_h3_vae.hpp" @@ -35,6 +37,43 @@ int64_t h3_resolve_seed(int64_t seed) { return seed >= 0 ? seed : static_cast(std::time(nullptr)); } +void h3_resize_for_vision(int source_width, int source_height, int* width, int* height) { + constexpr int factor = 32; + constexpr int min_pixels = 3136; + constexpr int max_pixels = 12845056; + int resized_width = std::max(factor, static_cast(std::round(static_cast(source_width) / factor)) * factor); + int resized_height = std::max(factor, static_cast(std::round(static_cast(source_height) / factor)) * factor); + const double area = static_cast(resized_width) * resized_height; + if (area > max_pixels) { + const double scale = std::sqrt(static_cast(source_width) * source_height / max_pixels); + resized_width = std::max(factor, static_cast(std::floor(source_width / scale / factor)) * factor); + resized_height = std::max(factor, static_cast(std::floor(source_height / scale / factor)) * factor); + } else if (area < min_pixels) { + const double scale = std::sqrt(static_cast(min_pixels) / (static_cast(source_width) * source_height)); + resized_width = static_cast(std::ceil(source_width * scale / factor)) * factor; + resized_height = static_cast(std::ceil(source_height * scale / factor)) * factor; + } + *width = resized_width; + *height = resized_height; +} + +void h3_reference_video_dimensions(const ed_image_t& image, int* width, int* height) { + const double ratio = static_cast(image.width) / image.height; + double nominal_width = ratio >= 1.0 ? 768.0 * ratio : 768.0; + double nominal_height = ratio >= 1.0 ? 768.0 : 768.0 / ratio; + if (nominal_width * nominal_height > 768.0 * 1344.0) { + const double scale = std::sqrt((768.0 * 1344.0) / (nominal_width * nominal_height)); + nominal_width *= scale; + nominal_height *= scale; + } + *width = std::max(32, static_cast(std::round(nominal_width / 32.0)) * 32); + *height = std::max(32, static_cast(std::round(nominal_height / 32.0)) * 32); + if (static_cast(image.width) * image.height < static_cast(*width) * *height) { + *width = std::max(32, static_cast(std::round(static_cast(image.width) / 32.0)) * 32); + *height = std::max(32, static_cast(std::round(static_cast(image.height) / 32.0)) * 32); + } +} + float h3_discrete_flow_sigma(int step, int steps, float shift) { if (steps <= 1) { return 1.0f; @@ -232,12 +271,12 @@ bool MiniMaxH3Pipeline::prepare(const ed_context_params_t& params, diffusion_->get_param_tensors(registry.tensors(), "model.diffusion_model"); const bool text_offload = runtime.clip_offload_params_to_cpu(); - conditioner_ = std::make_unique(LLM::LLMArch::QWEN3, + conditioner_ = std::make_unique(LLM::LLMArch::QWEN3_VL, runtime.clip_backend(), text_offload, loader.get_tensor_storage_map(), "text_encoders.llm", - false); + true); conditioner_->alloc_params_buffer(); conditioner_->get_param_tensors(registry.tensors(), "text_encoders.llm"); @@ -310,13 +349,110 @@ ed_status_t MiniMaxH3Pipeline::generate_image(const ed_image_generation_params_t } bool MiniMaxH3Pipeline::build_text_context(const char* prompt, + const ed_image_t* ref_images, + int ref_image_count, + const ed_ref_video_t* ref_videos, + int ref_video_count, + int ref_audio_count, + int max_video_frames, sd::Tensor* context, sd::Tensor* token_tags, std::string* error) { if (context == nullptr || token_tags == nullptr || conditioner_ == nullptr || runtime_ == nullptr) { return set_minimax_error(error, "MiniMax-H3 text conditioner is not initialized"); } - const std::string text = "<|im_start|>user\n" + + std::string presentations; + std::vector>> image_embeds; + std::vector image_embed_infos; + std::vector>>> deepstack_image_embeds(3); + for (int image_index = 0; image_index < ref_image_count; ++image_index) { + const ed_image_t& image = ref_images[image_index]; + int width = 0; + int height = 0; + h3_resize_for_vision(image.width, image.height, &width, &height); + sd::Tensor image_tensor = h3_image_to_tensor(image, width, height); + if (image_tensor.empty()) { + return set_minimax_error(error, "MiniMax-H3 Ref2VA image reference is invalid"); + } + image_tensor = image_tensor * 2.0f - 1.0f; + LOG_DEBUG("MiniMax-H3 Ref2VA vision image=%dx%d tensor=%s", + width, + height, + sd::tensor_shape_to_string(image_tensor.shape()).c_str()); + std::vector> image_outputs = conditioner_->model.encode_image_outputs(runtime_->n_threads(), image_tensor); + if (image_outputs.size() != 4 || image_outputs[0].empty()) { + return set_minimax_error(error, "MiniMax-H3 Ref2VA vision encoder failed"); + } + sd::Tensor image_embed = std::move(image_outputs[0]); + const int64_t image_tokens = image_embed.shape()[1]; + const std::string prefix = ": <|vision_start|>"; + const std::string prefix_text = "<|im_start|>user\n" + presentations + prefix; + const std::vector prefix_tokens = conditioner_->tokenizer->tokenize(prefix_text, nullptr, true, 0, 4096, false); + presentations += prefix; + for (int64_t token = 0; token < image_tokens; ++token) { + presentations += "<|image_pad|>"; + } + presentations += "<|vision_end|>\n"; + image_embeds.emplace_back(static_cast(prefix_tokens.size()), std::move(image_embed)); + for (size_t layer = 0; layer < deepstack_image_embeds.size(); ++layer) { + deepstack_image_embeds[layer].emplace_back(static_cast(prefix_tokens.size()), std::move(image_outputs[layer + 1])); + } + image_embed_infos.push_back({static_cast(prefix_tokens.size()), image_tokens, 1, height / 16, width / 16}); + } + int video_number = 0; + int audio_number = 0; + for (int video_index = 0; video_index < ref_video_count; ++video_index) { + const ed_ref_video_t& reference = ref_videos[video_index]; + if (reference.frames == nullptr || reference.frame_count <= 0) { + return set_minimax_error(error, "MiniMax-H3 Ref2VA video reference is invalid"); + } + const int source_fps = reference.fps > 0 ? reference.fps : 24; + int normalized_frames = static_cast(std::lround(static_cast(reference.frame_count) * 24.0 / source_fps)); + normalized_frames = std::min(normalized_frames, max_video_frames); + if (normalized_frames < 5) { + set_minimax_error(error, "MiniMax-H3 Ref2VA video reference needs at least 5 frames at 24 fps"); + return ED_STATUS_INVALID_ARGUMENT; + } + if (reference.audio.data != nullptr && reference.audio.sample_count > 0) { + presentations += "