Compare commits

...
Author SHA1 Message Date
stduhpf bfbef5b7e6 feat: trained Minimax VAE Latent2rgb proj (#1856) 2026-08-06 01:39:10 +08:00
leejet ea7f0c87cf feat: add minimax-h3 support (#1854) 2026-08-04 23:14:18 +08:00
Huang, Hong-Chang b4e67d1221 fix(cmake): only apply /MP to the MSVC compiler, not icx (#1846) 2026-08-04 22:37:30 +08:00
leejet db99efdd6d refactor: extract model loader initialization (#1844) 2026-08-02 17:24:10 +08:00
vmobilis eb7f35ca49 feat: add linear multi-step sampling method (#1843) 2026-08-02 16:19:47 +08:00
fszontagh 50062a4bba feat: add IP-Adapter Plus (Resampler image projection) support (#1839) 2026-08-02 16:15:28 +08:00
stduhpf 8457624101 feat: support more LoRA models (Kroma-v0.1 support) (#1842) 2026-08-02 16:08:47 +08:00
stduhpfandleejet 10378f42db fix: lora with split qkv compatibility check at runtime (#1836)
Co-authored-by: leejet <leejet714@gmail.com>
2026-08-02 16:08:08 +08:00
leejet e31a86ce91 refactor: centralize CLIP prefix conversion (#1837) 2026-07-30 22:44:31 +08:00
akleine 735a4ef520 fix(PhotoMaker): avoid GGML_ASSERT if trigger word 'img' was not found in prompt (#1835) 2026-07-30 21:26:57 +08:00
Wagner Bruna af92790ffc feat: allow customizing the alpha and beta parameters of the beta scheduler (#1834) 2026-07-30 21:24:51 +08:00
leejet e92e86fb11 fix: prevent torch checkpoint offset overflow (#1832) 2026-07-29 23:16:29 +08:00
vmobilis 9cfe2af8f9 feat: display number of tokens for SD models (#1831) 2026-07-29 22:21:54 +08:00
yzyyzyhhh 2993b7fb43 fix: make parameter loading backend-aware (#1828) 2026-07-29 22:18:44 +08:00
41 changed files with 5077 additions and 282 deletions
+3 -2
View File
@@ -11,10 +11,11 @@ endif()
if (MSVC)
add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
# /MP is MSVC-only: icx rejects it outright once offloading is enabled.
add_compile_options(
$<$<COMPILE_LANGUAGE:C>:/MP>
$<$<AND:$<COMPILE_LANGUAGE:C>,$<C_COMPILER_ID:MSVC>>:/MP>
$<$<COMPILE_LANGUAGE:C>:/utf-8>
$<$<COMPILE_LANGUAGE:CXX>:/MP>
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/MP>
$<$<COMPILE_LANGUAGE:CXX>:/utf-8>
)
endif()
+3 -1
View File
@@ -15,6 +15,7 @@ API and command-line option may change frequently.***
## 🔥Important News
* **2026/08/04** 🚀 stable-diffusion.cpp adds **Day-1 support for MiniMax-H3**
* **2026/06/25** 🚀 stable-diffusion.cpp now supports **Krea2**
* **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4**
* **2026/05/31** 🚀 stable-diffusion.cpp now supports **PiD**
@@ -66,11 +67,12 @@ API and command-line option may change frequently.***
- [Mage-Flow-Edit](./docs/mage_flow.md#image-editing)
- Video Models
- [Wan2.1/Wan2.2](./docs/wan.md)
- [MiniMax-H3](./docs/minimax_h3.md)
- [LTX-2.3](./docs/ltx2.md)
- [HunyuanVideo 1.5](./docs/hunyuan_video.md)
- [LingBot-Video](./docs/lingbot_video.md)
- [PhotoMaker](./docs/photo_maker.md) support.
- [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL)
- [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL, including Plus)
- Control Net support with SD 1.5
- [ADetailer](./docs/adetailer.md)
- LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
Binary file not shown.
Binary file not shown.
Binary file not shown.
+32
View File
@@ -11,6 +11,10 @@ through a decoupled cross-attention added to every attn2 layer of the
UNet. It composes with Control Net, so a reference image (appearance) and
an OpenPose hint (pose) can be combined in a single generation.
Both the classic adapters and the higher-fidelity **Plus** adapters are
supported; see [Plus variants](#plus-variants) below. The variant is
detected from the weight file, so the same options work for both.
## Required weights
1. A base SD 1.5 or SDXL model.
@@ -21,6 +25,11 @@ an OpenPose hint (pose) can be combined in a single generation.
[h94/IP-Adapter](https://huggingface.co/h94/IP-Adapter):
- SD 1.5: `models/ip-adapter_sd15.safetensors`
- SDXL: `sdxl_models/ip-adapter_sdxl_vit-h.safetensors`
- SD 1.5 Plus: `models/ip-adapter-plus_sd15.safetensors`
- SDXL Plus: `sdxl_models/ip-adapter-plus_sdxl_vit-h.safetensors`
The Plus files (`ip-adapter-plus_*`) are used exactly like the classic
ones; see [Plus variants](#plus-variants).
## Options
@@ -45,6 +54,29 @@ sd-cli -m ..\models\sdxl.safetensors --clip_vision ..\models\clip_vision_h.safet
The SDXL VAE decode at 1024x1024 is memory heavy; add `--vae-tiling` (and
`--offload-to-cpu`) on GPUs with limited VRAM.
## Plus variants
The Plus adapters (`ip-adapter-plus_sd15`, `ip-adapter-plus_sdxl_vit-h`)
replace the small linear image projection with a Resampler (a
Perceiver-style module with learned latent queries). Instead of pooling the
CLIP-Vision output into one vector, the Resampler attends over the full grid
of penultimate CLIP-Vision hidden states and emits more image tokens (16
instead of 4). The result transfers finer detail and layout from the
reference, at a small extra cost in the image-projection step.
No extra flags are needed. The variant is detected from the weight file (the
Resampler's `image_proj.latents` tensor), and every Resampler dimension is
read from the tensor shapes, so the same `--ip-adapter`,
`--ip-adapter-image`, and `--ip-adapter-strength` options apply. Plus
composes with Control Net in the same way as the classic adapters.
```
sd-cli -m ..\models\sd_v1.5.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter-plus_sd15.safetensors --ip-adapter-image ..\assets\reference.png --ip-adapter-strength 0.8 -p "a woman, best quality" -n "lowres, bad anatomy" --cfg-scale 7 --steps 30 --sampling-method dpm++2m --scheduler karras -W 512 -H 512
```
The startup log line `IP-Adapter: 16 image tokens` (versus `4` for the
classic adapters) confirms a Plus file was loaded.
## Combining with Control Net
Add the usual Control Net options to keep the reference appearance while
+96
View File
@@ -0,0 +1,96 @@
# MiniMax-H3
MiniMax-H3 jointly generates video and stereo audio with a packed diffusion
transformer. The implementation supports text-to-audio-video (T2VA), optional
first-frame conditioning (I2VA), first/last-frame conditioning (FL2VA), and
image/video/audio reference conditioning (Ref2VA).
## Model files
Pass the four MiniMax-H3 components separately:
- `--diffusion-model`: MiniMax-H3 diffusion transformer
- `--vae`: MiniMax-H3 video VAE
- `--audio-vae`: MiniMax-H3 audio VAE
- `--llm`: the MiniMax-H3 Qwen3-VL-32B text encoder checkpoint
The text encoder must be the MiniMax-H3 variant: Qwen3-VL-32B truncated to 50
language layers and exported without the final language-model normalization.
Its Qwen3-VL vision tower, including the three DeepStack mergers, must also be
present. If the vision tower is stored separately, pass it with `--llm_vision`.
Both the original time-embedder DiT and the smaller AdaLN curve-table variant
are detected from their weights.
### Download weights
- Download minimax_h3_fl2va/minimax_h3_ref2va
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/diffusion_models
- gguf: https://huggingface.co/leejet/MiniMax-H3-GGUF/tree/main
- Download qwen3vl_32b_minimax_h3
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/text_encoders
- gguf: https://huggingface.co/leejet/MiniMax-H3-GGUF/tree/main
- Download vae
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/vae
- Download audio vae
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/vae
## Text-to-audio-video
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\minimax_h3_fl2va-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "A cute American Shorthair silver tabby kitten surfs on a tropical ocean wave, riding a white surfboard with the clear text 'sd.cpp' on it. Cinematic tracking shot, realistic water, bright sunlight, smooth motion, and consistent character appearance. Add upbeat tropical surf-rock background music with cheerful drums and guitar, synchronized with the kittens energetic surfing." --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --fps 24 --video-frames 56
```
<video src=../assets/minimax-h3/t2av.mp4 controls="controls" muted="muted" type="video/mp4"></video>
Omitting `--audio-vae` still runs the joint diffusion model but produces video without a
decoded audio track.
## First/last-frame conditioning
Add `--init-img` for I2VA, or both `--init-img` and `--end-img` for FL2VA:
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\minimax_h3_fl2va-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "a lovely cat" -i ..\assets\ernie_image\turbo_example.png --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --fps 24 --video-frames 56
```
<video src=../assets/minimax-h3/i2av.mp4 controls="controls" muted="muted" type="video/mp4"></video>
## Reference-to-audio-video conditioning
Ref2VA accepts any combination of reference images, reference videos, paired
video soundtracks, and standalone audio references:
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\minimax_h3_ref2va_pruned-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "Use the cat from <Picture 1> as the main character. Keep the cats appearance, fur color, facial features, and identity consistent with the reference image. Create a 2-second cinematic video: start with an extreme close-up shot of the cats face, focusing on its cute expression and detailed fur texture. The camera slowly rotates around the cats head, creating a dynamic reveal. Then smoothly pull back and zoom out to reveal the full scene: the cat is standing confidently on a surfboard, riding ocean waves. Water splashes around the board, sea breeze gently moves the cats fur, and the cat maintains a cute and fearless expression while surfing. Smooth camera movement, cinematic orbit shot, seamless zoom-out transition, low-angle wide shot, realistic ocean environment, golden sunlight, dynamic waves, high-quality realistic style, natural motion, no distortion, keep the cats identity unchanged." -r ..\assets\ernie_image\turbo_example.png --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --fps 24 --video-frames 56
```
<video src=../assets/minimax-h3/r2av.mp4 controls="controls" muted="muted" type="video/mp4"></video>
`--ref-image`, `--ref-video`, and `--ref-audio` can each be repeated. A
reference video is a directory of image frames sorted lexicographically and is
treated as 24 fps. Repeated `--ref-video-audio` WAV files are paired by index
with repeated `--ref-video` inputs. WAV PCM (8/16/24/32-bit) and 32/64-bit
floating-point samples are accepted; audio is converted to stereo 32 kHz by the
pipeline.
Reference inputs are presented to Qwen3-VL in image, video, then audio order.
Videos are sampled at 2 fps for the Qwen presentation while their full 24 fps
latents condition the diffusion transformer. Paired video and audio references
share the same timeline. Ref2VA cannot be combined with `--init-img` or
`--end-img` in one request.
Reference images keep their aspect ratio and are only downscaled when their
pixel area exceeds the requested generation canvas.
The C API exposes the same inputs through `ref_images`, `ref_videos`, and
`ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own
frame rate and optional soundtrack; non-24-fps inputs are resampled internally.
## Shape and runtime notes
- Width and height are aligned upward to a multiple of 32.
- Frame count is aligned upward to the `17k + 5` grid, with a minimum of 5.
- MiniMax-H3 runs at 24 fps; another requested value is overridden.
- The default video flow shift is 12. The audio stream is mapped internally to
its shift of 3, so the regular samplers can operate on the packed AV latent.
+43
View File
@@ -754,6 +754,18 @@ int main(int argc, const char* argv[]) {
return true;
};
auto load_audio = [&](const std::string& path, SDAudioOwner& audio) -> bool {
std::vector<float> samples;
uint32_t sample_rate = 0;
uint32_t channels = 0;
if (!load_wav_from_file(path, samples, sample_rate, channels)) {
LOG_ERROR("load WAV audio from '%s' failed", path.c_str());
return false;
}
audio.reset(std::move(samples), sample_rate, channels);
return true;
};
if (gen_params.init_image_path.size() > 0) {
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
return 1;
@@ -777,6 +789,37 @@ int main(int argc, const char* argv[]) {
}
}
if (!gen_params.ref_video_paths.empty()) {
gen_params.ref_videos.clear();
gen_params.ref_videos.reserve(gen_params.ref_video_paths.size());
for (const auto& path : gen_params.ref_video_paths) {
std::vector<SDImageOwner> frames;
if (!load_images_from_dir(path, frames, 0, 0, 0, cli_params.verbose) || frames.empty()) {
LOG_ERROR("load reference video frames from '%s' failed", path.c_str());
return 1;
}
gen_params.ref_videos.push_back(std::move(frames));
}
gen_params.ref_video_audios.clear();
gen_params.ref_video_audios.resize(gen_params.ref_videos.size());
for (size_t i = 0; i < gen_params.ref_video_audio_paths.size(); ++i) {
if (!load_audio(gen_params.ref_video_audio_paths[i], gen_params.ref_video_audios[i])) {
return 1;
}
}
}
if (!gen_params.ref_audio_paths.empty()) {
gen_params.ref_audios.clear();
gen_params.ref_audios.resize(gen_params.ref_audio_paths.size());
for (size_t i = 0; i < gen_params.ref_audio_paths.size(); ++i) {
if (!load_audio(gen_params.ref_audio_paths[i], gen_params.ref_audios[i])) {
return 1;
}
}
}
if (gen_params.mask_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.mask_image.put(),
gen_params.mask_image_path.c_str(),
+88 -4
View File
@@ -1008,7 +1008,7 @@ ArgOptions SDGenerationParams::get_options() {
&hires_upscaler},
{"",
"--extra-sample-args",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_divisions",
(int)',',
&extra_sample_args},
{"",
@@ -1404,6 +1404,30 @@ ArgOptions SDGenerationParams::get_options() {
return 1;
};
auto on_ref_video_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
ref_video_paths.push_back(argv[index]);
return 1;
};
auto on_ref_video_audio_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
ref_video_audio_paths.push_back(argv[index]);
return 1;
};
auto on_ref_audio_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
ref_audio_paths.push_back(argv[index]);
return 1;
};
auto on_cache_mode_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
@@ -1538,12 +1562,12 @@ ArgOptions SDGenerationParams::get_options() {
on_seed_arg},
{"",
"--sampling-method",
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
"(default: euler for Flux/SD3/Wan, euler_a otherwise)",
on_sample_method_arg},
{"",
"--high-noise-sampling-method",
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
" default: euler for Flux/SD3/Wan, euler_a otherwise",
on_high_noise_sample_method_arg},
{"",
@@ -1568,8 +1592,20 @@ ArgOptions SDGenerationParams::get_options() {
on_high_noise_skip_layers_arg},
{"-r",
"--ref-image",
"reference image for Flux Kontext models (can be used multiple times)",
"reference image for Flux Kontext or MiniMax-H3 Ref2VA (can be used multiple times)",
on_ref_image_arg},
{"",
"--ref-video",
"MiniMax-H3 Ref2VA reference video frame directory at 24 fps (can be used multiple times)",
on_ref_video_arg},
{"",
"--ref-video-audio",
"WAV soundtrack paired by index with --ref-video (can be used multiple times)",
on_ref_video_audio_arg},
{"",
"--ref-audio",
"standalone WAV reference for MiniMax-H3 Ref2VA (can be used multiple times)",
on_ref_audio_arg},
{"",
"--cache-mode",
"caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)",
@@ -2366,6 +2402,16 @@ bool SDGenerationParams::validate(SDMode mode) {
return false;
}
if (ref_video_audio_paths.size() > ref_video_paths.size()) {
LOG_ERROR("error: each --ref-video-audio needs a corresponding --ref-video");
return false;
}
if (mode != VID_GEN && (!ref_video_paths.empty() || !ref_video_audio_paths.empty() || !ref_audio_paths.empty())) {
LOG_ERROR("error: reference video and audio inputs require vid_gen mode");
return false;
}
if (sample_params.shifted_timestep < 0 || sample_params.shifted_timestep > 1000) {
LOG_ERROR("error: shifted_timestep must be in range [0, 1000]");
return false;
@@ -2560,6 +2606,35 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
control_frame_views.push_back(frame.get());
}
ref_image_views.clear();
ref_image_views.reserve(ref_images.size());
for (auto& image : ref_images) {
ref_image_views.push_back(image.get());
}
ref_video_frame_views.clear();
ref_video_frame_views.resize(ref_videos.size());
ref_video_views.clear();
ref_video_views.reserve(ref_videos.size());
for (size_t i = 0; i < ref_videos.size(); ++i) {
auto& frame_views = ref_video_frame_views[i];
frame_views.reserve(ref_videos[i].size());
for (auto& frame : ref_videos[i]) {
frame_views.push_back(frame.get());
}
sd_audio_t audio = i < ref_video_audios.size() ? ref_video_audios[i].get() : sd_audio_t{};
ref_video_views.push_back({frame_views.empty() ? nullptr : frame_views.data(),
static_cast<int>(frame_views.size()),
24,
audio});
}
ref_audio_views.clear();
ref_audio_views.reserve(ref_audios.size());
for (auto& audio : ref_audios) {
ref_audio_views.push_back(audio.get());
}
sample_params.guidance.slg.layers = skip_layers.empty() ? nullptr : skip_layers.data();
sample_params.guidance.slg.layer_count = skip_layers.size();
high_noise_sample_params.guidance.slg.layers = high_noise_skip_layers.empty() ? nullptr : high_noise_skip_layers.data();
@@ -2578,6 +2653,12 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
params.clip_skip = clip_skip;
params.init_image = init_image.get();
params.end_image = end_image.get();
params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
params.ref_images_count = static_cast<int>(ref_image_views.size());
params.ref_videos = ref_video_views.empty() ? nullptr : ref_video_views.data();
params.ref_videos_count = static_cast<int>(ref_video_views.size());
params.ref_audios = ref_audio_views.empty() ? nullptr : ref_audio_views.data();
params.ref_audios_count = static_cast<int>(ref_audio_views.size());
params.control_frames = control_frame_views.empty() ? nullptr : control_frame_views.data();
params.control_frames_size = static_cast<int>(control_frame_views.size());
params.width = get_resolved_width();
@@ -2657,6 +2738,9 @@ std::string SDGenerationParams::to_string() const {
<< " mask_image_path: \"" << mask_image_path << "\",\n"
<< " control_image_path: \"" << control_image_path << "\",\n"
<< " ref_image_paths: " << vec_str_to_string(ref_image_paths) << ",\n"
<< " ref_video_paths: " << vec_str_to_string(ref_video_paths) << ",\n"
<< " ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n"
<< " ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n"
<< " control_video_path: \"" << control_video_path << "\",\n"
<< " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n"
<< " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n"
+9
View File
@@ -212,6 +212,9 @@ struct SDGenerationParams {
std::string control_image_path;
std::string ip_adapter_image_path;
std::vector<std::string> ref_image_paths;
std::vector<std::string> ref_video_paths;
std::vector<std::string> ref_video_audio_paths;
std::vector<std::string> ref_audio_paths;
std::string control_video_path;
sd_sample_params_t sample_params;
@@ -275,6 +278,9 @@ struct SDGenerationParams {
SDImageOwner init_image;
SDImageOwner end_image;
std::vector<SDImageOwner> ref_images;
std::vector<std::vector<SDImageOwner>> ref_videos;
std::vector<SDAudioOwner> ref_video_audios;
std::vector<SDAudioOwner> ref_audios;
SDImageOwner mask_image;
SDImageOwner control_image;
SDImageOwner ip_adapter_image;
@@ -283,6 +289,9 @@ struct SDGenerationParams {
// Backing storage for sd_img_gen_params_t view fields.
std::vector<sd_image_t> ref_image_views;
std::vector<std::vector<sd_image_t>> ref_video_frame_views;
std::vector<sd_ref_video_t> ref_video_views;
std::vector<sd_audio_t> ref_audio_views;
std::vector<sd_image_t> pm_id_image_views;
std::vector<sd_image_t> control_frame_views;
+129
View File
@@ -1374,3 +1374,132 @@ bool write_wav_to_file(const std::string& path,
file.write(reinterpret_cast<const char*>(pcm.data()), static_cast<std::streamsize>(pcm.size() * sizeof(int16_t)));
return file.good();
}
static uint16_t read_le16(const uint8_t* data) {
return static_cast<uint16_t>(data[0]) |
(static_cast<uint16_t>(data[1]) << 8);
}
static uint32_t read_le32(const uint8_t* data) {
return static_cast<uint32_t>(data[0]) |
(static_cast<uint32_t>(data[1]) << 8) |
(static_cast<uint32_t>(data[2]) << 16) |
(static_cast<uint32_t>(data[3]) << 24);
}
bool load_wav_from_file(const std::string& path,
std::vector<float>& interleaved_samples,
uint32_t& sample_rate,
uint32_t& channels) {
interleaved_samples.clear();
sample_rate = 0;
channels = 0;
std::ifstream file(path, std::ios::binary);
uint8_t riff_header[12];
if (!file.read(reinterpret_cast<char*>(riff_header), sizeof(riff_header)) ||
std::memcmp(riff_header, "RIFF", 4) != 0 ||
std::memcmp(riff_header + 8, "WAVE", 4) != 0) {
return false;
}
uint16_t audio_format = 0;
uint16_t bits_per_sample = 0;
uint16_t block_align = 0;
std::streampos data_pos = std::streampos(-1);
uint32_t data_size = 0;
while (file.good()) {
uint8_t chunk_header[8];
if (!file.read(reinterpret_cast<char*>(chunk_header), sizeof(chunk_header))) {
break;
}
uint32_t chunk_size = read_le32(chunk_header + 4);
std::streampos chunk_data_pos = file.tellg();
if (std::memcmp(chunk_header, "fmt ", 4) == 0) {
if (chunk_size < 16) {
return false;
}
std::vector<uint8_t> fmt(chunk_size);
if (!file.read(reinterpret_cast<char*>(fmt.data()), chunk_size)) {
return false;
}
audio_format = read_le16(fmt.data());
channels = read_le16(fmt.data() + 2);
sample_rate = read_le32(fmt.data() + 4);
block_align = read_le16(fmt.data() + 12);
bits_per_sample = read_le16(fmt.data() + 14);
if (audio_format == 0xfffe && chunk_size >= 40) {
audio_format = read_le16(fmt.data() + 24);
}
} else if (std::memcmp(chunk_header, "data", 4) == 0) {
data_pos = chunk_data_pos;
data_size = chunk_size;
file.seekg(chunk_size, std::ios::cur);
} else {
file.seekg(chunk_size, std::ios::cur);
}
if (!file.good()) {
break;
}
if ((chunk_size & 1) != 0) {
file.seekg(1, std::ios::cur);
}
}
const uint32_t bytes_per_sample = (bits_per_sample + 7) / 8;
if (data_pos == std::streampos(-1) || data_size == 0 || channels == 0 || sample_rate == 0 ||
block_align == 0 || bytes_per_sample == 0 || block_align < channels * bytes_per_sample ||
(audio_format != 1 && audio_format != 3)) {
return false;
}
const uint64_t frame_count = data_size / block_align;
if (frame_count == 0 || frame_count > SIZE_MAX / channels) {
return false;
}
std::vector<uint8_t> pcm(data_size);
file.clear();
file.seekg(data_pos);
if (!file.read(reinterpret_cast<char*>(pcm.data()), data_size)) {
return false;
}
interleaved_samples.resize(static_cast<size_t>(frame_count * channels));
for (uint64_t frame = 0; frame < frame_count; ++frame) {
const uint8_t* frame_data = pcm.data() + frame * block_align;
for (uint32_t channel = 0; channel < channels; ++channel) {
const uint8_t* sample_data = frame_data + channel * bytes_per_sample;
float sample = 0.0f;
if (audio_format == 3 && bits_per_sample == 32) {
std::memcpy(&sample, sample_data, sizeof(sample));
} else if (audio_format == 3 && bits_per_sample == 64) {
double value;
std::memcpy(&value, sample_data, sizeof(value));
sample = static_cast<float>(value);
} else if (audio_format == 1 && bits_per_sample == 8) {
sample = (static_cast<int>(sample_data[0]) - 128) / 128.0f;
} else if (audio_format == 1 && bits_per_sample == 16) {
sample = static_cast<int16_t>(read_le16(sample_data)) / 32768.0f;
} else if (audio_format == 1 && bits_per_sample == 24) {
int32_t value = static_cast<int32_t>(sample_data[0]) |
(static_cast<int32_t>(sample_data[1]) << 8) |
(static_cast<int32_t>(sample_data[2]) << 16);
if ((value & 0x800000) != 0) {
value |= ~0xffffff;
}
sample = value / 8388608.0f;
} else if (audio_format == 1 && bits_per_sample == 32) {
int32_t value = static_cast<int32_t>(read_le32(sample_data));
sample = value / 2147483648.0f;
} else {
interleaved_samples.clear();
return false;
}
interleaved_samples[static_cast<size_t>(frame * channels + channel)] = sample;
}
}
return true;
}
+5
View File
@@ -110,4 +110,9 @@ bool write_wav_to_file(const std::string& path,
uint32_t channels,
uint32_t sample_rate);
bool load_wav_from_file(const std::string& path,
std::vector<float>& interleaved_samples,
uint32_t& sample_rate,
uint32_t& channels);
#endif // __MEDIA_IO_H__
+31
View File
@@ -141,6 +141,37 @@ public:
}
};
class SDAudioOwner {
private:
uint32_t sample_rate_ = 0;
uint32_t channels_ = 0;
std::vector<float> samples_;
public:
SDAudioOwner() = default;
void reset(std::vector<float> samples = {}, uint32_t sample_rate = 0, uint32_t channels = 0) {
samples_ = std::move(samples);
sample_rate_ = sample_rate;
channels_ = channels;
}
bool empty() const {
return samples_.empty();
}
sd_audio_t get() {
return {sample_rate_,
channels_,
channels_ == 0 ? 0 : static_cast<uint64_t>(samples_.size() / channels_),
samples_.empty() ? nullptr : samples_.data()};
}
const std::vector<float>& samples() const {
return samples_;
}
};
class SDImageVec {
private:
std::vector<sd_image_t> images_;
+14
View File
@@ -56,6 +56,7 @@ enum sample_method_t {
EULER_GE_SAMPLE_METHOD,
DPMPP2M_SDE_SAMPLE_METHOD,
DPMPP2M_SDE_BT_SAMPLE_METHOD,
LMS_SAMPLE_METHOD,
SAMPLE_METHOD_COUNT
};
@@ -246,6 +247,13 @@ typedef struct {
uint8_t* data;
} sd_image_t;
typedef struct {
sd_image_t* frames;
int frame_count;
int fps;
sd_audio_t audio;
} sd_ref_video_t;
typedef struct {
int* layers;
size_t layer_count;
@@ -396,6 +404,12 @@ typedef struct {
int clip_skip;
sd_image_t init_image;
sd_image_t end_image;
sd_image_t* ref_images;
int ref_images_count;
sd_ref_video_t* ref_videos;
int ref_videos_count;
sd_audio_t* ref_audios;
int ref_audios_count;
sd_image_t* control_frames;
int control_frames_size;
int width;
+335
View File
@@ -0,0 +1,335 @@
#!/usr/bin/env python3
"""Merge selected tensors from multiple safetensors files without loading weights.
Edit ``OUTPUT_PATH`` and ``SOURCE_RULES`` below, then run:
python scripts/merge_safetensors.py
Each source rule uses regular expressions against complete tensor names.
``include`` is required and matches when any expression succeeds. ``exclude``
wins over ``include``. Expressions are evaluated with ``re.search``.
"""
import json
import os
import re
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO
# -----------------------------------------------------------------------------
# Configuration
# -----------------------------------------------------------------------------
OUTPUT_PATH = Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_bf16.safetensors")
SOURCE_RULES = [
{
"path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_bf16.safetensors"),
"include": [r".*"],
"exclude": [r".*adaln_proj\.linear.*", r"time_embedder.*"],
},
{
"path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_int8_convrot.safetensors"),
"include": [r"^.*adaln_proj\.linear.*", "adaln_t_table"],
"exclude": [],
},
]
# Safetensors metadata is optional. Set this to a dict[str, str] if needed.
OUTPUT_METADATA = None
# Refuse to replace an existing output unless explicitly enabled.
OVERWRITE_OUTPUT = False
# Only tensor headers and this fixed-size buffer are held in memory.
COPY_BUFFER_SIZE = 8 * 1024 * 1024
PROGRESS_INTERVAL = 1024 * 1024 * 1024
MAX_HEADER_SIZE = 256 * 1024 * 1024
@dataclass(frozen=True)
class TensorEntry:
name: str
source_path: Path
source_data_offset: int
source_start: int
source_end: int
dtype: str
shape: list[int]
@property
def size(self) -> int:
return self.source_end - self.source_start
def format_bytes(size: int) -> str:
value = float(size)
for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
if value < 1024.0 or unit == "TiB":
return f"{value:.2f} {unit}"
value /= 1024.0
raise AssertionError("unreachable")
def read_exact(file: BinaryIO, size: int, description: str) -> bytes:
data = file.read(size)
if len(data) != size:
raise ValueError(f"truncated {description}: expected {size} bytes, got {len(data)}")
return data
def read_safetensors_header(path: Path) -> tuple[dict, int, int]:
file_size = path.stat().st_size
with path.open("rb") as file:
header_size = struct.unpack("<Q", read_exact(file, 8, f"header size in {path}"))[0]
if header_size == 0 or header_size > MAX_HEADER_SIZE:
raise ValueError(
f"invalid header size in {path}: {header_size} "
f"(limit: {MAX_HEADER_SIZE})"
)
header_bytes = read_exact(file, header_size, f"header in {path}")
try:
header = json.loads(header_bytes)
except (UnicodeDecodeError, json.JSONDecodeError) as error:
raise ValueError(f"invalid safetensors JSON header in {path}: {error}") from error
if not isinstance(header, dict):
raise ValueError(f"safetensors header in {path} is not an object")
data_offset = 8 + header_size
if data_offset > file_size:
raise ValueError(f"safetensors data offset is past end of file: {path}")
return header, data_offset, file_size
def parse_tensor_entry(
name: str,
info: object,
source_path: Path,
source_data_offset: int,
source_file_size: int,
) -> TensorEntry:
if not isinstance(info, dict):
raise ValueError(f"{source_path}: tensor {name!r} has an invalid header entry")
dtype = info.get("dtype")
shape = info.get("shape")
offsets = info.get("data_offsets")
if not isinstance(dtype, str):
raise ValueError(f"{source_path}: tensor {name!r} has an invalid dtype")
if not isinstance(shape, list) or not all(
isinstance(dimension, int) and dimension >= 0 for dimension in shape
):
raise ValueError(f"{source_path}: tensor {name!r} has an invalid shape")
if (
not isinstance(offsets, list)
or len(offsets) != 2
or not all(isinstance(offset, int) for offset in offsets)
):
raise ValueError(f"{source_path}: tensor {name!r} has invalid data offsets")
start, end = offsets
if start < 0 or end < start or source_data_offset + end > source_file_size:
raise ValueError(
f"{source_path}: tensor {name!r} byte range [{start}, {end}) "
"is outside the file"
)
return TensorEntry(
name=name,
source_path=source_path,
source_data_offset=source_data_offset,
source_start=start,
source_end=end,
dtype=dtype,
shape=list(shape),
)
def compile_patterns(rule_index: int, field: str, values: object) -> list[re.Pattern[str]]:
if not isinstance(values, list) or not all(isinstance(value, str) for value in values):
raise TypeError(f"SOURCE_RULES[{rule_index}][{field!r}] must be a list of strings")
try:
return [re.compile(value) for value in values]
except re.error as error:
raise ValueError(
f"invalid regex in SOURCE_RULES[{rule_index}][{field!r}]: {error}"
) from error
def collect_entries() -> list[TensorEntry]:
if not SOURCE_RULES:
raise ValueError("SOURCE_RULES must contain at least one source")
entries: list[TensorEntry] = []
selected_by_name: dict[str, TensorEntry] = {}
header_cache: dict[Path, tuple[dict, int, int]] = {}
for rule_index, rule in enumerate(SOURCE_RULES):
if not isinstance(rule, dict) or "path" not in rule or "include" not in rule:
raise TypeError(
f"SOURCE_RULES[{rule_index}] must contain 'path' and 'include'"
)
source_path = Path(rule["path"])
if not source_path.is_file():
raise FileNotFoundError(f"source file does not exist: {source_path}")
source_path = source_path.resolve()
include = compile_patterns(rule_index, "include", rule["include"])
exclude = compile_patterns(rule_index, "exclude", rule.get("exclude", []))
if not include:
raise ValueError(f"SOURCE_RULES[{rule_index}]['include'] must not be empty")
if source_path not in header_cache:
header_cache[source_path] = read_safetensors_header(source_path)
header, data_offset, file_size = header_cache[source_path]
matched = 0
for name, info in header.items():
if name == "__metadata__":
continue
if not any(pattern.search(name) for pattern in include):
continue
if any(pattern.search(name) for pattern in exclude):
continue
entry = parse_tensor_entry(name, info, source_path, data_offset, file_size)
previous = selected_by_name.get(name)
if previous is not None:
raise ValueError(
f"tensor {name!r} was selected more than once:\n"
f" first: {previous.source_path}\n"
f" second: {source_path}"
)
selected_by_name[name] = entry
print(f"entry {entry}")
entries.append(entry)
matched += 1
print(f"Rule {rule_index}: selected {matched} tensors from {source_path}")
if matched == 0:
raise ValueError(
f"SOURCE_RULES[{rule_index}] did not select any tensors; check its regexes"
)
if not entries:
raise ValueError("no tensors were selected")
return entries
def build_output_header(entries: list[TensorEntry]) -> tuple[bytes, int]:
header: dict[str, object] = {}
if OUTPUT_METADATA is not None:
if not isinstance(OUTPUT_METADATA, dict) or not all(
isinstance(key, str) and isinstance(value, str)
for key, value in OUTPUT_METADATA.items()
):
raise TypeError("OUTPUT_METADATA must be None or a dict[str, str]")
header["__metadata__"] = OUTPUT_METADATA
output_offset = 0
for entry in entries:
header[entry.name] = {
"dtype": entry.dtype,
"shape": entry.shape,
"data_offsets": [output_offset, output_offset + entry.size],
}
output_offset += entry.size
header_bytes = json.dumps(header, separators=(",", ":"), ensure_ascii=False).encode(
"utf-8"
)
header_bytes += b" " * (-len(header_bytes) % 8)
return header_bytes, output_offset
def copy_tensor(source: BinaryIO, output: BinaryIO, entry: TensorEntry) -> None:
source.seek(entry.source_data_offset + entry.source_start)
remaining = entry.size
while remaining:
chunk = source.read(min(COPY_BUFFER_SIZE, remaining))
if not chunk:
raise OSError(
f"unexpected end of file while copying {entry.name!r} "
f"from {entry.source_path}"
)
output.write(chunk)
remaining -= len(chunk)
def write_output(entries: list[TensorEntry]) -> None:
if COPY_BUFFER_SIZE <= 0:
raise ValueError("COPY_BUFFER_SIZE must be positive")
output_path = OUTPUT_PATH.resolve()
source_paths = {entry.source_path.resolve() for entry in entries}
if output_path in source_paths:
raise ValueError("OUTPUT_PATH must not be one of the source files")
if output_path.exists() and not OVERWRITE_OUTPUT:
raise FileExistsError(
f"output already exists: {output_path}; set OVERWRITE_OUTPUT = True to replace it"
)
output_path.parent.mkdir(parents=True, exist_ok=True)
partial_path = output_path.with_name(output_path.name + ".partial")
if partial_path.exists():
raise FileExistsError(
f"partial output already exists: {partial_path}; remove it before retrying"
)
header_bytes, tensor_bytes = build_output_header(entries)
print(
f"Writing {len(entries)} tensors ({format_bytes(tensor_bytes)}) to {output_path}"
)
current_source_path: Path | None = None
current_source: BinaryIO | None = None
copied = 0
next_progress = PROGRESS_INTERVAL
try:
with partial_path.open("xb") as output:
output.write(struct.pack("<Q", len(header_bytes)))
output.write(header_bytes)
try:
for entry in entries:
if entry.source_path != current_source_path:
if current_source is not None:
current_source.close()
current_source = entry.source_path.open("rb")
current_source_path = entry.source_path
copy_tensor(current_source, output, entry)
copied += entry.size
if PROGRESS_INTERVAL > 0 and copied >= next_progress:
print(
f" copied {format_bytes(copied)} / "
f"{format_bytes(tensor_bytes)}"
)
while next_progress <= copied:
next_progress += PROGRESS_INTERVAL
finally:
if current_source is not None:
current_source.close()
if copied != tensor_bytes:
raise OSError(f"copied {copied} tensor bytes, expected {tensor_bytes}")
os.replace(partial_path, output_path)
except BaseException:
partial_path.unlink(missing_ok=True)
raise
print(f"Done: {output_path} ({format_bytes(output_path.stat().st_size)})")
def main() -> None:
entries = collect_entries()
write_output(entries)
if __name__ == "__main__":
main()
+190 -10
View File
@@ -2,8 +2,10 @@
#define __SD_CONDITIONING_CONDITIONER_HPP__
#include <cmath>
#include <iomanip>
#include <limits>
#include <optional>
#include <sstream>
#include "core/tensor_ggml.hpp"
#include "core/util.h"
@@ -25,6 +27,8 @@ struct SDCondition {
sd::Tensor<int32_t> c_vinput_mask;
std::vector<std::pair<int, sd::Tensor<float>>> c_image_embeds;
std::vector<sd::Tensor<float>> c_ref_images;
std::vector<sd::Tensor<float>> c_ref_audios;
std::vector<MiniMaxH3ReferenceBlock> c_reference_blocks;
std::vector<sd::Tensor<float>> extra_c_crossattns;
@@ -55,6 +59,12 @@ struct SDCondition {
}
}
for (const auto& tensor : c_ref_audios) {
if (!tensor.empty()) {
return false;
}
}
for (const auto& tensor : extra_c_crossattns) {
if (!tensor.empty()) {
return false;
@@ -65,6 +75,18 @@ struct SDCondition {
}
};
enum class MiniMaxH3PresentationKind {
IMAGE,
VIDEO,
AUDIO,
};
struct MiniMaxH3PresentationItem {
MiniMaxH3PresentationKind kind = MiniMaxH3PresentationKind::IMAGE;
std::vector<sd::Tensor<float>> frames;
std::vector<float> timestamps;
};
static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_states,
const std::vector<float>& weights) {
if (hidden_states.empty()) {
@@ -102,11 +124,12 @@ static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_sta
struct ConditionerParams {
std::string text;
int clip_skip = -1;
int width = -1;
int height = -1;
bool zero_out_masked = false;
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
int clip_skip = -1;
int width = -1;
int height = -1;
bool zero_out_masked = false;
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
const std::vector<MiniMaxH3PresentationItem>* minimax_h3_references = nullptr;
RefImageParams ref_image_params;
};
@@ -117,6 +140,7 @@ public:
virtual SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
@@ -1664,6 +1688,10 @@ struct AnimaConditioner : public Conditioner {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
@@ -1805,6 +1833,7 @@ struct LLMEmbedder : public Conditioner {
sd_version_is_boogu_image(version) ||
sd_version_is_sefi_image(version) ||
sd_version_is_krea2(version) ||
sd_version_is_minimax_h3(version) ||
sd_version_is_mage_flow(version)) {
arch = LLM::LLMArch::QWEN3_VL;
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
@@ -1847,6 +1876,10 @@ struct LLMEmbedder : public Conditioner {
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
if (byt5) {
@@ -1983,8 +2016,10 @@ struct LLMEmbedder : public Conditioner {
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
const std::set<int>& out_layers,
int prompt_template_encode_start_idx,
bool spell_quotes = false,
int max_length = 100000000) {
bool spell_quotes = false,
int max_length = 100000000,
const LLM::DeepStackImageEmbeds& deepstack_image_embeds = {},
const std::vector<LLM::ImageGrid>& image_grids = {}) {
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
auto& tokens = std::get<0>(tokens_weights_mask);
auto& weights = std::get<1>(tokens_weights_mask);
@@ -2017,7 +2052,9 @@ struct LLMEmbedder : public Conditioner {
false,
false,
true,
true);
true,
deepstack_image_embeds,
image_grids);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), weights);
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
@@ -2099,6 +2136,8 @@ struct LLMEmbedder : public Conditioner {
std::vector<std::string> extra_prompts;
std::vector<std::pair<int, int>> extra_prompts_attn_range;
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
LLM::DeepStackImageEmbeds deepstack_image_embeds;
std::vector<LLM::ImageGrid> image_grids;
int prompt_template_encode_start_idx = 34;
int min_length = 0; // pad tokens
int max_length = 100000000;
@@ -2109,7 +2148,131 @@ struct LLMEmbedder : public Conditioner {
int64_t t0 = ggml_time_ms();
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
if (sd_version_is_hunyuan_video(version)) {
if (sd_version_is_minimax_h3(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {50};
prompt_attn_range = {0, 0};
if (llm->enable_vision) {
const std::string placeholder = "<|image_pad|>";
const int patch_size = llm->config.vision.patch_size;
const int factor = patch_size * llm->config.vision.spatial_merge_size;
auto resize_for_vision = [&](const sd::Tensor<float>& image) {
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height,
width,
h_bar,
w_bar,
factor,
3136,
12845056,
RefImageResizeMode::AREA);
auto resized = sd::ops::interpolate(
image,
std::vector<int64_t>{w_bar, h_bar, image.shape()[2], image.shape()[3]});
for (int64_t i = 0; i < resized.numel(); ++i) {
resized[i] = std::clamp(resized[i], 0.f, 1.f) * 2.f - 1.f;
}
return resized;
};
auto add_vision_outputs = [&](std::vector<sd::Tensor<float>> image_outputs,
int grid_h,
int grid_w) {
GGML_ASSERT(image_outputs.size() == 4);
auto image_embed = std::move(image_outputs[0]);
prompt += "<|vision_start|>";
int image_embed_idx = static_cast<int>(tokenizer->encode(prompt, nullptr).size());
image_embeds.emplace_back(image_embed_idx, image_embed);
if (deepstack_image_embeds.empty()) {
deepstack_image_embeds.resize(image_outputs.size() - 1);
}
for (size_t layer = 0; layer < deepstack_image_embeds.size(); ++layer) {
deepstack_image_embeds[layer].emplace_back(image_embed_idx, std::move(image_outputs[layer + 1]));
}
image_grids.push_back({image_embed_idx,
static_cast<int>(image_embed.shape()[1]),
grid_h,
grid_w});
for (int64_t i = 0; i < image_embed.shape()[1]; ++i) {
prompt += placeholder;
}
prompt += "<|vision_end|>";
};
const auto* references = conditioner_params.minimax_h3_references;
if (references != nullptr && !references->empty()) {
int picture_index = 0;
int video_index = 0;
int audio_index = 0;
for (const auto& item : *references) {
if (item.kind == MiniMaxH3PresentationKind::AUDIO) {
prompt += "<Audio " + std::to_string(++audio_index) + ">: ";
continue;
}
if (item.kind == MiniMaxH3PresentationKind::IMAGE) {
GGML_ASSERT(item.frames.size() == 1);
auto resized = resize_for_vision(item.frames[0]);
prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
continue;
}
GGML_ASSERT(!item.frames.empty());
prompt += "<Video " + std::to_string(++video_index) + ">: ";
for (size_t frame = 0; frame < item.frames.size(); frame += 2) {
size_t next = std::min(frame + 1, item.frames.size() - 1);
float t0 = frame < item.timestamps.size() ? item.timestamps[frame] : frame / 2.f;
float t1 = next < item.timestamps.size() ? item.timestamps[next] : next / 2.f;
std::ostringstream timestamp;
timestamp << '<' << std::fixed << std::setprecision(1) << (t0 + t1) * 0.5f << " seconds>";
prompt += timestamp.str();
auto first = resize_for_vision(item.frames[frame]);
auto second = resize_for_vision(item.frames[next]);
if (first.shape()[0] != second.shape()[0] || first.shape()[1] != second.shape()[1]) {
second = sd::ops::interpolate(second,
std::vector<int64_t>{first.shape()[0],
first.shape()[1],
second.shape()[2],
second.shape()[3]});
}
auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
add_vision_outputs(llm->encode_video_block_outputs(n_threads,
pair,
false,
true,
true),
static_cast<int>(first.shape()[1]) / patch_size,
static_cast<int>(first.shape()[0]) / patch_size);
}
}
} else if (conditioner_params.ref_images != nullptr) {
for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
auto resized = resize_for_vision((*conditioner_params.ref_images)[i]);
prompt += "<Picture " + std::to_string(i + 1) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
}
}
}
prompt += conditioner_params.text;
} else if (sd_version_is_hunyuan_video(version)) {
prompt_template_encode_start_idx = 98;
out_layers = {26};
@@ -2657,7 +2820,9 @@ struct LLMEmbedder : public Conditioner {
out_layers,
prompt_template_encode_start_idx,
spell_quotes,
max_length);
max_length,
deepstack_image_embeds,
image_grids);
std::vector<sd::Tensor<float>> extra_hidden_states_vec;
if (sd_version_is_hunyuan_video(version) && byt5) {
std::vector<std::string> quoted_texts;
@@ -2718,6 +2883,17 @@ struct LLMEmbedder : public Conditioner {
SDCondition result;
result.c_crossattn = std::move(hidden_states);
result.extra_c_crossattns = std::move(extra_hidden_states_vec);
if (sd_version_is_minimax_h3(version)) {
std::vector<int32_t> tags(static_cast<size_t>(result.c_crossattn.shape()[1]), 1);
for (const auto& [index, image_embed] : image_embeds) {
int64_t begin = std::max<int64_t>(0, index - 1);
int64_t end = std::min<int64_t>(static_cast<int64_t>(tags.size()),
index + image_embed.shape()[1] + 1);
std::fill(tags.begin() + begin, tags.begin() + end, 0);
}
int64_t tag_count = static_cast<int64_t>(tags.size());
result.c_token_types = sd::Tensor<int32_t>({tag_count}, std::move(tags));
}
return result;
}
};
@@ -2828,6 +3004,10 @@ struct LTXAVEmbedder : public Conditioner {
projector->get_param_tensors(tensors, "text_embedding_projection");
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
projector->set_flash_attention_enabled(enabled);
+47 -11
View File
@@ -1753,7 +1753,7 @@ protected:
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
ggml_backend_sched_t sched = nullptr; // owned, multi-device only
ggml_backend_sched_t sched = nullptr; // owned
ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend
bool multi_device_eval_callback_warned = false;
@@ -2147,8 +2147,22 @@ protected:
return !extra_runtime_backends.empty();
}
bool graph_requires_backend_fallback(ggml_cgraph* gf) const {
if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) {
return false;
}
const int n_nodes = ggml_graph_n_nodes(gf);
for (int i = 0; i < n_nodes; ++i) {
ggml_tensor* node = ggml_graph_node(gf, i);
if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) {
return true;
}
}
return false;
}
bool alloc_compute_buffer(ggml_cgraph* gf) {
if (is_multi_device()) {
if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) {
// The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
// the graph here: reserve runs split_graph, which rewires the
// graph's src pointers to sched-internal copy tensors, and the
@@ -2156,6 +2170,10 @@ protected:
// rewired graph, silently corrupting every cross-backend input. A
// graph must be split at most once; the alloc in execute_graph
// performs the real allocation.
if (compute_allocr != nullptr) {
ggml_gallocr_free(compute_allocr);
compute_allocr = nullptr;
}
return ensure_sched(gf);
}
if (compute_allocr != nullptr) {
@@ -2753,7 +2771,7 @@ protected:
};
ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
if (is_multi_device()) {
if (sched != nullptr) {
ggml_backend_sched_reset(sched);
pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
@@ -2774,9 +2792,9 @@ protected:
}
ggml_status status;
if (is_multi_device()) {
if (sched != nullptr) {
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring",
get_desc().c_str());
multi_device_eval_callback_warned = true;
}
@@ -3018,12 +3036,9 @@ public:
// do copy after alloc graph
void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
if (is_multi_device()) {
// The sched only assigns a backend (and thus a buffer) to tensors
// that participate in the graph; flag standalone data tensors as
// inputs so they get one.
ggml_set_input(tensor);
}
// The scheduler only allocates standalone data tensors when they are
// marked as graph inputs. The flag is harmless for single-backend graphs.
ggml_set_input(tensor);
backend_tensor_data_map[tensor] = data;
}
@@ -3240,6 +3255,11 @@ protected:
virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
virtual enum ggml_op param_usage_op(const std::string& name) const {
(void)name;
return GGML_OP_NONE;
}
public:
void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
if (prefix.size() > 0) {
@@ -3290,6 +3310,18 @@ public:
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
for (auto& pair : blocks) {
pair.second->get_param_tensor_ops(tensor_ops);
}
for (auto& pair : params) {
enum ggml_op op = param_usage_op(pair.first);
if (op != GGML_OP_NONE) {
tensor_ops[pair.second] = op;
}
}
}
virtual std::string get_desc() {
return "GGMLBlock";
}
@@ -3417,6 +3449,10 @@ protected:
params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
}
enum ggml_op param_usage_op(const std::string& name) const override {
return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
}
public:
Embedding(int64_t num_embeddings, int64_t embedding_dim)
: embedding_dim(embedding_dim),
+1 -1
View File
@@ -56,7 +56,7 @@ tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
true);
std::vector<bool> class_token_mask;
for (int i = 0; i < tokens.size(); i++) {
class_token_mask.push_back(class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
class_token_mask.push_back(class_idx >= 0 && class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
}
return std::make_tuple(tokens, weights, class_token_mask);
+6
View File
@@ -43,6 +43,7 @@ enum SDVersion {
VERSION_FLUX2,
VERSION_FLUX2_KLEIN,
VERSION_LTXAV,
VERSION_MINIMAX_H3,
VERSION_HIDREAM_O1,
VERSION_Z_IMAGE,
VERSION_BOOGU_IMAGE,
@@ -123,6 +124,10 @@ static inline bool sd_version_is_ltxav(SDVersion version) {
return false;
}
static inline bool sd_version_is_minimax_h3(SDVersion version) {
return version == VERSION_MINIMAX_H3;
}
static inline bool sd_version_is_wan(SDVersion version) {
if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) {
return true;
@@ -272,6 +277,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
if (sd_version_is_flux(version) ||
sd_version_is_flux2(version) ||
sd_version_is_ltxav(version) ||
sd_version_is_minimax_h3(version) ||
sd_version_is_sd3(version) ||
sd_version_is_wan(version) ||
sd_version_is_lingbot_video(version) ||
+137 -16
View File
@@ -31,8 +31,92 @@ namespace IPAdapter {
}
};
struct Resampler : public GGMLBlock {
int64_t dim = 1280;
int64_t depth = 4;
int64_t num_queries = 16;
int64_t embed_dim = 1280;
int64_t output_dim = 2048;
int64_t ff_inner = 5120;
int64_t dim_head = 64;
int64_t heads = 20;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
params["latents"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, dim, num_queries, 1);
}
Resampler() {}
Resampler(int64_t dim, int64_t depth, int64_t num_queries, int64_t embed_dim, int64_t output_dim, int64_t ff_inner)
: dim(dim), depth(depth), num_queries(num_queries), embed_dim(embed_dim), output_dim(output_dim), ff_inner(ff_inner) {
heads = dim / dim_head;
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, dim, true));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(dim, output_dim, true));
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new LayerNorm(output_dim));
for (int64_t i = 0; i < depth; i++) {
std::string p = "layers." + std::to_string(i);
blocks[p + ".0.norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks[p + ".0.norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks[p + ".0.to_q"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, false));
blocks[p + ".0.to_kv"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * 2, false));
blocks[p + ".0.to_out"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, false));
blocks[p + ".1.0"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks[p + ".1.1"] = std::shared_ptr<GGMLBlock>(new Linear(dim, ff_inner, false));
blocks[p + ".1.3"] = std::shared_ptr<GGMLBlock>(new Linear(ff_inner, dim, false));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image_embeds) {
int64_t N = image_embeds->ne[2];
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
ggml_tensor* x = proj_in->forward(ctx, image_embeds);
ggml_tensor* latents = params["latents"];
if (N > 1) {
latents = ggml_repeat(ctx->ggml_ctx, latents, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, dim, num_queries, N));
}
for (int64_t i = 0; i < depth; i++) {
std::string p = "layers." + std::to_string(i);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".0.norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".0.norm2"]);
auto to_q = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_q"]);
auto to_kv = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_kv"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_out"]);
ggml_tensor* xn = norm1->forward(ctx, x);
ggml_tensor* ln = norm2->forward(ctx, latents);
ggml_tensor* q = to_q->forward(ctx, ln);
ggml_tensor* kv_in = ggml_concat(ctx->ggml_ctx, xn, ln, 1);
ggml_tensor* kv = to_kv->forward(ctx, kv_in);
int64_t L = kv->ne[1];
ggml_tensor* k = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], 0));
ggml_tensor* v = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], dim * kv->nb[0]));
ggml_tensor* attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, heads, nullptr, false, false);
attn = to_out->forward(ctx, attn);
latents = ggml_add(ctx->ggml_ctx, latents, attn);
auto ff_norm = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".1.0"]);
auto ff_fc1 = std::dynamic_pointer_cast<Linear>(blocks[p + ".1.1"]);
auto ff_fc2 = std::dynamic_pointer_cast<Linear>(blocks[p + ".1.3"]);
ggml_tensor* h = ff_norm->forward(ctx, latents);
h = ff_fc1->forward(ctx, h);
h = ggml_gelu_erf(ctx->ggml_ctx, h);
h = ff_fc2->forward(ctx, h);
latents = ggml_add(ctx->ggml_ctx, latents, h);
}
latents = proj_out->forward(ctx, latents);
latents = norm_out->forward(ctx, latents);
return latents;
}
};
struct IPAdapterRunner : public GGMLRunner {
ImageProjModel image_proj;
Resampler resampler;
bool is_plus = false;
int64_t num_tokens = 4;
std::string prefix;
@@ -41,21 +125,54 @@ namespace IPAdapter {
const std::string prefix,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager), prefix(prefix) {
int64_t ctx_dim = 768;
int64_t clip_dim = 1024;
int64_t out_dim = 3072;
auto norm_iter = tensor_storage_map.find(prefix + ".image_proj.norm.weight");
if (norm_iter != tensor_storage_map.end()) {
ctx_dim = norm_iter->second.ne[0];
is_plus = tensor_storage_map.find(prefix + ".image_proj.latents") != tensor_storage_map.end();
if (is_plus) {
int64_t dim = 1280;
int64_t num_queries = 16;
int64_t embed_dim = 1280;
int64_t output_dim = 2048;
int64_t ff_inner = 5120;
auto latents_iter = tensor_storage_map.find(prefix + ".image_proj.latents");
if (latents_iter != tensor_storage_map.end()) {
dim = latents_iter->second.ne[0];
num_queries = latents_iter->second.ne[1];
}
auto proj_in_iter = tensor_storage_map.find(prefix + ".image_proj.proj_in.weight");
if (proj_in_iter != tensor_storage_map.end()) {
embed_dim = proj_in_iter->second.ne[0];
}
auto proj_out_iter = tensor_storage_map.find(prefix + ".image_proj.proj_out.weight");
if (proj_out_iter != tensor_storage_map.end()) {
output_dim = proj_out_iter->second.ne[1];
}
auto ff_iter = tensor_storage_map.find(prefix + ".image_proj.layers.0.1.1.weight");
if (ff_iter != tensor_storage_map.end()) {
ff_inner = ff_iter->second.ne[1];
}
int64_t depth = 0;
while (tensor_storage_map.find(prefix + ".image_proj.layers." + std::to_string(depth) + ".0.to_q.weight") != tensor_storage_map.end()) {
depth++;
}
num_tokens = num_queries;
resampler = Resampler(dim, depth, num_queries, embed_dim, output_dim, ff_inner);
resampler.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
} else {
int64_t ctx_dim = 768;
int64_t clip_dim = 1024;
int64_t out_dim = 3072;
auto norm_iter = tensor_storage_map.find(prefix + ".image_proj.norm.weight");
if (norm_iter != tensor_storage_map.end()) {
ctx_dim = norm_iter->second.ne[0];
}
auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
if (proj_iter != tensor_storage_map.end()) {
clip_dim = proj_iter->second.ne[0];
out_dim = proj_iter->second.ne[1];
}
num_tokens = out_dim / ctx_dim;
image_proj = ImageProjModel(num_tokens, ctx_dim, clip_dim);
image_proj.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
}
auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
if (proj_iter != tensor_storage_map.end()) {
clip_dim = proj_iter->second.ne[0];
out_dim = proj_iter->second.ne[1];
}
num_tokens = out_dim / ctx_dim;
image_proj = ImageProjModel(num_tokens, ctx_dim, clip_dim);
image_proj.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
}
std::string get_desc() override {
@@ -63,14 +180,18 @@ namespace IPAdapter {
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string = "") {
image_proj.get_param_tensors(tensors, prefix + ".image_proj");
if (is_plus) {
resampler.get_param_tensors(tensors, prefix + ".image_proj");
} else {
image_proj.get_param_tensors(tensors, prefix + ".image_proj");
}
}
ggml_cgraph* build_graph(const sd::Tensor<float>& image_embeds_tensor) {
ggml_cgraph* gf = new_graph_custom(1024);
ggml_tensor* embeds = make_input(image_embeds_tensor);
auto runner_ctx = get_context();
ggml_tensor* out = image_proj.forward(&runner_ctx, embeds);
ggml_tensor* out = is_plus ? resampler.forward(&runner_ctx, embeds) : image_proj.forward(&runner_ctx, embeds);
ggml_build_forward_expand(gf, out);
return gf;
}
+74 -38
View File
@@ -588,6 +588,10 @@ struct LoraModel : public GGMLRunner {
const std::string& model_tensor_name) {
ggml_tensor* out_diff = nullptr;
int index = 0;
std::vector<std::string> used_tensors;
bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
while (true) {
std::string key;
if (index == 0) {
@@ -595,7 +599,6 @@ struct LoraModel : public GGMLRunner {
} else {
key = model_tensor_name + "." + std::to_string(index);
}
bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
std::string lokr_w1_name = "lora." + key + ".lokr_w1";
std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
@@ -663,7 +666,6 @@ struct LoraModel : public GGMLRunner {
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
applied_lora_tensors.insert(alpha_name);
}
if (rank == 1) {
@@ -678,19 +680,27 @@ struct LoraModel : public GGMLRunner {
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
}
if (lokr_w1)
applied_lora_tensors.insert(lokr_w1_name);
if (lokr_w1_a)
applied_lora_tensors.insert(lokr_w1_a_name);
if (lokr_w1_b)
applied_lora_tensors.insert(lokr_w1_b_name);
if (lokr_w2)
applied_lora_tensors.insert(lokr_w2_name);
if (lokr_w2_a)
applied_lora_tensors.insert(lokr_w2_a_name);
if (lokr_w2_b)
applied_lora_tensors.insert(lokr_w2_b_name);
applied_lora_tensors.insert(alpha_name);
if (lokr_w1) {
used_tensors.push_back(lokr_w1_name);
}
if (lokr_w1_a) {
used_tensors.push_back(lokr_w1_a_name);
}
if (lokr_w1_b) {
used_tensors.push_back(lokr_w1_b_name);
}
if (lokr_w2) {
used_tensors.push_back(lokr_w2_name);
}
if (lokr_w2_a) {
used_tensors.push_back(lokr_w2_a_name);
}
if (lokr_w2_b) {
used_tensors.push_back(lokr_w2_b_name);
}
if (iter != lora_tensors.end()) {
used_tensors.push_back(alpha_name);
}
index++;
continue;
@@ -740,10 +750,8 @@ struct LoraModel : public GGMLRunner {
const int64_t down_in = lora_down->ne[0];
const int64_t down_out = lora_down->ne[1];
const int64_t up_in = lora_up->ne[0];
const int64_t up_out = lora_up->ne[1];
bool compatible = down_in == model_weight->ne[0] &&
up_out == model_weight->ne[1];
bool compatible = down_in == model_weight->ne[0];
if (lora_mid != nullptr) {
compatible = compatible &&
lora_mid->ne[0] == down_out &&
@@ -755,45 +763,43 @@ struct LoraModel : public GGMLRunner {
if (!compatible) {
skipped_incompatible_lora_tensors.insert(lora_down_name);
skipped_incompatible_lora_tensors.insert(lora_up_name);
skipped_incompatible_lora_tensors.insert(lora_mid_name);
skipped_incompatible_lora_tensors.insert(scale_name);
skipped_incompatible_lora_tensors.insert(alpha_name);
if (lora_mid != nullptr) {
skipped_incompatible_lora_tensors.insert(lora_mid_name);
}
if (lora_tensors.find(scale_name) != lora_tensors.end()) {
skipped_incompatible_lora_tensors.insert(scale_name);
} else if (lora_tensors.find(alpha_name) != lora_tensors.end()) {
skipped_incompatible_lora_tensors.insert(alpha_name);
}
if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
LOG_WARN("skip incompatible LoRA tensor |%s|: model shape = [%lld, %lld], down shape = [%lld, %lld], up shape = [%lld, %lld]",
LOG_WARN("skip incompatible LoRA tensor |%s|: model input dim = %lld, down shape = [%lld, %lld], up shape = [%lld, %lld]",
model_tensor_name.c_str(),
static_cast<long long>(model_weight->ne[0]),
static_cast<long long>(model_weight->ne[1]),
static_cast<long long>(down_in),
static_cast<long long>(down_out),
static_cast<long long>(up_in),
static_cast<long long>(up_out));
static_cast<long long>(lora_up->ne[1]));
}
index++;
continue;
}
}
applied_lora_tensors.insert(lora_up_name);
applied_lora_tensors.insert(lora_down_name);
if (lora_mid) {
applied_lora_tensors.insert(lora_mid_name);
}
float scale_value = 1.0f;
std::string scale_tensor_name;
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
iter = lora_tensors.find(scale_name);
if (iter != lora_tensors.end()) {
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
applied_lora_tensors.insert(scale_name);
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
scale_tensor_name = scale_name;
} else {
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
scale_tensor_name = alpha_name;
// LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
applied_lora_tensors.insert(alpha_name);
}
}
scale_value *= multiplier;
@@ -853,15 +859,45 @@ struct LoraModel : public GGMLRunner {
}
auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
if (out_diff == nullptr) {
out_diff = curr_out_diff;
} else {
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, is_conv2d ? 2 : 0);
}
used_tensors.push_back(lora_up_name);
used_tensors.push_back(lora_down_name);
if (lora_mid) {
used_tensors.push_back(lora_mid_name);
}
if (!scale_tensor_name.empty()) {
used_tensors.push_back(scale_tensor_name);
}
index++;
}
if (out_diff == nullptr)
return nullptr;
int64_t expected_out_dim = is_conv2d ? model_weight->ne[3] : model_weight->ne[1];
int64_t actual_out_dim = out_diff->ne[is_conv2d ? 2 : 0];
if (actual_out_dim != expected_out_dim) {
for (const auto& name : used_tensors) {
skipped_incompatible_lora_tensors.insert(name);
}
if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
LOG_WARN("skip incompatible LoRA tensors for |%s|: output dim %lld != model dim %lld",
model_tensor_name.c_str(), actual_out_dim, expected_out_dim);
}
return nullptr;
}
for (const auto& name : used_tensors) {
applied_lora_tensors.insert(name);
}
return out_diff;
}
+20 -7
View File
@@ -180,9 +180,12 @@ namespace Krea2 {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* scale = params["scale"];
scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
x = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
if (ctx->weight_adapter) {
scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight");
}
scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
x = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
return x;
}
};
@@ -295,10 +298,11 @@ namespace Krea2 {
class KreaDoubleSharedModulation : public GGMLBlock {
protected:
int64_t dim;
std::string prefix;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
GGML_UNUSED(tensor_storage_map);
GGML_UNUSED(prefix);
this->prefix = prefix;
params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6);
}
@@ -307,7 +311,11 @@ namespace Krea2 {
: dim(dim) {}
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
auto lin = ggml_repeat(ctx->ggml_ctx, params["lin"], vec);
auto lin = params["lin"];
if (ctx->weight_adapter) {
lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
}
lin = ggml_repeat(ctx->ggml_ctx, lin, vec);
auto out = ggml_add(ctx->ggml_ctx, vec, lin);
return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
}
@@ -316,10 +324,11 @@ namespace Krea2 {
class KreaFinalModulation : public GGMLBlock {
protected:
int64_t dim;
std::string prefix;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
GGML_UNUSED(tensor_storage_map);
GGML_UNUSED(prefix);
this->prefix = prefix;
params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2);
}
@@ -328,7 +337,11 @@ namespace Krea2 {
: dim(dim) {}
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
auto out = ggml_add(ctx->ggml_ctx, params["lin"], vec);
auto lin = params["lin"];
if (ctx->weight_adapter) {
lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
}
auto out = ggml_add(ctx->ggml_ctx, lin, vec);
return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
}
};
File diff suppressed because it is too large Load Diff
+24
View File
@@ -87,6 +87,29 @@ struct LTXAVDiffusionExtra {
const sd::Tensor<float>* video_positions = nullptr;
};
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 MiniMaxH3DiffusionExtra {
const sd::Tensor<int32_t>* text_token_tags = nullptr;
const sd::Tensor<int32_t>* keyframe_indices = nullptr;
const std::vector<sd::Tensor<float>>* reference_audio_latents = nullptr;
const std::vector<MiniMaxH3ReferenceBlock>* reference_blocks = nullptr;
int audio_length = 0;
float video_sigma_shift = 12.f;
float audio_sigma_shift = 3.f;
};
struct MiniT2IDiffusionExtra {
const sd::Tensor<float>* mask = nullptr;
};
@@ -106,6 +129,7 @@ using DiffusionExtraParams = std::variant<std::monostate,
WanDiffusionExtra,
HiDreamO1DiffusionExtra,
LTXAVDiffusionExtra,
MiniMaxH3DiffusionExtra,
MiniT2IDiffusionExtra,
HunyuanVideoDiffusionExtra>;
+373 -26
View File
@@ -79,9 +79,20 @@ namespace LLM {
int window_size = 112;
int num_position_embeddings = 0;
std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
bool split_patch_embed = false;
std::vector<int> deepstack_visual_indexes;
bool split_patch_embed = false;
};
struct ImageGrid {
int index = 0;
int size = 0;
int grid_h = 0;
int grid_w = 0;
};
using ImageEmbeds = std::vector<std::pair<int, sd::Tensor<float>>>;
using DeepStackImageEmbeds = std::vector<ImageEmbeds>;
struct LLMConfig {
LLMArch arch = LLMArch::QWEN2_5_VL;
int64_t num_layers = 28;
@@ -93,6 +104,7 @@ namespace LLM {
bool qkv_bias = true;
bool attention_out_bias = false;
bool qk_norm = false;
bool final_norm = true;
bool rms_norm_add = false;
bool normalize_input = false;
int64_t vocab_size = 152064;
@@ -257,9 +269,20 @@ namespace LLM {
if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
config.num_heads = 16;
}
if (arch == LLMArch::QWEN3_VL && config.num_layers == 50 && config.hidden_size == 5120) {
config.num_heads = 64;
config.final_norm = false;
}
if (detected_vision_layers > 0) {
config.vision.num_layers = detected_vision_layers;
}
if (arch == LLMArch::QWEN3_VL) {
if (config.vision.num_layers == 24) {
config.vision.deepstack_visual_indexes = {5, 11, 17};
} else if (config.vision.num_layers == 27) {
config.vision.deepstack_visual_indexes = {8, 16, 24};
}
}
LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
config.num_layers,
config.vocab_size,
@@ -541,6 +564,37 @@ namespace LLM {
return input_embed;
}
static ggml_tensor* add_deepstack_image_embeds(GGMLRunnerContext* ctx,
ggml_tensor* x,
const std::vector<std::pair<int, ggml_tensor*>>& image_embeds) {
if (image_embeds.empty()) {
return x;
}
GGML_ASSERT(x->ne[2] == 1);
auto raw_x = ggml_cast(ctx->ggml_ctx, x, image_embeds[0].second->type);
int64_t token_start = 0;
ggml_tensor* output = nullptr;
for (const auto& [index, image_embed] : image_embeds) {
GGML_ASSERT(index >= token_start);
GGML_ASSERT(index + image_embed->ne[1] <= raw_x->ne[1]);
if (index > token_start) {
auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, index);
output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
}
auto visual_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, index, index + image_embed->ne[1]);
visual_embed = ggml_add(ctx->ggml_ctx, visual_embed, image_embed);
output = output == nullptr ? visual_embed : ggml_concat(ctx->ggml_ctx, output, visual_embed, 1);
token_start = index + image_embed->ne[1];
}
if (token_start < raw_x->ne[1]) {
auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, raw_x->ne[1]);
output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
}
GGML_ASSERT(output != nullptr && output->ne[1] == raw_x->ne[1]);
return output;
}
struct VisionMLP : public GGMLBlock {
protected:
LLMVisionArch arch_;
@@ -723,6 +777,33 @@ 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<LayerNorm>(merge_dim, 1e-6f);
blocks["linear_fc1"] = std::make_shared<Linear>(merge_dim, merge_dim, true);
blocks["linear_fc2"] = std::make_shared<Linear>(merge_dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
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 VisionAttention : public GGMLBlock {
protected:
bool llama_cpp_style;
@@ -844,6 +925,7 @@ namespace LLM {
int spatial_merge_size;
int num_grid_per_side;
std::set<int> fullatt_block_indexes;
std::vector<int> deepstack_visual_indexes;
public:
VisionModel(bool llama_cpp_style,
@@ -853,7 +935,8 @@ namespace LLM {
num_layers(vision_params.num_layers),
spatial_merge_size(vision_params.spatial_merge_size),
num_grid_per_side(vision_params.num_position_embeddings > 0 ? static_cast<int>(std::sqrt(vision_params.num_position_embeddings)) : 0),
fullatt_block_indexes(vision_params.fullatt_block_indexes) {
fullatt_block_indexes(vision_params.fullatt_block_indexes),
deepstack_visual_indexes(vision_params.deepstack_visual_indexes) {
blocks["patch_embed"] = std::shared_ptr<GGMLBlock>(new VisionPatchEmbed(vision_params.split_patch_embed,
arch_,
vision_params.patch_size,
@@ -875,6 +958,11 @@ namespace LLM {
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<Qwen3VLDeepStackMerger>(vision_params.out_hidden_size,
vision_params.hidden_size,
spatial_merge_size);
}
}
std::shared_ptr<Embedding> pos_embedder() {
@@ -893,13 +981,13 @@ namespace LLM {
return spatial_merge_size;
}
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,
ggml_tensor* pos_embeds = nullptr) {
std::vector<ggml_tensor*> 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_tensor* pos_embeds = nullptr) {
// pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
// window_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
// window_inverse_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
@@ -919,6 +1007,7 @@ namespace LLM {
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]);
}
std::vector<ggml_tensor*> deepstack_outputs;
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<VisionBlock>(blocks["blocks." + std::to_string(i)]);
@@ -926,8 +1015,12 @@ namespace LLM {
if (fullatt_block_indexes.find(i) != fullatt_block_indexes.end()) {
mask = nullptr;
}
x = block->forward(ctx, x, pe, mask);
if (i == 0) {
x = block->forward(ctx, x, pe, mask);
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<size_t>(std::distance(deepstack_visual_indexes.begin(), deepstack_it));
auto deepstack_merger = std::dynamic_pointer_cast<Qwen3VLDeepStackMerger>(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");
}
@@ -939,7 +1032,19 @@ namespace LLM {
x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index);
}
return x;
std::vector<ggml_tensor*> outputs = {x};
outputs.insert(outputs.end(), deepstack_outputs.begin(), deepstack_outputs.end());
return outputs;
}
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,
ggml_tensor* pos_embeds = nullptr) {
return forward_outputs(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask, pos_embeds)[0];
}
};
@@ -1263,7 +1368,9 @@ namespace LLM {
for (int i = 0; i < num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, i));
}
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
if (config.final_norm) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
}
}
ggml_tensor* embed(GGMLRunnerContext* ctx,
@@ -1278,9 +1385,11 @@ namespace LLM {
ggml_tensor* input_pos,
ggml_tensor* attention_mask,
std::set<int> out_layers,
ggml_tensor* sliding_attention_mask = nullptr,
bool return_all_hidden_states = false) {
auto norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"]);
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds = {},
ggml_tensor* sliding_attention_mask = nullptr,
bool return_all_hidden_states = false) {
auto norm = config.final_norm ? std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])
: nullptr;
std::vector<ggml_tensor*> intermediate_outputs;
if (config.normalize_input) {
@@ -1295,6 +1404,9 @@ namespace LLM {
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask);
if (i < static_cast<int>(deepstack_image_embeds.size())) {
x = add_deepstack_image_embeds(ctx, x, deepstack_image_embeds[static_cast<size_t>(i)]);
}
if (return_all_hidden_states || out_layers.size() > 1) {
x = ggml_cont(ctx->ggml_ctx, x);
}
@@ -1308,7 +1420,7 @@ namespace LLM {
}
}
auto normed_x = norm->forward(ctx, x);
auto normed_x = norm == nullptr ? x : norm->forward(ctx, x);
if (return_all_hidden_states) {
intermediate_outputs.push_back(normed_x);
x = intermediate_outputs[0];
@@ -1336,6 +1448,7 @@ namespace LLM {
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
// input_ids: [N, n_token]
@@ -1347,6 +1460,7 @@ namespace LLM {
input_pos,
attention_mask,
std::move(out_layers),
deepstack_image_embeds,
sliding_attention_mask,
return_all_hidden_states);
}
@@ -1372,6 +1486,7 @@ namespace LLM {
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
// input_ids: [N, n_token]
@@ -1383,6 +1498,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states);
return x;
@@ -1524,7 +1640,8 @@ namespace LLM {
std::vector<float>& window_mask_vec,
std::vector<float>& pe_vec,
std::array<std::vector<int32_t>, 4>& pos_embed_idx_data,
std::array<std::vector<float>, 4>& pos_embed_weight_data) {
std::array<std::vector<float>, 4>& pos_embed_weight_data,
std::vector<ggml_tensor*>* output_tensors = nullptr) {
GGML_ASSERT(image->ne[1] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
GGML_ASSERT(image->ne[0] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
@@ -1556,7 +1673,11 @@ namespace LLM {
int pos_len = static_cast<int>(pe_vec.size() / head_dim / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
runner->set_backend_tensor_data(pe, pe_vec.data());
return vision_model->forward(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds);
auto outputs = vision_model->forward_outputs(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds);
if (output_tensors != nullptr) {
*output_tensors = outputs;
}
return outputs[0];
}
int llm_grid_h = grid_h / vision_params.spatial_merge_size;
@@ -1622,7 +1743,11 @@ namespace LLM {
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
runner->set_backend_tensor_data(pe, pe_vec.data());
return vision_model->forward(runner_ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
auto output = vision_model->forward(runner_ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
if (output_tensors != nullptr) {
*output_tensors = {output};
}
return output;
}
public:
@@ -1657,12 +1782,17 @@ namespace LLM {
model.get_param_tensors(tensors, prefix);
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
model.get_param_tensor_ops(tensor_ops);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* input_pos,
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
auto hidden_states = model.forward(ctx,
@@ -1671,6 +1801,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states); // [N, n_token, hidden_size]
return hidden_states;
@@ -1689,7 +1820,9 @@ namespace LLM {
ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
const sd::Tensor<float>& attention_mask_tensor,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
const ImageEmbeds& image_embeds_tensor,
const DeepStackImageEmbeds& deepstack_image_embeds_tensor,
const std::vector<ImageGrid>& image_grids,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
@@ -1700,6 +1833,13 @@ namespace LLM {
ggml_tensor* embed = make_input(embed_tensor);
image_embeds.emplace_back(idx, embed);
}
std::vector<std::vector<std::pair<int, ggml_tensor*>>> deepstack_image_embeds(deepstack_image_embeds_tensor.size());
for (size_t layer = 0; layer < deepstack_image_embeds_tensor.size(); ++layer) {
deepstack_image_embeds[layer].reserve(deepstack_image_embeds_tensor[layer].size());
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 (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
@@ -1720,6 +1860,30 @@ namespace LLM {
input_pos_vec[2 * n_tokens + i] = i;
input_pos_vec[3 * n_tokens + i] = 0;
}
if (config.arch == LLMArch::QWEN3_VL && !image_grids.empty()) {
int offset = 0;
for (const auto& grid : image_grids) {
int end = grid.index + grid.size;
int grid_h = grid.grid_h / config.vision.spatial_merge_size;
int grid_w = grid.grid_w / config.vision.spatial_merge_size;
int len_max = std::max(grid_h, grid_w);
int next_pos = grid.index + len_max + offset;
GGML_ASSERT(grid.index >= 0 && end <= n_tokens);
GGML_ASSERT(grid_h > 0 && grid_w > 0 && grid.size == grid_h * grid_w);
for (int token = end; token < n_tokens; ++token) {
int pos = next_pos + token - end;
input_pos_vec[token] = pos;
input_pos_vec[n_tokens + token] = pos;
input_pos_vec[2 * n_tokens + token] = pos;
}
for (int token = 0; token < grid.size; ++token) {
input_pos_vec[grid.index + token] = grid.index + offset;
input_pos_vec[n_tokens + grid.index + token] = grid.index + offset + token / grid_w;
input_pos_vec[2 * n_tokens + grid.index + token] = grid.index + offset + token % grid_w;
}
offset += len_max - grid.size;
}
}
}
auto input_pos = ggml_new_tensor_1d(compute_ctx,
@@ -1777,6 +1941,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states);
@@ -1788,16 +1953,20 @@ namespace LLM {
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<float>& attention_mask,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
const ImageEmbeds& image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true,
const DeepStackImageEmbeds& deepstack_image_embeds = {},
const std::vector<ImageGrid>& image_grids = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids,
attention_mask,
image_embeds,
deepstack_image_embeds,
image_grids,
out_layers,
return_all_hidden_states);
};
@@ -1847,6 +2016,24 @@ namespace LLM {
pos_embed_weight_data_);
}
std::vector<ggml_tensor*> encode_image_outputs(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
std::vector<ggml_tensor*> outputs;
encode_image_common(this,
compute_ctx,
runner_ctx,
image,
config.vision,
model.vision_model(),
window_index_vec,
window_inverse_index_vec,
window_mask_vec,
pe_vec,
pos_embed_idx_data_,
pos_embed_weight_data_,
&outputs);
return outputs;
}
ggml_cgraph* build_encode_image_graph(const sd::Tensor<float>& image_tensor) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
@@ -1871,6 +2058,166 @@ namespace LLM {
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
}
ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
auto runner_ctx = get_context();
auto outputs = encode_image_outputs(&runner_ctx, image);
GGML_ASSERT(!outputs.empty());
auto combined = outputs[0];
for (size_t i = 1; i < outputs.size(); ++i) {
combined = ggml_concat(compute_ctx, combined, outputs[i], 0);
}
ggml_build_forward_expand(gf, combined);
return gf;
}
static sd::Tensor<float> process_video_block_tensor(const sd::Tensor<float>& frames,
const LLMVisionConfig& vision_params) {
GGML_ASSERT(frames.dim() == 5);
GGML_ASSERT(frames.shape()[2] == vision_params.temporal_patch_size);
GGML_ASSERT(frames.shape()[3] == vision_params.in_channels);
GGML_ASSERT(frames.shape()[4] == 1);
int64_t width = frames.shape()[0];
int64_t height = frames.shape()[1];
int64_t temporal = frames.shape()[2];
int64_t channels = frames.shape()[3];
int64_t patch = vision_params.patch_size;
int64_t merge = vision_params.spatial_merge_size;
int64_t grid_w = width / patch;
int64_t grid_h = height / patch;
int64_t feature = channels * temporal * patch * patch;
int64_t token_count = grid_h * grid_w;
sd::Tensor<float> output({feature, token_count});
int64_t token = 0;
for (int64_t block_h = 0; block_h < grid_h / merge; ++block_h) {
for (int64_t block_w = 0; block_w < grid_w / merge; ++block_w) {
for (int64_t inner_h = 0; inner_h < merge; ++inner_h) {
for (int64_t inner_w = 0; inner_w < merge; ++inner_w) {
int64_t patch_h = block_h * merge + inner_h;
int64_t patch_w = block_w * merge + inner_w;
int64_t offset = 0;
for (int64_t c = 0; c < channels; ++c) {
for (int64_t t = 0; t < temporal; ++t) {
for (int64_t y = 0; y < patch; ++y) {
for (int64_t x = 0; x < patch; ++x) {
output.index(offset++, token) =
frames.index(patch_w * patch + x,
patch_h * patch + y,
t,
c,
0);
}
}
}
}
++token;
}
}
}
}
return output;
}
ggml_cgraph* build_encode_video_block_outputs_graph(const sd::Tensor<float>& pixel_values_tensor,
int grid_h,
int grid_w) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
auto pixel_values = make_input(pixel_values_tensor);
auto runner_ctx = get_context();
auto vision = model.vision_model();
int head_dim = static_cast<int>(config.vision.hidden_size / config.vision.num_heads);
auto pos_embeds = build_patch_pos_embeds(&runner_ctx, vision, grid_h, grid_w);
window_index_vec.resize(static_cast<size_t>((grid_h / config.vision.spatial_merge_size) *
(grid_w / config.vision.spatial_merge_size)));
for (int i = 0; i < static_cast<int>(window_index_vec.size()); ++i) {
window_index_vec[static_cast<size_t>(i)] = i;
}
pe_vec = Rope::gen_qwen2vl_pe(grid_h,
grid_w,
config.vision.spatial_merge_size,
window_index_vec,
10000,
{head_dim / 2, head_dim / 2});
int pos_len = static_cast<int>(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 outputs = vision->forward_outputs(&runner_ctx,
pixel_values,
pe,
nullptr,
nullptr,
nullptr,
pos_embeds);
GGML_ASSERT(!outputs.empty());
auto combined = outputs[0];
for (size_t i = 1; i < outputs.size(); ++i) {
combined = ggml_concat(compute_ctx, combined, outputs[i], 0);
}
ggml_build_forward_expand(gf, combined);
return gf;
}
std::vector<sd::Tensor<float>> encode_image_outputs(const int n_threads,
const sd::Tensor<float>& image,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_image_outputs_graph(image);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
if (combined.empty()) {
return {};
}
size_t output_count = config.vision.deepstack_visual_indexes.size() + 1;
GGML_ASSERT(combined.shape()[0] == config.hidden_size * static_cast<int64_t>(output_count));
std::vector<sd::Tensor<float>> outputs;
outputs.reserve(output_count);
for (size_t i = 0; i < output_count; ++i) {
outputs.push_back(sd::ops::slice(combined,
0,
static_cast<int64_t>(i) * config.hidden_size,
static_cast<int64_t>(i + 1) * config.hidden_size));
}
return outputs;
}
std::vector<sd::Tensor<float>> encode_video_block_outputs(const int n_threads,
const sd::Tensor<float>& frames,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
int grid_h = static_cast<int>(frames.shape()[1] / config.vision.patch_size);
int grid_w = static_cast<int>(frames.shape()[0] / config.vision.patch_size);
auto pixel_values = process_video_block_tensor(frames, config.vision);
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph,
n_threads,
auto_free,
free_compute_buffer,
free_compute_params));
if (combined.empty()) {
return {};
}
size_t output_count = config.vision.deepstack_visual_indexes.size() + 1;
GGML_ASSERT(combined.shape()[0] == config.hidden_size * static_cast<int64_t>(output_count));
std::vector<sd::Tensor<float>> outputs;
outputs.reserve(output_count);
for (size_t i = 0; i < output_count; ++i) {
outputs.push_back(sd::ops::slice(combined,
0,
static_cast<int64_t>(i) * config.hidden_size,
static_cast<int64_t>(i + 1) * config.hidden_size));
}
return outputs;
}
};
struct LLMEmbedder {
+28
View File
@@ -0,0 +1,28 @@
#ifndef __SD_MODEL_VAE_AUDIO_VAE_HPP__
#define __SD_MODEL_VAE_AUDIO_VAE_HPP__
#include "core/ggml_extend.hpp"
struct AudioVAERunner : public GGMLRunner {
AudioVAERunner(ggml_backend_t backend,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {}
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual size_t get_params_mem_size() = 0;
virtual std::string get_desc() = 0;
virtual sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& waveform) {
SD_UNUSED(n_threads);
SD_UNUSED(waveform);
return {};
}
virtual sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) = 0;
virtual int input_sample_rate() const {
return output_sample_rate();
}
virtual int output_sample_rate() const = 0;
};
#endif // __SD_MODEL_VAE_AUDIO_VAE_HPP__
+11 -6
View File
@@ -8,6 +8,7 @@
#include <vector>
#include "core/ggml_extend.hpp"
#include "model/vae/audio_vae.hpp"
#include "model_loader.h"
#include "model_manager.h"
@@ -996,7 +997,7 @@ namespace LTXV {
}
};
struct LTXAudioVAERunner : public GGMLRunner {
struct LTXAudioVAERunner : public AudioVAERunner {
LTXAudioVAEConfig config;
LTXAudioVAE model;
std::string weight_prefix;
@@ -1006,7 +1007,7 @@ namespace LTXV {
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
: AudioVAERunner(backend, weight_manager),
weight_prefix(prefix),
config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
model(config) {
@@ -1017,20 +1018,20 @@ namespace LTXV {
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
size_t get_params_mem_size() {
size_t get_params_mem_size() override {
return model.get_params_mem_size();
}
std::string get_desc() {
std::string get_desc() override {
return "ltx_audio_vae";
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) {
const sd::Tensor<float>& latent_tensor) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto latent = make_input(latent_tensor);
@@ -1047,6 +1048,10 @@ namespace LTXV {
return result;
}
int output_sample_rate() const override {
return config.output_sample_rate();
}
void test(const std::string& input_path) {
auto z = sd::load_tensor_from_file_as_tensor<float>(input_path);
GGML_ASSERT(!z.empty());
+497
View File
@@ -0,0 +1,497 @@
#ifndef __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
#define __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
#include <array>
#include <string>
#include <vector>
#include "model/vae/audio_vae.hpp"
#include "model/vae/ltx_audio_vae.hpp"
namespace MiniMaxH3 {
struct AudioSnake1D : public UnaryBlock {
int64_t channels;
explicit AudioSnake1D(int64_t channels)
: channels(channels) {}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_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 {
int64_t channels;
AudioEncoderResidualUnit(int64_t channels, int dilation)
: channels(channels) {
blocks["block.0"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.1"] = std::make_shared<LTXV::Conv1D>(channels,
channels,
7,
1,
3 * dilation,
dilation);
blocks["block.2"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.3"] = std::make_shared<LTXV::Conv1D>(channels, channels, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto act1 = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.0"]);
auto conv1 = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.1"]);
auto act2 = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.2"]);
auto conv2 = std::dynamic_pointer_cast<LTXV::Conv1D>(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]) {
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 {
int64_t out_channels;
AudioEncoderBlock(int64_t out_channels, int stride)
: out_channels(out_channels) {
int64_t in_channels = out_channels / 2;
blocks["block.0"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 1);
blocks["block.1"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 3);
blocks["block.2"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 9);
blocks["block.3"] = std::make_shared<AudioSnake1D>(in_channels);
blocks["block.4"] = std::make_shared<LTXV::Conv1D>(in_channels,
out_channels,
2 * stride,
stride,
static_cast<int>(std::ceil(stride / 2.f)));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
for (int i = 0; i < 3; ++i) {
auto unit = std::dynamic_pointer_cast<AudioEncoderResidualUnit>(blocks["block." + std::to_string(i)]);
x = unit->forward(ctx, x);
}
auto act = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.3"]);
auto conv = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.4"]);
return conv->forward(ctx, act->forward(ctx, x));
}
};
struct AudioEncoder : public GGMLBlock {
static constexpr std::array<int, 5> strides = {2, 4, 4, 5, 5};
AudioEncoder() {
int64_t channels = 64;
blocks["block.0"] = std::make_shared<LTXV::Conv1D>(1, channels, 7, 1, 3);
for (size_t i = 0; i < strides.size(); ++i) {
channels *= 2;
blocks["block." + std::to_string(i + 1)] = std::make_shared<AudioEncoderBlock>(channels, strides[i]);
}
blocks["block.6"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.7"] = std::make_shared<LTXV::Conv1D>(channels, 2048, 3, 1, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto input = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.0"]);
x = input->forward(ctx, x);
for (size_t i = 0; i < strides.size(); ++i) {
auto block = std::dynamic_pointer_cast<AudioEncoderBlock>(blocks["block." + std::to_string(i + 1)]);
x = block->forward(ctx, x);
}
auto act = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.6"]);
auto out = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.7"]);
return out->forward(ctx, act->forward(ctx, x));
}
};
struct AudioGeGLUMLP : public GGMLBlock {
AudioGeGLUMLP(int64_t hidden_size, int64_t intermediate_size) {
blocks["norm"] = std::make_shared<LayerNorm>(hidden_size);
blocks["w0"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
blocks["w1"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
blocks["w2"] = std::make_shared<Linear>(intermediate_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto w0 = std::dynamic_pointer_cast<Linear>(blocks["w0"]);
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
x = norm->forward(ctx, x);
auto gate = ggml_ext_gelu(ctx->ggml_ctx, w0->forward(ctx, x), true);
return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, w1->forward(ctx, x)));
}
};
struct AudioCausalAttention : public GGMLBlock {
static constexpr int64_t in_channels = 2048;
static constexpr int64_t out_channels = 32;
static constexpr int64_t num_head = 8;
static constexpr int64_t head_dim = in_channels / num_head;
AudioCausalAttention() {
blocks["qkv"] = std::make_shared<Linear>(in_channels, in_channels * 3, false);
blocks["proj"] = std::make_shared<Linear>(out_channels, out_channels, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
params["q_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
params["v_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto qkv_layer = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_layer->forward(ctx, x), 3, 0);
auto bias_shape = [&](ggml_tensor* bias) {
return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1);
};
auto q = ggml_add(ctx->ggml_ctx, qkv[0], bias_shape(params["q_bias"]));
auto k = qkv[1];
auto v = ggml_add(ctx->ggml_ctx, qkv[2], bias_shape(params["v_bias"]));
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 attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
ctx->backend,
q,
k,
v,
num_head,
mask,
false,
ctx->flash_attn_enabled);
int64_t batch = attn_out->ne[2] * attn_out->ne[3];
attn_out = ggml_reshape_4d(ctx->ggml_ctx, attn_out, head_dim, num_head, sequence, batch);
attn_out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, attn_out, 1, 0, 2, 3));
attn_out = ggml_mean(ctx->ggml_ctx, attn_out);
attn_out = ggml_reshape_3d(ctx->ggml_ctx, attn_out, head_dim, sequence, batch);
constexpr int64_t pool = head_dim / out_channels;
attn_out = ggml_reshape_4d(ctx->ggml_ctx, attn_out, pool, out_channels, sequence, batch);
attn_out = ggml_mean(ctx->ggml_ctx, attn_out);
attn_out = ggml_reshape_3d(ctx->ggml_ctx, attn_out, out_channels, sequence, batch);
return proj->forward(ctx, attn_out);
}
};
struct AudioAttentionProjection : public GGMLBlock {
AudioAttentionProjection() {
blocks["norm1"] = std::make_shared<LayerNorm>(2048);
blocks["attn"] = std::make_shared<AudioCausalAttention>();
blocks["proj"] = std::make_shared<Linear>(2048, 32, true);
blocks["norm3"] = std::make_shared<LayerNorm>(2048);
blocks["norm2"] = std::make_shared<LayerNorm>(32);
blocks["mlp"] = std::make_shared<AudioGeGLUMLP>(32, 64);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<AudioCausalAttention>(blocks["attn"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto norm3 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm3"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto mlp = std::dynamic_pointer_cast<AudioGeGLUMLP>(blocks["mlp"]);
x = ggml_add(ctx->ggml_ctx,
proj->forward(ctx, norm3->forward(ctx, x)),
attn->forward(ctx, norm1->forward(ctx, x)));
return ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, norm2->forward(ctx, x)));
}
};
struct AudioAMPBlock : public GGMLBlock {
int channels;
AudioAMPBlock(int channels,
int kernel_size,
const std::array<int, 3>& dilations)
: channels(channels) {
for (int i = 0; i < 3; ++i) {
blocks["activations." + std::to_string(i * 2)] =
std::make_shared<LTXV::Activation1D>(channels);
blocks["activations." + std::to_string(i * 2 + 1)] =
std::make_shared<LTXV::Activation1D>(channels);
blocks["convs1." + std::to_string(i)] =
std::make_shared<LTXV::Conv1D>(channels,
channels,
kernel_size,
1,
(kernel_size * dilations[i] - dilations[i]) / 2,
dilations[i]);
blocks["convs2." + std::to_string(i)] =
std::make_shared<LTXV::Conv1D>(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<LTXV::Activation1D>(
blocks["activations." + std::to_string(i * 2)]);
auto act2 = std::dynamic_pointer_cast<LTXV::Activation1D>(
blocks["activations." + std::to_string(i * 2 + 1)]);
auto conv1 = std::dynamic_pointer_cast<LTXV::Conv1D>(
blocks["convs1." + std::to_string(i)]);
auto conv2 = std::dynamic_pointer_cast<LTXV::Conv1D>(
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<int, num_upsamples> rates = {5, 5, 2, 2, 2, 2, 2};
static constexpr std::array<int, num_upsamples> kernels = {9, 9, 4, 4, 4, 4, 4};
static constexpr std::array<int, num_kernels> res_kernels = {3, 7, 11};
BigVGAN() {
blocks["conv_pre"] = std::make_shared<LTXV::Conv1D>(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<LTXV::ConvTranspose1D>(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<AudioAMPBlock>(next_channels,
res_kernels[j],
std::array<int, 3>{1, 3, 5});
}
channels = next_channels;
}
blocks["activation_post"] = std::make_shared<LTXV::Activation1D>(channels);
blocks["conv_post"] = std::make_shared<LTXV::Conv1D>(channels,
1,
7,
1,
3,
1,
false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv_pre = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_pre"]);
x = conv_pre->forward(ctx, x);
for (int i = 0; i < num_upsamples; ++i) {
auto up = std::dynamic_pointer_cast<LTXV::ConvTranspose1D>(
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<AudioAMPBlock>(
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<LTXV::Activation1D>(blocks["activation_post"]);
auto conv_post = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_post"]);
return ggml_clamp(ctx->ggml_ctx,
conv_post->forward(ctx, activation->forward(ctx, x)),
-1.f,
1.f);
}
};
struct AudioVAE : public GGMLBlock {
static constexpr int kLatentChannels = 32;
AudioVAE() {
blocks["encoder"] = std::make_shared<AudioEncoder>();
blocks["pre_block"] = std::make_shared<AudioAttentionProjection>();
blocks["mean_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels, kLatentChannels, 1);
blocks["logs_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels, kLatentChannels, 1);
blocks["dec_in_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels,
2048,
1);
blocks["decoder"] = std::make_shared<BigVGAN>();
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_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* encode(GGMLRunnerContext* ctx, ggml_tensor* waveform) {
GGML_ASSERT(waveform->ne[1] == 2);
auto encoder = std::dynamic_pointer_cast<AudioEncoder>(blocks["encoder"]);
auto pre = std::dynamic_pointer_cast<AudioAttentionProjection>(blocks["pre_block"]);
auto mean_proj = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["mean_proj"]);
waveform = ggml_reshape_3d(ctx->ggml_ctx, waveform, waveform->ne[0], 1, waveform->ne[1]);
auto x = encoder->forward(ctx, waveform); // [B*S, 2048, T]
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = pre->forward(ctx, x);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
auto z = mean_proj->forward(ctx, x);
auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
auto std = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
z = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
}
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<LTXV::Conv1D>(blocks["dec_in_proj"]);
auto decoder = std::dynamic_pointer_cast<BigVGAN>(blocks["decoder"]);
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 stream_latent = ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1);
auto stream_waveform = decoder->forward(ctx, dec_in->forward(ctx, stream_latent));
waveform = waveform == nullptr
? stream_waveform
: ggml_concat(ctx->ggml_ctx, waveform, stream_waveform, 2);
}
return ggml_reshape_4d(ctx->ggml_ctx,
waveform,
waveform->ne[0],
streams,
1,
1);
}
};
struct AudioVAERunner : public ::AudioVAERunner {
AudioVAE model;
std::string weight_prefix;
AudioVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: ::AudioVAERunner(backend, weight_manager),
weight_prefix(prefix) {
model.init(params_ctx, tensor_storage_map, prefix);
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
size_t get_params_mem_size() override {
return model.get_params_mem_size();
}
std::string get_desc() override {
return "minimax_h3_audio_vae";
}
int output_sample_rate() const override {
return 32000;
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& waveform) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto input = make_input(waveform);
auto runner_ctx = get_context();
auto latent = model.encode(&runner_ctx, input);
auto graph = new_graph_custom(655360);
ggml_build_forward_expand(graph, latent);
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
(t1 - t0) / 1000.f);
return result;
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto latent = make_input(latent_tensor);
auto runner_ctx = get_context();
auto waveform = model.decode(&runner_ctx, latent);
auto graph = new_graph_custom(655360);
ggml_build_forward_expand(graph, waveform);
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",
(t1 - t0) / 1000.f);
return result;
}
};
} // namespace MiniMaxH3
#endif // __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
+805
View File
@@ -0,0 +1,805 @@
#ifndef __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
#define __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
#include <algorithm>
#include <array>
#include <cmath>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/vae/vae.hpp"
namespace MiniMaxH3VAE {
constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144;
struct CausalConv3d : public Conv3d {
std::tuple<int, int, int> temporal_padding;
CausalConv3d(int64_t in_channels,
int64_t out_channels,
std::tuple<int, int, int> kernel_size,
std::tuple<int, int, int> stride = {1, 1, 1},
std::tuple<int, int, int> 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,
ctx->backend,
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<CausalConv3d>(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<CausalConv3d>(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<TemporalGroupNorm>(in_channels);
blocks["norm2"] = std::make_shared<TemporalGroupNorm>(out_channels);
blocks["conv1"] = std::make_shared<CausalConv3d>(in_channels,
out_channels,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
blocks["conv2"] = std::make_shared<CausalConv3d>(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<CausalConv3d>(in_channels,
out_channels,
std::tuple{1, 1, 1});
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm2"]);
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(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<CausalConv3d>(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<int, levels> multipliers = {1, 2, 2, 4, 4, 8};
static constexpr std::array<int, levels> spatial_down = {2, 2, 2, 2, 1, 1};
static constexpr std::array<int, levels> temporal_down = {1, 2, 2, 1, 1, 1};
Encoder() {
constexpr int ch = 128;
blocks["conv_in"] = std::make_shared<CausalConv3d>(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<ResnetBlock3D>(block == 0 ? previous : current,
current);
}
if (spatial_down[level] * temporal_down[level] > 1) {
blocks["down." + std::to_string(level) + ".downsample"] =
std::make_shared<Downsample3D>(current,
current,
temporal_down[level],
spatial_down[level]);
}
previous = current;
}
blocks["norm_out"] = std::make_shared<TemporalGroupNorm>(previous);
blocks["conv_out"] = std::make_shared<CausalConv3d>(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<CausalConv3d>(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<ResnetBlock3D>(
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<Downsample3D>(downsample->second)->forward(ctx, x);
}
}
auto norm = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm_out"]);
auto conv = std::dynamic_pointer_cast<CausalConv3d>(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<Linear>(dim, dim * 3, true);
blocks["to_out"] = std::make_shared<Linear>(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto to_qkv = std::dynamic_pointer_cast<Linear>(blocks["to_qkv"]);
auto to_out = std::dynamic_pointer_cast<Linear>(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<Linear>(dim, kInnerDim * 2, true);
blocks["w2"] = std::make_shared<Linear>(kInnerDim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(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<RMSNorm>(dim, 1e-5f);
blocks["attn"] = std::make_shared<DecoderAttention>();
blocks["norm2"] = std::make_shared<RMSNorm>(dim, 1e-5f);
blocks["ff"] = std::make_shared<DecoderFeedForward>();
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_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<RMSNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<DecoderAttention>(blocks["attn"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto ff = std::dynamic_pointer_cast<DecoderFeedForward>(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<Linear>(24, dim, true);
for (int i = 0; i < num_layers; ++i) {
blocks["transformer_blocks." + std::to_string(i)] =
std::make_shared<DecoderBlock>();
}
blocks["norm_out"] = std::make_shared<LayerNorm>(dim, 1e-5f, true, true);
blocks["proj_out"] = std::make_shared<Linear>(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 {
SD_UNUSED(tensor_storage_map);
SD_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<Linear>(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<DecoderBlock>(
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<LayerNorm>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(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<Encoder>();
blocks["quant_conv"] = std::make_shared<Conv3d>(48,
48,
std::tuple{1, 1, 1});
blocks["post_quant_conv"] = std::make_shared<Conv3d>(24,
24,
std::tuple{1, 1, 1});
blocks["decoder"] = std::make_shared<Decoder>();
}
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<Encoder>(blocks["encoder"]);
auto quant = std::dynamic_pointer_cast<Conv3d>(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<Conv3d>(blocks["post_quant_conv"]);
auto decoder = std::dynamic_pointer_cast<Decoder>(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<float> pixel_mean;
sd::Tensor<float> pixel_std;
sd::Tensor<float> latents_mean;
sd::Tensor<float> latents_std;
sd::Tensor<float> rope_cache;
MiniMaxH3VideoVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "first_stage_model",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(VERSION_MINIMAX_H3, backend, prefix, weight_manager),
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 {
SD_UNUSED(input_channels);
return 24;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output,
std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents * latents_std + latents_mean;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return (latents - latents_mean) / latents_std;
}
static sd::Tensor<float> ensure_video_shape(const sd::Tensor<float>& 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 sd_tiling_params_t h3_tiling(sd_tiling_params_t params) {
params.enabled = true;
params.tile_size_x = 16;
params.tile_size_y = 16;
params.target_overlap = 0.25f;
return params;
}
static sd::Tensor<float> repeat_last_frame(const sd::Tensor<float>& 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<float> blend_temporal(const sd::Tensor<float>& previous,
const sd::Tensor<float>& 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<float>(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<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) override {
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<float> result;
for (int64_t start = 0; start < input.shape()[2]; start += 17) {
auto chunk = sd::ops::slice(input, 2, start, start + 17);
auto encoded = VAE::encode(n_threads, chunk, tiling, circular_x, circular_y);
if (encoded.empty()) {
return {};
}
result = result.empty() ? std::move(encoded)
: sd::ops::concat(result, encoded, 2);
}
if (result.shape()[2] > 3) {
result = sd::ops::slice(result, 2, 0, result.shape()[2] - 3);
}
return result;
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) override {
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<float> result;
sd::Tensor<float> overlap;
for (int64_t i = 0; i < num_chunks; ++i) {
int64_t start = i * tokens_per_chunk;
int64_t end = std::min(start + tokens_per_chunk + token_overlap,
input.shape()[2]);
auto chunk = sd::ops::slice(input, 2, start, end);
auto decoded = VAE::decode(n_threads,
chunk,
tiling,
true,
circular_x,
circular_y,
silent);
if (decoded.empty()) {
return {};
}
int64_t first_end = std::min<int64_t>(frames_per_chunk, decoded.shape()[2]);
auto first = sd::ops::slice(decoded,
2,
std::min<int64_t>(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<int64_t>(1, expected_frames);
if (result.shape()[2] > expected_frames) {
result = sd::ops::slice(result, 2, 0, expected_frames);
}
return result;
}
sd::Tensor<float> build_rope(int64_t width,
int64_t height,
int64_t num_frames) {
std::vector<std::vector<float>> ids;
ids.reserve(static_cast<size_t>(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<int>{16, 16, 16});
return sd::Tensor<float>({2,
2,
24,
static_cast<int64_t>(ids.size())},
std::move(values));
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& 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<float>(get_graph,
n_threads,
false,
false,
false),
5);
}
};
} // namespace MiniMaxH3VAE
#endif // __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
+15 -14
View File
@@ -74,7 +74,7 @@ public:
int scale_factor = 8;
if (version == VERSION_LTXAV) {
scale_factor = 32;
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16;
@@ -115,11 +115,11 @@ public:
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) {
virtual sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) {
int64_t t0 = ggml_time_ms();
sd::Tensor<float> input = x;
sd::Tensor<float> output;
@@ -136,7 +136,8 @@ public:
// Image VAE encode is more sensitive to tile boundary context than decode.
// Keep the smaller legacy factor for video VAEs, but default image encode
// tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f
: 2.0f;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
output = tiled_compute(input,
@@ -166,13 +167,13 @@ public:
return std::move(output);
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) {
virtual sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) {
int64_t t0 = ggml_time_ms();
sd::Tensor<float> input = x;
sd::Tensor<float> output;
+60 -5
View File
@@ -2,6 +2,7 @@
#include <cstdlib>
#include <cstring>
#include <limits>
#include <string>
#include <unordered_map>
#include <utility>
@@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
return false;
}
static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
if (tensor.tensor_storage.nelements() == 0) {
static bool checked_pickle_byte_count(int64_t element_count,
uint64_t element_nbytes,
uint64_t* byte_count) {
if (element_count < 0 || element_nbytes == 0) {
return false;
}
uint64_t count = static_cast<uint64_t>(element_count);
if (count > std::numeric_limits<uint64_t>::max() / element_nbytes) {
return false;
}
*byte_count = count * element_nbytes;
return true;
}
static bool tensor_layout_is_valid(const PickleTensorInfo& tensor, uint64_t raw_element_nbytes) {
if (raw_element_nbytes == 0) {
return false;
}
bool has_zero_dimension = false;
uint64_t element_count = 1;
for (int i = 0; i < tensor.tensor_storage.n_dims; ++i) {
int64_t dimension = tensor.tensor_storage.ne[i];
if (dimension < 0) {
return false;
}
if (dimension == 0) {
has_zero_dimension = true;
continue;
}
uint64_t size = static_cast<uint64_t>(dimension);
if (element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / size) {
return false;
}
element_count *= size;
}
if (!has_zero_dimension &&
element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / raw_element_nbytes) {
return false;
}
if (has_zero_dimension) {
return true;
}
if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
@@ -932,7 +976,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
return false;
}
storage.nbytes = (uint64_t)pid.items[4].int_value * storage.raw_element_nbytes;
if (!checked_pickle_byte_count(pid.items[4].int_value,
storage.raw_element_nbytes,
&storage.nbytes)) {
set_error(error, "invalid storage size in torch pickle");
return false;
}
storage_nbytes[storage.key] = storage.nbytes;
stack.push_back(make_storage_value(storage));
} break;
@@ -963,7 +1012,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.tensor_storage.is_f64 = args.items[0].storage.is_f64;
tensor.tensor_storage.is_i64 = args.items[0].storage.is_i64;
tensor.tensor_storage.storage_key = args.items[0].storage.key;
tensor.tensor_storage.offset = (uint64_t)args.items[1].int_value * args.items[0].storage.raw_element_nbytes;
if (!checked_pickle_byte_count(args.items[1].int_value,
args.items[0].storage.raw_element_nbytes,
&tensor.tensor_storage.offset)) {
set_error(error, "invalid tensor storage offset in torch pickle");
return false;
}
for (const auto& item : args.items[2].items) {
if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
@@ -979,7 +1033,8 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.stride[tensor.stride_n_dims++] = item.int_value;
}
if (!tensor_is_contiguous(tensor)) {
if (!tensor_layout_is_valid(tensor, args.items[0].storage.raw_element_nbytes)) {
set_error(error, "invalid tensor shape or stride in torch pickle");
return false;
}
stack.push_back(make_tensor_value(tensor));
+12 -5
View File
@@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
if (it == legacy_storage_map.end()) {
return false;
}
if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
if (current_offset > file_size ||
LEGACY_STORAGE_HEADER_SIZE > file_size - current_offset) {
return false;
}
storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
if (it->second > file_size - storage_offset) {
return false;
}
storage_offsets[storage_key] = storage_offset;
current_offset = storage_offset + it->second;
}
for (auto& tensor_storage : tensor_storages) {
@@ -159,8 +164,10 @@ bool read_torch_legacy_file(const std::string& file_path,
uint64_t base_offset = it_offset->second;
uint64_t storage_nbytes = it_size->second;
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > storage_nbytes ||
static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
return false;
}
+4 -2
View File
@@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
return false;
}
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > entry_size) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > entry_size ||
static_cast<uint64_t>(tensor_nbytes) > entry_size - tensor_storage.offset) {
set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
return false;
}
+46 -21
View File
@@ -537,6 +537,10 @@ SDVersion ModelLoader::get_sd_version() {
if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) {
return VERSION_LTXAV;
}
if (tensor_storage.name.find("model.diffusion_model.video_patch_proj.weight") != std::string::npos &&
tensor_storage_map.find("model.diffusion_model.audio_patch_proj.weight") != tensor_storage_map.end()) {
return VERSION_MINIMAX_H3;
}
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
is_wan = true;
}
@@ -1053,7 +1057,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
std::atomic<size_t> tensor_idx(0);
std::atomic<bool> failed(false);
std::vector<std::thread> workers;
std::mutex rpc_backend_mutex;
std::mutex backend_tensor_set_mutex;
for (int i = 0; i < n_threads; ++i) {
workers.emplace_back([&, file_path, is_zip]() {
@@ -1077,6 +1081,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
std::vector<uint8_t> read_buffer;
std::vector<uint8_t> convert_buffer;
std::vector<uint8_t> zip_entry_buffer;
while (true) {
int64_t t0, t1;
@@ -1115,34 +1120,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
auto read_data = [&](char* buf, size_t n) {
auto read_data = [&](char* buf, size_t n) -> bool {
if (zip != nullptr) {
zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
return false;
}
size_t entry_size = zip_entry_size(zip);
if (tensor_storage.offset > entry_size) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
if (n > entry_size - tensor_offset) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
if (entry_size != n) {
int64_t t_memcpy_start;
read_buffer.resize(entry_size);
zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
zip_entry_buffer.resize(entry_size);
auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
t_memcpy_start = ggml_time_ms();
memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
} else {
zip_entry_noallocread(zip, (void*)buf, n);
auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
}
zip_entry_close(zip);
} else if (mmapped) {
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
} else {
file.seekg(tensor_storage.offset);
file.read(buf, n);
if (!file) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
}
return true;
};
char* read_buf = nullptr;
@@ -1176,7 +1207,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
}
t0 = ggml_time_ms();
read_data(read_buf, nbytes_to_read);
if (!read_data(read_buf, nbytes_to_read)) {
failed = true;
break;
}
t1 = ggml_time_ms();
read_time_ms.fetch_add(t1 - t0);
@@ -1214,17 +1248,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
t0 = ggml_time_ms();
// RPC backends require serialized access to prevent concurrency issues
const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
bool is_rpc_buffer = buffer_type_name != nullptr &&
std::string(buffer_type_name).find("RPC") != std::string::npos;
if (is_rpc_buffer) {
std::lock_guard<std::mutex> lock(rpc_backend_mutex);
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
} else {
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
}
std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
t1 = ggml_time_ms();
copy_to_backend_time_ms.fetch_add(t1 - t0);
+66 -1
View File
@@ -53,6 +53,48 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) {
return props.caps.buffer_from_host_ptr;
}
static bool device_supports_param_op(ggml_backend_dev_t device,
ggml_tensor* weight,
enum ggml_op op,
ggml_backend_buffer_type_t buft) {
if (op == GGML_OP_NONE) {
return true;
}
if (device == nullptr || weight == nullptr || buft == nullptr || weight->buffer != nullptr) {
return false;
}
ggml_init_params params;
params.mem_size = ggml_tensor_overhead() * 2;
params.mem_buffer = nullptr;
params.no_alloc = true;
ggml_context* ctx = ggml_init(params);
if (ctx == nullptr) {
return false;
}
ggml_tensor* op_tensor = nullptr;
if (op == GGML_OP_GET_ROWS) {
ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
op_tensor = ggml_get_rows(ctx, weight, indices);
}
if (op_tensor == nullptr) {
ggml_free(ctx);
return false;
}
weight->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
if (weight->buffer == nullptr) {
ggml_free(ctx);
return false;
}
bool supported = ggml_backend_dev_supports_op(device, op_tensor);
ggml_backend_buffer_free(weight->buffer);
weight->buffer = nullptr;
ggml_free(ctx);
return supported;
}
ModelManager::~ModelManager() {
release_all();
}
@@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
ggml_backend_t params_backend,
size_t* registered_tensor_size,
bool allow_split_buffer,
bool params_follow_compute_backend) {
bool params_follow_compute_backend,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
if (desc.empty()) {
LOG_ERROR("model manager tensor desc is empty");
return false;
@@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc,
state->params_backend = params_backend;
state->allow_split_buffer = allow_split_buffer;
state->params_follow_compute_backend = params_follow_compute_backend;
if (tensor_ops != nullptr) {
auto op_it = tensor_ops->find(tensor);
if (op_it != tensor_ops->end()) {
state->usage_op = op_it->second;
}
}
new_states.push_back(std::move(state));
}
@@ -844,6 +893,22 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
if (params_buft == nullptr) {
params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
}
if (state.usage_op != GGML_OP_NONE &&
state.compute_backend != nullptr) {
ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend);
if (device_supports_param_op(compute_dev, state.tensor, state.usage_op, params_buft)) {
return params_buft;
}
ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
params_buft = cpu_dev != nullptr ? ggml_backend_dev_buffer_type(cpu_dev) : nullptr;
if (!device_supports_param_op(cpu_dev, state.tensor, state.usage_op, params_buft)) {
LOG_ERROR("model manager has no compatible buffer for tensor '%s' used by %s",
state.name.c_str(),
ggml_op_name(state.usage_op));
return nullptr;
}
}
return params_buft;
}
+5 -3
View File
@@ -39,6 +39,7 @@ private:
bool allow_split_buffer = false;
bool params_follow_compute_backend = false;
bool metadata_validated = false;
enum ggml_op usage_op = GGML_OP_NONE;
int active_prepare_count = 0;
@@ -130,9 +131,10 @@ public:
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
ggml_backend_t params_backend,
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false);
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
bool unregister_param_tensors(const std::string& desc,
size_t* registered_tensor_size = nullptr);
+13 -2
View File
@@ -1048,7 +1048,7 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) {
}
std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) {
if (sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
if (sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
return name;
}
if (sd_version_uses_wan_vae(version)) {
@@ -1384,6 +1384,8 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
{".lora_B.weight", ".weight.lora_up"},
{".lora_A.default.weight", ".weight.lora_down"},
{".lora_B.default.weight", ".weight.lora_up"},
{".lora_A", ".weight.lora_down"},
{".lora_B", ".weight.lora_up"},
{".lora_linear", ".weight.alpha"},
{".alpha", ".weight.alpha"},
{".scale", ".weight.scale"},
@@ -1449,15 +1451,24 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
{"te2.", "cond_stage_model.1.transformer."},
{"te1.", "cond_stage_model.transformer."},
{"te3.", "text_encoders.t5xxl.transformer."},
{"clip_vision.", "cond_stage_model.transformer."},
};
if (sd_version_is_flux(version)) {
prefix_map["te1."] = "text_encoders.clip_l.transformer.";
}
if (sd_version_is_unet(version)) {
prefix_map["clip_l."] = "cond_stage_model.transformer.";
prefix_map["clip_g."] = "cond_stage_model.1.transformer.";
} else {
prefix_map["clip_l."] = "text_encoders.clip_l.transformer.";
prefix_map["clip_g."] = "text_encoders.clip_g.transformer.";
}
replace_with_prefix_map(name, prefix_map);
if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version)) {
if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_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());
+112 -3
View File
@@ -306,8 +306,33 @@ struct KarrasScheduler : SigmaScheduler {
};
struct BetaScheduler : SigmaScheduler {
static constexpr double alpha = 0.6;
static constexpr double beta = 0.6;
double alpha = 0.6;
double beta = 0.6;
explicit BetaScheduler(const char* extra_sample_args = nullptr) {
parse_extra_sample_args(extra_sample_args);
LOG_DEBUG("Beta scheduler: alpha=%.4f, beta=%.4f", alpha, beta);
}
void parse_extra_sample_args(const char* extra_sample_args) {
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "beta scheduler arg")) {
if (key == "alpha") {
float parsed;
if (!parse_strict_float(value, parsed) || parsed <= 0.0) {
LOG_WARN("ignoring invalid beta scheduler arg '%s=%s'", key.c_str(), value.c_str());
} else {
alpha = static_cast<double>(parsed);
}
} else if (key == "beta") {
float parsed;
if (!parse_strict_float(value, parsed) || parsed <= 0.0) {
LOG_WARN("ignoring invalid beta scheduler arg '%s=%s'", key.c_str(), value.c_str());
} else {
beta = static_cast<double>(parsed);
}
}
}
}
static double log_beta(double a, double b) {
return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
@@ -1032,7 +1057,7 @@ struct Denoiser {
break;
case BETA_SCHEDULER:
LOG_INFO("get_sigmas with Beta scheduler");
scheduler = std::make_shared<BetaScheduler>();
scheduler = std::make_shared<BetaScheduler>(extra_sample_args);
break;
case EXPONENTIAL_SCHEDULER:
LOG_INFO("get_sigmas exponential scheduler");
@@ -2553,6 +2578,88 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
return x;
}
static sd::Tensor<float> sample_lms(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
const SamplerExtraArgs& extra_sample_args) {
// Linear Multi-Step from https://github.com/crowsonkb/k-diffusion
int divisions = 1000;
for (const auto& [key, value] : extra_sample_args) {
int parsed = 0;
if (key == "lms_divisions") {
if (!parse_strict_int(value, parsed)) {
LOG_WARN("ignoring invalid lms extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
divisions = parsed; // std::max(1, parsed);
// values above 35M produce noise, can be fixed by double precision
// values < 1 always produce noise
}
}
LOG_DEBUG("linear multi-step sampler: integrating using %i division%s", divisions, (divisions == 1) ? "" : "s");
auto linear_multistep_coeff = [=](const int order, const int m, const int j) -> float {
if (!divisions)
return sigmas[m + 1] - sigmas[m]; // delta / 0 * 0
#define LMS_PRECISION float // double
const LMS_PRECISION a = sigmas[m], dx = (sigmas[m + 1] - a) / divisions, s = sigmas[m - j];
const LMS_PRECISION b0 = a + 0.5f * dx; // using Riemann middle integral
LMS_PRECISION sum = 0.0f;
for (int h = 0; h < divisions; h++) {
const LMS_PRECISION b = h * dx + b0;
LMS_PRECISION prod = 1.0f;
for (int k = 0; k < j; k++) {
const LMS_PRECISION t = sigmas[m - k];
prod *= (b - t) / (s - t);
}
for (int k = j + 1; k < order; k++) {
const LMS_PRECISION t = sigmas[m - k];
prod *= (b - t) / (s - t);
}
sum += prod;
}
return sum * dx;
};
const int max_order = 4;
float lms_coeff[max_order];
std::vector<sd::Tensor<float>> hist = {};
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
const float sigma = sigmas[i];
auto denoised_opt = model(x, sigma, i + 1);
if (denoised_opt.pred.empty()) {
return {};
}
sd::Tensor<float> denoised = std::move(denoised_opt.pred);
const int order = std::min(max_order, i + 1);
for (int c = 0; c < order; c++) // computing coefficients
lms_coeff[c] = linear_multistep_coeff(order, i, c);
sd::Tensor<float> d_cur = (x - denoised) / sigma;
switch (order) {
case 4: // derivative + 3 history points
x += hist[hist.size() - 2] * lms_coeff[3];
case 3:
x += hist[hist.size() - 1] * lms_coeff[2];
case 2:
x += hist.back() * lms_coeff[1];
case 1:
x += d_cur * lms_coeff[0];
}
if (hist.size() == static_cast<size_t>(max_order - 1)) {
hist.erase(hist.begin());
}
hist.push_back(std::move(d_cur));
}
return x;
}
static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
@@ -2714,6 +2821,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
case TCD_SAMPLE_METHOD:
return sample_tcd(model, std::move(x), sigmas, rng, eta);
case LMS_SAMPLE_METHOD:
return sample_lms(model, std::move(x), sigmas, extra_args);
case EULER_CFG_PP_SAMPLE_METHOD:
return sample_euler_cfg_pp(model, std::move(x), sigmas);
case EULER_A_CFG_PP_SAMPLE_METHOD:
+27
View File
@@ -4,6 +4,33 @@
#include "core/tensor.hpp"
#include "ggml.h"
const float minimax_latent_rgb_proj[24][3] = {
{0.19819857f, 0.11584999f, 0.07929777f},
{-0.16047224f, -0.10601170f, -0.15996324f},
{0.47391951f, 0.37602475f, 0.20267826f},
{-0.09857441f, -0.27435449f, -0.51681751f},
{-0.18930605f, -0.10512278f, -0.28571478f},
{-0.15639569f, -0.18000929f, -0.25432852f},
{-0.07176921f, -0.10901598f, -0.06654253f},
{-0.05014077f, -0.05839826f, -0.05516087f},
{-0.05201424f, -0.04351913f, -0.01507579f},
{0.24750438f, 0.13307422f, 0.17684120f},
{0.07377446f, 0.10235858f, 0.11707827f},
{0.02908304f, 0.06587022f, 0.10643690f},
{-0.00670531f, -0.03857879f, 0.01750151f},
{-0.07119107f, -0.03083323f, -0.01995450f},
{-0.08612627f, -0.07253841f, -0.01442890f},
{0.08793202f, 0.08681750f, 0.02994647f},
{0.00876893f, 0.02721868f, 0.00091178f},
{-0.03484412f, -0.02711262f, -0.00110101f},
{-0.00679772f, -0.01844275f, -0.01683359f},
{0.04287028f, 0.01601068f, 0.04037397f},
{-0.00493432f, -0.00230528f, 0.00353911f},
{0.01495088f, 0.00292306f, 0.00416671f},
{0.00495307f, 0.05066542f, 0.05210543f},
{-0.02154842f, -0.01518524f, 0.00442402f}};
float minimax_latent_rgb_bias[3] = {0.07776964f, -0.01580954f, -0.06561434f};
const float ltxav_latent_rgb_proj[128][3] = {
{-0.0293802f, -0.0362516f, -0.0291386f},
{0.0117735f, 0.0223435f, 0.018856f},
+537 -103
View File
@@ -38,6 +38,7 @@
#include "model/diffusion/lingbot_video.hpp"
#include "model/diffusion/ltxv.hpp"
#include "model/diffusion/mage_flow.hpp"
#include "model/diffusion/minimax_h3.hpp"
#include "model/diffusion/minit2i.hpp"
#include "model/diffusion/mmdit.hpp"
#include "model/diffusion/model.hpp"
@@ -53,6 +54,8 @@
#include "model/vae/ltx_audio_vae.hpp"
#include "model/vae/ltx_vae.hpp"
#include "model/vae/mage_vae.hpp"
#include "model/vae/minimax_h3_audio_vae.hpp"
#include "model/vae/minimax_h3_vae.hpp"
#include "model/vae/tae.hpp"
#include "model/vae/vae.hpp"
#include "model/vae/wan_vae.hpp"
@@ -106,6 +109,7 @@ const char* model_version_to_str[] = {
"Flux.2",
"Flux.2 klein",
"LTXAV",
"MiniMax-H3",
"HiDream O1",
"Z-Image",
"Boogu Image",
@@ -143,6 +147,7 @@ const char* sampling_methods_str[] = {
"Euler GE",
"DPM++ (2M) SDE",
"DPM++ (2M) SDE BT",
"LMS",
};
/*================================================== Helper Functions ================================================*/
@@ -215,6 +220,7 @@ public:
std::shared_ptr<RNG> sampler_rng = nullptr;
int n_threads = -1;
float default_flow_shift = INFINITY;
float active_flow_shift = INFINITY;
std::shared_ptr<Conditioner> cond_stage_model;
std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision; // for svd or wan2.1 i2v
@@ -222,7 +228,7 @@ public:
std::shared_ptr<DiffusionModelRunner> high_noise_diffusion_model;
std::shared_ptr<VAE> first_stage_model;
std::shared_ptr<VAE> preview_vae;
std::shared_ptr<LTXV::LTXAudioVAERunner> audio_vae_model;
std::shared_ptr<AudioVAERunner> audio_vae_model;
std::shared_ptr<ControlNet> control_net;
std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
sd::Tensor<float> ip_adapter_tokens;
@@ -320,7 +326,11 @@ public:
return true;
}
std::map<std::string, ggml_tensor*> group_tensors;
std::map<ggml_tensor*, enum ggml_op> tensor_ops;
model->get_param_tensors(group_tensors);
if constexpr (std::is_base_of_v<Conditioner, T>) {
model->get_param_tensor_ops(tensor_ops);
}
if (model_manager == nullptr) {
return true;
}
@@ -337,6 +347,7 @@ public:
module,
module_backends,
std::move(group_tensors),
tensor_ops,
residency_mode,
params_mem_size);
}
@@ -345,6 +356,7 @@ public:
module,
module_backends,
std::move(group_tensors),
tensor_ops,
residency_mode,
params_mem_size);
}
@@ -358,7 +370,10 @@ public:
residency_mode,
backend_for(module),
params_backend_for(module),
params_mem_size);
params_mem_size,
false,
false,
&tensor_ops);
}
template <typename T>
@@ -367,6 +382,7 @@ public:
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size) {
ggml_backend_t main_backend = module_backends[0];
@@ -378,6 +394,7 @@ public:
module,
module_backends,
std::move(group_tensors),
tensor_ops,
residency_mode,
params_mem_size);
};
@@ -452,7 +469,9 @@ public:
main_backend,
params_backend_for(module),
params_mem_size,
/*allow_split_buffer=*/true)) {
/*allow_split_buffer=*/true,
false,
&tensor_ops)) {
return false;
}
return model_manager->register_param_tensors(desc,
@@ -460,7 +479,10 @@ public:
residency_mode,
main_backend,
params_backend_for(module),
params_mem_size);
params_mem_size,
false,
false,
&tensor_ops);
}
// Register graph-cut layer-split tensors on the primary backend first.
@@ -472,6 +494,7 @@ public:
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size) {
bool has_cpu_device = false;
@@ -493,7 +516,10 @@ public:
residency_mode,
module_backends[0],
params_backend_for(module),
params_mem_size);
params_mem_size,
false,
false,
&tensor_ops);
}
model->set_runtime_backends(module_backends);
@@ -518,7 +544,8 @@ public:
initial_params_backend,
params_mem_size,
false,
params_follow_runtime);
params_follow_runtime,
&tensor_ops);
}
bool unload_control_net() {
@@ -674,45 +701,11 @@ public:
LOG_DEBUG("loaded alphas_cumprod from model file");
}
bool init(const sd_ctx_params_t* sd_ctx_params) {
n_threads = sd_ctx_params->n_threads;
enable_mmap = sd_ctx_params->enable_mmap;
stream_layers = sd_ctx_params->stream_layers;
eager_load = sd_ctx_params->eager_load;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
split_mode_spec = SAFE_STR(sd_ctx_params->split_mode);
auto_fit_enabled = sd_ctx_params->auto_fit;
max_vram_assignment.reset(0.f);
{
std::string error;
if (!max_vram_assignment.parse(SAFE_STR(sd_ctx_params->max_vram), &error)) {
LOG_ERROR("%s", error.c_str());
return false;
}
}
std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers);
add_rpc_devices(rpc_servers_spec);
bool use_tae = false;
bool use_audio_vae = false;
bool use_control_net = false;
rng = get_rng(sd_ctx_params->rng_type);
if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) {
sampler_rng = get_rng(sd_ctx_params->sampler_rng_type);
} else {
sampler_rng = rng;
}
ggml_log_set(ggml_log_callback_default, nullptr);
model_manager = std::make_shared<ModelManager>();
model_manager->set_n_threads(n_threads);
model_manager->set_enable_mmap(enable_mmap);
ModelLoader& model_loader = model_manager->loader();
bool init_model_loader(ModelLoader& model_loader,
const sd_ctx_params_t* sd_ctx_params,
bool& use_tae,
bool& use_audio_vae,
bool& use_control_net) {
if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) {
LOG_INFO("loading model from '%s'", sd_ctx_params->model_path);
if (!model_loader.init_from_file(sd_ctx_params->model_path)) {
@@ -741,28 +734,23 @@ public:
}
}
bool is_unet = sd_version_is_unet(model_loader.get_sd_version());
if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) {
LOG_INFO("loading clip_l from '%s'", sd_ctx_params->clip_l_path);
std::string prefix = is_unet ? "cond_stage_model.transformer." : "text_encoders.clip_l.transformer.";
if (!model_loader.init_from_file(sd_ctx_params->clip_l_path, prefix)) {
if (!model_loader.init_from_file(sd_ctx_params->clip_l_path, "clip_l.")) {
LOG_WARN("loading clip_l from '%s' failed", sd_ctx_params->clip_l_path);
}
}
if (strlen(SAFE_STR(sd_ctx_params->clip_g_path)) > 0) {
LOG_INFO("loading clip_g from '%s'", sd_ctx_params->clip_g_path);
std::string prefix = is_unet ? "cond_stage_model.1.transformer." : "text_encoders.clip_g.transformer.";
if (!model_loader.init_from_file(sd_ctx_params->clip_g_path, prefix)) {
if (!model_loader.init_from_file(sd_ctx_params->clip_g_path, "clip_g.")) {
LOG_WARN("loading clip_g from '%s' failed", sd_ctx_params->clip_g_path);
}
}
if (strlen(SAFE_STR(sd_ctx_params->clip_vision_path)) > 0) {
LOG_INFO("loading clip_vision from '%s'", sd_ctx_params->clip_vision_path);
std::string prefix = "cond_stage_model.transformer.";
if (!model_loader.init_from_file(sd_ctx_params->clip_vision_path, prefix)) {
if (!model_loader.init_from_file(sd_ctx_params->clip_vision_path, "clip_vision.")) {
LOG_WARN("loading clip_vision from '%s' failed", sd_ctx_params->clip_vision_path);
}
}
@@ -821,9 +809,9 @@ public:
}
if (strlen(SAFE_STR(sd_ctx_params->audio_vae_path)) > 0) {
LOG_INFO("loading LTX audio VAE from '%s'", sd_ctx_params->audio_vae_path);
LOG_INFO("loading audio VAE from '%s'", sd_ctx_params->audio_vae_path);
if (!model_loader.init_from_file(sd_ctx_params->audio_vae_path)) {
LOG_WARN("loading LTX audio VAE weights from '%s' failed", sd_ctx_params->audio_vae_path);
LOG_WARN("loading audio VAE weights from '%s' failed", sd_ctx_params->audio_vae_path);
} else {
use_audio_vae = true;
}
@@ -857,24 +845,69 @@ public:
model_loader.convert_tensors_name();
version = model_loader.get_sd_version();
if (version == VERSION_COUNT) {
LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path));
return false;
}
auto& tensor_storage_map = model_loader.get_tensor_storage_map();
LOG_INFO("Version: %s ", model_version_to_str[version]);
ggml_type wtype = sd_type_to_ggml_type(sd_ctx_params->wtype);
std::string tensor_type_rules = SAFE_STR(sd_ctx_params->tensor_type_rules);
if (wtype != GGML_TYPE_COUNT || tensor_type_rules.size() > 0) {
model_loader.set_wtype_override(wtype, tensor_type_rules);
}
return true;
}
bool init(const sd_ctx_params_t* sd_ctx_params) {
n_threads = sd_ctx_params->n_threads;
enable_mmap = sd_ctx_params->enable_mmap;
stream_layers = sd_ctx_params->stream_layers;
eager_load = sd_ctx_params->eager_load;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
split_mode_spec = SAFE_STR(sd_ctx_params->split_mode);
auto_fit_enabled = sd_ctx_params->auto_fit;
max_vram_assignment.reset(0.f);
{
std::string error;
if (!max_vram_assignment.parse(SAFE_STR(sd_ctx_params->max_vram), &error)) {
LOG_ERROR("%s", error.c_str());
return false;
}
}
std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers);
add_rpc_devices(rpc_servers_spec);
bool use_tae = false;
bool use_audio_vae = false;
bool use_control_net = false;
rng = get_rng(sd_ctx_params->rng_type);
if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) {
sampler_rng = get_rng(sd_ctx_params->sampler_rng_type);
} else {
sampler_rng = rng;
}
ggml_log_set(ggml_log_callback_default, nullptr);
model_manager = std::make_shared<ModelManager>();
model_manager->set_n_threads(n_threads);
model_manager->set_enable_mmap(enable_mmap);
ModelLoader& model_loader = model_manager->loader();
if (!init_model_loader(model_loader, sd_ctx_params, use_tae, use_audio_vae, use_control_net)) {
return false;
}
version = model_loader.get_sd_version();
if (version == VERSION_COUNT) {
LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path));
return false;
} else {
LOG_INFO("Version: %s ", model_version_to_str[version]);
}
if (auto_fit_enabled) {
if (!sd::backend_fit::derive_backend_specs(model_loader,
wtype,
sd_type_to_ggml_type(sd_ctx_params->wtype),
max_vram_assignment,
backend_spec,
params_backend_spec)) {
@@ -929,14 +962,10 @@ public:
if (sd_ctx_params->lora_apply_mode == LORA_APPLY_AUTO) {
bool have_quantized_weight = false;
if (wtype != GGML_TYPE_COUNT && ggml_is_quantized(wtype)) {
have_quantized_weight = true;
} else {
for (const auto& [type, _] : wtype_stat) {
if (ggml_is_quantized(type)) {
have_quantized_weight = true;
break;
}
for (const auto& [type, _] : wtype_stat) {
if (ggml_is_quantized(type)) {
have_quantized_weight = true;
break;
}
}
// Avoid full-model LoRA merge buffers on constrained setups.
@@ -979,6 +1008,13 @@ public:
tae_preview_only = false;
use_tae = true;
}
if (sd_version_is_minimax_h3(version) && use_tae) {
LOG_WARN("MiniMax-H3 does not have a compatible TAE; ignoring --taesd");
tae_preview_only = false;
use_tae = false;
}
auto& tensor_storage_map = model_loader.get_tensor_storage_map();
{
if (!ensure_backend_pair(SDBackendModule::TE) ||
@@ -1085,6 +1121,17 @@ public:
tensor_storage_map,
"model.diffusion_model",
model_manager);
} else if (sd_version_is_minimax_h3(version)) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
model_manager);
diffusion_model = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
model_manager);
} else if (sd_version_is_hunyuan_video(version)) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@@ -1377,6 +1424,11 @@ public:
false,
version,
model_manager);
} else if (sd_version_is_minimax_h3(version)) {
return std::make_shared<MiniMaxH3VAE::MiniMaxH3VideoVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
model_manager);
} else if (sd_version_is_mage_flow(vae_version)) {
return std::make_shared<MageVAE::MageVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
@@ -1462,11 +1514,18 @@ public:
}
if (use_audio_vae) {
audio_vae_model = std::make_shared<LTXV::LTXAudioVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"",
model_manager);
if (!register_runner_params("LTX audio VAE",
if (sd_version_is_minimax_h3(version)) {
audio_vae_model = std::make_shared<MiniMaxH3::AudioVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"",
model_manager);
} else {
audio_vae_model = std::make_shared<LTXV::LTXAudioVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"",
model_manager);
}
if (!register_runner_params("audio VAE",
audio_vae_model,
SDBackendModule::VAE,
&vae_params_mem_size)) {
@@ -1579,7 +1638,9 @@ public:
ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__");
if (audio_vae_model) {
ignore_tensors.insert("audio_vae.encoder");
if (!sd_version_is_minimax_h3(version)) {
ignore_tensors.insert("audio_vae.encoder");
}
}
if (version == VERSION_OVIS_IMAGE) {
ignore_tensors.insert("text_encoders.llm.vision_model.");
@@ -1702,6 +1763,7 @@ public:
sd_version_is_wan(version) ||
sd_version_is_hunyuan_video(version) ||
sd_version_is_lingbot_video(version) ||
sd_version_is_minimax_h3(version) ||
sd_version_is_qwen_image(version) ||
sd_version_is_mage_flow(version) ||
version == VERSION_HIDREAM_O1 ||
@@ -1716,6 +1778,8 @@ public:
default_flow_shift = 5.f;
} else if (sd_version_is_hunyuan_video(version)) {
default_flow_shift = 7.f;
} else if (sd_version_is_minimax_h3(version)) {
default_flow_shift = 12.f;
} else if (sd_version_is_ernie_image(version)) {
default_flow_shift = 4.f;
} else if (sd_version_is_pid(version)) {
@@ -2110,7 +2174,9 @@ public:
return;
}
auto image_tensor = sd_image_to_tensor(image);
auto embed = get_clip_vision_output(image_tensor, true, -1);
auto embed = ip_adapter->is_plus
? get_clip_vision_output(image_tensor, false, 2)
: get_clip_vision_output(image_tensor, true, -1);
if (embed.empty()) {
return;
}
@@ -2230,6 +2296,14 @@ public:
LOG_WARN("No latent to RGB projection known for this model");
return;
}
} else if (channels == 24) {
if(sd_version_is_minimax_h3(version)){
latent_rgb_proj = minimax_latent_rgb_proj;
latent_rgb_bias = minimax_latent_rgb_bias;
} else {
LOG_WARN("No latent to RGB projection known for this model");
return;
}
} else if (channels == 16) {
if (sd_version_is_sd3(version)) {
latent_rgb_proj = sd3_latent_rgb_proj;
@@ -2661,6 +2735,15 @@ public:
condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask,
condition.c_image_embeds.empty() ? nullptr : &condition.c_image_embeds};
} else if (sd_version_is_minimax_h3(version)) {
diffusion_params.extra = MiniMaxH3DiffusionExtra{
condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids,
condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios,
condition.c_reference_blocks.empty() ? nullptr : &condition.c_reference_blocks,
audio_length,
std::isfinite(active_flow_shift) ? active_flow_shift : 12.f,
3.f};
} else if (sd_version_is_ltxav(version)) {
diffusion_params.extra = LTXAVDiffusionExtra{
nullptr,
@@ -2834,7 +2917,7 @@ public:
int get_diffusion_model_down_factor() {
int down_factor = 8; // unet
if (sd_version_is_dit(version)) {
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) {
down_factor = 2;
} else {
down_factor = 1;
@@ -2848,6 +2931,8 @@ public:
if (sd_version_is_dit(version)) {
if (sd_version_is_ltxav(version)) {
latent_channel = 128;
} else if (sd_version_is_minimax_h3(version)) {
latent_channel = 24;
} else if (version == VERSION_WAN2_2_TI2V) {
latent_channel = 48;
} else if (sd_version_is_hunyuan_video(version)) {
@@ -2901,6 +2986,8 @@ public:
int latent_frames = frames;
if (sd_version_is_ltxav(version)) {
latent_frames = ((frames - 1) / 8) + 1;
} else if (sd_version_is_minimax_h3(version)) {
latent_frames = frames <= 5 ? 2 : ((frames - 5) / 17) * 5 + 2;
} else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
latent_frames = ((frames - 1) / 4) + 1;
}
@@ -2914,6 +3001,9 @@ public:
if (sd_version_is_ltxav(version)) {
return (latent_frames - 1) * 8 + 1;
}
if (sd_version_is_minimax_h3(version)) {
return latent_frames <= 2 ? 5 : ((latent_frames - 2) / 5) * 17 + 5;
}
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
return (latent_frames - 1) * 4 + 1;
}
@@ -2921,6 +3011,13 @@ public:
}
int align_video_frames(int frames) {
if (sd_version_is_minimax_h3(version)) {
frames = std::max(frames, 5);
while (frames % 17 != 5) {
++frames;
}
return frames;
}
return latent_frames_to_video_frames(video_frames_to_latent_frames(frames));
}
@@ -2995,6 +3092,7 @@ public:
flow_shift = default_flow_shift;
}
flow_denoiser->set_shift(flow_shift);
active_flow_shift = flow_shift;
}
}
@@ -3190,6 +3288,7 @@ const char* sample_method_to_str[] = {
"euler_ge",
"dpm++2m_sde",
"dpm++2m_sde_bt",
"lms",
};
const char* sd_sample_method_name(enum sample_method_t sample_method) {
@@ -3714,7 +3813,7 @@ struct sd_ctx_t {
};
static bool sd_version_supports_video_generation(SDVersion version) {
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version);
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version);
}
static bool sd_version_supports_image_generation(SDVersion version) {
@@ -3776,7 +3875,7 @@ static sd_audio_t* waveform_to_sd_audio(const StableDiffusionGGML* sd,
return nullptr;
}
audio->sample_rate = static_cast<uint32_t>(sd->audio_vae_model != nullptr ? sd->audio_vae_model->config.output_sample_rate() : 0);
audio->sample_rate = static_cast<uint32_t>(sd->audio_vae_model != nullptr ? sd->audio_vae_model->output_sample_rate() : 0);
audio->channels = static_cast<uint32_t>(channels);
audio->sample_count = static_cast<uint64_t>(sample_count);
size_t sample_bytes = waveform.numel() * sizeof(float);
@@ -3981,14 +4080,18 @@ struct GenerationRequest {
}
GenerationRequest(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params) {
prompt = SAFE_STR(sd_vid_gen_params->prompt);
negative_prompt = SAFE_STR(sd_vid_gen_params->negative_prompt);
width = sd_vid_gen_params->width;
height = sd_vid_gen_params->height;
requested_frames = std::max(1, sd_vid_gen_params->video_frames);
frames = sd_ctx->sd->align_video_frames(requested_frames);
clip_skip = sd_vid_gen_params->clip_skip;
fps = std::max(1, sd_vid_gen_params->fps);
prompt = SAFE_STR(sd_vid_gen_params->prompt);
negative_prompt = SAFE_STR(sd_vid_gen_params->negative_prompt);
width = sd_vid_gen_params->width;
height = sd_vid_gen_params->height;
requested_frames = std::max(1, sd_vid_gen_params->video_frames);
frames = sd_ctx->sd->align_video_frames(requested_frames);
clip_skip = sd_vid_gen_params->clip_skip;
fps = std::max(1, sd_vid_gen_params->fps);
if (sd_version_is_minimax_h3(sd_ctx->sd->version) && fps != 24) {
LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps);
fps = 24;
}
vae_scale_factor = sd_ctx->sd->get_vae_scale_factor();
diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor();
seed = sd_vid_gen_params->seed;
@@ -4228,6 +4331,8 @@ struct SamplePlan {
if (sd_version_is_ltxav(sd_ctx->sd->version) && request->frames > 0) {
int latent_frames = ((request->frames - 1) / 8) + 1;
sample_seq_len *= latent_frames;
} else if (sd_version_is_minimax_h3(sd_ctx->sd->version) && request->frames > 0) {
sample_seq_len *= sd_ctx->sd->video_frames_to_latent_frames(request->frames);
}
sigmas = sd_ctx->sd->denoiser->get_sigmas(total_steps,
sample_seq_len,
@@ -4267,6 +4372,10 @@ struct ImageGenerationLatents {
sd::Tensor<float> control_image;
std::vector<sd::Tensor<float>> ref_images;
std::vector<sd::Tensor<float>> ref_latents;
std::vector<sd::Tensor<float>> reference_audio_latents;
std::vector<MiniMaxH3ReferenceBlock> minimax_reference_blocks;
std::vector<MiniMaxH3PresentationItem> minimax_presentation_refs;
std::vector<int32_t> keyframe_indices;
sd::Tensor<float> denoise_mask;
sd::Tensor<float> clip_vision_output;
sd::Tensor<float> vace_context;
@@ -4581,6 +4690,67 @@ static int get_ltxav_num_audio_latents(int frames, int fps) {
return static_cast<int>(std::ceil((static_cast<float>(frames) / static_cast<float>(fps)) * kLatentsPerSecond));
}
static int get_minimax_h3_num_audio_latents(int frames, int fps) {
GGML_ASSERT(frames > 0 && fps > 0);
return std::max(1,
static_cast<int>(std::lround(
static_cast<double>(frames) * 40.0 / fps)));
}
static sd::Tensor<float> make_minimax_h3_empty_audio_latent(int audio_length) {
if (audio_length <= 0) {
return {};
}
return sd::zeros<float>({audio_length, 2, 32, 1});
}
static sd::Tensor<float> prepare_minimax_h3_reference_waveform(const sd_audio_t& audio,
int target_sample_rate = 32000) {
if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) {
return {};
}
uint64_t output_samples = static_cast<uint64_t>(std::llround(
static_cast<long double>(audio.sample_count) * target_sample_rate / audio.sample_rate));
output_samples = std::max<uint64_t>(1, output_samples);
uint64_t padded_samples = (output_samples + 799) / 800 * 800;
sd::Tensor<float> waveform({static_cast<int64_t>(padded_samples), 2, 1, 1});
for (uint64_t i = 0; i < output_samples; ++i) {
long double source_pos = static_cast<long double>(i) * audio.sample_rate / target_sample_rate;
uint64_t source0 = std::min<uint64_t>(static_cast<uint64_t>(source_pos), audio.sample_count - 1);
uint64_t source1 = std::min<uint64_t>(source0 + 1, audio.sample_count - 1);
float fraction = static_cast<float>(source_pos - source0);
for (uint32_t channel = 0; channel < 2; ++channel) {
uint32_t source_channel = audio.channels == 1 ? 0 : std::min<uint32_t>(channel, audio.channels - 1);
float a = audio.data[source0 * audio.channels + source_channel];
float b = audio.data[source1 * audio.channels + source_channel];
waveform.index(static_cast<int64_t>(i), channel, 0, 0) =
std::clamp(a + (b - a) * fraction, -1.f, 1.f);
}
}
return waveform;
}
static sd::Tensor<float> unpack_minimax_h3_audio_latent(const sd::Tensor<float>& packed_latent,
int audio_length,
int video_channels) {
if (packed_latent.empty() || audio_length <= 0) {
return {};
}
GGML_ASSERT(packed_latent.dim() == 4 || packed_latent.dim() == 5);
int64_t spatial_size = packed_latent.shape()[0] * packed_latent.shape()[1] * packed_latent.shape()[2];
int64_t required = static_cast<int64_t>(audio_length) * 2 * 32;
int64_t available = (packed_latent.shape()[3] - video_channels) * spatial_size;
if (available < required) {
return {};
}
sd::Tensor<float> audio({audio_length, 2, 32, 1});
const float* source = packed_latent.data() +
static_cast<size_t>(video_channels) * static_cast<size_t>(spatial_size);
std::copy_n(source, static_cast<size_t>(required), audio.data());
return audio;
}
struct ImageGenerationEmbeds {
SDCondition cond;
SDCondition uncond;
@@ -5681,6 +5851,248 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height);
}
if (sd_version_is_minimax_h3(sd_ctx->sd->version)) {
if (sd_vid_gen_params->ref_images_count < 0 || sd_vid_gen_params->ref_videos_count < 0 ||
sd_vid_gen_params->ref_audios_count < 0 ||
(sd_vid_gen_params->ref_images_count > 0 && sd_vid_gen_params->ref_images == nullptr) ||
(sd_vid_gen_params->ref_videos_count > 0 && sd_vid_gen_params->ref_videos == nullptr) ||
(sd_vid_gen_params->ref_audios_count > 0 && sd_vid_gen_params->ref_audios == nullptr)) {
LOG_ERROR("invalid MiniMax-H3 Ref2VA input arrays");
return std::nullopt;
}
latents.audio_length = get_minimax_h3_num_audio_latents(request->frames,
request->fps);
latents.audio_latent = make_minimax_h3_empty_audio_latent(latents.audio_length);
bool has_references = sd_vid_gen_params->ref_images_count > 0 ||
sd_vid_gen_params->ref_videos_count > 0 ||
sd_vid_gen_params->ref_audios_count > 0;
if (has_references && (!start_image.empty() || !end_image.empty())) {
LOG_ERROR("MiniMax-H3 keyframes and Ref2VA references cannot be used together");
return std::nullopt;
}
if (sd_vid_gen_params->control_frames_size > 0) {
LOG_ERROR("MiniMax-H3 control_frames are not implemented");
return std::nullopt;
}
auto add_visual_noise = [&](sd::Tensor<float> latent) {
auto condition_rng = std::make_shared<PhiloxRNG>();
condition_rng->manual_seed(static_cast<uint64_t>(request->seed));
return latent * MiniMaxH3::VISUAL_COND_TIMESTEP +
sd::Tensor<float>::randn_like(latent, condition_rng) *
(1.f - MiniMaxH3::VISUAL_COND_TIMESTEP);
};
auto add_keyframe = [&](const sd::Tensor<float>& image,
int32_t frame_index,
const char* name) -> bool {
if (image.empty()) {
return true;
}
auto video_image = image.reshape({image.shape()[0],
image.shape()[1],
1,
image.shape()[2],
image.shape()[3]});
auto latent = sd_ctx->sd->encode_first_stage(video_image);
if (latent.empty()) {
LOG_ERROR("failed to encode MiniMax-H3 %s keyframe", name);
return false;
}
latents.ref_images.push_back(image);
latents.ref_latents.push_back(add_visual_noise(std::move(latent)));
latents.keyframe_indices.push_back(frame_index);
return true;
};
auto resize_reference = [&](const sd::Tensor<float>& image,
int width,
int height) {
return sd::ops::interpolate(
image,
std::vector<int64_t>{width, height, image.shape()[2], image.shape()[3]});
};
auto encode_reference_audio = [&](const sd_audio_t& audio,
int32_t* audio_index) -> bool {
if (sd_ctx->sd->audio_vae_model == nullptr) {
LOG_ERROR("MiniMax-H3 Ref2VA audio requires --audio-vae with encoder weights");
return false;
}
auto waveform = prepare_minimax_h3_reference_waveform(
audio,
sd_ctx->sd->audio_vae_model->input_sample_rate());
if (waveform.empty()) {
LOG_ERROR("invalid MiniMax-H3 reference audio");
return false;
}
auto encoded = sd_ctx->sd->audio_vae_model->encode(sd_ctx->sd->n_threads, waveform);
if (encoded.empty()) {
LOG_ERROR("failed to encode MiniMax-H3 reference audio");
return false;
}
*audio_index = static_cast<int32_t>(latents.reference_audio_latents.size());
latents.reference_audio_latents.push_back(std::move(encoded));
return true;
};
if (has_references) {
LOG_INFO("MiniMax-H3 Ref2VA: %d image(s), %d video(s), %d audio clip(s)",
sd_vid_gen_params->ref_images_count,
sd_vid_gen_params->ref_videos_count,
sd_vid_gen_params->ref_audios_count);
for (int i = 0; i < sd_vid_gen_params->ref_images_count; ++i) {
auto image = ensure_image_tensor_channels(
sd_image_to_tensor(sd_vid_gen_params->ref_images[i]),
3);
if (image.empty()) {
LOG_ERROR("failed to load MiniMax-H3 reference image %d", i + 1);
return std::nullopt;
}
int source_w = static_cast<int>(image.shape()[0]);
int source_h = static_cast<int>(image.shape()[1]);
double source_area = static_cast<double>(source_w) * source_h;
double target_area = static_cast<double>(request->width) * request->height;
double scale = std::min(1.0, std::sqrt(target_area / source_area));
int width = std::max(32, static_cast<int>(std::round(source_w * scale / 32.f)) * 32);
int height = std::max(32, static_cast<int>(std::round(source_h * scale / 32.f)) * 32);
image = resize_reference(image, width, height);
auto latent = sd_ctx->sd->encode_first_stage(image);
if (latent.empty()) {
LOG_ERROR("failed to encode MiniMax-H3 reference image %d", i + 1);
return std::nullopt;
}
int32_t video_index = static_cast<int32_t>(latents.ref_latents.size());
latents.ref_latents.push_back(add_visual_noise(std::move(latent)));
latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::IMAGE,
video_index,
-1});
MiniMaxH3PresentationItem item;
item.kind = MiniMaxH3PresentationKind::IMAGE;
item.frames.push_back(std::move(image));
latents.minimax_presentation_refs.push_back(std::move(item));
}
for (int video_idx = 0; video_idx < sd_vid_gen_params->ref_videos_count; ++video_idx) {
const auto& reference = sd_vid_gen_params->ref_videos[video_idx];
if (reference.frames == nullptr || reference.frame_count < 1) {
LOG_ERROR("invalid MiniMax-H3 reference video %d", video_idx + 1);
return std::nullopt;
}
int source_fps = reference.fps > 0 ? reference.fps : 24;
int normalized_frames = static_cast<int>(std::lround(
static_cast<double>(reference.frame_count) * 24.0 / source_fps));
normalized_frames = std::min(normalized_frames, request->frames);
if (normalized_frames < 5) {
LOG_ERROR("MiniMax-H3 reference video %d needs at least 5 frames at 24 fps",
video_idx + 1);
return std::nullopt;
}
while (normalized_frames % 17 != 5) {
--normalized_frames;
}
auto first = ensure_image_tensor_channels(sd_image_to_tensor(reference.frames[0]), 3);
if (first.empty()) {
LOG_ERROR("invalid first frame in MiniMax-H3 reference video %d", video_idx + 1);
return std::nullopt;
}
int source_w = static_cast<int>(first.shape()[0]);
int source_h = static_cast<int>(first.shape()[1]);
double ratio = static_cast<double>(source_w) / source_h;
double nominal_w = ratio >= 1.0 ? 768.0 * ratio : 768.0;
double nominal_h = ratio >= 1.0 ? 768.0 : 768.0 / ratio;
if (nominal_w * nominal_h > 768.0 * 1344.0) {
double scale = std::sqrt((768.0 * 1344.0) / (nominal_w * nominal_h));
nominal_w *= scale;
nominal_h *= scale;
}
int width = std::max(32, static_cast<int>(std::round(nominal_w / 32.0)) * 32);
int height = std::max(32, static_cast<int>(std::round(nominal_h / 32.0)) * 32);
if (source_w * source_h < width * height) {
width = std::max(32, static_cast<int>(std::round(source_w / 32.0)) * 32);
height = std::max(32, static_cast<int>(std::round(source_h / 32.0)) * 32);
}
sd::Tensor<float> video({width, height, normalized_frames, 3, 1});
for (int frame = 0; frame < normalized_frames; ++frame) {
int source_index = std::min(reference.frame_count - 1,
static_cast<int>(std::floor(frame * source_fps / 24.0)));
auto source = ensure_image_tensor_channels(
sd_image_to_tensor(reference.frames[source_index]),
3);
if (source.empty()) {
LOG_ERROR("invalid frame %d in MiniMax-H3 reference video %d",
source_index + 1,
video_idx + 1);
return std::nullopt;
}
source = resize_reference(source, width, height);
sd::ops::slice_assign(&video, 2, frame, frame + 1, source.unsqueeze(2));
}
auto video_latent = sd_ctx->sd->encode_first_stage(video);
if (video_latent.empty()) {
LOG_ERROR("failed to encode MiniMax-H3 reference video %d", video_idx + 1);
return std::nullopt;
}
int32_t audio_index = -1;
bool has_audio = reference.audio.data != nullptr && reference.audio.sample_count > 0;
if (has_audio) {
if (!encode_reference_audio(reference.audio, &audio_index)) {
return std::nullopt;
}
MiniMaxH3PresentationItem audio_item;
audio_item.kind = MiniMaxH3PresentationKind::AUDIO;
latents.minimax_presentation_refs.push_back(std::move(audio_item));
}
MiniMaxH3PresentationItem video_item;
video_item.kind = MiniMaxH3PresentationKind::VIDEO;
for (int frame = 0; frame < normalized_frames; frame += 12) {
auto sampled = sd::ops::slice(video, 2, frame, frame + 1)
.reshape({width, height, 3, 1});
video_item.frames.push_back(std::move(sampled));
video_item.timestamps.push_back(frame / 24.f);
}
latents.minimax_presentation_refs.push_back(std::move(video_item));
int32_t video_index = static_cast<int32_t>(latents.ref_latents.size());
latents.ref_latents.push_back(add_visual_noise(std::move(video_latent)));
latents.minimax_reference_blocks.push_back({has_audio ? MiniMaxH3ReferenceKind::VIDEO_AUDIO
: MiniMaxH3ReferenceKind::VIDEO,
video_index,
audio_index});
}
for (int audio_idx = 0; audio_idx < sd_vid_gen_params->ref_audios_count; ++audio_idx) {
int32_t encoded_index = -1;
if (!encode_reference_audio(sd_vid_gen_params->ref_audios[audio_idx], &encoded_index)) {
return std::nullopt;
}
MiniMaxH3PresentationItem item;
item.kind = MiniMaxH3PresentationKind::AUDIO;
latents.minimax_presentation_refs.push_back(std::move(item));
latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::AUDIO,
-1,
encoded_index});
}
}
if (!has_references && (!start_image.empty() || !end_image.empty())) {
LOG_INFO(!start_image.empty() && !end_image.empty() ? "MiniMax-H3 FL2VA" : !start_image.empty() ? "MiniMax-H3 I2VA"
: "MiniMax-H3 end-frame conditioning");
}
if (!has_references &&
(!add_keyframe(start_image, 0, "start") ||
!add_keyframe(end_image, request->frames - 1, "end"))) {
return std::nullopt;
}
}
if (sd_version_is_ltxav(sd_ctx->sd->version)) {
latents.audio_length = get_ltxav_num_audio_latents(request->frames, request->fps);
latents.audio_latent = make_ltxav_empty_audio_latent(latents.audio_length);
@@ -6041,7 +6453,8 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true);
}
if (sd_version_is_ltxav(sd_ctx->sd->version) && !latents.audio_latent.empty()) {
if ((sd_version_is_ltxav(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) &&
!latents.audio_latent.empty()) {
if (!latents.denoise_mask.empty()) {
latents.denoise_mask = pack_ltxav_audio_and_video_denoise_mask(latents.denoise_mask,
latents.init_latent,
@@ -6061,11 +6474,12 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
ImageGenerationEmbeds embeds;
ConditionerParams condition_params;
condition_params.clip_skip = request.clip_skip;
condition_params.text = request.prompt;
condition_params.zero_out_masked = true;
condition_params.ref_images = &latents.ref_images;
if (sd_version_is_lingbot_video(sd_ctx->sd->version)) {
condition_params.clip_skip = request.clip_skip;
condition_params.text = request.prompt;
condition_params.zero_out_masked = true;
condition_params.ref_images = &latents.ref_images;
condition_params.minimax_h3_references = &latents.minimax_presentation_refs;
if (sd_version_is_lingbot_video(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) {
condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA;
}
@@ -6074,12 +6488,28 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
condition_params);
embeds.cond.c_concat = latents.concat_latent;
embeds.cond.c_vector = latents.clip_vision_output;
if (sd_version_is_minimax_h3(sd_ctx->sd->version)) {
embeds.cond.c_ref_images = latents.ref_latents;
embeds.cond.c_ref_audios = latents.reference_audio_latents;
embeds.cond.c_reference_blocks = latents.minimax_reference_blocks;
if (!latents.keyframe_indices.empty()) {
embeds.cond.c_position_ids = sd::Tensor<int32_t>(
{static_cast<int64_t>(latents.keyframe_indices.size())},
latents.keyframe_indices);
}
}
if (request.use_uncond) {
condition_params.text = request.negative_prompt;
embeds.uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params);
embeds.uncond.c_concat = latents.concat_latent;
embeds.uncond.c_vector = latents.clip_vision_output;
if (sd_version_is_minimax_h3(sd_ctx->sd->version)) {
embeds.uncond.c_ref_images = latents.ref_latents;
embeds.uncond.c_ref_audios = latents.reference_audio_latents;
embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks;
embeds.uncond.c_position_ids = embeds.cond.c_position_ids;
}
}
int64_t t1 = ggml_time_ms();
@@ -6101,7 +6531,7 @@ static sd_image_t* decode_video_outputs(sd_ctx_t* sd_ctx,
return nullptr;
}
sd::Tensor<float> video_latent = final_latent;
if (sd_version_is_ltxav(sd_ctx->sd->version) &&
if ((sd_version_is_ltxav(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) &&
video_latent.shape()[3] > sd_ctx->sd->get_latent_channel()) {
video_latent = sd::ops::slice(video_latent, 3, 0, sd_ctx->sd->get_latent_channel());
}
@@ -6689,7 +7119,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000);
sd_audio_t* generated_audio = nullptr;
if (sd_version_is_ltxav(sd_ctx->sd->version) &&
if ((sd_version_is_ltxav(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) &&
latents.audio_length > 0 &&
sd_ctx->sd->audio_vae_model != nullptr) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
@@ -6698,9 +7128,13 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
}
int64_t audio_latent_decode_start = ggml_time_ms();
auto audio_latent = unpack_ltxav_audio_latent(final_latent,
latents.audio_length,
sd_ctx->sd->get_latent_channel());
auto audio_latent = sd_version_is_minimax_h3(sd_ctx->sd->version)
? unpack_minimax_h3_audio_latent(final_latent,
latents.audio_length,
sd_ctx->sd->get_latent_channel())
: unpack_ltxav_audio_latent(final_latent,
latents.audio_length,
sd_ctx->sd->get_latent_channel());
if (!audio_latent.empty()) {
LOG_DEBUG("decode audio latent %dx%dx%dx%d",
(int)audio_latent.shape()[0],
@@ -6711,7 +7145,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
if (!waveform.empty()) {
generated_audio = waveform_to_sd_audio(sd_ctx->sd, waveform);
} else {
LOG_WARN("LTX audio latent decode failed; continuing with silent video output");
LOG_WARN("audio latent decode failed; continuing with silent video output");
}
}
int64_t audio_latent_decode_end = ggml_time_ms();
+1 -1
View File
@@ -205,7 +205,7 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
ss << "\"" << token << "\", ";
}
ss << "]";
LOG_DEBUG("split prompt \"%s\" to tokens %s", text.c_str(), ss.str().c_str());
LOG_DEBUG("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str());
return bpe_tokens;
}