mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-21 21:47:49 -05:00
Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ca37fad89a | ||
|
|
0bd72f075a | ||
|
|
4a7da26b73 | ||
|
|
9a977388a8 | ||
|
|
44dd13716d | ||
|
|
7f410a3793 | ||
|
|
5ebce93342 | ||
|
|
7f986a9d73 | ||
|
|
e06b205384 | ||
|
|
3191b23d4b | ||
|
|
e95ab96997 | ||
|
|
b68d58624d |
@@ -6,7 +6,9 @@ body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Please use this template and include as many details as possible to help us reproduce and fix the issue.
|
||||
Before submitting a bug report, please read the [Troubleshooting guide](https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/troubleshooting.md) and try the steps relevant to your problem.
|
||||
|
||||
If the problem persists, complete this form and include what you tried and the results, along with enough details to help us reproduce and fix the issue.
|
||||
- type: textarea
|
||||
id: commit
|
||||
attributes:
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
contact_links:
|
||||
- name: Troubleshooting
|
||||
url: https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/troubleshooting.md
|
||||
about: Read the troubleshooting guide first. If the problem persists, submit a bug report.
|
||||
@@ -294,6 +294,8 @@ endif()
|
||||
|
||||
if(MSVC)
|
||||
target_compile_options(${SD_LIB} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>)
|
||||
# ggml backends can throw C++ exceptions through their C API.
|
||||
target_compile_options(${SD_LIB} PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/EHsc->)
|
||||
endif()
|
||||
|
||||
if(APPLE)
|
||||
|
||||
@@ -51,6 +51,7 @@ API and command-line option may change frequently.***
|
||||
- [LongCat Image](./docs/longcat_image.md)
|
||||
- [Z-Image](./docs/z_image.md)
|
||||
- [MiniT2I](./docs/minit2i.md)
|
||||
- [SenseNova U1.5](./docs/sensenova_u1.md)
|
||||
- [Ovis-Image](./docs/ovis_image.md)
|
||||
- [Anima](./docs/anima.md)
|
||||
- [ERNIE-Image](./docs/ernie_image.md)
|
||||
@@ -147,6 +148,7 @@ For runtime and parameter backend placement, see the [backend selection guide](.
|
||||
|
||||
## More Guides
|
||||
|
||||
- [Troubleshooting](./docs/troubleshooting.md)
|
||||
- [Backend selection](./docs/backend.md)
|
||||
- [RPC](./docs/rpc.md)
|
||||
- [LoRA](./docs/lora.md)
|
||||
|
||||
+33
-14
@@ -5,7 +5,8 @@
|
||||
- `--backend` selects the runtime backend used to execute model graphs.
|
||||
- `--params-backend` selects where model parameters are kept.
|
||||
|
||||
If `--params-backend` is not set, parameters use the same backend as their module runtime backend.
|
||||
If `--params-backend` is not set, auto-fit chooses parameter placement. With
|
||||
`--auto-fit off`, parameters use the same backend as their module runtime backend.
|
||||
|
||||
## Syntax
|
||||
|
||||
@@ -129,17 +130,21 @@ warning.
|
||||
## Automatic placement (`--auto-fit on|off`)
|
||||
|
||||
`--auto-fit` requires `on` or `off` and defaults to `on` when omitted.
|
||||
Explicit `--backend` or `--params-backend` assignments disable auto-fit,
|
||||
Explicit `--params-backend` assignments disable auto-fit,
|
||||
regardless of argument order, even with `--auto-fit on`.
|
||||
|
||||
When enabled, auto-fit uses one GPU for `diffusion` / `te` / `vae` computation. It chooses
|
||||
the GPU with the largest available memory budget (the first device on a tie),
|
||||
then derives parameter placements from the model metadata and the remaining
|
||||
memory budgets. The chosen backend specifications are printed.
|
||||
Auto-fit preserves explicit `--backend` assignments, including per-module
|
||||
assignments and device lists. For modules without a runtime assignment, it chooses
|
||||
the GPU with the largest available memory budget (the first device on a tie).
|
||||
It then derives parameter placements from the model metadata, each module's
|
||||
compute devices, and the remaining memory budgets. The chosen backend
|
||||
specifications are printed.
|
||||
|
||||
```shell
|
||||
sd-cli -m model.safetensors -p "a cat" --auto-fit on
|
||||
sd-cli -m model.safetensors -p "a cat" --auto-fit on --max-vram cuda0=8,cuda1=14
|
||||
sd-cli -m model.safetensors -p "a cat" --backend cuda0
|
||||
sd-cli -m model.safetensors -p "a cat" --backend diffusion=cuda0,te=cpu,vae=cuda1
|
||||
sd-cli -m model.safetensors -p "a cat" --auto-fit off
|
||||
```
|
||||
|
||||
@@ -149,11 +154,16 @@ GiB", and with no budget set each device's free memory minus a 512 MiB margin
|
||||
is used. These resolved GPU budgets, including the safety margin, also drive
|
||||
the runner's graph-cut capacity checks.
|
||||
|
||||
Runtime capacity checks also leave 512 MiB of currently free device memory for
|
||||
backend scratch buffers and pipelines, including with explicit backend assignments.
|
||||
They cap stale free-memory reports by the device's total memory minus tracked
|
||||
resident allocations and reject reports that exceed the device's total memory.
|
||||
|
||||
Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
|
||||
used diffusion weights have priority. Each component's weights use the first
|
||||
storage location with enough remaining budget:
|
||||
|
||||
1. The main GPU, leaving estimated space for computation and weight staging.
|
||||
1. The component's compute GPU, leaving estimated space for computation and weight staging.
|
||||
2. CPU RAM, reserving the larger of 2 GiB or 10% of available RAM for other work.
|
||||
3. Another GPU, choosing the one with the largest remaining budget that fits.
|
||||
4. Disk, reloading weights on demand.
|
||||
@@ -170,10 +180,17 @@ weight to be copied again at every step.
|
||||
RAM and GPU budgets are shared across components. Each component uses a single
|
||||
parameter backend; several other GPUs' capacities are not combined to store
|
||||
one component. If available RAM cannot be queried, RAM residency is skipped.
|
||||
Other GPUs store weights only: weights are copied to the main GPU for execution.
|
||||
Auto-fit does not select multi-GPU layer/row computation, so `--split-mode` does
|
||||
not change its placements. Use explicit backend assignments for multi-GPU
|
||||
computation.
|
||||
Weights stored on another GPU are copied to the component's compute devices for
|
||||
execution. CPU modules use RAM or disk. Compute reserves and cache priority are
|
||||
accounted for separately on each device, so a CPU module does not reserve GPU
|
||||
space. Storage on another module's GPU also leaves room for that module's work.
|
||||
|
||||
Auto-fit does not select multi-GPU layer/row computation itself. Explicit device
|
||||
lists and `--split-mode` still control that computation. Before the runners have
|
||||
built their split plans, auto-fit conservatively counts the full component size
|
||||
on each listed GPU when checking residency and cache space. This can offload
|
||||
parameters even when a split layout would fit; use `--auto-fit off` to keep the
|
||||
default split-device parameter placement.
|
||||
|
||||
For example, a diffusion model whose full weights exceed the main GPU's budget
|
||||
can use `--backend diffusion=cuda0 --params-backend diffusion=cpu` when RAM is
|
||||
@@ -188,8 +205,9 @@ weights, compute buffers and caches must
|
||||
still fit the runner's capacity checks. Offloading weights does not guarantee
|
||||
that every resolution or frame count will fit, and auto-fit does not change a
|
||||
component to CPU computation solely because its full weights exceed VRAM.
|
||||
If a VAE decode fails, auto-fit retries with spatial tiling; supported video
|
||||
decoders try temporal tiling first and can then add spatial tiling.
|
||||
If a VAE decode fails, decoding retries with spatial tiling even when `--auto-fit`
|
||||
is off; supported video decoders try temporal tiling first and can then add
|
||||
spatial tiling. Spatial retries use half-size tiles along each latent dimension.
|
||||
|
||||
## Modules
|
||||
|
||||
@@ -291,6 +309,7 @@ The example CLI/server still accepts these older CPU placement flags as compatib
|
||||
Because this default is inserted first, later explicit `--params-backend` entries can still override it, for example `--offload-to-cpu --params-backend te=disk` keeps non-TE parameters on CPU and reloads TE parameters from disk.
|
||||
|
||||
Library callers should set `backend` and `params_backend` directly. `sd_ctx_params_init()`
|
||||
enables `auto_fit` by default; nonempty `backend` or `params_backend` assignments disable it.
|
||||
enables `auto_fit` by default; a nonempty `params_backend` assignment disables it.
|
||||
The `backend` assignment constrains auto-fit's compute placement.
|
||||
The old CPU/offload fields are no longer part of the C API. Explicit `--backend` and
|
||||
`--params-backend` assignments are preferred for new commands.
|
||||
|
||||
+1
-1
@@ -27,7 +27,7 @@ Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM w
|
||||
|
||||
## Use params backend to reduce VRAM or RAM usage.
|
||||
|
||||
`--params-backend` controls where model parameters are kept. If it is not set, parameters use the same backend as `--backend`, so a GPU runtime backend also keeps parameters in VRAM.
|
||||
`--params-backend` controls where model parameters are kept. If it is not set, auto-fit chooses parameter placement while preserving `--backend`. With `--auto-fit off`, parameters use the same backend as `--backend`, so a GPU runtime backend also keeps parameters in VRAM.
|
||||
|
||||
Use CPU params to reduce VRAM usage:
|
||||
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
# How to Use
|
||||
|
||||
SenseNova U1.5 is an 8B MoT model that performs diffusion directly in RGB pixel
|
||||
space. It does not require a separate text encoder or VAE.
|
||||
|
||||
## Download weights
|
||||
|
||||
- Download SenseNova U1.5 8B MoT
|
||||
- safetensors: https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT
|
||||
|
||||
Pass the complete downloaded repository directory to `--model`. The directory
|
||||
must contain `model.safetensors.index.json`, every referenced Safetensors shard,
|
||||
and the tokenizer files.
|
||||
|
||||
## Examples
|
||||
|
||||
### CUDA
|
||||
|
||||
```bash
|
||||
./bin/sd-cli \
|
||||
--model /path/to/SenseNova-U1.5-8B-MoT \
|
||||
--prompt "a red cube on a white background" \
|
||||
--width 2048 \
|
||||
--height 2048 \
|
||||
--steps 50 \
|
||||
--cfg-scale 4 \
|
||||
--flow-shift 3 \
|
||||
--seed 42 \
|
||||
--sampling-method euler \
|
||||
--rng cuda \
|
||||
--fa \
|
||||
--output output.png
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- To match the official non-thinking text-to-image pipeline, use 50 Euler
|
||||
steps, CFG 4, flow shift 3, seed 42, CUDA RNG, and an empty negative prompt.
|
||||
- Width and height must be multiples of 32. The trained 1:1 resolution is
|
||||
2048x2048; lower resolutions are useful for smoke tests but are outside the
|
||||
training buckets.
|
||||
- The SenseNova prompt template and unconditional prompt are built
|
||||
automatically.
|
||||
- This implementation supports non-thinking text-to-image generation. Image
|
||||
editing, visual understanding, interleaved generation, and thinking-mode
|
||||
prompt expansion are not implemented.
|
||||
@@ -0,0 +1,45 @@
|
||||
# Troubleshooting
|
||||
|
||||
## Completely black or white images or videos / NaNs
|
||||
|
||||
Some ggml backends can encounter numerical overflow during inference, producing
|
||||
NaN (not-a-number) values. This can result in completely black or white images or videos.
|
||||
Whether it happens can depend on the backend, device, model, and weight format.
|
||||
|
||||
Known overflow issues have been addressed as far as possible, but the maintainer
|
||||
has limited hardware and cannot test every combination. Some cases may therefore
|
||||
still need a manual workaround.
|
||||
|
||||
These options are supported by both `sd-cli` and `sd-server`. If you encounter
|
||||
this problem, add them to your CLI generation command or server startup command:
|
||||
|
||||
```sh
|
||||
--linear-scale 0.0078125 --attn-scale 0.0078125
|
||||
```
|
||||
|
||||
For `sd-server`, restart the server after changing these startup options. Run the
|
||||
same prompt and seed again to see whether the output recovers. If the problem
|
||||
persists, try smaller positive values, for example:
|
||||
|
||||
```sh
|
||||
--linear-scale 0.00390625 --attn-scale 0.00390625
|
||||
```
|
||||
|
||||
These options reduce intermediate values and compensate afterwards to preserve
|
||||
the intended output scale:
|
||||
|
||||
- `--linear-scale` scales Linear inputs before matrix multiplication and rescales
|
||||
the result.
|
||||
- `--attn-scale` scales attention keys and values (K/V). It takes effect only in
|
||||
the Flash Attention path, where `--fa` or `--diffusion-fa` is enabled and the
|
||||
backend supports it.
|
||||
|
||||
The two values can be set independently and apply across model components. The
|
||||
default `0` preserves each model's built-in settings; `1` explicitly disables the
|
||||
corresponding scaling. Overrides must be finite positive values. C API users can
|
||||
set `linear_scale` and `attn_scale` in `sd_ctx_params_t`.
|
||||
|
||||
If the problem persists after trying the relevant steps above,
|
||||
[submit a bug report](https://github.com/leejet/stable-diffusion.cpp/issues/new?template=bug_report.yml).
|
||||
Include your full command, backend and hardware, model and weight format, logs,
|
||||
and the scale values you tried with their results.
|
||||
+49
@@ -34,6 +34,10 @@
|
||||
- Wan2.2 I2V A14B
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models
|
||||
- gguf: https://huggingface.co/QuantStack/Wan2.2-I2V-A14B-GGUF/tree/main
|
||||
- Wan2.2 S2V 14B
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models
|
||||
- gguf: https://huggingface.co/QuantStack/Wan2.2-S2V-14B-GGUF/tree/main
|
||||
- int8_convrot safetensors: https://huggingface.co/noctrex/Wan2.2-S2V-14B-int8_convrot
|
||||
- Download vae
|
||||
- wan_2.1_vae (for all the wan model except Wan2.2 TI2V 5B)
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
|
||||
@@ -49,6 +53,9 @@
|
||||
- Download clip_vison_h (for Wan2.1 I2V/FLF2V only)
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/clip_vision/clip_vision_h.safetensors
|
||||
|
||||
- Download audio_encoder (for Wan2.2 S2V only)
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/blob/main/split_files/audio_encoders/wav2vec2_large_english_fp16.safetensors
|
||||
|
||||
|
||||
## Examples
|
||||
|
||||
@@ -94,6 +101,48 @@
|
||||
|
||||
<video src=../assets/wan/Wan2.2_14B_i2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
|
||||
|
||||
### Wan2.2 S2V 14B
|
||||
|
||||
Audio-driven video (speech-to-video). The reference image (`-i`) is the speaker
|
||||
portrait, `--audio` is the driving audio track and `--audio-encoder` is the
|
||||
wav2vec2 audio encoder. Wan2.2 S2V requires the wan_2.1 vae (16 channel), not
|
||||
the wan2.2 vae.
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\wan2.2_s2v-14B-Q8_0.gguf --audio-encoder ..\models\audio_encoders\wav2vec2_large_english_fp16.safetensors --vae ..\models\vae\wan_2.1_vae.safetensors --t5xxl ..\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf -p "a person is talking" --cfg-scale 6.0 --steps 20 --sampling-method euler -v -n "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" -W 832 -H 480 --diffusion-fa --offload-to-cpu --vae-tiling --video-frames 81 -i ..\assets\cat_with_sd_cpp_42.png --audio .\input\speech.wav --flow-shift 3.0
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- Recommended settings: `--sampling-method euler --steps 20 --cfg-scale 6.0`.
|
||||
`dpm++2m` produces heavy artifacts on S2V. 4 steps with the lightning LoRA
|
||||
(below) is the fast option.
|
||||
- Resolutions: width and height must be multiples of 16; the examples use
|
||||
multiples of 64. 832x480 is a fast starting point; generation cost scales
|
||||
with pixel area.
|
||||
- `--audio` accepts a WAV file; it is downmixed to mono and resampled to 16 kHz
|
||||
internally. Audio longer than the video is truncated, video longer than the
|
||||
audio is padded with silence. Pick `--video-frames` to match the audio:
|
||||
roughly `audio_seconds * 16` frames, capped at one chunk (77-81 frames,
|
||||
~5 s at the model's 16 fps). 33, 77 and 81 map to clean latent frame counts.
|
||||
- S2V always uses 16 fps. Other requested frame rates are automatically
|
||||
changed to 16 with a warning, including the CLI and server video output.
|
||||
`generate_video()` returns the actual frame rate through `fps_out`; C API
|
||||
callers should use that value when encoding the output video.
|
||||
- One generation covers the first S2V chunk window (`--video-frames` frames).
|
||||
Long-video chunked extend mode is not implemented yet.
|
||||
- Speed: the lightx2v lightning LoRA works with S2V at 4 steps and
|
||||
`--cfg-scale 1.0`. Use the **low_noise** variant;
|
||||
the high_noise variant produces artifacts on S2V:
|
||||
|
||||
```
|
||||
--lora-model-dir ..\models\loras
|
||||
-p "...<lora:lightx2v-Wan2.2-T2V-A14B-4steps-lora-rank64-Seko-V2.0-low_noise:1.0>"
|
||||
--cfg-scale 1.0 --steps 4
|
||||
```
|
||||
|
||||
Expect some quality/dynamics loss compared to the full 20-step run.
|
||||
|
||||
### Wan2.2 T2V A14B T2I
|
||||
|
||||
```
|
||||
|
||||
@@ -22,3 +22,6 @@ Metadata mode inspects PNG/JPEG container metadata without loading any model:
|
||||
./bin/sd-cli -M metadata --image ./output.png --metadata-raw
|
||||
./bin/sd-cli -M metadata --image ./output.png --metadata-all
|
||||
```
|
||||
|
||||
For completely black or white images or videos, NaNs, and the `--linear-scale` /
|
||||
`--attn-scale` workaround, see [Troubleshooting](../../docs/troubleshooting.md).
|
||||
|
||||
@@ -419,7 +419,8 @@ void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy,
|
||||
LOG_ERROR("save preview image to '%s' failed", path.string().c_str());
|
||||
}
|
||||
} else {
|
||||
if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps, cli_params->compression_quality) != 0) {
|
||||
int fps = cli_params->preview_method == PREVIEW_PROJ ? cli_params->preview_fps / 4 : cli_params->preview_fps;
|
||||
if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, fps, cli_params->compression_quality) != 0) {
|
||||
LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
|
||||
}
|
||||
}
|
||||
@@ -687,8 +688,6 @@ int main(int argc, const char* argv[]) {
|
||||
}
|
||||
}
|
||||
cli_params.preview_fps = gen_params.fps;
|
||||
if (cli_params.preview_method == PREVIEW_PROJ)
|
||||
cli_params.preview_fps /= 4;
|
||||
|
||||
sd_set_preview_callback(step_callback,
|
||||
cli_params.preview_method,
|
||||
@@ -951,9 +950,10 @@ int main(int argc, const char* argv[]) {
|
||||
} else if (cli_params.mode == VID_GEN) {
|
||||
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
|
||||
sd_image_t* generated_video = nullptr;
|
||||
if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio)) {
|
||||
if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio, &cli_params.preview_fps)) {
|
||||
generated_video = nullptr;
|
||||
}
|
||||
gen_params.fps = cli_params.preview_fps;
|
||||
results.adopt(generated_video, num_results);
|
||||
}
|
||||
|
||||
|
||||
@@ -302,8 +302,12 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
|
||||
invalid_arg = true;
|
||||
return;
|
||||
}
|
||||
*option.target = std::stoi(argv[i]);
|
||||
found_arg = true;
|
||||
try {
|
||||
*option.target = std::stoi(argv[i]);
|
||||
} catch (const std::invalid_argument&) {
|
||||
invalid_arg = true;
|
||||
}
|
||||
found_arg = true;
|
||||
}))
|
||||
break;
|
||||
|
||||
@@ -312,8 +316,12 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
|
||||
invalid_arg = true;
|
||||
return;
|
||||
}
|
||||
*option.target = std::stof(argv[i]);
|
||||
found_arg = true;
|
||||
try {
|
||||
*option.target = std::stof(argv[i]);
|
||||
} catch (const std::invalid_argument&) {
|
||||
invalid_arg = true;
|
||||
}
|
||||
found_arg = true;
|
||||
}))
|
||||
break;
|
||||
|
||||
@@ -337,7 +345,8 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
|
||||
|
||||
if (invalid_arg) {
|
||||
if (!valid) {
|
||||
LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
|
||||
LOG_ERROR("error: invalid parameter for argument \"%s\": \"%s\"",
|
||||
arg.c_str(), (i >= argc) ? "" : argv[i]);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
@@ -350,6 +359,25 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
|
||||
return true;
|
||||
}
|
||||
|
||||
static int parse_scale_override(int argc, const char** argv, int index, float& scale) {
|
||||
if (++index >= argc) {
|
||||
return -1;
|
||||
}
|
||||
try {
|
||||
size_t end = 0;
|
||||
const std::string value = argv[index];
|
||||
float parsed = std::stof(value, &end);
|
||||
if (end != value.size() || !std::isfinite(parsed) || parsed < 0.f ||
|
||||
(parsed > 0.f && !std::isfinite(1.f / parsed))) {
|
||||
return -1;
|
||||
}
|
||||
scale = parsed;
|
||||
} catch (const std::exception&) {
|
||||
return -1;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
ArgOptions SDContextParams::get_options() {
|
||||
ArgOptions options;
|
||||
options.string_options = {
|
||||
@@ -432,6 +460,11 @@ ArgOptions SDContextParams::get_options() {
|
||||
"path to standalone LTX audio vae model",
|
||||
0,
|
||||
&audio_vae_path},
|
||||
{"",
|
||||
"--audio-encoder",
|
||||
"path to wav2vec2 audio encoder model (Wan2.2 S2V)",
|
||||
0,
|
||||
&audio_encoder_path},
|
||||
{"",
|
||||
"--taesd",
|
||||
"path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)",
|
||||
@@ -678,11 +711,23 @@ ArgOptions SDContextParams::get_options() {
|
||||
};
|
||||
|
||||
options.manual_options = {
|
||||
{"",
|
||||
"--linear-scale",
|
||||
"linear input scale override (float, default: 0 = model default, 1 = no scaling)",
|
||||
[this](int argc, const char** argv, int index) {
|
||||
return parse_scale_override(argc, argv, index, linear_scale);
|
||||
}},
|
||||
{"",
|
||||
"--attn-scale",
|
||||
"flash-attention K/V scale override (float, default: 0 = model default, 1 = no scaling); requires --fa or --diffusion-fa",
|
||||
[this](int argc, const char** argv, int index) {
|
||||
return parse_scale_override(argc, argv, index, attn_scale);
|
||||
}},
|
||||
{"",
|
||||
"--auto-fit",
|
||||
"on|off (default: on). Use one GPU for diffusion/te/vae computation and place weights on that GPU, "
|
||||
"on|off (default: on). Preserve --backend (otherwise select one GPU) and place weights on the compute GPU, "
|
||||
"RAM, another GPU, or disk in that order, according to available memory (--max-vram limits GPU budgets). "
|
||||
"Disabled by explicit --backend or --params-backend; uses automatic graph segmentation when needed",
|
||||
"Disabled by explicit --params-backend; uses automatic graph segmentation when needed",
|
||||
on_auto_fit_arg},
|
||||
{"",
|
||||
"--type",
|
||||
@@ -858,6 +903,7 @@ std::string SDContextParams::to_string() const {
|
||||
<< " vae_path: \"" << vae_path << "\",\n"
|
||||
<< " vae_format: \"" << vae_format << "\",\n"
|
||||
<< " audio_vae_path: \"" << audio_vae_path << "\",\n"
|
||||
<< " audio_encoder_path: \"" << audio_encoder_path << "\",\n"
|
||||
<< " taesd_path: \"" << taesd_path << "\",\n"
|
||||
<< " esrgan_path: \"" << esrgan_path << "\",\n"
|
||||
<< " control_net_path: \"" << control_net_path << "\",\n"
|
||||
@@ -886,6 +932,8 @@ std::string SDContextParams::to_string() const {
|
||||
<< " vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n"
|
||||
<< " flash_attn: " << (flash_attn ? "true" : "false") << ",\n"
|
||||
<< " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
|
||||
<< " linear_scale: " << linear_scale << ",\n"
|
||||
<< " attn_scale: " << attn_scale << ",\n"
|
||||
<< " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
|
||||
<< " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n"
|
||||
<< " prediction: " << sd_prediction_name(prediction) << ",\n"
|
||||
@@ -921,6 +969,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
|
||||
sd_ctx_params.embeddings_connectors_path = embeddings_connectors_path.c_str();
|
||||
sd_ctx_params.vae_path = vae_path.c_str();
|
||||
sd_ctx_params.audio_vae_path = audio_vae_path.c_str();
|
||||
sd_ctx_params.audio_encoder_path = audio_encoder_path.c_str();
|
||||
sd_ctx_params.taesd_path = taesd_path.c_str();
|
||||
sd_ctx_params.control_net_path = control_net_path.c_str();
|
||||
sd_ctx_params.ip_adapter_path = ip_adapter_path.c_str();
|
||||
@@ -939,6 +988,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
|
||||
sd_ctx_params.enable_mmap = enable_mmap;
|
||||
sd_ctx_params.flash_attn = flash_attn;
|
||||
sd_ctx_params.diffusion_flash_attn = diffusion_flash_attn;
|
||||
sd_ctx_params.linear_scale = linear_scale;
|
||||
sd_ctx_params.attn_scale = attn_scale;
|
||||
sd_ctx_params.tae_preview_only = taesd_preview;
|
||||
sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct;
|
||||
sd_ctx_params.vae_conv_direct = vae_conv_direct;
|
||||
@@ -1471,6 +1522,14 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
return 1;
|
||||
};
|
||||
|
||||
auto on_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;
|
||||
@@ -1660,6 +1719,10 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"--ref-audio",
|
||||
"standalone WAV reference for MiniMax-H3 Ref2VA (can be used multiple times)",
|
||||
on_ref_audio_arg},
|
||||
{"",
|
||||
"--audio",
|
||||
"driving audio track (Wan2.2 S2V; can be used once)",
|
||||
on_audio_arg},
|
||||
{"",
|
||||
"--cache-mode",
|
||||
"caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)",
|
||||
|
||||
@@ -131,6 +131,7 @@ struct SDContextParams {
|
||||
std::string vae_path;
|
||||
std::string vae_format = "auto";
|
||||
std::string audio_vae_path;
|
||||
std::string audio_encoder_path;
|
||||
std::string taesd_path;
|
||||
std::string esrgan_path;
|
||||
std::string control_net_path;
|
||||
@@ -175,6 +176,8 @@ struct SDContextParams {
|
||||
lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
|
||||
|
||||
bool force_sdxl_vae_conv_scale = false;
|
||||
float linear_scale = 0.f;
|
||||
float attn_scale = 0.f;
|
||||
|
||||
float flow_shift = INFINITY;
|
||||
ArgOptions get_options();
|
||||
|
||||
@@ -129,3 +129,6 @@ For detailed command-line arguments, run:
|
||||
```bash
|
||||
./bin/sd-server -h
|
||||
```
|
||||
|
||||
For completely black or white images or videos, NaNs, and the `--linear-scale` /
|
||||
`--attn-scale` startup options, see [Troubleshooting](../../docs/troubleshooting.md).
|
||||
|
||||
@@ -245,7 +245,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
|
||||
sd_image_t* raw_results = nullptr;
|
||||
if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio)) {
|
||||
if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio, &output_fps)) {
|
||||
raw_results = nullptr;
|
||||
}
|
||||
results.adopt(raw_results, num_results);
|
||||
@@ -261,7 +261,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
|
||||
std::vector<uint8_t> video_bytes = create_video_from_sd_images_to_vector(job.vid_gen.output_format,
|
||||
results.data(),
|
||||
num_results,
|
||||
job.vid_gen.gen_params.fps,
|
||||
output_fps,
|
||||
job.vid_gen.output_compression,
|
||||
generated_audio);
|
||||
free_sd_audio(generated_audio);
|
||||
@@ -273,7 +273,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
|
||||
output_media_b64 = base64_encode(video_bytes);
|
||||
output_media_mime_type = video_mime_type(job.vid_gen.output_format);
|
||||
output_frame_count = num_results;
|
||||
output_fps = job.vid_gen.gen_params.fps;
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
@@ -92,6 +92,7 @@ enum prediction_t {
|
||||
FLUX_FLOW_PRED,
|
||||
SEFI_FLOW_PRED,
|
||||
MINIT2I_FLOW_PRED,
|
||||
SENSENOVA_U1_FLOW_PRED,
|
||||
PREDICTION_COUNT
|
||||
};
|
||||
|
||||
@@ -207,6 +208,7 @@ typedef struct {
|
||||
const char* embeddings_connectors_path;
|
||||
const char* vae_path;
|
||||
const char* audio_vae_path;
|
||||
const char* audio_encoder_path;
|
||||
const char* taesd_path;
|
||||
const char* control_net_path;
|
||||
const char* ip_adapter_path;
|
||||
@@ -240,6 +242,8 @@ typedef struct {
|
||||
const char* rpc_servers;
|
||||
const char* model_args;
|
||||
bool disable_segmented_compute; // Force monolithic graph execution even when automatic graph cutting would fit memory better
|
||||
float linear_scale; // Override linear input scaling; 0 keeps the model default
|
||||
float attn_scale; // Override flash-attention K/V scaling; 0 keeps the model default
|
||||
} sd_ctx_params_t;
|
||||
|
||||
typedef struct {
|
||||
@@ -493,6 +497,9 @@ SD_API void free_sd_audio(sd_audio_t* audio);
|
||||
SD_API void sd_sample_params_init(sd_sample_params_t* sample_params);
|
||||
SD_API char* sd_sample_params_to_str(const sd_sample_params_t* sample_params);
|
||||
|
||||
// Requires a loaded context; returns a static string owned by the library, or "Unknown".
|
||||
SD_API const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx);
|
||||
|
||||
SD_API enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx);
|
||||
SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_method_t sample_method);
|
||||
|
||||
@@ -515,11 +522,13 @@ enum sd_cancel_mode_t {
|
||||
SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode);
|
||||
|
||||
SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params);
|
||||
// If non-NULL, fps_out receives the effective encoding frame rate before preview callbacks.
|
||||
SD_API bool generate_video(sd_ctx_t* sd_ctx,
|
||||
const sd_vid_gen_params_t* sd_vid_gen_params,
|
||||
sd_image_t** frames_out,
|
||||
int* num_frames_out,
|
||||
sd_audio_t** audio_out);
|
||||
sd_audio_t** audio_out,
|
||||
int* fps_out);
|
||||
|
||||
typedef struct upscaler_ctx_t upscaler_ctx_t;
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
#include "model/te/llm.hpp"
|
||||
#include "model/te/t5.hpp"
|
||||
#include "model_loader.h"
|
||||
#include "tokenizers/sensenova_u1_tokenizer.h"
|
||||
|
||||
struct SDCondition {
|
||||
sd::Tensor<float> c_crossattn;
|
||||
@@ -149,6 +150,7 @@ public:
|
||||
virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
|
||||
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
|
||||
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
||||
virtual void set_scale_overrides(float linear_scale, float attn_scale) {}
|
||||
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
|
||||
virtual void runner_end() {}
|
||||
};
|
||||
@@ -231,6 +233,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
text_model->set_scale_overrides(linear_scale, attn_scale);
|
||||
if (sd_version_is_sdxl(version)) {
|
||||
text_model2->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
text_model->set_weight_adapter(adapter);
|
||||
if (sd_version_is_sdxl(version)) {
|
||||
@@ -736,6 +745,18 @@ struct SD3CLIPEmbedder : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
if (clip_g) {
|
||||
clip_g->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_weight_adapter(adapter);
|
||||
@@ -1106,6 +1127,15 @@ struct FluxCLIPEmbedder : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_weight_adapter(adapter);
|
||||
@@ -1224,7 +1254,10 @@ struct FluxCLIPEmbedder : public Conditioner {
|
||||
true,
|
||||
clip_skip,
|
||||
false);
|
||||
GGML_ASSERT(!pooled.empty());
|
||||
if (pooled.empty()) {
|
||||
LOG_ERROR("Flux CLIP-L encoding failed");
|
||||
return {};
|
||||
}
|
||||
} else {
|
||||
pooled = sd::Tensor<float>::zeros({768});
|
||||
}
|
||||
@@ -1243,7 +1276,10 @@ struct FluxCLIPEmbedder : public Conditioner {
|
||||
input_ids,
|
||||
sd::Tensor<float>(),
|
||||
false);
|
||||
GGML_ASSERT(!chunk_hidden_states.empty());
|
||||
if (chunk_hidden_states.empty()) {
|
||||
LOG_ERROR("Flux T5 encoding failed at chunk %d/%zu", chunk_idx + 1, chunk_count);
|
||||
return {};
|
||||
}
|
||||
chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
|
||||
if (zero_out_masked) {
|
||||
chunk_hidden_states.fill_(0.0f);
|
||||
@@ -1368,6 +1404,12 @@ struct T5CLIPEmbedder : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
if (t5) {
|
||||
t5->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
if (t5) {
|
||||
t5->set_weight_adapter(adapter);
|
||||
@@ -1576,6 +1618,12 @@ struct MiniT2IConditioner : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
if (t5) {
|
||||
t5->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
if (t5) {
|
||||
t5->set_weight_adapter(adapter);
|
||||
@@ -1623,6 +1671,71 @@ struct MiniT2IConditioner : public Conditioner {
|
||||
}
|
||||
};
|
||||
|
||||
struct SenseNovaU1Conditioner : public Conditioner {
|
||||
static constexpr size_t kMaxPromptTokens = 12288;
|
||||
SenseNovaU1Tokenizer tokenizer;
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
SD_UNUSED(tensors);
|
||||
}
|
||||
|
||||
void set_flash_attention_enabled(bool enabled) override {
|
||||
SD_UNUSED(enabled);
|
||||
}
|
||||
|
||||
static std::string build_query(const std::string& text, bool is_negative) {
|
||||
static const std::string kSystemMessage =
|
||||
"You are an image generation and editing assistant that accurately understands and executes user intent.\n\n"
|
||||
"You support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you MUST start with a "
|
||||
"<think></think> block. Put all reasoning inside the block using plain text. DO NOT include any image tags. "
|
||||
"Keep it reasonable and directly useful for producing the final image.\n\n2. Non-Think Mode:\nIf no reasoning "
|
||||
"is needed, directly produce the final image.\n\nTask Types:\n\nA. Text-to-Image Generation:\n- Generate a "
|
||||
"high-quality image based on the user's description.\n- Ensure visual clarity, semantic consistency, and "
|
||||
"completeness.\n- DO NOT introduce elements that contradict or override the user's intent.\n\nB. Image Editing:\n"
|
||||
"- Use the provided image(s) as input or reference for modification or transformation.\n- The result can be an "
|
||||
"edited image or a new image based on the reference(s).\n- Preserve all unspecified attributes unless explicitly "
|
||||
"changed.\n\nGeneral Rules:\n- For any visible text in the image, follow the language specified for the rendered "
|
||||
"text in the user's description, not the language of the prompt. If no language is specified, use the user's input "
|
||||
"language.";
|
||||
|
||||
std::string query;
|
||||
if (!is_negative) {
|
||||
query += "<|im_start|>system\n";
|
||||
query += kSystemMessage;
|
||||
query += "<|im_end|>\n";
|
||||
}
|
||||
query += "<|im_start|>user\n";
|
||||
query += text;
|
||||
query += "<|im_end|>\n<|im_start|>assistant\n";
|
||||
query += is_negative ? "<img>" : "<think>\n\n</think>\n\n<img>";
|
||||
return query;
|
||||
}
|
||||
|
||||
SDCondition tokenize_condition(const std::string& text, bool is_negative) {
|
||||
auto tokens = tokenizer.encode(build_query(text, is_negative));
|
||||
if (tokens.empty() || tokens.size() > kMaxPromptTokens) {
|
||||
LOG_ERROR("SenseNova U1.5 prompt token count %zu is outside [1, %zu]",
|
||||
tokens.size(),
|
||||
kMaxPromptTokens);
|
||||
return {};
|
||||
}
|
||||
|
||||
SDCondition result;
|
||||
result.c_input_ids = sd::Tensor<int32_t>({static_cast<int64_t>(tokens.size())}, tokens);
|
||||
return result;
|
||||
}
|
||||
|
||||
SDCondition get_learned_condition(int n_threads,
|
||||
const ConditionerParams& conditioner_params) override {
|
||||
SD_UNUSED(n_threads);
|
||||
return tokenize_condition(conditioner_params.text, false);
|
||||
}
|
||||
|
||||
SDCondition get_unconditional_condition(const std::string& text) {
|
||||
return tokenize_condition(text, true);
|
||||
}
|
||||
};
|
||||
|
||||
struct AnimaConditioner : public Conditioner {
|
||||
std::shared_ptr<BPETokenizer> qwen_tokenizer;
|
||||
T5UniGramTokenizer t5_tokenizer;
|
||||
@@ -1672,6 +1785,10 @@ struct AnimaConditioner : public Conditioner {
|
||||
llm->set_flash_attention_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
llm->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
llm->set_weight_adapter(adapter);
|
||||
}
|
||||
@@ -1876,6 +1993,13 @@ struct LLMEmbedder : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
llm->set_scale_overrides(linear_scale, attn_scale);
|
||||
if (byt5) {
|
||||
byt5->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
if (llm) {
|
||||
llm->set_weight_adapter(adapter);
|
||||
@@ -2965,6 +3089,11 @@ struct LTXAVEmbedder : public Conditioner {
|
||||
projector->set_flash_attention_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
llm->set_scale_overrides(linear_scale, attn_scale);
|
||||
projector->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
|
||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
||||
llm->set_max_graph_vram_bytes(max_vram_bytes);
|
||||
projector->set_max_graph_vram_bytes(max_vram_bytes);
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
#include "wan_audio.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstddef>
|
||||
|
||||
namespace sd::wan_audio {
|
||||
|
||||
static BucketPlan plan_buckets(int audio_frames, int batch_frames, int video_rate, int fps) {
|
||||
BucketPlan plan;
|
||||
plan.audio_frames = audio_frames;
|
||||
plan.batch_frames = batch_frames;
|
||||
plan.video_rate = video_rate;
|
||||
plan.fps = fps;
|
||||
const double scale = static_cast<double>(video_rate) / fps;
|
||||
// Keep a trailing chunk even when audio ends on a chunk boundary.
|
||||
plan.num_chunks = static_cast<int>(audio_frames / (batch_frames * scale)) + 1;
|
||||
plan.bucket_frames = plan.num_chunks * batch_frames;
|
||||
plan.padded_audio_frames = static_cast<int>(
|
||||
std::ceil(plan.bucket_frames / static_cast<double>(fps) * video_rate));
|
||||
return plan;
|
||||
}
|
||||
|
||||
// Match NumPy's round-half-even sampling.
|
||||
static int bucket_source_frame(int bucket_frame, int video_rate, int fps) {
|
||||
return static_cast<int>(std::nearbyint(static_cast<double>(bucket_frame) * video_rate / fps));
|
||||
}
|
||||
|
||||
static int interpolated_frame_count(int in_frames, int input_fps, int output_fps) {
|
||||
return static_cast<int>(in_frames / static_cast<double>(input_fps) * output_fps);
|
||||
}
|
||||
|
||||
// Match PyTorch linear interpolation with align_corners=True.
|
||||
static std::vector<float> linear_interpolate_frames(const std::vector<float>& in,
|
||||
int num_layers,
|
||||
int in_frames,
|
||||
int dim,
|
||||
int out_frames) {
|
||||
std::vector<float> out(static_cast<size_t>(num_layers) * out_frames * dim, 0.0f);
|
||||
if (in.empty() || in_frames <= 0 || out_frames <= 0 || num_layers <= 0 || dim <= 0) {
|
||||
return out;
|
||||
}
|
||||
const double scale = out_frames > 1 ? static_cast<double>(in_frames - 1) / (out_frames - 1) : 0.0;
|
||||
for (int layer = 0; layer < num_layers; ++layer) {
|
||||
for (int out_i = 0; out_i < out_frames; ++out_i) {
|
||||
const double pos = out_i * scale;
|
||||
const int src0 = static_cast<int>(pos);
|
||||
const int src1 = std::min(src0 + 1, in_frames - 1);
|
||||
const float frac = static_cast<float>(pos - src0);
|
||||
const float* in_row = &in[(static_cast<size_t>(layer) * in_frames + src0) * dim];
|
||||
const float* in_next = &in[(static_cast<size_t>(layer) * in_frames + src1) * dim];
|
||||
float* out_row = &out[(static_cast<size_t>(layer) * out_frames + out_i) * dim];
|
||||
for (int d = 0; d < dim; ++d) {
|
||||
out_row[d] = in_row[d] * (1.0f - frac) + in_next[d] * frac;
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<float> build_audio_buckets(const float* stacked_states,
|
||||
int num_layers,
|
||||
int in_frames,
|
||||
int dim,
|
||||
int batch_frames,
|
||||
BucketPlan* plan_out,
|
||||
int input_fps,
|
||||
int video_rate,
|
||||
int fps) {
|
||||
if (stacked_states == nullptr || num_layers <= 0 || in_frames <= 0 || dim <= 0 || batch_frames <= 0) {
|
||||
return {};
|
||||
}
|
||||
const int audio_frames = interpolated_frame_count(in_frames, input_fps, video_rate);
|
||||
if (audio_frames <= 0) {
|
||||
return {};
|
||||
}
|
||||
const std::vector<float> interpolated =
|
||||
linear_interpolate_frames(std::vector<float>(stacked_states,
|
||||
stacked_states + static_cast<size_t>(num_layers) * in_frames * dim),
|
||||
num_layers,
|
||||
in_frames,
|
||||
dim,
|
||||
audio_frames);
|
||||
const BucketPlan plan = plan_buckets(audio_frames, batch_frames, video_rate, fps);
|
||||
if (plan_out != nullptr) {
|
||||
*plan_out = plan;
|
||||
}
|
||||
std::vector<float> buckets(static_cast<size_t>(plan.bucket_frames) * num_layers * dim, 0.0f);
|
||||
for (int frame = 0; frame < plan.bucket_frames; ++frame) {
|
||||
const int src = bucket_source_frame(frame, video_rate, fps);
|
||||
if (src >= plan.audio_frames) {
|
||||
continue;
|
||||
}
|
||||
for (int layer = 0; layer < num_layers; ++layer) {
|
||||
std::copy_n(interpolated.data() + (static_cast<size_t>(layer) * audio_frames + src) * dim,
|
||||
static_cast<size_t>(dim),
|
||||
buckets.data() + (static_cast<size_t>(frame) * num_layers + layer) * dim);
|
||||
}
|
||||
}
|
||||
return buckets;
|
||||
}
|
||||
|
||||
} // namespace sd::wan_audio
|
||||
@@ -0,0 +1,32 @@
|
||||
#ifndef __SD_CONDITIONING_WAN_AUDIO_H__
|
||||
#define __SD_CONDITIONING_WAN_AUDIO_H__
|
||||
|
||||
#include <vector>
|
||||
|
||||
namespace sd::wan_audio {
|
||||
|
||||
struct BucketPlan {
|
||||
int audio_frames; // frames at video_rate
|
||||
int batch_frames; // latent_t * 4
|
||||
int video_rate;
|
||||
int fps; // bucket frame rate
|
||||
int num_chunks; // includes trailing padding
|
||||
int bucket_frames;
|
||||
int padded_audio_frames;
|
||||
};
|
||||
|
||||
// [layers, frames, dim] at input_fps -> [bucket_frames, layers, dim] at fps.
|
||||
// Pads past the audio end; returns an empty vector on invalid input.
|
||||
std::vector<float> build_audio_buckets(const float* stacked_states,
|
||||
int num_layers,
|
||||
int in_frames,
|
||||
int dim,
|
||||
int batch_frames,
|
||||
BucketPlan* plan_out = nullptr,
|
||||
int input_fps = 50,
|
||||
int video_rate = 30,
|
||||
int fps = 16);
|
||||
|
||||
} // namespace sd::wan_audio
|
||||
|
||||
#endif // __SD_CONDITIONING_WAN_AUDIO_H__
|
||||
+137
-40
@@ -58,9 +58,15 @@ namespace sd::backend_fit {
|
||||
size_t params_device = SIZE_MAX;
|
||||
};
|
||||
|
||||
struct Runtime {
|
||||
std::string name;
|
||||
std::vector<size_t> devices;
|
||||
};
|
||||
|
||||
struct Plan {
|
||||
bool valid = false;
|
||||
size_t main_device = SIZE_MAX;
|
||||
std::vector<Runtime> runtimes;
|
||||
std::vector<Decision> decisions;
|
||||
};
|
||||
|
||||
@@ -121,11 +127,14 @@ namespace sd::backend_fit {
|
||||
return name;
|
||||
}
|
||||
|
||||
static std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
|
||||
static std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets,
|
||||
bool include_other_devices) {
|
||||
std::vector<Device> out;
|
||||
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
|
||||
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
|
||||
if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) {
|
||||
const auto type = ggml_backend_dev_type(dev);
|
||||
if (type != GGML_BACKEND_DEVICE_TYPE_GPU &&
|
||||
(!include_other_devices || type == GGML_BACKEND_DEVICE_TYPE_CPU)) {
|
||||
continue;
|
||||
}
|
||||
Device device;
|
||||
@@ -183,18 +192,29 @@ namespace sd::backend_fit {
|
||||
return -1;
|
||||
}
|
||||
|
||||
static Plan compute_plan(const std::vector<Component>& components,
|
||||
const std::vector<Device>& devices,
|
||||
int64_t ram_budget_bytes) {
|
||||
Plan plan;
|
||||
static size_t select_main_device(const std::vector<Device>& devices) {
|
||||
size_t main_device = SIZE_MAX;
|
||||
for (size_t di = 0; di < devices.size(); ++di) {
|
||||
if (devices[di].budget_bytes > 0 &&
|
||||
(plan.main_device == SIZE_MAX || devices[di].budget_bytes > devices[plan.main_device].budget_bytes)) {
|
||||
plan.main_device = di;
|
||||
(main_device == SIZE_MAX || devices[di].budget_bytes > devices[main_device].budget_bytes)) {
|
||||
main_device = di;
|
||||
}
|
||||
}
|
||||
if (plan.main_device == SIZE_MAX) {
|
||||
return plan;
|
||||
return main_device;
|
||||
}
|
||||
|
||||
static Plan compute_plan(const std::vector<Component>& components,
|
||||
const std::vector<Device>& devices,
|
||||
int64_t ram_budget_bytes,
|
||||
const std::vector<Runtime>& runtimes = {}) {
|
||||
Plan plan;
|
||||
plan.main_device = select_main_device(devices);
|
||||
plan.runtimes = runtimes;
|
||||
if (plan.runtimes.empty()) {
|
||||
if (plan.main_device == SIZE_MAX) {
|
||||
return plan;
|
||||
}
|
||||
plan.runtimes.resize(components.size(), {devices[plan.main_device].name, {plan.main_device}});
|
||||
}
|
||||
|
||||
std::vector<size_t> order(components.size());
|
||||
@@ -212,6 +232,27 @@ namespace sd::backend_fit {
|
||||
ram_budget_bytes = std::max<int64_t>(ram_budget_bytes, 0);
|
||||
plan.decisions.resize(components.size());
|
||||
|
||||
auto uses_device = [&](size_t ci, size_t di) {
|
||||
const auto& runtime_devices = plan.runtimes[ci].devices;
|
||||
return std::find(runtime_devices.begin(), runtime_devices.end(), di) != runtime_devices.end();
|
||||
};
|
||||
auto headroom_for = [&](size_t ci, size_t di) {
|
||||
// Higher-priority offloaded weights need cache space on their compute devices.
|
||||
int64_t headroom = 0;
|
||||
for (size_t other = 0; other < components.size(); ++other) {
|
||||
if (components[other].params_bytes == 0 || !uses_device(other, di)) {
|
||||
continue;
|
||||
}
|
||||
const bool resident = other == ci || plan.decisions[other].params_location == ParamsLocation::MAIN_GPU;
|
||||
const int64_t cached_weights = components[other].kind < components[ci].kind
|
||||
? components[other].params_bytes
|
||||
: components[other].staging_bytes;
|
||||
headroom = std::max(headroom, components[other].reserve_bytes +
|
||||
(resident ? 0 : cached_weights));
|
||||
}
|
||||
return headroom;
|
||||
};
|
||||
|
||||
for (size_t ci : order) {
|
||||
const Component& comp = components[ci];
|
||||
Decision& decision = plan.decisions[ci];
|
||||
@@ -219,24 +260,19 @@ namespace sd::backend_fit {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Higher-priority offloaded weights need GPU cache space across graph runs.
|
||||
int64_t headroom = 0;
|
||||
for (size_t other = 0; other < components.size(); ++other) {
|
||||
if (components[other].params_bytes == 0) {
|
||||
continue;
|
||||
}
|
||||
const bool resident = other == ci || plan.decisions[other].params_location == ParamsLocation::MAIN_GPU;
|
||||
const int64_t cached_weights = components[other].kind < comp.kind
|
||||
? components[other].params_bytes
|
||||
: components[other].staging_bytes;
|
||||
headroom = std::max(headroom, components[other].reserve_bytes +
|
||||
(resident ? 0 : cached_weights));
|
||||
}
|
||||
int64_t& main_remaining = remaining[plan.main_device];
|
||||
if (headroom <= main_remaining && comp.params_bytes <= main_remaining - headroom) {
|
||||
const auto& runtime_devices = plan.runtimes[ci].devices;
|
||||
const bool fits_runtime = !runtime_devices.empty() &&
|
||||
std::all_of(runtime_devices.begin(), runtime_devices.end(), [&](size_t di) {
|
||||
const int64_t headroom = headroom_for(ci, di);
|
||||
return headroom <= remaining[di] && comp.params_bytes <= remaining[di] - headroom;
|
||||
});
|
||||
if (fits_runtime) {
|
||||
decision.params_location = ParamsLocation::MAIN_GPU;
|
||||
decision.params_device = plan.main_device;
|
||||
main_remaining -= comp.params_bytes;
|
||||
decision.params_device = runtime_devices.front();
|
||||
// Exact split allocations are unavailable until the runners build their plans.
|
||||
for (size_t di : runtime_devices) {
|
||||
remaining[di] -= comp.params_bytes;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (comp.params_bytes <= ram_budget_bytes) {
|
||||
@@ -244,10 +280,14 @@ namespace sd::backend_fit {
|
||||
ram_budget_bytes -= comp.params_bytes;
|
||||
continue;
|
||||
}
|
||||
if (runtime_devices.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t best = SIZE_MAX;
|
||||
for (size_t di = 0; di < devices.size(); ++di) {
|
||||
if (di != plan.main_device && comp.params_bytes <= remaining[di] &&
|
||||
const int64_t headroom = headroom_for(ci, di);
|
||||
if (!uses_device(ci, di) && headroom <= remaining[di] && comp.params_bytes <= remaining[di] - headroom &&
|
||||
(best == SIZE_MAX || remaining[di] > remaining[best])) {
|
||||
best = di;
|
||||
}
|
||||
@@ -280,7 +320,7 @@ namespace sd::backend_fit {
|
||||
const std::vector<Device>& devices,
|
||||
int64_t free_ram,
|
||||
int64_t ram_budget) {
|
||||
LOG_INFO("auto-fit plan (single-GPU compute on %s):", devices[plan.main_device].name.c_str());
|
||||
LOG_INFO("auto-fit plan:");
|
||||
LOG_INFO(" devices:");
|
||||
for (const Device& device : devices) {
|
||||
LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB",
|
||||
@@ -293,17 +333,19 @@ namespace sd::backend_fit {
|
||||
LOG_INFO(" RAM free %6lld MiB, params budget %6lld MiB",
|
||||
(long long)(free_ram / MiB), (long long)(ram_budget / MiB));
|
||||
}
|
||||
LOG_INFO(" main-GPU weight cache priority: diffusion > te > vae");
|
||||
LOG_INFO(" components (params: main GPU -> RAM -> other GPU -> disk):");
|
||||
LOG_INFO(" compute-device weight cache priority: diffusion > te > vae");
|
||||
LOG_INFO(" components (params: compute device -> RAM -> other GPU -> disk):");
|
||||
for (size_t ci = 0; ci < components.size(); ++ci) {
|
||||
const Component& comp = components[ci];
|
||||
if (comp.params_bytes == 0) {
|
||||
continue;
|
||||
}
|
||||
const std::string params = params_backend_name(plan.decisions[ci], devices);
|
||||
const std::string params = plan.decisions[ci].params_location == ParamsLocation::MAIN_GPU
|
||||
? plan.runtimes[ci].name
|
||||
: params_backend_name(plan.decisions[ci], devices);
|
||||
LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> compute %s, params %s",
|
||||
comp.name, (long long)(comp.params_bytes / MiB), (long long)(comp.reserve_bytes / MiB),
|
||||
devices[plan.main_device].name.c_str(), params.c_str());
|
||||
plan.runtimes[ci].name.c_str(), params.c_str());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -328,6 +370,51 @@ namespace sd::backend_fit {
|
||||
return "";
|
||||
}
|
||||
|
||||
static bool resolve_runtimes(const std::vector<Component>& components,
|
||||
const std::vector<Device>& devices,
|
||||
std::string& runtime_spec,
|
||||
std::vector<Runtime>& runtimes,
|
||||
std::string& error) {
|
||||
SDBackendAssignment assignment;
|
||||
if (!sd_parse_backend_assignment(runtime_spec, &assignment, &error)) {
|
||||
return false;
|
||||
}
|
||||
const size_t main_device = select_main_device(devices);
|
||||
const SDBackendModule modules[] = {SDBackendModule::DIFFUSION, SDBackendModule::TE, SDBackendModule::VAE};
|
||||
for (const Component& comp : components) {
|
||||
std::string name = assignment.get(modules[int(comp.kind)]);
|
||||
if (name.empty()) {
|
||||
name = main_device == SIZE_MAX ? "cpu" : devices[main_device].name;
|
||||
if (comp.params_bytes > 0) {
|
||||
append_assignment(runtime_spec, module_key(comp.kind), name);
|
||||
}
|
||||
}
|
||||
Runtime runtime;
|
||||
for (const std::string& part : split_string(name, '&')) {
|
||||
if (trim(part).empty()) {
|
||||
continue;
|
||||
}
|
||||
const std::string resolved = sd_backend_resolve_name(part);
|
||||
if (resolved.empty()) {
|
||||
error = "backend '" + part + "' was not found";
|
||||
return false;
|
||||
}
|
||||
if (!runtime.name.empty()) {
|
||||
runtime.name += "&";
|
||||
}
|
||||
runtime.name += resolved;
|
||||
for (size_t di = 0; di < devices.size(); ++di) {
|
||||
if (devices[di].name == resolved &&
|
||||
std::find(runtime.devices.begin(), runtime.devices.end(), di) == runtime.devices.end()) {
|
||||
runtime.devices.push_back(di);
|
||||
}
|
||||
}
|
||||
}
|
||||
runtimes.push_back(std::move(runtime));
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool derive_backend_specs(ModelLoader& loader,
|
||||
ggml_type override_wtype,
|
||||
sd::ggml_graph_cut::MaxVramAssignment& budgets,
|
||||
@@ -339,12 +426,18 @@ namespace sd::backend_fit {
|
||||
return false;
|
||||
}
|
||||
|
||||
const auto components = estimate_components(loader, override_wtype);
|
||||
const auto devices = enumerate_gpu_devices(budgets);
|
||||
// Resolve once to ensure dynamic backends are loaded before enumerating devices.
|
||||
sd_backend_resolve_name("");
|
||||
const auto components = estimate_components(loader, override_wtype);
|
||||
const auto devices = enumerate_gpu_devices(budgets, !runtime_spec.empty());
|
||||
std::vector<Runtime> runtimes;
|
||||
if (!runtime_spec.empty() && !resolve_runtimes(components, devices, runtime_spec, runtimes, error)) {
|
||||
LOG_ERROR("%s", error.c_str());
|
||||
return false;
|
||||
}
|
||||
const int64_t free_ram = available_ram_bytes();
|
||||
const int64_t ram_budget = std::max<int64_t>(free_ram - std::max<int64_t>(2048 * MiB, free_ram / 10), 0);
|
||||
const auto plan = compute_plan(components, devices, ram_budget);
|
||||
runtime_spec.clear();
|
||||
const auto plan = compute_plan(components, devices, ram_budget, runtimes);
|
||||
params_spec.clear();
|
||||
if (!plan.valid) {
|
||||
if (devices.empty()) {
|
||||
@@ -362,7 +455,9 @@ namespace sd::backend_fit {
|
||||
continue;
|
||||
}
|
||||
const char* key = module_key(components[ci].kind);
|
||||
append_assignment(runtime_spec, key, devices[plan.main_device].name);
|
||||
if (runtimes.empty()) {
|
||||
append_assignment(runtime_spec, key, plan.runtimes[ci].name);
|
||||
}
|
||||
if (plan.decisions[ci].params_location != ParamsLocation::MAIN_GPU) {
|
||||
append_assignment(params_spec, key, params_backend_name(plan.decisions[ci], devices));
|
||||
}
|
||||
@@ -389,7 +484,9 @@ namespace sd::backend_fit {
|
||||
tiling_params.temporal_tiling = true;
|
||||
retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal";
|
||||
} else if (!tiling_params.enabled) {
|
||||
tiling_params.enabled = true;
|
||||
tiling_params.enabled = true;
|
||||
tiling_params.rel_size_x = 0.5f;
|
||||
tiling_params.rel_size_y = 0.5f;
|
||||
if (tiling_params.tile_size_x <= 0) {
|
||||
tiling_params.tile_size_x = 256;
|
||||
}
|
||||
@@ -401,7 +498,7 @@ namespace sd::backend_fit {
|
||||
return false;
|
||||
}
|
||||
|
||||
LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling",
|
||||
LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling",
|
||||
retry_mode);
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -2,12 +2,14 @@
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
#include <exception>
|
||||
#include <map>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
|
||||
#include "core/ggml_extend_backend.h"
|
||||
#include "core/ggml_graph_cut.h"
|
||||
#include "core/util.h"
|
||||
#include "ggml-cpu.h"
|
||||
#include "ggml/src/ggml-impl.h"
|
||||
|
||||
@@ -228,11 +230,23 @@ namespace sd {
|
||||
}
|
||||
}
|
||||
|
||||
void ComputeWorkspace::segment_end() {
|
||||
if (active_) {
|
||||
synchronize();
|
||||
active_ = false;
|
||||
bool ComputeWorkspace::segment_end() noexcept {
|
||||
if (!active_) {
|
||||
return true;
|
||||
}
|
||||
// Outer cleanup guards must not retry a failed backend submission.
|
||||
active_ = false;
|
||||
try {
|
||||
synchronize();
|
||||
return true;
|
||||
} catch (const std::exception& error) {
|
||||
LOG_ERROR("%s workspace synchronization failed during segment cleanup: %s",
|
||||
ggml_backend_name(backend_), error.what());
|
||||
} catch (...) {
|
||||
LOG_ERROR("%s workspace synchronization failed during segment cleanup: unknown exception",
|
||||
ggml_backend_name(backend_));
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool ComputeWorkspace::release() {
|
||||
|
||||
@@ -51,7 +51,7 @@ namespace sd {
|
||||
const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
|
||||
const AssignNodes& assign_nodes);
|
||||
void synchronize() const;
|
||||
void segment_end();
|
||||
bool segment_end() noexcept;
|
||||
bool release();
|
||||
bool active() const { return active_; }
|
||||
ggml_backend_sched_t scheduler() const { return scheduler_; }
|
||||
|
||||
@@ -325,6 +325,76 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
|
||||
return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
|
||||
}
|
||||
|
||||
static ggml_tensor* conv_1d(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, int s0, int p0, int d0, bool force_prec_f32) {
|
||||
ggml_tensor* result;
|
||||
if (force_prec_f32) {
|
||||
ggml_tensor* patches = ggml_im2col(ctx, w, x, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F32);
|
||||
result = ggml_mul_mat(ctx,
|
||||
ggml_reshape_2d(ctx, patches, patches->ne[0], patches->ne[2] * patches->ne[1]),
|
||||
ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1], w->ne[2]));
|
||||
result = ggml_reshape_3d(ctx, result, patches->ne[1], w->ne[2], patches->ne[2]);
|
||||
} else {
|
||||
result = ggml_conv_1d(ctx, w, x, s0, p0, d0);
|
||||
}
|
||||
if (x->ne[2] > 1) {
|
||||
// mul_mat packs positions and batches before output channels: [OL, N, OC].
|
||||
result = ggml_reshape_3d(ctx, result, result->ne[0], x->ne[2], w->ne[2]);
|
||||
result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 2, 1, 3));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* w,
|
||||
ggml_tensor* b,
|
||||
int s0,
|
||||
int p0,
|
||||
int d0,
|
||||
int64_t groups,
|
||||
bool force_prec_f32) {
|
||||
GGML_ASSERT(s0 > 0 && p0 >= 0 && d0 > 0 && groups > 0);
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32 && x->ne[3] == 1 && w->ne[3] == 1);
|
||||
GGML_ASSERT(x->ne[1] % groups == 0 && w->ne[2] % groups == 0);
|
||||
GGML_ASSERT(w->ne[1] == x->ne[1] / groups);
|
||||
GGML_ASSERT(b == nullptr || (b->type == GGML_TYPE_F32 && ggml_is_vector(b) && b->ne[0] == w->ne[2]));
|
||||
|
||||
// im2col requires contiguous time rows; group views must retain the real channel and batch strides.
|
||||
if (!ggml_is_contiguous(x)) {
|
||||
x = ggml_cont(ctx, x);
|
||||
}
|
||||
if (force_prec_f32 && w->type != GGML_TYPE_F32) {
|
||||
w = ggml_cast(ctx, w, GGML_TYPE_F32);
|
||||
}
|
||||
if (!ggml_is_contiguous(w)) {
|
||||
w = ggml_cont(ctx, w);
|
||||
}
|
||||
|
||||
ggml_tensor* result = nullptr;
|
||||
if (groups == 1) {
|
||||
result = conv_1d(ctx, x, w, s0, p0, d0, force_prec_f32);
|
||||
} else {
|
||||
const int64_t ic_g = x->ne[1] / groups;
|
||||
const int64_t oc_g = w->ne[2] / groups;
|
||||
std::vector<ggml_tensor*> outputs;
|
||||
outputs.reserve(groups);
|
||||
for (int64_t group = 0; group < groups; ++group) {
|
||||
ggml_tensor* x_i = ggml_view_3d(ctx, x, x->ne[0], ic_g, x->ne[2], x->nb[1], x->nb[2], group * ic_g * x->nb[1]);
|
||||
ggml_tensor* w_i = ggml_view_3d(ctx, w, w->ne[0], ic_g, oc_g, w->nb[1], w->nb[2], group * oc_g * w->nb[2]);
|
||||
outputs.push_back(conv_1d(ctx, x_i, w_i, s0, p0, d0, force_prec_f32));
|
||||
}
|
||||
result = ggml_ext_vec_concat(ctx, outputs, 1);
|
||||
}
|
||||
if (b != nullptr) {
|
||||
if (!ggml_is_contiguous(b)) {
|
||||
b = ggml_cont(ctx, b);
|
||||
}
|
||||
b = ggml_reshape_3d(ctx, b, 1, w->ne[2], 1);
|
||||
result = ggml_add_inplace(ctx, result, b);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* w,
|
||||
@@ -683,17 +753,16 @@ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* w,
|
||||
ggml_tensor* b,
|
||||
int num_groups) {
|
||||
int num_groups,
|
||||
float eps) {
|
||||
if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) {
|
||||
w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
|
||||
b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
|
||||
}
|
||||
|
||||
const float eps = 1e-6f; // default eps parameter
|
||||
x = ggml_group_norm(ctx, x, num_groups, eps);
|
||||
x = ggml_group_norm(ctx, x, num_groups, eps);
|
||||
if (w != nullptr && b != nullptr) {
|
||||
x = ggml_mul_inplace(ctx, x, w);
|
||||
// b = ggml_repeat(ctx, b, x);
|
||||
x = ggml_add_inplace(ctx, x, b);
|
||||
}
|
||||
return x;
|
||||
|
||||
+14
-1
@@ -103,6 +103,18 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
|
||||
bool circular_x = false,
|
||||
bool circular_y = false);
|
||||
|
||||
// ggml layout: x [L, IC, N], w [K, IC/groups, OC], b [OC], result [OL, OC, N].
|
||||
// force_prec_f32 keeps both input patches and weights in F32.
|
||||
ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* w,
|
||||
ggml_tensor* b,
|
||||
int s0 = 1,
|
||||
int p0 = 0,
|
||||
int d0 = 1,
|
||||
int64_t groups = 1,
|
||||
bool force_prec_f32 = false);
|
||||
|
||||
// w: [OC,IC, KH, KW]
|
||||
// x: [N, IC, IH, IW]
|
||||
// b: [OC,]
|
||||
@@ -219,7 +231,8 @@ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* w,
|
||||
ggml_tensor* b,
|
||||
int num_groups = 32);
|
||||
int num_groups = 32,
|
||||
float eps = 1e-6f);
|
||||
|
||||
ggml_tensor* ggml_ext_timestep_embedding(
|
||||
ggml_context* ctx,
|
||||
|
||||
@@ -87,6 +87,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m
|
||||
*module = SDBackendModule::DETECTOR;
|
||||
return true;
|
||||
}
|
||||
if (name == "audioencoder" || name == "audio") {
|
||||
*module = SDBackendModule::AUDIO_ENCODER;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -593,7 +597,7 @@ static ggml_backend_t sd_get_default_backend() {
|
||||
return backend;
|
||||
}
|
||||
|
||||
static bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
|
||||
bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
|
||||
if (assignment == nullptr) {
|
||||
return false;
|
||||
}
|
||||
@@ -660,7 +664,13 @@ void SDBackendAssignment::set_module(SDBackendModule module, const std::string&
|
||||
}
|
||||
|
||||
void SDBackendHandleDeleter::operator()(ggml_backend_t backend) const {
|
||||
ggml_backend_free(backend);
|
||||
try {
|
||||
ggml_backend_free(backend);
|
||||
} catch (const std::exception& error) {
|
||||
LOG_ERROR("backend cleanup failed: %s", error.what());
|
||||
} catch (...) {
|
||||
LOG_ERROR("backend cleanup failed: unknown exception");
|
||||
}
|
||||
}
|
||||
|
||||
SDBackendManager::~SDBackendManager() {
|
||||
@@ -962,6 +972,8 @@ const char* sd_backend_module_name(SDBackendModule module) {
|
||||
return "upscaler";
|
||||
case SDBackendModule::DETECTOR:
|
||||
return "detector";
|
||||
case SDBackendModule::AUDIO_ENCODER:
|
||||
return "audio_encoder";
|
||||
}
|
||||
return "unknown";
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ enum class SDBackendModule {
|
||||
PHOTOMAKER,
|
||||
UPSCALER,
|
||||
DETECTOR,
|
||||
AUDIO_ENCODER,
|
||||
};
|
||||
|
||||
struct SDBackendAssignment {
|
||||
@@ -93,6 +94,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
|
||||
sd_graph_eval_callback_t callback_eval,
|
||||
void* callback_eval_user_data);
|
||||
std::string sd_backend_resolve_name(const std::string& name);
|
||||
bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
|
||||
const char* sd_backend_module_name(SDBackendModule module);
|
||||
void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
|
||||
bool add_rpc_devices(const std::string& servers);
|
||||
|
||||
@@ -482,6 +482,16 @@ namespace sd::ggml_graph_cut {
|
||||
return ggml_nbytes(cache_src);
|
||||
}
|
||||
|
||||
static bool can_ignore_op_params(ggml_op op) {
|
||||
// Exempt only parameters that cannot affect graph layout or backend allocation size.
|
||||
switch (op) {
|
||||
case GGML_OP_SCALE:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
|
||||
std::vector<const ggml_tensor*> tensors;
|
||||
std::unordered_map<const ggml_tensor*, size_t> indices;
|
||||
@@ -530,8 +540,10 @@ namespace sd::ggml_graph_cut {
|
||||
for (auto source : tensor->src) {
|
||||
signature.push_back(source == nullptr ? 0 : indices.at(source));
|
||||
}
|
||||
for (int value : tensor->op_params) {
|
||||
signature.push_back(static_cast<uint32_t>(value));
|
||||
if (!can_ignore_op_params(tensor->op)) {
|
||||
for (int value : tensor->op_params) {
|
||||
signature.push_back(static_cast<uint32_t>(value));
|
||||
}
|
||||
}
|
||||
}
|
||||
return signature;
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
#include <algorithm>
|
||||
#include <exception>
|
||||
#include <map>
|
||||
#include <utility>
|
||||
|
||||
#include "core/ggml_extend.h"
|
||||
#include "core/ggml_extend_backend.h"
|
||||
#include "core/ggml_runner.h"
|
||||
#include "core/ggml_tensor_utils.h"
|
||||
@@ -11,6 +13,21 @@
|
||||
|
||||
using namespace sd;
|
||||
|
||||
ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* q,
|
||||
ggml_tensor* k,
|
||||
ggml_tensor* v,
|
||||
int64_t n_head,
|
||||
ggml_tensor* mask,
|
||||
bool skip_reshape,
|
||||
bool flash_attn,
|
||||
float kv_scale) {
|
||||
if (ctx->attn_scale > 0.f) {
|
||||
kv_scale = ctx->attn_scale;
|
||||
}
|
||||
return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale);
|
||||
}
|
||||
|
||||
void GGMLRunner::alloc_params_ctx() {
|
||||
ggml_init_params params;
|
||||
params.mem_size = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
|
||||
@@ -510,6 +527,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
|
||||
runner_ctx.ggml_ctx = compute_ctx;
|
||||
runner_ctx.backend = runtime_backend;
|
||||
runner_ctx.flash_attn_enabled = flash_attn_enabled;
|
||||
runner_ctx.linear_scale = linear_scale;
|
||||
runner_ctx.attn_scale = attn_scale;
|
||||
runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
|
||||
runner_ctx.circular_x_enabled = circular_x_enabled;
|
||||
runner_ctx.circular_y_enabled = circular_y_enabled;
|
||||
@@ -624,8 +643,15 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
|
||||
params_tensor_set_.insert(parameter);
|
||||
}
|
||||
}
|
||||
auto output = execute_graph(graph, n_threads, no_return, read_outputs);
|
||||
success = output.has_value();
|
||||
std::optional<sd::Tensor<float>> output;
|
||||
try {
|
||||
output = execute_graph(graph, n_threads, no_return, read_outputs);
|
||||
} catch (const std::exception& error) {
|
||||
LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
|
||||
ggml_backend_name(runtime_backend), error.what());
|
||||
return std::nullopt;
|
||||
}
|
||||
success = output.has_value();
|
||||
if (success) {
|
||||
cache_.graph_end(true);
|
||||
}
|
||||
@@ -938,6 +964,9 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!workspace_.segment_end()) {
|
||||
return fail_segment("workspace synchronization");
|
||||
}
|
||||
// Final outputs and their callbacks may still be views of consumed cuts.
|
||||
cut_cache_.prune(segment.future_cut_names);
|
||||
}
|
||||
|
||||
@@ -68,6 +68,8 @@ struct GGMLRunnerContext {
|
||||
ggml_backend_t backend = nullptr;
|
||||
ggml_context* ggml_ctx = nullptr;
|
||||
bool flash_attn_enabled = false;
|
||||
float linear_scale = 0.f;
|
||||
float attn_scale = 0.f;
|
||||
bool conv2d_direct_enabled = false;
|
||||
bool circular_x_enabled = false;
|
||||
bool circular_y_enabled = false;
|
||||
@@ -113,6 +115,16 @@ struct GGMLRunnerContext {
|
||||
}
|
||||
};
|
||||
|
||||
ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* q,
|
||||
ggml_tensor* k,
|
||||
ggml_tensor* v,
|
||||
int64_t n_head,
|
||||
ggml_tensor* mask = nullptr,
|
||||
bool skip_reshape = false,
|
||||
bool flash_attn = false,
|
||||
float kv_scale = 1.f);
|
||||
|
||||
struct GGMLRunner {
|
||||
private:
|
||||
std::map<ggml_backend_t, size_t> logged_compute_bytes_;
|
||||
@@ -163,6 +175,8 @@ protected:
|
||||
const std::string final_result_name = "ggml_runner_final_result_tensor";
|
||||
|
||||
bool flash_attn_enabled = false;
|
||||
float linear_scale = 0.f;
|
||||
float attn_scale = 0.f;
|
||||
bool conv2d_direct_enabled = false;
|
||||
bool circular_x_enabled = false;
|
||||
bool circular_y_enabled = false;
|
||||
@@ -323,6 +337,11 @@ public:
|
||||
flash_attn_enabled = enabled;
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) {
|
||||
this->linear_scale = linear_scale;
|
||||
this->attn_scale = attn_scale;
|
||||
}
|
||||
|
||||
void set_conv2d_direct_enabled(bool enabled) {
|
||||
conv2d_direct_enabled = enabled;
|
||||
}
|
||||
|
||||
+45
-8
@@ -821,7 +821,11 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
|
||||
float round_bracket_multiplier = 1.1f;
|
||||
float square_bracket_multiplier = 1 / 1.1f;
|
||||
|
||||
std::regex re_attention(R"(\\\(|\\\)|\\\[|\\\]|\\\\|\\|\(|\[|:([+-]?[.\d]+)\)|\)|\]|\bBREAK\b|[^\\()\[\]:B]+|:|\bB)");
|
||||
// libstdc++ std::regex recurses per matched character, so unbounded runs
|
||||
// overflow the stack. Split runs are merged back below.
|
||||
const int max_plain_text_run = 1024;
|
||||
std::regex re_attention(R"(\\\(|\\\)|\\\[|\\\]|\\\\|\\|\(|\[|\)|\]|\bBREAK\b|[^\\()\[\]:B]{1,)" +
|
||||
std::to_string(max_plain_text_run) + R"(}|:|\bB)");
|
||||
std::regex re_break(R"(\s*\bBREAK\b\s*)");
|
||||
|
||||
auto multiply_range = [&](int start_position, float multiplier) {
|
||||
@@ -830,22 +834,55 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
|
||||
}
|
||||
};
|
||||
|
||||
// Kept out of the regex: bounding the repetition rejects valid long weights,
|
||||
// leaving it unbounded overflows the stack.
|
||||
auto lex_weight = [](const std::string& s, float& value) -> size_t {
|
||||
size_t end = 0;
|
||||
if (end < s.size() && (s[end] == '+' || s[end] == '-')) {
|
||||
++end;
|
||||
}
|
||||
while (end < s.size() && (std::isdigit((unsigned char)s[end]) || s[end] == '.')) {
|
||||
++end;
|
||||
}
|
||||
if (end >= s.size() || s[end] != ')') {
|
||||
return 0;
|
||||
}
|
||||
std::string number = s.substr(0, end);
|
||||
char* number_end = nullptr;
|
||||
float parsed = std::strtof(number.c_str(), &number_end);
|
||||
const char* expected = number.c_str() + number.size();
|
||||
// Without this ".", "+." and "1.2.3" would silently become weights.
|
||||
if (number.empty() || number_end != expected || !std::isfinite(parsed)) {
|
||||
return 0;
|
||||
}
|
||||
value = parsed;
|
||||
return end + 1;
|
||||
};
|
||||
|
||||
std::smatch m, m2;
|
||||
std::string remaining_text = text;
|
||||
|
||||
while (std::regex_search(remaining_text, m, re_attention)) {
|
||||
std::string text = m[0];
|
||||
std::string weight = m[1];
|
||||
std::string suffix = m.suffix();
|
||||
|
||||
if (text == ":") {
|
||||
float weight_value = 1.0f;
|
||||
size_t weight_length = lex_weight(suffix, weight_value);
|
||||
if (weight_length > 0) {
|
||||
if (!round_brackets.empty()) {
|
||||
multiply_range(round_brackets.back(), weight_value);
|
||||
round_brackets.pop_back();
|
||||
}
|
||||
remaining_text = suffix.substr(weight_length);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if (text == "(") {
|
||||
round_brackets.push_back((int)res.size());
|
||||
} else if (text == "[") {
|
||||
square_brackets.push_back((int)res.size());
|
||||
} else if (!weight.empty()) {
|
||||
if (!round_brackets.empty()) {
|
||||
multiply_range(round_brackets.back(), std::stof(weight));
|
||||
round_brackets.pop_back();
|
||||
}
|
||||
} else if (text == ")" && !round_brackets.empty()) {
|
||||
multiply_range(round_brackets.back(), round_bracket_multiplier);
|
||||
round_brackets.pop_back();
|
||||
@@ -860,7 +897,7 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
|
||||
res.push_back({text, 1.0f});
|
||||
}
|
||||
|
||||
remaining_text = m.suffix();
|
||||
remaining_text = suffix;
|
||||
}
|
||||
|
||||
for (int pos : round_brackets) {
|
||||
|
||||
@@ -135,6 +135,7 @@ struct PhotoMakerExtension : public GenerationExtension {
|
||||
pm_version,
|
||||
20.f,
|
||||
ctx.model_manager);
|
||||
pmid_model->set_scale_overrides(ctx.params->linear_scale, ctx.params->attn_scale);
|
||||
if (pm_version == PM_VERSION_2) {
|
||||
LOG_INFO("using PhotoMaker Version 2");
|
||||
}
|
||||
|
||||
+9
-2
@@ -35,6 +35,7 @@ enum SDVersion {
|
||||
VERSION_WAN2,
|
||||
VERSION_WAN2_2_I2V,
|
||||
VERSION_WAN2_2_TI2V,
|
||||
VERSION_WAN2_2_S2V,
|
||||
VERSION_LINGBOT_VIDEO,
|
||||
VERSION_QWEN_IMAGE,
|
||||
VERSION_QWEN_IMAGE_LAYERED,
|
||||
@@ -57,6 +58,7 @@ enum SDVersion {
|
||||
VERSION_SEFI_IMAGE,
|
||||
VERSION_KREA2,
|
||||
VERSION_MAGE_FLOW,
|
||||
VERSION_SENSENOVA_U1_5,
|
||||
VERSION_ESRGAN,
|
||||
VERSION_COUNT,
|
||||
};
|
||||
@@ -129,7 +131,7 @@ static inline bool sd_version_is_minimax_h3(SDVersion version) {
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_wan(SDVersion version) {
|
||||
if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) {
|
||||
if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V || version == VERSION_WAN2_2_S2V) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
@@ -237,6 +239,10 @@ static inline bool sd_version_is_mage_flow(SDVersion version) {
|
||||
return version == VERSION_MAGE_FLOW;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_sensenova_u1(SDVersion version) {
|
||||
return version == VERSION_SENSENOVA_U1_5;
|
||||
}
|
||||
|
||||
static inline bool sd_version_uses_flux_vae(SDVersion version) {
|
||||
if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
|
||||
return true;
|
||||
@@ -295,7 +301,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
|
||||
sd_version_is_ideogram4(version) ||
|
||||
sd_version_is_sefi_image(version) ||
|
||||
sd_version_is_krea2(version) ||
|
||||
sd_version_is_mage_flow(version)) {
|
||||
sd_version_is_mage_flow(version) ||
|
||||
sd_version_is_sensenova_u1(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
|
||||
@@ -95,7 +95,7 @@ namespace IPAdapter {
|
||||
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);
|
||||
ggml_tensor* attn = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, false, false);
|
||||
attn = to_out->forward(ctx, attn);
|
||||
latents = ggml_add(ctx->ggml_ctx, latents, attn);
|
||||
|
||||
|
||||
@@ -63,12 +63,11 @@ public:
|
||||
k = ggml_cont(ctx->ggml_ctx, k);
|
||||
v = ggml_cont(ctx->ggml_ctx, v);
|
||||
|
||||
ggml_tensor* attn_out = ggml_ext_attention_ext(
|
||||
ctx->ggml_ctx, ctx->backend,
|
||||
q, k, v,
|
||||
heads,
|
||||
/*mask=*/nullptr,
|
||||
/*diag_mask_inf=*/false);
|
||||
ggml_tensor* attn_out = ggml_ext_attention_ext(ctx,
|
||||
q, k, v,
|
||||
heads,
|
||||
/*mask=*/nullptr,
|
||||
/*diag_mask_inf=*/false);
|
||||
|
||||
ggml_tensor* out = to_out->forward(ctx, attn_out);
|
||||
return out;
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
#ifndef __SD_MODEL_AUDIO_WAV2VEC2_HPP__
|
||||
#define __SD_MODEL_AUDIO_WAV2VEC2_HPP__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "core/ggml_extend.h"
|
||||
#include "core/ggml_runner.h"
|
||||
#include "model.h"
|
||||
#include "model/common/ggml_block.hpp"
|
||||
|
||||
namespace Wav2Vec2 {
|
||||
|
||||
struct Wav2Vec2Config {
|
||||
int64_t embed_dim = 1024;
|
||||
int64_t conv_dim = 512;
|
||||
int num_heads = 16;
|
||||
int num_layers = 24;
|
||||
std::string feat_extract_norm = "layer";
|
||||
bool conv_bias = true;
|
||||
bool do_normalize = true;
|
||||
bool do_stable_layer_norm = true;
|
||||
|
||||
static Wav2Vec2Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
|
||||
Wav2Vec2Config config;
|
||||
auto it = tensor_storage_map.find(prefix + "encoder.layer_norm.bias");
|
||||
if (it == tensor_storage_map.end()) {
|
||||
LOG_WARN("wav2vec2: %sencoder.layer_norm.bias not found, using large defaults", prefix.c_str());
|
||||
return config;
|
||||
}
|
||||
config.embed_dim = it->second.ne[0];
|
||||
if (config.embed_dim == 1024) {
|
||||
config.embed_dim = 1024;
|
||||
config.num_heads = 16;
|
||||
config.num_layers = 24;
|
||||
config.feat_extract_norm = "layer";
|
||||
config.conv_bias = true;
|
||||
config.do_normalize = true;
|
||||
config.do_stable_layer_norm = true;
|
||||
} else if (config.embed_dim == 768) {
|
||||
config.embed_dim = 768;
|
||||
config.num_heads = 12;
|
||||
config.num_layers = 12;
|
||||
config.feat_extract_norm = "group";
|
||||
config.conv_bias = false;
|
||||
config.do_normalize = false;
|
||||
config.do_stable_layer_norm = false;
|
||||
} else {
|
||||
LOG_WARN("wav2vec2: unsupported embed_dim %" PRId64 ", using large defaults", config.embed_dim);
|
||||
config.embed_dim = 1024;
|
||||
}
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2NoLayerNormConvLayer : public UnaryBlock {
|
||||
Wav2Vec2NoLayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
|
||||
blocks["conv"] = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
|
||||
x = conv->forward(ctx, x);
|
||||
return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2LayerNormConvLayer : public UnaryBlock {
|
||||
Wav2Vec2LayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
|
||||
blocks["conv"] = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
|
||||
blocks["layer_norm"] = std::make_shared<LayerNorm>(out_channels);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
|
||||
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
|
||||
x = conv->forward(ctx, x);
|
||||
// LayerNorm normalizes channels: [N, C, L] -> [N, L, C].
|
||||
x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
|
||||
x = layer_norm->forward(ctx, x);
|
||||
x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
|
||||
return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2GroupNormConvLayer : public UnaryBlock {
|
||||
Wav2Vec2GroupNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
|
||||
blocks["conv"] = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
|
||||
blocks["layer_norm"] = std::make_shared<GroupNorm>((int)out_channels, out_channels, 1e-05f);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
|
||||
auto layer_norm = std::dynamic_pointer_cast<GroupNorm>(blocks["layer_norm"]);
|
||||
x = conv->forward(ctx, x);
|
||||
// ggml GroupNorm needs [N, C, H, W], with H=1 for audio.
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
|
||||
x = layer_norm->forward(ctx, x);
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->ne[3]);
|
||||
return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2FeatureEncoder : public UnaryBlock {
|
||||
Wav2Vec2FeatureEncoder(const Wav2Vec2Config& config) {
|
||||
GGML_ASSERT(config.feat_extract_norm == "layer" || config.feat_extract_norm == "group");
|
||||
const int kernels[7] = {10, 3, 3, 3, 3, 2, 2};
|
||||
const int strides[7] = {5, 2, 2, 2, 2, 2, 2};
|
||||
int64_t in_channels = 1;
|
||||
for (int i = 0; i < 7; ++i) {
|
||||
const std::string name = "conv_layers." + std::to_string(i);
|
||||
if (config.feat_extract_norm == "layer") {
|
||||
blocks[name] = std::make_shared<Wav2Vec2LayerNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
|
||||
} else if (i == 0) {
|
||||
blocks[name] = std::make_shared<Wav2Vec2GroupNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
|
||||
} else {
|
||||
blocks[name] = std::make_shared<Wav2Vec2NoLayerNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
|
||||
}
|
||||
in_channels = config.conv_dim;
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
for (int i = 0; i < 7; ++i) {
|
||||
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv_layers." + std::to_string(i)]);
|
||||
x = conv->forward(ctx, x);
|
||||
}
|
||||
return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2FeatureProjection : public UnaryBlock {
|
||||
Wav2Vec2FeatureProjection(const Wav2Vec2Config& config) {
|
||||
blocks["layer_norm"] = std::make_shared<LayerNorm>(config.conv_dim);
|
||||
blocks["projection"] = std::make_shared<Linear>(config.conv_dim, config.embed_dim);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto ln = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
|
||||
auto projection = std::dynamic_pointer_cast<Linear>(blocks["projection"]);
|
||||
x = ln->forward(ctx, x);
|
||||
x = projection->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class Wav2Vec2PositionalConvEmbedding : public UnaryBlock {
|
||||
private:
|
||||
int64_t embed_dim_;
|
||||
static constexpr int groups_ = 16;
|
||||
static constexpr int kernel_size_ = 128;
|
||||
std::string weight_g_name_;
|
||||
std::string weight_v_name_;
|
||||
|
||||
ggml_tensor* weight(GGMLRunnerContext* ctx) {
|
||||
auto g = params[weight_g_name_];
|
||||
auto v = ggml_cast(ctx->ggml_ctx, params[weight_v_name_], GGML_TYPE_F32);
|
||||
auto squared = ggml_mul(ctx->ggml_ctx, v, v);
|
||||
// PyTorch weight_norm(dim=2) reduces both channel axes, retaining each kernel tap.
|
||||
squared = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, squared, 2, 0, 1, 3));
|
||||
squared = ggml_reshape_2d(ctx->ggml_ctx, squared, embed_dim_ / groups_ * embed_dim_, kernel_size_);
|
||||
auto norm = ggml_sqrt(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, squared));
|
||||
norm = ggml_reshape_3d(ctx->ggml_ctx, norm, kernel_size_, 1, 1);
|
||||
return ggml_mul(ctx->ggml_ctx, v, ggml_div(ctx->ggml_ctx, g, norm));
|
||||
}
|
||||
|
||||
public:
|
||||
Wav2Vec2PositionalConvEmbedding(const Wav2Vec2Config& config)
|
||||
: embed_dim_(config.embed_dim) {
|
||||
GGML_ASSERT(embed_dim_ > 0 && embed_dim_ % groups_ == 0);
|
||||
}
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
bool legacy = tensor_storage_map.count(prefix + "conv.weight_g") > 0;
|
||||
weight_g_name_ = legacy ? "conv.weight_g" : "conv.parametrizations.weight.original0";
|
||||
weight_v_name_ = legacy ? "conv.weight_v" : "conv.parametrizations.weight.original1";
|
||||
auto g = tensor_storage_map.find(prefix + weight_g_name_);
|
||||
auto v = tensor_storage_map.find(prefix + weight_v_name_);
|
||||
GGML_ASSERT(g != tensor_storage_map.end() && v != tensor_storage_map.end());
|
||||
GGML_ASSERT(g->second.ne[0] == kernel_size_ && g->second.ne[1] == 1 && g->second.ne[2] == 1 && g->second.ne[3] == 1);
|
||||
GGML_ASSERT(v->second.ne[0] == kernel_size_ && v->second.ne[1] == embed_dim_ / groups_ && v->second.ne[2] == embed_dim_ && v->second.ne[3] == 1);
|
||||
|
||||
params[weight_g_name_] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kernel_size_, 1, 1);
|
||||
params[weight_v_name_] = ggml_new_tensor_3d(ctx, get_type(prefix + weight_v_name_, tensor_storage_map, GGML_TYPE_F16),
|
||||
kernel_size_, embed_dim_ / groups_, embed_dim_);
|
||||
if (tensor_storage_map.count(prefix + "conv.bias") > 0) {
|
||||
params["conv.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim_);
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto w = weight(ctx);
|
||||
auto b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr;
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
|
||||
x = ggml_ext_conv_1d(ctx->ggml_ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true);
|
||||
// Apply GELU out of place before cropping to keep graph buffer reuse safe.
|
||||
x = ggml_gelu_erf(ctx->ggml_ctx, x);
|
||||
x = ggml_view_3d(ctx->ggml_ctx, x, x->ne[0] - 1, x->ne[1], x->ne[2], x->nb[1], x->nb[2], 0);
|
||||
return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2FeedForward : public UnaryBlock {
|
||||
Wav2Vec2FeedForward(const Wav2Vec2Config& config) {
|
||||
blocks["intermediate_dense"] = std::make_shared<Linear>(config.embed_dim, config.embed_dim * 4);
|
||||
blocks["output_dense"] = std::make_shared<Linear>(config.embed_dim * 4, config.embed_dim);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto intermediate_dense = std::dynamic_pointer_cast<Linear>(blocks["intermediate_dense"]);
|
||||
auto output_dense = std::dynamic_pointer_cast<Linear>(blocks["output_dense"]);
|
||||
x = intermediate_dense->forward(ctx, x);
|
||||
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
|
||||
x = output_dense->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2EncoderLayer : public UnaryBlock {
|
||||
bool do_stable_layer_norm;
|
||||
|
||||
Wav2Vec2EncoderLayer(const Wav2Vec2Config& config)
|
||||
: do_stable_layer_norm(config.do_stable_layer_norm) {
|
||||
blocks["attention"] = std::make_shared<MultiheadAttention>(config.embed_dim, config.num_heads, true, true);
|
||||
blocks["layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
|
||||
blocks["feed_forward"] = std::make_shared<Wav2Vec2FeedForward>(config);
|
||||
blocks["final_layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto attention = std::dynamic_pointer_cast<MultiheadAttention>(blocks["attention"]);
|
||||
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
|
||||
auto feed_forward = std::dynamic_pointer_cast<Wav2Vec2FeedForward>(blocks["feed_forward"]);
|
||||
auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
|
||||
|
||||
ggml_tensor* residual = x;
|
||||
if (do_stable_layer_norm) {
|
||||
x = layer_norm->forward(ctx, x);
|
||||
x = attention->forward(ctx, x);
|
||||
x = ggml_add(ctx->ggml_ctx, residual, x);
|
||||
x = ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, final_layer_norm->forward(ctx, x)));
|
||||
} else {
|
||||
x = attention->forward(ctx, x);
|
||||
x = ggml_add(ctx->ggml_ctx, residual, x);
|
||||
x = layer_norm->forward(ctx, x);
|
||||
x = final_layer_norm->forward(ctx, ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, x)));
|
||||
}
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2Encoder : public GGMLBlock {
|
||||
int num_layers;
|
||||
bool do_stable_layer_norm;
|
||||
|
||||
Wav2Vec2Encoder(const Wav2Vec2Config& config)
|
||||
: num_layers(config.num_layers), do_stable_layer_norm(config.do_stable_layer_norm) {
|
||||
blocks["pos_conv_embed"] = std::make_shared<Wav2Vec2PositionalConvEmbedding>(config);
|
||||
for (int i = 0; i < config.num_layers; ++i) {
|
||||
blocks["layers." + std::to_string(i)] = std::make_shared<Wav2Vec2EncoderLayer>(config);
|
||||
}
|
||||
blocks["layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
|
||||
}
|
||||
|
||||
// For N == 1, all_layers stacks pre-layer states and the final state as [embed_dim, L, num_layers + 1].
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) {
|
||||
auto pos_conv_embed = std::dynamic_pointer_cast<Wav2Vec2PositionalConvEmbedding>(blocks["pos_conv_embed"]);
|
||||
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
|
||||
|
||||
std::vector<ggml_tensor*> collected;
|
||||
if (all_layers != nullptr) {
|
||||
collected.reserve(num_layers + 1);
|
||||
}
|
||||
|
||||
x = ggml_add(ctx->ggml_ctx, x, pos_conv_embed->forward(ctx, x));
|
||||
if (!do_stable_layer_norm) {
|
||||
x = layer_norm->forward(ctx, x);
|
||||
}
|
||||
for (int i = 0; i < num_layers; ++i) {
|
||||
if (all_layers != nullptr) {
|
||||
collected.push_back(x);
|
||||
}
|
||||
auto layer = std::dynamic_pointer_cast<Wav2Vec2EncoderLayer>(blocks["layers." + std::to_string(i)]);
|
||||
x = layer->forward(ctx, x);
|
||||
}
|
||||
if (do_stable_layer_norm) {
|
||||
x = layer_norm->forward(ctx, x);
|
||||
}
|
||||
if (all_layers != nullptr) {
|
||||
collected.push_back(x);
|
||||
ggml_tensor* stack = collected[0];
|
||||
for (size_t i = 1; i < collected.size(); ++i) {
|
||||
stack = ggml_concat(ctx->ggml_ctx, stack, collected[i], 2);
|
||||
}
|
||||
*all_layers = stack;
|
||||
}
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct Wav2Vec2Model : public GGMLBlock {
|
||||
Wav2Vec2Config config;
|
||||
|
||||
Wav2Vec2Model() = default;
|
||||
Wav2Vec2Model(const Wav2Vec2Config& config_)
|
||||
: config(config_) {
|
||||
blocks["feature_extractor"] = std::make_shared<Wav2Vec2FeatureEncoder>(config);
|
||||
blocks["feature_projection"] = std::make_shared<Wav2Vec2FeatureProjection>(config);
|
||||
blocks["encoder"] = std::make_shared<Wav2Vec2Encoder>(config);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) {
|
||||
auto feature_extractor = std::dynamic_pointer_cast<Wav2Vec2FeatureEncoder>(blocks["feature_extractor"]);
|
||||
auto feature_projection = std::dynamic_pointer_cast<Wav2Vec2FeatureProjection>(blocks["feature_projection"]);
|
||||
auto encoder = std::dynamic_pointer_cast<Wav2Vec2Encoder>(blocks["encoder"]);
|
||||
|
||||
x = feature_extractor->forward(ctx, x);
|
||||
x = feature_projection->forward(ctx, x);
|
||||
x = encoder->forward(ctx, x, all_layers);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class Wav2Vec2ModelRunner : public GGMLRunner {
|
||||
private:
|
||||
Wav2Vec2Config config;
|
||||
|
||||
public:
|
||||
Wav2Vec2Model model;
|
||||
std::string weight_prefix;
|
||||
|
||||
Wav2Vec2ModelRunner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "wav2vec2.",
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
: GGMLRunner(backend, weight_manager),
|
||||
config(Wav2Vec2Config::detect_from_weights(tensor_storage_map, prefix)),
|
||||
model(config),
|
||||
weight_prefix(prefix) {
|
||||
// GGMLBlock appends its own separator; loader prefixes already include one.
|
||||
std::string block_prefix = weight_prefix;
|
||||
if (!block_prefix.empty() && block_prefix.back() == '.') {
|
||||
block_prefix.pop_back();
|
||||
}
|
||||
model.init(params_ctx, tensor_storage_map, block_prefix);
|
||||
LOG_INFO("%s", get_desc().c_str());
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "wav2vec2";
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
|
||||
std::string block_prefix = weight_prefix;
|
||||
if (!block_prefix.empty() && block_prefix.back() == '.') {
|
||||
block_prefix.pop_back();
|
||||
}
|
||||
model.get_param_tensors(tensors, block_prefix);
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& waveform_tensor) {
|
||||
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||
ggml_tensor* waveform = make_input(waveform_tensor);
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* all_layers = nullptr;
|
||||
model.forward(&runner_ctx, waveform, &all_layers);
|
||||
GGML_ASSERT(all_layers != nullptr);
|
||||
ggml_build_forward_expand(gf, all_layers);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(const int n_threads, const std::vector<float>& mono_waveform) {
|
||||
GGML_ASSERT(!mono_waveform.empty());
|
||||
const int64_t num_samples = (int64_t)mono_waveform.size();
|
||||
sd::Tensor<float> waveform({num_samples, 1, 1});
|
||||
std::copy(mono_waveform.begin(), mono_waveform.end(), waveform.data());
|
||||
normalize(waveform.data(), num_samples);
|
||||
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(waveform);
|
||||
};
|
||||
return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
|
||||
}
|
||||
|
||||
private:
|
||||
static void normalize(float* x, int64_t n) {
|
||||
double mean = 0.0;
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
mean += x[i];
|
||||
}
|
||||
mean /= n;
|
||||
double var = 0.0;
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
const double d = x[i] - mean;
|
||||
var += d * d;
|
||||
}
|
||||
var /= n;
|
||||
const float scale = (float)(1.0 / std::sqrt(var + 1e-7));
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
x[i] = (float)((x[i] - mean) * scale);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace Wav2Vec2
|
||||
|
||||
#endif // __SD_MODEL_AUDIO_WAV2VEC2_HPP__
|
||||
@@ -380,14 +380,14 @@ public:
|
||||
if (xtra_dim) {
|
||||
context->ne[0] = 320; // reset dim to orig
|
||||
}
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim]
|
||||
|
||||
if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) {
|
||||
auto to_k_ip = std::dynamic_pointer_cast<Linear>(blocks["to_k_ip"]);
|
||||
auto to_v_ip = std::dynamic_pointer_cast<Linear>(blocks["to_v_ip"]);
|
||||
auto k_ip = to_k_ip->forward(ctx, ctx->ip_context);
|
||||
auto v_ip = to_v_ip->forward(ctx, ctx->ip_context);
|
||||
auto x_ip = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
|
||||
auto x_ip = ggml_ext_attention_ext(ctx, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale));
|
||||
}
|
||||
|
||||
|
||||
@@ -206,6 +206,7 @@ public:
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = params["weight"];
|
||||
const float scale = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;
|
||||
ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
|
||||
if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
|
||||
bool supports_fp8_matmul = false;
|
||||
@@ -309,6 +310,7 @@ public:
|
||||
__STATIC_INLINE__ bool support_get_rows(ggml_type wtype) {
|
||||
switch (wtype) {
|
||||
case GGML_TYPE_F16:
|
||||
case GGML_TYPE_BF16:
|
||||
case GGML_TYPE_Q8_0:
|
||||
case GGML_TYPE_Q5_1:
|
||||
case GGML_TYPE_Q5_0:
|
||||
@@ -366,6 +368,61 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
class Conv1d : public UnaryBlock {
|
||||
protected:
|
||||
int64_t in_channels;
|
||||
int64_t out_channels;
|
||||
int64_t groups;
|
||||
int kernel_size;
|
||||
int stride;
|
||||
int padding;
|
||||
int dilation;
|
||||
bool bias;
|
||||
bool force_prec_f32;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16);
|
||||
params["weight"] = ggml_new_tensor_3d(ctx, wtype, kernel_size, in_channels / groups, out_channels);
|
||||
if (bias) {
|
||||
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
Conv1d(int64_t in_channels,
|
||||
int64_t out_channels,
|
||||
int kernel_size,
|
||||
int stride = 1,
|
||||
int padding = 0,
|
||||
int dilation = 1,
|
||||
int64_t groups = 1,
|
||||
bool bias = true,
|
||||
bool force_prec_f32 = false)
|
||||
: in_channels(in_channels),
|
||||
out_channels(out_channels),
|
||||
groups(groups),
|
||||
kernel_size(kernel_size),
|
||||
stride(stride),
|
||||
padding(padding),
|
||||
dilation(dilation),
|
||||
bias(bias),
|
||||
force_prec_f32(force_prec_f32) {
|
||||
GGML_ASSERT(in_channels > 0 && out_channels > 0 && groups > 0);
|
||||
GGML_ASSERT(in_channels % groups == 0 && out_channels % groups == 0);
|
||||
GGML_ASSERT(kernel_size > 0 && stride > 0 && padding >= 0 && dilation > 0);
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "Conv1d";
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
GGML_ASSERT(x->ne[1] == in_channels);
|
||||
return ggml_ext_conv_1d(ctx->ggml_ctx, x, params["weight"], bias ? params["bias"] : nullptr,
|
||||
stride, padding, dilation, groups, force_prec_f32);
|
||||
}
|
||||
};
|
||||
|
||||
class Conv2d : public UnaryBlock {
|
||||
protected:
|
||||
int64_t in_channels;
|
||||
@@ -764,7 +821,7 @@ public:
|
||||
b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
|
||||
}
|
||||
}
|
||||
return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups);
|
||||
return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups, eps);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -869,7 +926,7 @@ public:
|
||||
v = v_proj->forward(ctx, x);
|
||||
}
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false); // [N, n_token, embed_dim]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, false); // [N, n_token, embed_dim]
|
||||
|
||||
x = out_proj->forward(ctx, x); // [N, n_token, embed_dim]
|
||||
return x;
|
||||
|
||||
@@ -818,8 +818,9 @@ namespace Rope {
|
||||
int pw,
|
||||
int bs,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs);
|
||||
const std::vector<int>& axes_dim,
|
||||
int t_offset = 0) {
|
||||
std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs, t_offset);
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
@@ -1024,7 +1025,7 @@ namespace Rope {
|
||||
q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head]
|
||||
k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head]
|
||||
|
||||
auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head]
|
||||
auto x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head]
|
||||
return x;
|
||||
}
|
||||
}; // namespace Rope
|
||||
|
||||
@@ -237,8 +237,7 @@ namespace Anima {
|
||||
}
|
||||
auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
|
||||
auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
|
||||
attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
attn_out = ggml_ext_attention_ext(ctx,
|
||||
q_rope,
|
||||
k_rope,
|
||||
v4,
|
||||
@@ -249,8 +248,7 @@ namespace Anima {
|
||||
} else {
|
||||
auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
|
||||
auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
|
||||
attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
attn_out = ggml_ext_attention_ext(ctx,
|
||||
q_flat,
|
||||
k_flat,
|
||||
v,
|
||||
|
||||
@@ -61,7 +61,7 @@ namespace AnimateDiff {
|
||||
auto k = to_k->forward(ctx, x_pe);
|
||||
auto v = to_v->forward(ctx, x_pe);
|
||||
|
||||
auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false);
|
||||
auto a = ggml_ext_attention_ext(ctx, q, k, v, (int)num_heads, nullptr, false);
|
||||
return to_out->forward(ctx, a);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -183,7 +183,7 @@ namespace ErnieImage {
|
||||
k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); // [N, heads, S, head_dim]
|
||||
k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size]
|
||||
x = to_out_0->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
|
||||
@@ -504,6 +504,10 @@ namespace HiDreamO1 {
|
||||
vision_runner->set_flash_attention_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
vision_runner->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
vision_runner->set_weight_adapter(adapter);
|
||||
}
|
||||
|
||||
@@ -54,7 +54,7 @@ namespace Hunyuan {
|
||||
auto k = qkv_vec[1];
|
||||
auto v = qkv_vec[2];
|
||||
|
||||
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
|
||||
auto attn_out = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
|
||||
attn_out = self_attn_proj->forward(ctx, attn_out);
|
||||
|
||||
// adaLN_modulation
|
||||
|
||||
@@ -232,8 +232,7 @@ namespace Krea2 {
|
||||
q = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, q), head_dim_ * heads, Lq, N);
|
||||
k = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, k), head_dim_ * kv_heads, Lk, N);
|
||||
v = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, v), head_dim_ * kv_heads, Lk, N);
|
||||
return ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
return ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -709,8 +709,7 @@ namespace LTXV {
|
||||
k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved);
|
||||
}
|
||||
|
||||
auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
auto out = ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -215,8 +215,7 @@ namespace MiniMaxH3 {
|
||||
q = attention_layout(ctx->ggml_ctx, q);
|
||||
k = attention_layout(ctx->ggml_ctx, k);
|
||||
}
|
||||
auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
auto out = ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -365,8 +365,8 @@ public:
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x) {
|
||||
auto qkv = pre_attention(ctx, x);
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention(ctx, x); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention(ctx, x); // [N, n_token, dim]
|
||||
return x;
|
||||
}
|
||||
};
|
||||
@@ -587,8 +587,8 @@ public:
|
||||
auto qkv2 = std::get<1>(qkv_intermediates);
|
||||
auto intermediates = std::get<2>(qkv_intermediates);
|
||||
|
||||
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn2_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn2_out = ggml_ext_attention_ext(ctx, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention_x(ctx,
|
||||
attn_out,
|
||||
attn2_out,
|
||||
@@ -604,7 +604,7 @@ public:
|
||||
auto qkv = qkv_intermediates.first;
|
||||
auto intermediates = qkv_intermediates.second;
|
||||
|
||||
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention(ctx,
|
||||
attn_out,
|
||||
intermediates[0],
|
||||
@@ -648,7 +648,7 @@ block_mixing(GGMLRunnerContext* ctx,
|
||||
qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
|
||||
}
|
||||
|
||||
auto attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size]
|
||||
auto attn = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size]
|
||||
|
||||
auto context_attn = ggml_view_3d(ctx->ggml_ctx,
|
||||
attn,
|
||||
@@ -680,7 +680,7 @@ block_mixing(GGMLRunnerContext* ctx,
|
||||
}
|
||||
|
||||
if (x_block->self_attn) {
|
||||
auto attn2 = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, hidden_size]
|
||||
auto attn2 = ggml_ext_attention_ext(ctx, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, hidden_size]
|
||||
|
||||
x = x_block->post_attention_x(ctx,
|
||||
x_attn,
|
||||
|
||||
@@ -69,6 +69,8 @@ struct AnimaDiffusionExtra {
|
||||
struct WanDiffusionExtra {
|
||||
const sd::Tensor<float>* vace_context = nullptr;
|
||||
float vace_strength = 1.f;
|
||||
// S2V audio, sd::Tensor layout: [dim, T_latent*4, layers].
|
||||
const sd::Tensor<float>* audio_embed = nullptr;
|
||||
};
|
||||
|
||||
struct HiDreamO1DiffusionExtra {
|
||||
@@ -114,6 +116,10 @@ struct MiniT2IDiffusionExtra {
|
||||
const sd::Tensor<float>* mask = nullptr;
|
||||
};
|
||||
|
||||
struct SenseNovaU1DiffusionExtra {
|
||||
const sd::Tensor<int32_t>* input_ids = nullptr;
|
||||
};
|
||||
|
||||
struct HunyuanVideoDiffusionExtra {
|
||||
const sd::Tensor<float>* guidance = nullptr;
|
||||
const sd::Tensor<float>* byt5 = nullptr;
|
||||
@@ -131,6 +137,7 @@ using DiffusionExtraParams = std::variant<std::monostate,
|
||||
LTXAVDiffusionExtra,
|
||||
MiniMaxH3DiffusionExtra,
|
||||
MiniT2IDiffusionExtra,
|
||||
SenseNovaU1DiffusionExtra,
|
||||
HunyuanVideoDiffusionExtra>;
|
||||
|
||||
struct DiffusionParams {
|
||||
|
||||
@@ -0,0 +1,846 @@
|
||||
#ifndef __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__
|
||||
#define __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "core/ggml_extend.h"
|
||||
#include "model/diffusion/dit.hpp"
|
||||
#include "model/diffusion/model.hpp"
|
||||
#include "model/te/llm.hpp"
|
||||
#include "model_loader.h"
|
||||
|
||||
namespace SenseNovaU1 {
|
||||
constexpr int SENSENOVA_U1_GRAPH_SIZE = 327680;
|
||||
|
||||
struct SenseNovaU1Config {
|
||||
int64_t hidden_size = 4096;
|
||||
int64_t intermediate_size = 12288;
|
||||
int64_t num_layers = 42;
|
||||
int64_t num_heads = 32;
|
||||
int64_t num_kv_heads = 8;
|
||||
int64_t head_dim = 128;
|
||||
int64_t vocab_size = 151936;
|
||||
int64_t max_position_embeddings = 262144;
|
||||
int64_t max_position_embeddings_hw = 10000;
|
||||
int64_t vision_hidden_size = 1024;
|
||||
int64_t patch_size = 16;
|
||||
int64_t vision_downsample_factor = 2;
|
||||
int64_t in_channels = 3;
|
||||
int64_t timestep_embedding_size = 256;
|
||||
float rms_norm_eps = 1e-6f;
|
||||
float rope_theta = 5000000.f;
|
||||
float rope_theta_hw = 10000.f;
|
||||
float noise_scale_base_image_seq_len = 64.f;
|
||||
float noise_scale_max_value = 16.f;
|
||||
float t_eps = 0.02f;
|
||||
bool add_noise_scale_embedding = true;
|
||||
|
||||
int64_t image_token_stride() const {
|
||||
return patch_size * vision_downsample_factor;
|
||||
}
|
||||
|
||||
static SenseNovaU1Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& prefix) {
|
||||
SenseNovaU1Config config;
|
||||
config.num_layers = 0;
|
||||
const std::string root = prefix.empty() ? "" : prefix + ".";
|
||||
|
||||
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!starts_with(name, root)) {
|
||||
continue;
|
||||
}
|
||||
if (ends_with(name, "language_model.model.embed_tokens.weight") && tensor_storage.n_dims == 2) {
|
||||
config.hidden_size = tensor_storage.ne[0];
|
||||
config.vocab_size = tensor_storage.ne[1];
|
||||
} else if (ends_with(name, "language_model.model.layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
|
||||
config.intermediate_size = tensor_storage.ne[1];
|
||||
} else if (ends_with(name, "language_model.model.layers.0.self_attn.q_proj.weight") && tensor_storage.n_dims == 2) {
|
||||
config.num_heads = tensor_storage.ne[1] / config.head_dim;
|
||||
} else if (ends_with(name, "language_model.model.layers.0.self_attn.k_proj.weight") && tensor_storage.n_dims == 2) {
|
||||
config.num_kv_heads = tensor_storage.ne[1] / config.head_dim;
|
||||
} else if (ends_with(name, "fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight") && tensor_storage.n_dims == 4) {
|
||||
config.patch_size = tensor_storage.ne[0];
|
||||
config.in_channels = tensor_storage.ne[2];
|
||||
config.vision_hidden_size = tensor_storage.ne[3];
|
||||
} else if (ends_with(name, "fm_modules.vision_model_mot_gen.embeddings.dense_embedding.weight") && tensor_storage.n_dims == 4) {
|
||||
config.vision_downsample_factor = tensor_storage.ne[0];
|
||||
}
|
||||
|
||||
const std::string layer_prefix = root + "language_model.model.layers.";
|
||||
if (starts_with(name, layer_prefix)) {
|
||||
const char* index_begin = name.c_str() + layer_prefix.size();
|
||||
config.num_layers = std::max<int64_t>(config.num_layers, std::strtoll(index_begin, nullptr, 10) + 1);
|
||||
}
|
||||
}
|
||||
|
||||
if (config.num_layers == 0) {
|
||||
config.num_layers = 42;
|
||||
}
|
||||
config.add_noise_scale_embedding = tensor_storage_map.find(root + "fm_modules.noise_scale_embedder.mlp.0.weight") != tensor_storage_map.end();
|
||||
|
||||
LOG_DEBUG("sensenova-u1.5: layers=%" PRId64 ", hidden=%" PRId64 ", intermediate=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", patch=%" PRId64 "x%" PRId64,
|
||||
config.num_layers,
|
||||
config.hidden_size,
|
||||
config.intermediate_size,
|
||||
config.num_heads,
|
||||
config.num_kv_heads,
|
||||
config.patch_size,
|
||||
config.vision_downsample_factor);
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
class StorageConv2d : public Conv2d {
|
||||
protected:
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
this->prefix = prefix;
|
||||
ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16);
|
||||
params["weight"] = ggml_new_tensor_4d(ctx,
|
||||
wtype,
|
||||
kernel_size.second,
|
||||
kernel_size.first,
|
||||
in_channels,
|
||||
out_channels);
|
||||
if (bias) {
|
||||
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
StorageConv2d(int64_t in_channels,
|
||||
int64_t out_channels,
|
||||
std::pair<int, int> kernel_size,
|
||||
std::pair<int, int> stride = {1, 1},
|
||||
std::pair<int, int> padding = {0, 0},
|
||||
bool bias = true)
|
||||
: Conv2d(in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
{1, 1},
|
||||
bias) {}
|
||||
};
|
||||
|
||||
struct TimestepEmbedder : public GGMLBlock {
|
||||
int64_t frequency_embedding_size;
|
||||
|
||||
TimestepEmbedder(int64_t hidden_size, int64_t frequency_embedding_size = 256)
|
||||
: frequency_embedding_size(frequency_embedding_size) {
|
||||
blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true);
|
||||
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
|
||||
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
|
||||
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
|
||||
auto x = ggml_ext_timestep_embedding(ctx->ggml_ctx,
|
||||
timesteps,
|
||||
static_cast<int>(frequency_embedding_size),
|
||||
10000.f,
|
||||
1.f);
|
||||
x = mlp_0->forward(ctx, x);
|
||||
x = ggml_silu_inplace(ctx->ggml_ctx, x);
|
||||
return mlp_2->forward(ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
inline ggml_tensor* apply_vision_rope(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* position_x,
|
||||
ggml_tensor* position_y,
|
||||
float theta,
|
||||
int max_position) {
|
||||
GGML_ASSERT(x->ne[0] % 2 == 0);
|
||||
// ggml_rope_ext addresses positions through ne[2]. The vision
|
||||
// embeddings arrive as [hidden, tokens, batch], so add the singleton
|
||||
// head axis used by the RoPE kernel: [hidden, 1, tokens, batch].
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
|
||||
const int64_t half = x->ne[0] / 2;
|
||||
auto x_part = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, half);
|
||||
auto y_part = ggml_ext_slice(ctx->ggml_ctx, x, 0, half, x->ne[0]);
|
||||
x_part = ggml_rope_ext(ctx->ggml_ctx,
|
||||
x_part,
|
||||
position_x,
|
||||
nullptr,
|
||||
static_cast<int>(half),
|
||||
GGML_ROPE_TYPE_NORMAL,
|
||||
max_position,
|
||||
theta,
|
||||
1.f,
|
||||
0.f,
|
||||
1.f,
|
||||
32.f,
|
||||
1.f);
|
||||
y_part = ggml_rope_ext(ctx->ggml_ctx,
|
||||
y_part,
|
||||
position_y,
|
||||
nullptr,
|
||||
static_cast<int>(half),
|
||||
GGML_ROPE_TYPE_NORMAL,
|
||||
max_position,
|
||||
theta,
|
||||
1.f,
|
||||
0.f,
|
||||
1.f,
|
||||
32.f,
|
||||
1.f);
|
||||
return ggml_concat(ctx->ggml_ctx, x_part, y_part, 0);
|
||||
}
|
||||
|
||||
struct VisionEmbeddings : public GGMLBlock {
|
||||
SenseNovaU1Config config;
|
||||
|
||||
explicit VisionEmbeddings(const SenseNovaU1Config& config)
|
||||
: config(config) {
|
||||
blocks["patch_embedding"] = std::make_shared<StorageConv2d>(config.in_channels,
|
||||
config.vision_hidden_size,
|
||||
std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)},
|
||||
std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)},
|
||||
std::pair<int, int>{0, 0},
|
||||
true);
|
||||
blocks["dense_embedding"] = std::make_shared<StorageConv2d>(config.vision_hidden_size,
|
||||
config.hidden_size,
|
||||
std::pair<int, int>{static_cast<int>(config.vision_downsample_factor), static_cast<int>(config.vision_downsample_factor)},
|
||||
std::pair<int, int>{static_cast<int>(config.vision_downsample_factor), static_cast<int>(config.vision_downsample_factor)},
|
||||
std::pair<int, int>{0, 0},
|
||||
true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* image,
|
||||
ggml_tensor* position_x,
|
||||
ggml_tensor* position_y) {
|
||||
auto patch_embedding = std::dynamic_pointer_cast<StorageConv2d>(blocks["patch_embedding"]);
|
||||
auto dense_embedding = std::dynamic_pointer_cast<StorageConv2d>(blocks["dense_embedding"]);
|
||||
|
||||
auto x = patch_embedding->forward(ctx, image);
|
||||
x = ggml_gelu_erf(ctx->ggml_ctx, x);
|
||||
|
||||
const int64_t grid_w = x->ne[0];
|
||||
const int64_t grid_h = x->ne[1];
|
||||
const int64_t batch = x->ne[3];
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, grid_w * grid_h, x->ne[2], batch);
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
|
||||
x = apply_vision_rope(ctx,
|
||||
x,
|
||||
position_x,
|
||||
position_y,
|
||||
config.rope_theta_hw,
|
||||
static_cast<int>(config.max_position_embeddings_hw));
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, config.vision_hidden_size, grid_w, grid_h, batch);
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
|
||||
x = dense_embedding->forward(ctx, x);
|
||||
|
||||
const int64_t token_w = x->ne[0];
|
||||
const int64_t token_h = x->ne[1];
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, token_w * token_h, x->ne[2], x->ne[3]);
|
||||
return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
|
||||
}
|
||||
};
|
||||
|
||||
inline ggml_tensor* pixel_shuffle(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
int upscale_factor) {
|
||||
GGML_ASSERT(upscale_factor > 0);
|
||||
const int64_t h = x->ne[1];
|
||||
const int64_t w = x->ne[0];
|
||||
GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
|
||||
x = ggml_ext_cont(ctx->ggml_ctx,
|
||||
ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);
|
||||
return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
|
||||
}
|
||||
|
||||
struct PixelDecoder : public GGMLBlock {
|
||||
explicit PixelDecoder(const SenseNovaU1Config& config) {
|
||||
blocks["conv1"] = std::make_shared<StorageConv2d>(config.hidden_size / 4,
|
||||
1024,
|
||||
std::pair<int, int>{3, 3},
|
||||
std::pair<int, int>{1, 1},
|
||||
std::pair<int, int>{1, 1},
|
||||
true);
|
||||
blocks["conv2"] = std::make_shared<StorageConv2d>(256,
|
||||
192,
|
||||
std::pair<int, int>{3, 3},
|
||||
std::pair<int, int>{1, 1},
|
||||
std::pair<int, int>{1, 1},
|
||||
true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto conv1 = std::dynamic_pointer_cast<StorageConv2d>(blocks["conv1"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<StorageConv2d>(blocks["conv2"]);
|
||||
x = pixel_shuffle(ctx, x, 2);
|
||||
x = conv1->forward(ctx, x);
|
||||
x = ggml_gelu_erf(ctx->ggml_ctx, x);
|
||||
x = pixel_shuffle(ctx, x, 2);
|
||||
x = conv2->forward(ctx, x);
|
||||
return pixel_shuffle(ctx, x, 8);
|
||||
}
|
||||
};
|
||||
|
||||
enum class Branch {
|
||||
UNDERSTANDING,
|
||||
GENERATION,
|
||||
};
|
||||
|
||||
struct Attention : public GGMLBlock {
|
||||
SenseNovaU1Config config;
|
||||
int layer_index;
|
||||
|
||||
Attention(const SenseNovaU1Config& config, int layer_index)
|
||||
: config(config), layer_index(layer_index) {
|
||||
blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, config.num_heads * config.head_dim, false);
|
||||
blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
|
||||
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
|
||||
blocks["o_proj"] = std::make_shared<Linear>(config.num_heads * config.head_dim, config.hidden_size, false);
|
||||
blocks["q_proj_mot_gen"] = std::make_shared<Linear>(config.hidden_size, config.num_heads * config.head_dim, false);
|
||||
blocks["k_proj_mot_gen"] = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
|
||||
blocks["v_proj_mot_gen"] = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
|
||||
blocks["o_proj_mot_gen"] = std::make_shared<Linear>(config.num_heads * config.head_dim, config.hidden_size, false);
|
||||
|
||||
const int64_t axis_dim = config.head_dim / 2;
|
||||
blocks["q_norm"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
blocks["k_norm"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
blocks["q_norm_hw"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
blocks["k_norm_hw"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
blocks["q_norm_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
blocks["k_norm_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
blocks["q_norm_hw_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
blocks["k_norm_hw_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
|
||||
}
|
||||
|
||||
ggml_tensor* apply_axis_rope(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* positions,
|
||||
int dimensions,
|
||||
float theta,
|
||||
int max_position) {
|
||||
return ggml_rope_ext(ctx->ggml_ctx,
|
||||
x,
|
||||
positions,
|
||||
nullptr,
|
||||
dimensions,
|
||||
GGML_ROPE_TYPE_NEOX,
|
||||
max_position,
|
||||
theta,
|
||||
1.f,
|
||||
0.f,
|
||||
1.f,
|
||||
32.f,
|
||||
1.f);
|
||||
}
|
||||
|
||||
ggml_tensor* normalize_and_rotate(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* position_t,
|
||||
ggml_tensor* position_h,
|
||||
ggml_tensor* position_w,
|
||||
const std::string& norm_name,
|
||||
const std::string& norm_hw_name) {
|
||||
const int64_t temporal_dim = config.head_dim / 2;
|
||||
const int64_t spatial_dim = config.head_dim - temporal_dim;
|
||||
const int64_t axis_dim = spatial_dim / 2;
|
||||
|
||||
auto temporal = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, temporal_dim);
|
||||
auto spatial = ggml_ext_slice(ctx->ggml_ctx, x, 0, temporal_dim, config.head_dim);
|
||||
temporal = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_name])->forward(ctx, temporal);
|
||||
spatial = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_hw_name])->forward(ctx, spatial);
|
||||
|
||||
auto height = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, 0, axis_dim);
|
||||
auto width = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, axis_dim, spatial_dim);
|
||||
temporal = apply_axis_rope(ctx,
|
||||
temporal,
|
||||
position_t,
|
||||
static_cast<int>(temporal_dim),
|
||||
config.rope_theta,
|
||||
static_cast<int>(config.max_position_embeddings));
|
||||
height = apply_axis_rope(ctx,
|
||||
height,
|
||||
position_h,
|
||||
static_cast<int>(axis_dim),
|
||||
config.rope_theta_hw,
|
||||
static_cast<int>(config.max_position_embeddings_hw));
|
||||
width = apply_axis_rope(ctx,
|
||||
width,
|
||||
position_w,
|
||||
static_cast<int>(axis_dim),
|
||||
config.rope_theta_hw,
|
||||
static_cast<int>(config.max_position_embeddings_hw));
|
||||
return ggml_concat(ctx->ggml_ctx,
|
||||
ggml_concat(ctx->ggml_ctx, temporal, height, 0),
|
||||
width,
|
||||
0);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* position_t,
|
||||
ggml_tensor* position_h,
|
||||
ggml_tensor* position_w,
|
||||
ggml_tensor* attention_mask,
|
||||
Branch branch,
|
||||
const std::string& cache_prefix) {
|
||||
const bool generation = branch == Branch::GENERATION;
|
||||
const std::string suffix = generation ? "_mot_gen" : "";
|
||||
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj" + suffix]);
|
||||
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj" + suffix]);
|
||||
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v_proj" + suffix]);
|
||||
auto o_proj = std::dynamic_pointer_cast<Linear>(blocks["o_proj" + suffix]);
|
||||
|
||||
const int64_t n_tokens = x->ne[1];
|
||||
const int64_t batch = x->ne[2];
|
||||
auto q = ggml_reshape_4d(ctx->ggml_ctx,
|
||||
q_proj->forward(ctx, x),
|
||||
config.head_dim,
|
||||
config.num_heads,
|
||||
n_tokens,
|
||||
batch);
|
||||
auto k = ggml_reshape_4d(ctx->ggml_ctx,
|
||||
k_proj->forward(ctx, x),
|
||||
config.head_dim,
|
||||
config.num_kv_heads,
|
||||
n_tokens,
|
||||
batch);
|
||||
auto v = ggml_reshape_4d(ctx->ggml_ctx,
|
||||
v_proj->forward(ctx, x),
|
||||
config.head_dim,
|
||||
config.num_kv_heads,
|
||||
n_tokens,
|
||||
batch);
|
||||
|
||||
q = normalize_and_rotate(ctx,
|
||||
q,
|
||||
position_t,
|
||||
position_h,
|
||||
position_w,
|
||||
"q_norm" + suffix,
|
||||
"q_norm_hw" + suffix);
|
||||
k = normalize_and_rotate(ctx,
|
||||
k,
|
||||
position_t,
|
||||
position_h,
|
||||
position_w,
|
||||
"k_norm" + suffix,
|
||||
"k_norm_hw" + suffix);
|
||||
|
||||
const std::string layer_cache = cache_prefix + "." + std::to_string(layer_index);
|
||||
if (generation) {
|
||||
auto prefix_k = ctx->load_cache_tensor(layer_cache + ".k");
|
||||
auto prefix_v = ctx->load_cache_tensor(layer_cache + ".v");
|
||||
GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
|
||||
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2);
|
||||
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
|
||||
} else {
|
||||
// Keep dedicated graph outputs alive until the runner copies them
|
||||
// into its persistent cache buffer after graph execution.
|
||||
auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k);
|
||||
cache_k = ggml_cpy(ctx->ggml_ctx, k, cache_k);
|
||||
ggml_set_output(cache_k);
|
||||
auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v);
|
||||
cache_v = ggml_cpy(ctx->ggml_ctx, v, cache_v);
|
||||
ggml_set_output(cache_v);
|
||||
ctx->persist_cache_tensor(layer_cache + ".k", cache_k);
|
||||
ctx->persist_cache_tensor(layer_cache + ".v", cache_v);
|
||||
}
|
||||
|
||||
q = ggml_cont(ctx->ggml_ctx,
|
||||
ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));
|
||||
q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
|
||||
k = ggml_cont(ctx->ggml_ctx,
|
||||
ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));
|
||||
k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
|
||||
|
||||
auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
config.num_heads,
|
||||
attention_mask,
|
||||
true,
|
||||
ctx->flash_attn_enabled);
|
||||
return o_proj->forward(ctx, out);
|
||||
}
|
||||
};
|
||||
|
||||
struct TransformerBlock : public GGMLBlock {
|
||||
TransformerBlock(const SenseNovaU1Config& config, int layer_index) {
|
||||
blocks["self_attn"] = std::make_shared<Attention>(config, layer_index);
|
||||
blocks["mlp"] = std::make_shared<LLM::MLP>(config.hidden_size, config.intermediate_size, false);
|
||||
blocks["mlp_mot_gen"] = std::make_shared<LLM::MLP>(config.hidden_size, config.intermediate_size, false);
|
||||
blocks["input_layernorm"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
|
||||
blocks["input_layernorm_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
|
||||
blocks["post_attention_layernorm"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
|
||||
blocks["post_attention_layernorm_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* position_t,
|
||||
ggml_tensor* position_h,
|
||||
ggml_tensor* position_w,
|
||||
ggml_tensor* attention_mask,
|
||||
Branch branch,
|
||||
const std::string& cache_prefix) {
|
||||
const bool generation = branch == Branch::GENERATION;
|
||||
auto input_norm = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(
|
||||
blocks[generation ? "input_layernorm_mot_gen" : "input_layernorm"]);
|
||||
auto post_norm = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(
|
||||
blocks[generation ? "post_attention_layernorm_mot_gen" : "post_attention_layernorm"]);
|
||||
auto attention = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
|
||||
auto mlp = std::dynamic_pointer_cast<LLM::MLP>(blocks[generation ? "mlp_mot_gen" : "mlp"]);
|
||||
|
||||
auto residual = x;
|
||||
x = input_norm->forward(ctx, x);
|
||||
x = attention->forward(ctx,
|
||||
x,
|
||||
position_t,
|
||||
position_h,
|
||||
position_w,
|
||||
attention_mask,
|
||||
branch,
|
||||
cache_prefix);
|
||||
x = ggml_add_inplace(ctx->ggml_ctx, x, residual);
|
||||
|
||||
residual = x;
|
||||
x = post_norm->forward(ctx, x);
|
||||
x = mlp->forward(ctx, x);
|
||||
return ggml_add_inplace(ctx->ggml_ctx, x, residual);
|
||||
}
|
||||
};
|
||||
|
||||
struct TextModel : public GGMLBlock {
|
||||
SenseNovaU1Config config;
|
||||
|
||||
explicit TextModel(const SenseNovaU1Config& config)
|
||||
: config(config) {
|
||||
blocks["embed_tokens"] = std::make_shared<Embedding>(config.vocab_size, config.hidden_size);
|
||||
for (int i = 0; i < config.num_layers; ++i) {
|
||||
blocks["layers." + std::to_string(i)] = std::make_shared<TransformerBlock>(config, i);
|
||||
}
|
||||
blocks["norm"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
|
||||
blocks["norm_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
|
||||
}
|
||||
|
||||
ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) {
|
||||
return std::dynamic_pointer_cast<Embedding>(blocks["embed_tokens"])->forward(ctx, input_ids);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* position_t,
|
||||
ggml_tensor* position_h,
|
||||
ggml_tensor* position_w,
|
||||
ggml_tensor* attention_mask,
|
||||
Branch branch,
|
||||
const std::string& cache_prefix) {
|
||||
for (int i = 0; i < config.num_layers; ++i) {
|
||||
auto layer = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
|
||||
x = layer->forward(ctx,
|
||||
x,
|
||||
position_t,
|
||||
position_h,
|
||||
position_w,
|
||||
attention_mask,
|
||||
branch,
|
||||
cache_prefix);
|
||||
}
|
||||
auto norm = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(
|
||||
blocks[branch == Branch::GENERATION ? "norm_mot_gen" : "norm"]);
|
||||
return norm->forward(ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
struct SenseNovaU1Model : public GGMLBlock {
|
||||
SenseNovaU1Config config;
|
||||
|
||||
explicit SenseNovaU1Model(const SenseNovaU1Config& config)
|
||||
: config(config) {
|
||||
blocks["language_model.model"] = std::make_shared<TextModel>(config);
|
||||
blocks["fm_modules.vision_model_mot_gen.embeddings"] = std::make_shared<VisionEmbeddings>(config);
|
||||
blocks["fm_modules.timestep_embedder"] = std::make_shared<TimestepEmbedder>(config.hidden_size,
|
||||
config.timestep_embedding_size);
|
||||
if (config.add_noise_scale_embedding) {
|
||||
blocks["fm_modules.noise_scale_embedder"] = std::make_shared<TimestepEmbedder>(config.hidden_size,
|
||||
config.timestep_embedding_size);
|
||||
}
|
||||
blocks["fm_modules.fm_head"] = std::make_shared<PixelDecoder>(config);
|
||||
}
|
||||
|
||||
std::shared_ptr<TextModel> text_model() {
|
||||
return std::dynamic_pointer_cast<TextModel>(blocks["language_model.model"]);
|
||||
}
|
||||
|
||||
std::shared_ptr<VisionEmbeddings> vision_embeddings() {
|
||||
return std::dynamic_pointer_cast<VisionEmbeddings>(blocks["fm_modules.vision_model_mot_gen.embeddings"]);
|
||||
}
|
||||
|
||||
std::shared_ptr<TimestepEmbedder> timestep_embedder() {
|
||||
return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["fm_modules.timestep_embedder"]);
|
||||
}
|
||||
|
||||
std::shared_ptr<TimestepEmbedder> noise_scale_embedder() {
|
||||
if (!config.add_noise_scale_embedding) {
|
||||
return nullptr;
|
||||
}
|
||||
return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["fm_modules.noise_scale_embedder"]);
|
||||
}
|
||||
|
||||
std::shared_ptr<PixelDecoder> pixel_decoder() {
|
||||
return std::dynamic_pointer_cast<PixelDecoder>(blocks["fm_modules.fm_head"]);
|
||||
}
|
||||
};
|
||||
|
||||
struct SenseNovaU1Runner : public DiffusionModelRunner {
|
||||
SenseNovaU1Config config;
|
||||
SenseNovaU1Model model;
|
||||
std::unordered_set<uint64_t> cached_prefix_hashes;
|
||||
std::vector<int32_t> position_t_vec;
|
||||
std::vector<int32_t> position_h_vec;
|
||||
std::vector<int32_t> position_w_vec;
|
||||
std::vector<float> attention_mask_vec;
|
||||
std::vector<float> noise_scale_vec;
|
||||
|
||||
SenseNovaU1Runner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string& prefix = "",
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
: DiffusionModelRunner(backend, prefix, weight_manager),
|
||||
config(SenseNovaU1Config::detect_from_weights(tensor_storage_map, prefix)),
|
||||
model(config) {
|
||||
model.init(params_ctx, tensor_storage_map, prefix);
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "SenseNova U1.5";
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors,
|
||||
const std::string& prefix) override {
|
||||
model.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
static uint64_t hash_input_ids(const sd::Tensor<int32_t>& input_ids) {
|
||||
uint64_t hash = 1469598103934665603ULL;
|
||||
for (int32_t token : input_ids.values()) {
|
||||
uint32_t value = static_cast<uint32_t>(token);
|
||||
for (int byte = 0; byte < 4; ++byte) {
|
||||
hash ^= static_cast<uint8_t>(value & 0xffU);
|
||||
hash *= 1099511628211ULL;
|
||||
value >>= 8;
|
||||
}
|
||||
}
|
||||
hash ^= static_cast<uint64_t>(input_ids.numel());
|
||||
hash *= 1099511628211ULL;
|
||||
return hash;
|
||||
}
|
||||
|
||||
static std::string cache_prefix(uint64_t hash) {
|
||||
return "snu15." + std::to_string(hash);
|
||||
}
|
||||
|
||||
ggml_tensor* make_position_tensor(const std::vector<int32_t>& values,
|
||||
const std::string& name) {
|
||||
auto tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, values.size());
|
||||
ggml_set_name(tensor, name.c_str());
|
||||
set_backend_tensor_data(tensor, values.data());
|
||||
return tensor;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_prefix_graph(const sd::Tensor<int32_t>& input_ids_tensor,
|
||||
const std::string& prefix_cache) {
|
||||
ggml_cgraph* graph = new_graph_custom(SENSENOVA_U1_GRAPH_SIZE);
|
||||
ggml_tensor* ids = make_input(input_ids_tensor);
|
||||
const int64_t length = input_ids_tensor.numel();
|
||||
|
||||
position_t_vec.resize(length);
|
||||
position_h_vec.assign(length, 0);
|
||||
position_w_vec.assign(length, 0);
|
||||
for (int64_t i = 0; i < length; ++i) {
|
||||
position_t_vec[i] = static_cast<int32_t>(i);
|
||||
}
|
||||
auto position_t = make_position_tensor(position_t_vec, "snu15.prefix.position_t");
|
||||
auto position_h = make_position_tensor(position_h_vec, "snu15.prefix.position_h");
|
||||
auto position_w = make_position_tensor(position_w_vec, "snu15.prefix.position_w");
|
||||
|
||||
attention_mask_vec.assign(static_cast<size_t>(length * length), 0.f);
|
||||
for (int64_t query = 0; query < length; ++query) {
|
||||
for (int64_t key = query + 1; key < length; ++key) {
|
||||
attention_mask_vec[static_cast<size_t>(query * length + key)] = -INFINITY;
|
||||
}
|
||||
}
|
||||
auto attention_mask = ggml_new_tensor_2d(compute_ctx,
|
||||
GGML_TYPE_F32,
|
||||
length,
|
||||
length);
|
||||
ggml_set_name(attention_mask, "snu15.prefix.attention_mask");
|
||||
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
auto text_model = model.text_model();
|
||||
auto hidden = text_model->embed(&runner_ctx, ids);
|
||||
hidden = text_model->forward(&runner_ctx,
|
||||
hidden,
|
||||
position_t,
|
||||
position_h,
|
||||
position_w,
|
||||
attention_mask,
|
||||
Branch::UNDERSTANDING,
|
||||
prefix_cache);
|
||||
ggml_build_forward_expand(graph, hidden);
|
||||
return graph;
|
||||
}
|
||||
|
||||
bool ensure_prefix_cache(int n_threads,
|
||||
const sd::Tensor<int32_t>& input_ids,
|
||||
std::string* prefix_cache) {
|
||||
const uint64_t hash = hash_input_ids(input_ids);
|
||||
*prefix_cache = cache_prefix(hash);
|
||||
if (cached_prefix_hashes.find(hash) != cached_prefix_hashes.end() &&
|
||||
get_cache_tensor_by_name(*prefix_cache + ".0.k") != nullptr) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (cached_prefix_hashes.size() >= 2) {
|
||||
free_cache_ctx_and_buffer();
|
||||
cached_prefix_hashes.clear();
|
||||
}
|
||||
auto get_graph = [&]() {
|
||||
return build_prefix_graph(input_ids, *prefix_cache);
|
||||
};
|
||||
auto result = GGMLRunner::compute(get_graph, n_threads, false, true);
|
||||
if (!result.has_value()) {
|
||||
LOG_ERROR("SenseNova U1.5 prefix cache computation failed");
|
||||
return false;
|
||||
}
|
||||
cached_prefix_hashes.insert(hash);
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
const sd::Tensor<float>& timestep_tensor,
|
||||
const std::string& prefix_cache,
|
||||
int64_t prefix_length) {
|
||||
ggml_cgraph* graph = new_graph_custom(SENSENOVA_U1_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* t = make_input(timestep_tensor);
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
GGML_ASSERT(x->ne[0] % config.image_token_stride() == 0);
|
||||
GGML_ASSERT(x->ne[1] % config.image_token_stride() == 0);
|
||||
|
||||
const int64_t grid_w = x->ne[0] / config.patch_size;
|
||||
const int64_t grid_h = x->ne[1] / config.patch_size;
|
||||
const int64_t token_w = grid_w / config.vision_downsample_factor;
|
||||
const int64_t token_h = grid_h / config.vision_downsample_factor;
|
||||
const int64_t tokens = token_w * token_h;
|
||||
|
||||
position_h_vec.resize(grid_w * grid_h);
|
||||
position_w_vec.resize(grid_w * grid_h);
|
||||
for (int64_t index = 0; index < grid_w * grid_h; ++index) {
|
||||
position_h_vec[index] = static_cast<int32_t>(index / grid_w);
|
||||
position_w_vec[index] = static_cast<int32_t>(index % grid_w);
|
||||
}
|
||||
auto vision_position_x = make_position_tensor(position_w_vec, "snu15.vision.position_x");
|
||||
auto vision_position_y = make_position_tensor(position_h_vec, "snu15.vision.position_y");
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
auto hidden = model.vision_embeddings()->forward(&runner_ctx,
|
||||
x,
|
||||
vision_position_x,
|
||||
vision_position_y);
|
||||
auto time_embedding = model.timestep_embedder()->forward(&runner_ctx, t);
|
||||
time_embedding = ggml_reshape_3d(compute_ctx, time_embedding, config.hidden_size, 1, 1);
|
||||
hidden = ggml_add(compute_ctx, hidden, time_embedding);
|
||||
|
||||
if (config.add_noise_scale_embedding) {
|
||||
const float image_tokens = static_cast<float>(tokens);
|
||||
const float noise_scale = std::min(config.noise_scale_max_value,
|
||||
std::sqrt(image_tokens / config.noise_scale_base_image_seq_len));
|
||||
noise_scale_vec = {noise_scale / config.noise_scale_max_value};
|
||||
auto noise_scale_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1);
|
||||
ggml_set_name(noise_scale_tensor, "snu15.noise_scale");
|
||||
set_backend_tensor_data(noise_scale_tensor, noise_scale_vec.data());
|
||||
auto noise_embedding = model.noise_scale_embedder()->forward(&runner_ctx, noise_scale_tensor);
|
||||
noise_embedding = ggml_reshape_3d(compute_ctx, noise_embedding, config.hidden_size, 1, 1);
|
||||
hidden = ggml_add(compute_ctx, hidden, noise_embedding);
|
||||
}
|
||||
|
||||
position_t_vec.assign(tokens, static_cast<int32_t>(prefix_length));
|
||||
position_h_vec.resize(tokens);
|
||||
position_w_vec.resize(tokens);
|
||||
for (int64_t index = 0; index < tokens; ++index) {
|
||||
position_h_vec[index] = static_cast<int32_t>(index / token_w);
|
||||
position_w_vec[index] = static_cast<int32_t>(index % token_w);
|
||||
}
|
||||
auto position_t = make_position_tensor(position_t_vec, "snu15.image.position_t");
|
||||
auto position_h = make_position_tensor(position_h_vec, "snu15.image.position_h");
|
||||
auto position_w = make_position_tensor(position_w_vec, "snu15.image.position_w");
|
||||
|
||||
hidden = model.text_model()->forward(&runner_ctx,
|
||||
hidden,
|
||||
position_t,
|
||||
position_h,
|
||||
position_w,
|
||||
nullptr,
|
||||
Branch::GENERATION,
|
||||
prefix_cache);
|
||||
hidden = ggml_reshape_4d(compute_ctx,
|
||||
hidden,
|
||||
config.hidden_size,
|
||||
token_w,
|
||||
token_h,
|
||||
x->ne[3]);
|
||||
hidden = ggml_cont(compute_ctx, ggml_permute(compute_ctx, hidden, 2, 0, 1, 3));
|
||||
auto x_prediction = model.pixel_decoder()->forward(&runner_ctx, hidden);
|
||||
|
||||
const float timestep = timestep_tensor.values()[0];
|
||||
const float denom = std::max(1.f - timestep, config.t_eps);
|
||||
auto velocity = ggml_scale(compute_ctx,
|
||||
ggml_sub(compute_ctx, x_prediction, x),
|
||||
1.f / denom);
|
||||
ggml_build_forward_expand(graph, velocity);
|
||||
return graph;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
const sd::Tensor<float>& timestep,
|
||||
const sd::Tensor<int32_t>& input_ids) {
|
||||
std::string prefix_cache;
|
||||
if (!ensure_prefix_cache(n_threads, input_ids, &prefix_cache)) {
|
||||
return {};
|
||||
}
|
||||
auto get_graph = [&]() {
|
||||
return build_graph(x, timestep, prefix_cache, input_ids.numel());
|
||||
};
|
||||
return restore_trailing_singleton_dims(
|
||||
GGMLRunner::compute(get_graph, n_threads, false),
|
||||
x.dim());
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const DiffusionParams& diffusion_params) override {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
const auto* extra = diffusion_extra_as<SenseNovaU1DiffusionExtra>(diffusion_params);
|
||||
GGML_ASSERT(extra->input_ids != nullptr);
|
||||
return compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
*extra->input_ids);
|
||||
}
|
||||
};
|
||||
} // namespace SenseNovaU1
|
||||
|
||||
#endif // __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__
|
||||
+161
-35
@@ -1,6 +1,7 @@
|
||||
#ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
|
||||
#define __SD_MODEL_DIFFUSION_WAN_HPP__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
@@ -33,11 +34,16 @@ namespace WAN {
|
||||
int vace_layers = 0;
|
||||
int64_t vace_in_dim = 96;
|
||||
std::map<int, int> vace_layers_mapping = {};
|
||||
bool qk_norm = true;
|
||||
bool cross_attn_norm = true;
|
||||
float eps = 1e-6f;
|
||||
int64_t flf_pos_embed_token_number = 0;
|
||||
int theta = 10000;
|
||||
int64_t audio_dim = 1024;
|
||||
int num_audio_token = 4; // excludes the learned padding token
|
||||
std::vector<int> audio_inject_layers = {};
|
||||
std::map<int, int> audio_inject_mapping = {}; // block index -> injector index
|
||||
std::string adain_mode = "attn_norm";
|
||||
bool qk_norm = true;
|
||||
bool cross_attn_norm = true;
|
||||
float eps = 1e-6f;
|
||||
int64_t flf_pos_embed_token_number = 0;
|
||||
int theta = 10000;
|
||||
// wan2.1 1.3B: 1536/12, wan2.1/2.2 14B: 5120/40, wan2.2 5B: 3074/24
|
||||
std::vector<int> axes_dim = {44, 42, 42};
|
||||
int64_t axes_dim_sum = 128;
|
||||
@@ -74,6 +80,10 @@ namespace WAN {
|
||||
if (name.find("img_emb") != std::string::npos) {
|
||||
config.model_type = "i2v";
|
||||
}
|
||||
if (name.find("audio_injector") != std::string::npos || name.find("casual_audio_encoder") != std::string::npos) {
|
||||
config.model_type = "s2v";
|
||||
config.audio_inject_layers = {0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39};
|
||||
}
|
||||
if (name.find("img_emb.emb_pos") != std::string::npos) {
|
||||
config.flf_pos_embed_token_number = 514;
|
||||
}
|
||||
@@ -193,7 +203,7 @@ namespace WAN {
|
||||
k = norm_k->forward(ctx, k);
|
||||
auto v = v_proj->forward(ctx, context); // [N, n_context, dim]
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
|
||||
x = o_proj->forward(ctx, x); // [N, n_token, dim]
|
||||
return x;
|
||||
@@ -255,8 +265,8 @@ namespace WAN {
|
||||
k_img = norm_k_img->forward(ctx, k_img);
|
||||
auto v_img = v_img_proj->forward(ctx, context_img); // [N, context_img_len, dim]
|
||||
|
||||
auto img_x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto img_x = ggml_ext_attention_ext(ctx, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
|
||||
x = ggml_add(ctx->ggml_ctx, x, img_x);
|
||||
|
||||
@@ -265,6 +275,13 @@ namespace WAN {
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace WAN
|
||||
|
||||
// Audio injection reuses WanT2VCrossAttention defined above.
|
||||
#include "model/diffusion/wan_audio.hpp"
|
||||
|
||||
namespace WAN {
|
||||
|
||||
static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) {
|
||||
// x: [N, n_token, dim]
|
||||
// e: [N, 1, dim] or [N, T, 1, dim]
|
||||
@@ -532,6 +549,13 @@ namespace WAN {
|
||||
protected:
|
||||
WanConfig config;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
if (config.model_type == "s2v") {
|
||||
enum ggml_type wtype = GGML_TYPE_F32; // elementwise add vs F32 activations
|
||||
params["trainable_cond_mask.weight"] = ggml_new_tensor_2d(ctx, wtype, config.dim, 3);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
Wan() {}
|
||||
Wan(WanConfig config)
|
||||
@@ -554,7 +578,7 @@ namespace WAN {
|
||||
|
||||
// blocks
|
||||
for (int i = 0; i < config.num_layers; i++) {
|
||||
auto block = std::shared_ptr<GGMLBlock>(new WanAttentionBlock(config.model_type == "t2v",
|
||||
auto block = std::shared_ptr<GGMLBlock>(new WanAttentionBlock(config.model_type != "i2v",
|
||||
config.dim,
|
||||
config.ffn_dim,
|
||||
config.num_heads,
|
||||
@@ -595,6 +619,14 @@ namespace WAN {
|
||||
|
||||
blocks["vace_patch_embedding"] = std::shared_ptr<GGMLBlock>(new Conv3d(config.vace_in_dim, config.dim, config.patch_size, config.patch_size));
|
||||
}
|
||||
|
||||
if (config.model_type == "s2v") {
|
||||
blocks["casual_audio_encoder"] = std::make_shared<WanCausalAudioEncoder>(config.audio_dim, config.dim, config.num_audio_token);
|
||||
blocks["audio_injector"] = std::make_shared<WanAudioInjector>(config.dim, config.num_heads, (int)config.audio_inject_layers.size(), config.qk_norm, config.eps);
|
||||
for (size_t i = 0; i < config.audio_inject_layers.size(); i++) {
|
||||
config.audio_inject_mapping[config.audio_inject_layers[i]] = (int)i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
|
||||
@@ -642,18 +674,24 @@ namespace WAN {
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* clip_fea = nullptr,
|
||||
ggml_tensor* vace_context = nullptr,
|
||||
float vace_strength = 1.f,
|
||||
int64_t N = 1) {
|
||||
ggml_tensor* clip_fea = nullptr,
|
||||
ggml_tensor* vace_context = nullptr,
|
||||
float vace_strength = 1.f,
|
||||
int64_t N = 1,
|
||||
ggml_tensor* audio_embed = nullptr,
|
||||
ggml_tensor* reference_latent = nullptr) {
|
||||
// x: [N*C, T, H, W], C => in_dim
|
||||
// vace_context: [N*vace_in_dim, T, H, W]
|
||||
// timestep: [N,] or [T]
|
||||
// context: [N, L, text_dim]
|
||||
// return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw]
|
||||
// audio_embed: [layers, T*4, audio_dim]
|
||||
// reference_latent: [N*C, T_ref, H, W]
|
||||
// return: [N, (t_len [+ t_ref_len]) * h_len*w_len, out_dim*pt*ph*pw]
|
||||
|
||||
GGML_ASSERT(N == 1);
|
||||
|
||||
int64_t T = x->ne[2];
|
||||
|
||||
auto patch_embedding = std::dynamic_pointer_cast<Conv3d>(blocks["patch_embedding"]);
|
||||
|
||||
auto text_embedding_0 = std::dynamic_pointer_cast<Linear>(blocks["text_embedding.0"]);
|
||||
@@ -670,6 +708,40 @@ namespace WAN {
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N); // [N, dim, t_len*h_len*w_len]
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim]
|
||||
|
||||
ggml_tensor* audio_local = nullptr;
|
||||
ggml_tensor* audio_global = nullptr;
|
||||
int64_t seq_len = x->ne[1];
|
||||
int64_t t_ref_len = 0;
|
||||
if (config.model_type == "s2v") {
|
||||
if (audio_embed != nullptr) {
|
||||
GGML_ASSERT(audio_embed->ne[1] == T * 4);
|
||||
auto audio_encoder = std::dynamic_pointer_cast<WanCausalAudioEncoder>(blocks["casual_audio_encoder"]);
|
||||
auto audio_emb = audio_encoder->forward(ctx, audio_embed);
|
||||
audio_local = audio_emb.first;
|
||||
audio_global = audio_emb.second;
|
||||
GGML_ASSERT(audio_local->ne[2] == T);
|
||||
}
|
||||
|
||||
// video tokens get cond_mask[0], reference tokens cond_mask[1]
|
||||
auto cond_mask = params["trainable_cond_mask.weight"];
|
||||
auto cm0 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 0, 1), config.dim, 1, 1);
|
||||
x = ggml_add(ctx->ggml_ctx, x, cm0);
|
||||
|
||||
if (reference_latent != nullptr) {
|
||||
t_ref_len = reference_latent->ne[2];
|
||||
auto ref = patch_embedding->forward(ctx, reference_latent);
|
||||
ref = ggml_reshape_3d(ctx->ggml_ctx, ref, ref->ne[0] * ref->ne[1] * ref->ne[2], ref->ne[3] / N, N);
|
||||
ref = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ref, 1, 0, 2, 3)); // [N, t_ref*h_len*w_len, dim]
|
||||
auto cm1 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 1, 2), config.dim, 1, 1);
|
||||
ref = ggml_add(ctx->ggml_ctx, ref, cm1);
|
||||
x = ggml_concat(ctx->ggml_ctx, x, ref, 1);
|
||||
|
||||
// Reference tokens use timestep 0.
|
||||
GGML_ASSERT(timestep->ne[0] == T);
|
||||
timestep = ggml_ext_pad(ctx->ggml_ctx, timestep, (int)t_ref_len, 0, 0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
// time_embedding
|
||||
auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim);
|
||||
e = time_embedding_0->forward(ctx, e);
|
||||
@@ -714,6 +786,11 @@ namespace WAN {
|
||||
|
||||
auto x_orig = x;
|
||||
|
||||
std::shared_ptr<WanAudioInjector> audio_injector;
|
||||
if (audio_local != nullptr) {
|
||||
audio_injector = std::dynamic_pointer_cast<WanAudioInjector>(blocks["audio_injector"]);
|
||||
}
|
||||
|
||||
for (int i = 0; i < config.num_layers; i++) {
|
||||
auto block = std::dynamic_pointer_cast<WanAttentionBlock>(blocks["blocks." + std::to_string(i)]);
|
||||
|
||||
@@ -731,6 +808,13 @@ namespace WAN {
|
||||
c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
|
||||
x = ggml_add(ctx->ggml_ctx, x, c_skip);
|
||||
}
|
||||
|
||||
if (audio_injector != nullptr) {
|
||||
auto inject_iter = config.audio_inject_mapping.find(i);
|
||||
if (inject_iter != config.audio_inject_mapping.end()) {
|
||||
x = audio_injector->forward(ctx, x, seq_len, T, inject_iter->second, audio_local, audio_global);
|
||||
}
|
||||
}
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x");
|
||||
if (c != nullptr) {
|
||||
sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c");
|
||||
@@ -747,11 +831,13 @@ namespace WAN {
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* clip_fea = nullptr,
|
||||
ggml_tensor* time_dim_concat = nullptr,
|
||||
ggml_tensor* vace_context = nullptr,
|
||||
float vace_strength = 1.f,
|
||||
int64_t N = 1) {
|
||||
ggml_tensor* clip_fea = nullptr,
|
||||
ggml_tensor* time_dim_concat = nullptr,
|
||||
ggml_tensor* vace_context = nullptr,
|
||||
float vace_strength = 1.f,
|
||||
int64_t N = 1,
|
||||
ggml_tensor* audio_embed = nullptr,
|
||||
ggml_tensor* reference_latent = nullptr) {
|
||||
// Forward pass of DiT.
|
||||
// x: [N*C, T, H, W]
|
||||
// timestep: [N,]
|
||||
@@ -779,7 +865,12 @@ namespace WAN {
|
||||
t_len = ((x->ne[2] + (std::get<0>(config.patch_size) / 2)) / std::get<0>(config.patch_size));
|
||||
}
|
||||
|
||||
auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N); // [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||
auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N, audio_embed, reference_latent); // [N, (t_len [+t_ref]) *h_len*w_len, pt*ph*pw*C]
|
||||
|
||||
if (reference_latent != nullptr) {
|
||||
// Exclude reference tokens from the generated video.
|
||||
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len);
|
||||
}
|
||||
|
||||
out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
|
||||
|
||||
@@ -839,7 +930,10 @@ namespace WAN {
|
||||
config.text_len = 512;
|
||||
}
|
||||
} else if (config.num_layers == 40) {
|
||||
if (config.model_type == "t2v") {
|
||||
if (version == VERSION_WAN2_2_S2V) {
|
||||
desc = "Wan2.2-S2V-14B";
|
||||
config.in_dim = 16;
|
||||
} else if (config.model_type == "t2v") {
|
||||
if (version == VERSION_WAN2_2_I2V) {
|
||||
desc = "Wan2.2-I2V-14B";
|
||||
config.in_dim = 36;
|
||||
@@ -891,7 +985,9 @@ namespace WAN {
|
||||
const sd::Tensor<float>& c_concat_tensor = {},
|
||||
const sd::Tensor<float>& time_dim_concat_tensor = {},
|
||||
const sd::Tensor<float>& vace_context_tensor = {},
|
||||
float vace_strength = 1.f) {
|
||||
float vace_strength = 1.f,
|
||||
const sd::Tensor<float>& audio_embed_tensor = {},
|
||||
const sd::Tensor<float>& ref_latent_tensor = {}) {
|
||||
ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE);
|
||||
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
@@ -901,16 +997,33 @@ namespace WAN {
|
||||
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
|
||||
ggml_tensor* time_dim_concat = make_optional_input(time_dim_concat_tensor);
|
||||
ggml_tensor* vace_context = make_optional_input(vace_context_tensor);
|
||||
ggml_tensor* audio_embed = make_optional_input(audio_embed_tensor);
|
||||
ggml_tensor* ref_latent = make_optional_input(ref_latent_tensor);
|
||||
|
||||
pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
|
||||
static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
std::get<0>(config.patch_size),
|
||||
std::get<1>(config.patch_size),
|
||||
std::get<2>(config.patch_size),
|
||||
1,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
|
||||
static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
std::get<0>(config.patch_size),
|
||||
std::get<1>(config.patch_size),
|
||||
std::get<2>(config.patch_size),
|
||||
1,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
if (ref_latent != nullptr) {
|
||||
// Match S2V's reference-frame temporal offset.
|
||||
int t_start = std::max(30, static_cast<int>(x->ne[2]) + 9);
|
||||
auto ref_pe = Rope::gen_wan_pe(static_cast<int>(ref_latent->ne[2]),
|
||||
static_cast<int>(ref_latent->ne[1]),
|
||||
static_cast<int>(ref_latent->ne[0]),
|
||||
std::get<0>(config.patch_size),
|
||||
std::get<1>(config.patch_size),
|
||||
std::get<2>(config.patch_size),
|
||||
1,
|
||||
config.theta,
|
||||
config.axes_dim,
|
||||
t_start);
|
||||
pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end());
|
||||
}
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
// LOG_VERBOSE("pos_len %d", pos_len);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
@@ -933,7 +1046,10 @@ namespace WAN {
|
||||
clip_fea,
|
||||
time_dim_concat,
|
||||
vace_context,
|
||||
vace_strength);
|
||||
vace_strength,
|
||||
1,
|
||||
audio_embed,
|
||||
ref_latent);
|
||||
|
||||
ggml_build_forward_expand(gf, out);
|
||||
|
||||
@@ -948,9 +1064,11 @@ namespace WAN {
|
||||
const sd::Tensor<float>& c_concat = {},
|
||||
const sd::Tensor<float>& time_dim_concat = {},
|
||||
const sd::Tensor<float>& vace_context = {},
|
||||
float vace_strength = 1.f) {
|
||||
float vace_strength = 1.f,
|
||||
const sd::Tensor<float>& audio_embed = {},
|
||||
const sd::Tensor<float>& ref_latent = {}) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
|
||||
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength, audio_embed, ref_latent);
|
||||
};
|
||||
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
|
||||
@@ -961,6 +1079,12 @@ namespace WAN {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
const auto* extra = diffusion_extra_as<WanDiffusionExtra>(diffusion_params);
|
||||
static const std::vector<sd::Tensor<float>> no_ref_latents;
|
||||
const auto& ref_latents = config.model_type == "s2v" && diffusion_params.ref_latents != nullptr
|
||||
? *diffusion_params.ref_latents
|
||||
: no_ref_latents;
|
||||
const sd::Tensor<float> empty_tensor;
|
||||
const sd::Tensor<float>& ref_latent = ref_latents.empty() ? empty_tensor : ref_latents[0];
|
||||
return compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
@@ -969,7 +1093,9 @@ namespace WAN {
|
||||
tensor_or_empty(diffusion_params.c_concat),
|
||||
sd::Tensor<float>(),
|
||||
tensor_or_empty(extra->vace_context),
|
||||
extra->vace_strength);
|
||||
extra->vace_strength,
|
||||
tensor_or_empty(extra->audio_embed),
|
||||
ref_latent);
|
||||
}
|
||||
|
||||
void test() {
|
||||
|
||||
@@ -0,0 +1,215 @@
|
||||
#ifndef __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__
|
||||
#define __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__
|
||||
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "model/common/ggml_block.hpp"
|
||||
|
||||
namespace WAN {
|
||||
|
||||
class WanCausalConv1d : public UnaryBlock {
|
||||
private:
|
||||
int kernel_size_;
|
||||
|
||||
public:
|
||||
WanCausalConv1d(int64_t in_dim,
|
||||
int64_t out_dim,
|
||||
int kernel_size = 3,
|
||||
int stride = 1)
|
||||
: kernel_size_(kernel_size) {
|
||||
blocks["conv"] = std::make_shared<Conv1d>(in_dim, out_dim, kernel_size, stride, 0, 1, 1, true, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
// Replicate the first sample for causal left padding.
|
||||
if (kernel_size_ > 1) {
|
||||
auto first = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, 1);
|
||||
for (int i = 0; i < kernel_size_ - 1; i++) {
|
||||
x = ggml_concat(ctx->ggml_ctx, first, x, 0);
|
||||
}
|
||||
}
|
||||
return std::dynamic_pointer_cast<Conv1d>(blocks["conv"])->forward(ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
class WanMotionEncoder : public GGMLBlock {
|
||||
private:
|
||||
int64_t hidden_dim_;
|
||||
int num_token_;
|
||||
bool need_global_;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
// The padding token is combined with F32 activations.
|
||||
params["padding_tokens"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_dim_);
|
||||
}
|
||||
|
||||
ggml_tensor* conv_norm_silu(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
const std::string& conv_key,
|
||||
const std::string& norm_key,
|
||||
bool to_conv_layout) {
|
||||
x = std::dynamic_pointer_cast<WanCausalConv1d>(blocks[conv_key])->forward(ctx, x);
|
||||
x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
|
||||
x = std::dynamic_pointer_cast<LayerNorm>(blocks[norm_key])->forward(ctx, x);
|
||||
x = ggml_silu(ctx->ggml_ctx, x);
|
||||
if (to_conv_layout) {
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
||||
public:
|
||||
WanMotionEncoder(int64_t in_dim,
|
||||
int64_t hidden_dim,
|
||||
int num_token,
|
||||
bool need_global = true)
|
||||
: hidden_dim_(hidden_dim), num_token_(num_token), need_global_(need_global) {
|
||||
blocks["conv1_local"] = std::make_shared<WanCausalConv1d>(in_dim, hidden_dim / 4 * num_token);
|
||||
if (need_global) {
|
||||
blocks["conv1_global"] = std::make_shared<WanCausalConv1d>(in_dim, hidden_dim / 4);
|
||||
}
|
||||
blocks["norm1"] = std::make_shared<LayerNorm>(hidden_dim / 4, 1e-6f, false);
|
||||
blocks["conv2"] = std::make_shared<WanCausalConv1d>(hidden_dim / 4, hidden_dim / 2, 3, 2);
|
||||
blocks["norm2"] = std::make_shared<LayerNorm>(hidden_dim / 2, 1e-6f, false);
|
||||
blocks["conv3"] = std::make_shared<WanCausalConv1d>(hidden_dim / 2, hidden_dim, 3, 2);
|
||||
blocks["norm3"] = std::make_shared<LayerNorm>(hidden_dim, 1e-6f, false);
|
||||
if (need_global) {
|
||||
blocks["final_linear"] = std::make_shared<Linear>(hidden_dim, hidden_dim);
|
||||
}
|
||||
}
|
||||
|
||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto local = std::dynamic_pointer_cast<WanCausalConv1d>(blocks["conv1_local"])->forward(ctx, x);
|
||||
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
|
||||
std::vector<ggml_tensor*> tokens;
|
||||
// Each token group is normalized independently over channels.
|
||||
for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token_, 1)) {
|
||||
ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3);
|
||||
s = norm1->forward(ctx, s);
|
||||
s = ggml_silu(ctx->ggml_ctx, s);
|
||||
s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3));
|
||||
s = conv_norm_silu(ctx, s, "conv2", "norm2", true);
|
||||
s = conv_norm_silu(ctx, s, "conv3", "norm3", false);
|
||||
tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1]));
|
||||
}
|
||||
auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim_, 1, 1);
|
||||
padding = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]);
|
||||
tokens.push_back(padding);
|
||||
ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1);
|
||||
|
||||
if (!need_global_) {
|
||||
return {local_out, nullptr};
|
||||
}
|
||||
ggml_tensor* g = conv_norm_silu(ctx, x, "conv1_global", "norm1", true);
|
||||
g = conv_norm_silu(ctx, g, "conv2", "norm2", true);
|
||||
g = conv_norm_silu(ctx, g, "conv3", "norm3", false);
|
||||
g = std::dynamic_pointer_cast<Linear>(blocks["final_linear"])->forward(ctx, g);
|
||||
return {local_out, g};
|
||||
}
|
||||
};
|
||||
|
||||
class WanCausalAudioEncoder : public GGMLBlock {
|
||||
private:
|
||||
int num_layers_;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
// Preserve the checkpoint shape for loading; layer mixing requires F32.
|
||||
auto it = tensor_storage_map.find(prefix + "weights");
|
||||
if (it != tensor_storage_map.end()) {
|
||||
params["weights"] = ggml_new_tensor(ctx, GGML_TYPE_F32, it->second.n_dims, it->second.ne);
|
||||
} else {
|
||||
params["weights"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_layers_);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
WanCausalAudioEncoder(int64_t audio_dim,
|
||||
int64_t dim,
|
||||
int num_token,
|
||||
int num_layers = 25)
|
||||
: num_layers_(num_layers) {
|
||||
blocks["encoder"] = std::make_shared<WanMotionEncoder>(audio_dim, dim, num_token, true);
|
||||
}
|
||||
|
||||
// features: [layers, frames, audio_dim]; outputs: [T, tokens+1, dim] and [T, dim].
|
||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* features) {
|
||||
auto weights = ggml_silu(ctx->ggml_ctx, params["weights"]);
|
||||
auto x = ggml_mul(ctx->ggml_ctx, features, ggml_reshape_3d(ctx->ggml_ctx, weights, 1, 1, num_layers_));
|
||||
x = ggml_div(ctx->ggml_ctx, x, ggml_sum(ctx->ggml_ctx, weights));
|
||||
// Move the layer axis to ggml dimension 0 for reduction.
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
|
||||
x = ggml_sum_rows(ctx->ggml_ctx, x);
|
||||
x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]);
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
|
||||
return std::dynamic_pointer_cast<WanMotionEncoder>(blocks["encoder"])->forward(ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
class WanAudioInjector : public GGMLBlock {
|
||||
private:
|
||||
int64_t dim_;
|
||||
|
||||
public:
|
||||
WanAudioInjector(int64_t dim,
|
||||
int64_t num_heads,
|
||||
int count,
|
||||
bool qk_norm = true,
|
||||
float eps = 1e-6f)
|
||||
: dim_(dim) {
|
||||
for (int i = 0; i < count; i++) {
|
||||
blocks["injector." + std::to_string(i)] =
|
||||
std::make_shared<WanT2VCrossAttention>(dim, num_heads, qk_norm, eps);
|
||||
blocks["injector_adain_layers." + std::to_string(i) + ".linear"] =
|
||||
std::make_shared<Linear>(dim, dim * 2);
|
||||
}
|
||||
// S2V AdaLayerNorm uses its own epsilon, independent of attention norms.
|
||||
blocks["adain_norm"] = std::make_shared<LayerNorm>(dim, 1e-5f, false);
|
||||
}
|
||||
|
||||
// Inject into the video prefix; trailing reference tokens pass through unchanged.
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
int64_t seq_len,
|
||||
int64_t T,
|
||||
int injector_id,
|
||||
ggml_tensor* audio_local,
|
||||
ggml_tensor* audio_global) {
|
||||
int64_t n_tok = seq_len / T;
|
||||
int64_t n_token = x->ne[1];
|
||||
|
||||
auto adain_linear = std::dynamic_pointer_cast<Linear>(blocks["injector_adain_layers." + std::to_string(injector_id) + ".linear"]);
|
||||
auto injector = std::dynamic_pointer_cast<WanT2VCrossAttention>(blocks["injector." + std::to_string(injector_id)]);
|
||||
auto adain_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["adain_norm"]);
|
||||
|
||||
auto temb = ggml_silu(ctx->ggml_ctx, audio_global);
|
||||
temb = adain_linear->forward(ctx, temb);
|
||||
auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim_);
|
||||
auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim_, dim_ * 2);
|
||||
shift = ggml_reshape_3d(ctx->ggml_ctx, shift, dim_, 1, T);
|
||||
scale = ggml_reshape_3d(ctx->ggml_ctx, scale, dim_, 1, T);
|
||||
|
||||
auto x_vid = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, seq_len);
|
||||
auto h = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim_, n_tok, T);
|
||||
h = adain_norm->forward(ctx, h);
|
||||
h = ggml_add(ctx->ggml_ctx, h, ggml_mul(ctx->ggml_ctx, h, scale));
|
||||
h = ggml_add(ctx->ggml_ctx, h, shift);
|
||||
|
||||
auto res = injector->forward(ctx, h, audio_local, 0);
|
||||
res = ggml_reshape_2d(ctx->ggml_ctx, res, dim_, seq_len);
|
||||
|
||||
auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res);
|
||||
if (seq_len < n_token) {
|
||||
auto x_tail = ggml_ext_slice(ctx->ggml_ctx, x, 1, seq_len, n_token);
|
||||
return ggml_concat(ctx->ggml_ctx, x_head, x_tail, 1);
|
||||
}
|
||||
return x_head;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace WAN
|
||||
|
||||
#endif // __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__
|
||||
@@ -145,6 +145,10 @@ protected:
|
||||
params["position_embedding.weight"] = ggml_new_tensor_2d(ctx, position_wtype, embed_dim, num_positions);
|
||||
}
|
||||
|
||||
enum ggml_op param_usage_op(const std::string& name) const override {
|
||||
return name == "token_embedding.weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
|
||||
}
|
||||
|
||||
public:
|
||||
CLIPEmbeddings(int64_t embed_dim,
|
||||
int64_t vocab_size = 49408,
|
||||
|
||||
+21
-9
@@ -139,7 +139,8 @@ namespace LLM {
|
||||
|
||||
static LLMConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& prefix,
|
||||
LLMArch arch) {
|
||||
LLMArch arch,
|
||||
bool& enable_vision) {
|
||||
LLMConfig config;
|
||||
config.arch = arch;
|
||||
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
|
||||
@@ -230,8 +231,9 @@ namespace LLM {
|
||||
config.num_experts_per_tok = 4;
|
||||
}
|
||||
|
||||
config.num_layers = 0;
|
||||
int detected_vision_layers = 0;
|
||||
config.num_layers = 0;
|
||||
int detected_vision_layers = 0;
|
||||
bool out_hidden_size_detected = false;
|
||||
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!starts_with(name, prefix)) {
|
||||
continue;
|
||||
@@ -277,6 +279,7 @@ namespace LLM {
|
||||
if (ends_with(name, "visual.merger.linear_fc2.weight") ||
|
||||
ends_with(name, "visual.merger.mlp.2.weight")) {
|
||||
config.vision.out_hidden_size = tensor_storage.ne[1];
|
||||
out_hidden_size_detected = true;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
@@ -330,6 +333,19 @@ namespace LLM {
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
config.intermediate_size);
|
||||
if (enable_vision && !config.have_vision_weight) {
|
||||
LOG_WARN("no vision weights detected, vision disabled");
|
||||
enable_vision = false;
|
||||
}
|
||||
// The default would reject valid models, so only compare a detected dim.
|
||||
if (enable_vision && out_hidden_size_detected &&
|
||||
config.vision.out_hidden_size != config.hidden_size) {
|
||||
LOG_ERROR("vision projector output size (%" PRId64 ") does not match LLM hidden size (%" PRId64 "), "
|
||||
"the vision weights (mmproj) likely belong to a different LLM variant, vision disabled",
|
||||
config.vision.out_hidden_size,
|
||||
config.hidden_size);
|
||||
enable_vision = false;
|
||||
}
|
||||
return config;
|
||||
}
|
||||
};
|
||||
@@ -1359,7 +1375,7 @@ namespace LLM {
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, kqv);
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
|
||||
} else {
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false); // [N, n_token, hidden_size]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, false); // [N, n_token, hidden_size]
|
||||
}
|
||||
|
||||
x = out_proj->forward(ctx, x); // [N, n_token, hidden_size]
|
||||
@@ -1886,12 +1902,8 @@ namespace LLM {
|
||||
bool enable_vision_ = false,
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
: GGMLRunner(backend, weight_manager),
|
||||
config(LLMConfig::detect_from_weights(tensor_storage_map, prefix, arch)),
|
||||
config(LLMConfig::detect_from_weights(tensor_storage_map, prefix, arch, enable_vision_)),
|
||||
enable_vision(enable_vision_) {
|
||||
if (enable_vision && !config.have_vision_weight) {
|
||||
LOG_WARN("no vision weights detected, vision disabled");
|
||||
enable_vision = false;
|
||||
}
|
||||
if (enable_vision) {
|
||||
LOG_VERBOSE("enable llm vision");
|
||||
if (config.llama_cpp_style) {
|
||||
|
||||
+1
-1
@@ -251,7 +251,7 @@ public:
|
||||
|
||||
k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask); // [N, n_token, d_head * n_head]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, n_token, d_head * n_head]
|
||||
|
||||
x = out_proj->forward(ctx, x); // [N, n_token, model_dim]
|
||||
return {x, past_bias};
|
||||
|
||||
@@ -142,7 +142,7 @@ public:
|
||||
v = ggml_reshape_3d(ctx->ggml_ctx, v, c, h * w, n); // [N, h * w, in_channels]
|
||||
}
|
||||
|
||||
h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
|
||||
h_ = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
|
||||
|
||||
if (use_linear) {
|
||||
h_ = proj_out->forward(ctx, h_); // [N, h * w, in_channels]
|
||||
|
||||
@@ -193,7 +193,7 @@ namespace Hunyuan {
|
||||
v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b); // [b, c, t*h*w]
|
||||
v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [b, t*h*w, c]
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c]
|
||||
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [b, c, t*h*w]
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b); // [b*c, t, h, w]
|
||||
|
||||
@@ -253,7 +253,7 @@ namespace MageVAE {
|
||||
q = to_patches(ctx->ggml_ctx, q);
|
||||
k = to_patches(ctx->ggml_ctx, k);
|
||||
v = to_patches(ctx->ggml_ctx, v);
|
||||
h = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
|
||||
h = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
|
||||
h = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
|
||||
if (pad_h > 0) {
|
||||
h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);
|
||||
|
||||
@@ -174,8 +174,7 @@ namespace MiniMaxH3 {
|
||||
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,
|
||||
auto attn_out = ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -291,8 +291,7 @@ namespace MiniMaxH3VAE {
|
||||
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,
|
||||
auto out = ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -166,7 +166,7 @@ public:
|
||||
scale_factor = 16;
|
||||
} else if (sd_version_uses_flux2_vae(version)) {
|
||||
scale_factor = 16;
|
||||
} else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
|
||||
} else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) {
|
||||
scale_factor = 1;
|
||||
}
|
||||
return scale_factor;
|
||||
|
||||
@@ -615,8 +615,8 @@ namespace WAN {
|
||||
auto v = qkv_vec[2];
|
||||
v = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n); // [t, c, h * w]
|
||||
|
||||
v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [t, h * w, c]
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [t, h * w, c]
|
||||
v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [t, h * w, c]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [t, h * w, c]
|
||||
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [t, c, h * w]
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n); // [t, c, h, w]
|
||||
|
||||
@@ -10,6 +10,7 @@ enum class ModelComponent {
|
||||
VAE,
|
||||
PreviewVAE,
|
||||
AudioVAE,
|
||||
AudioEncoder,
|
||||
ControlNet,
|
||||
PhotoMaker,
|
||||
PuLID,
|
||||
@@ -38,6 +39,8 @@ inline const char* model_component_name(ModelComponent component) {
|
||||
return "preview VAE";
|
||||
case ModelComponent::AudioVAE:
|
||||
return "audio VAE";
|
||||
case ModelComponent::AudioEncoder:
|
||||
return "audio encoder";
|
||||
case ModelComponent::ControlNet:
|
||||
return "ControlNet";
|
||||
case ModelComponent::PhotoMaker:
|
||||
|
||||
@@ -247,6 +247,11 @@ bool read_safetensors_file(const std::string& file_path,
|
||||
std::string dtype = tensor_info["dtype"];
|
||||
nlohmann::json shape = tensor_info["shape"];
|
||||
|
||||
// ComfyUI FP8 activation scales cancel when inference uses F16/F32 activations.
|
||||
if (ends_with(name, ".scale_input")) {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t begin = tensor_info["data_offsets"][0].get<size_t>();
|
||||
size_t end = tensor_info["data_offsets"][1].get<size_t>();
|
||||
if (begin > end || end > file_size_ - data_start) {
|
||||
|
||||
@@ -69,6 +69,8 @@ const char* unused_tensors[] = {
|
||||
// "v_pred", // Used to detect SDXL vpred models
|
||||
"text_encoders.llm.output.weight",
|
||||
"text_encoders.llm.lm_head.",
|
||||
"language_model.lm_head.",
|
||||
"vision_model.",
|
||||
};
|
||||
|
||||
bool is_unused_tensor(const std::string& name) {
|
||||
@@ -185,6 +187,15 @@ bool ModelLoader::parse_file(const std::string& file_path, const std::string& pr
|
||||
}
|
||||
parsed_dependencies_.push_back(stamp);
|
||||
if (is_directory(file_path)) {
|
||||
const std::string diffusers_index_path = path_join(file_path, "model_index.json");
|
||||
const std::string diffusers_unet_path = path_join(file_path, "unet/diffusion_pytorch_model.safetensors");
|
||||
const bool has_diffusers_layout = file_exists(diffusers_index_path) || file_exists(diffusers_unet_path);
|
||||
|
||||
const std::string safetensors_index_path = path_join(file_path, "model.safetensors.index.json");
|
||||
if (!has_diffusers_layout && file_exists(safetensors_index_path)) {
|
||||
LOG_INFO("load %s using root safetensors index", file_path.c_str());
|
||||
return parse_file(safetensors_index_path, prefix);
|
||||
}
|
||||
LOG_INFO("load %s using diffusers format", file_path.c_str());
|
||||
return init_from_diffusers_file(file_path, prefix);
|
||||
} else if (is_gguf_file(file_path)) {
|
||||
@@ -424,6 +435,7 @@ SDVersion ModelLoader::get_sd_version() const {
|
||||
bool is_flux2 = false;
|
||||
bool has_single_block_47 = false;
|
||||
bool is_wan = false;
|
||||
bool is_s2v = false;
|
||||
int64_t patch_embedding_channels = 0;
|
||||
bool has_img_emb = false;
|
||||
bool has_middle_block_1 = false;
|
||||
@@ -463,6 +475,9 @@ SDVersion ModelLoader::get_sd_version() const {
|
||||
if (tensor_storage.name.find("net.img_embedder.proj1.weight") != std::string::npos) {
|
||||
return VERSION_MINIT2I;
|
||||
}
|
||||
if (tensor_storage.name.find("language_model.model.layers.0.self_attn.q_proj_mot_gen.weight") != std::string::npos) {
|
||||
return VERSION_SENSENOVA_U1_5;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
|
||||
auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
|
||||
if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
|
||||
@@ -510,6 +525,11 @@ SDVersion ModelLoader::get_sd_version() const {
|
||||
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
|
||||
is_wan = true;
|
||||
}
|
||||
if (tensor_storage.name.find("casual_audio_encoder.weights") != std::string::npos ||
|
||||
tensor_storage.name.find("audio_injector.injector.0.q.weight") != std::string::npos) {
|
||||
// S2V and T2V-14B share patch_embedding shapes.
|
||||
is_s2v = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) {
|
||||
return VERSION_LINGBOT_VIDEO;
|
||||
}
|
||||
@@ -573,6 +593,9 @@ SDVersion ModelLoader::get_sd_version() const {
|
||||
}
|
||||
if (is_wan) {
|
||||
LOG_VERBOSE("patch_embedding_channels %d", patch_embedding_channels);
|
||||
if (is_s2v) {
|
||||
return VERSION_WAN2_2_S2V;
|
||||
}
|
||||
if (patch_embedding_channels == 184320 && !has_img_emb) {
|
||||
return VERSION_WAN2_2_I2V;
|
||||
}
|
||||
|
||||
+37
-18
@@ -1584,18 +1584,35 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
|
||||
if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) {
|
||||
return result;
|
||||
}
|
||||
auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
|
||||
const size_t missing = compute_backend_alloc_size(states, true);
|
||||
result.required_device_bytes = add(request.pending_allocation_bytes, missing);
|
||||
result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
|
||||
auto device = ggml_backend_get_device(request.compute_backend);
|
||||
if (device != nullptr) {
|
||||
auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
|
||||
const size_t missing = compute_backend_alloc_size(states, true);
|
||||
// Backend scratch buffers and pipelines are not included in graph measurements.
|
||||
constexpr size_t safety_margin = 512ULL * 1024ULL * 1024ULL;
|
||||
result.required_device_bytes = add(add(request.pending_allocation_bytes, missing), safety_margin);
|
||||
result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
|
||||
auto available_device_bytes = [&](ggml_backend_t backend) {
|
||||
auto device = ggml_backend_get_device(backend);
|
||||
if (device == nullptr) {
|
||||
return SIZE_MAX;
|
||||
}
|
||||
size_t free_bytes = 0, total_bytes = 0;
|
||||
ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
|
||||
if (free_bytes != 0 || total_bytes != 0) {
|
||||
result.available_device_bytes = free_bytes;
|
||||
if (free_bytes == 0 && total_bytes == 0) {
|
||||
return SIZE_MAX;
|
||||
}
|
||||
}
|
||||
// Vulkan's heap budget subtraction can underflow when usage exceeds the budget.
|
||||
if (total_bytes > 0 && free_bytes > total_bytes) {
|
||||
return size_t{0};
|
||||
}
|
||||
const size_t resident = add(compute_backend_resident_bytes(backend),
|
||||
add(other_runtime_resident_bytes(request.owner_id, backend),
|
||||
request.runtime_resident_bytes));
|
||||
if (total_bytes > 0) {
|
||||
free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0);
|
||||
}
|
||||
return free_bytes;
|
||||
};
|
||||
result.available_device_bytes = available_device_bytes(request.compute_backend);
|
||||
if (request.max_backend_bytes > 0) {
|
||||
const size_t resident = add(compute_backend_resident_bytes(request.compute_backend),
|
||||
other_runtime_resident_bytes(request.owner_id, request.compute_backend));
|
||||
@@ -1619,11 +1636,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
|
||||
// GGML exposes only a split buffer's total size, not per-device allocations.
|
||||
// Charge that upper bound on every participant instead of undercounting a shard.
|
||||
for (const auto& entry : split_devices) {
|
||||
size_t free_bytes = 0, total_bytes = 0;
|
||||
ggml_backend_dev_memory(ggml_backend_get_device(entry.first), &free_bytes, &total_bytes);
|
||||
if (free_bytes != 0 || total_bytes != 0) {
|
||||
result.available_device_bytes = std::min(result.available_device_bytes, free_bytes);
|
||||
}
|
||||
result.available_device_bytes = std::min(result.available_device_bytes, available_device_bytes(entry.first));
|
||||
if (entry.second > 0) {
|
||||
const size_t resident = add(compute_backend_resident_bytes(entry.first),
|
||||
other_runtime_resident_bytes(request.owner_id, entry.first));
|
||||
@@ -1739,12 +1752,18 @@ bool ModelManager::ensure_compute_backend_capacity(
|
||||
}
|
||||
}
|
||||
|
||||
const auto capacity = check_capacity(request, required_states);
|
||||
LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %.2f MB device / %.2f MB budget",
|
||||
const auto capacity = check_capacity(request, required_states);
|
||||
const std::string available_device = capacity.available_device_bytes == SIZE_MAX
|
||||
? "unknown"
|
||||
: sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));
|
||||
const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX
|
||||
? "unlimited"
|
||||
: sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0));
|
||||
LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
|
||||
ggml_backend_name(compute_backend),
|
||||
capacity.required_device_bytes / (1024.0 * 1024.0),
|
||||
capacity.required_budget_bytes / (1024.0 * 1024.0),
|
||||
capacity.available_device_bytes / (1024.0 * 1024.0),
|
||||
capacity.available_budget_bytes / (1024.0 * 1024.0));
|
||||
available_device.c_str(),
|
||||
available_budget.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
+1
-1
@@ -31,7 +31,7 @@ public:
|
||||
};
|
||||
|
||||
private:
|
||||
static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 64ULL * 1024ULL * 1024ULL;
|
||||
static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 1024ULL * 1024ULL * 1024ULL;
|
||||
|
||||
struct TensorState {
|
||||
std::string name;
|
||||
|
||||
@@ -33,12 +33,14 @@
|
||||
#include "extensions/generation_extension.h"
|
||||
#include "model/adapter/ip_adapter.hpp"
|
||||
#include "model/adapter/lora.hpp"
|
||||
#include "model/audio/wav2vec2.hpp"
|
||||
#include "model/diffusion/animatediff.hpp"
|
||||
#include "model/diffusion/control.hpp"
|
||||
#include "model/diffusion/model.hpp"
|
||||
#include "model/vae/audio_vae.hpp"
|
||||
#include "model/vae/ltx_vae.hpp"
|
||||
#include "model/vae/vae.hpp"
|
||||
#include "runtime/audio_processing.h"
|
||||
#include "runtime/denoiser.hpp"
|
||||
#include "runtime/guidance.h"
|
||||
#include "runtime/preview_interval.h"
|
||||
@@ -74,6 +76,7 @@ const char* model_version_to_str[] = {
|
||||
"Wan 2.x",
|
||||
"Wan 2.2 I2V",
|
||||
"Wan 2.2 TI2V",
|
||||
"Wan 2.2 S2V",
|
||||
"LingBot Video",
|
||||
"Qwen Image",
|
||||
"Qwen Image Layered",
|
||||
@@ -96,9 +99,13 @@ const char* model_version_to_str[] = {
|
||||
"SeFi-Image",
|
||||
"Krea2",
|
||||
"Mage Flow",
|
||||
"SenseNova U1.5",
|
||||
"ESRGAN",
|
||||
};
|
||||
|
||||
static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]),
|
||||
"\nnumber of elements in model_version_to_str[] != VERSION_COUNT");
|
||||
|
||||
void calculate_alphas_cumprod(float* alphas_cumprod,
|
||||
float linear_start = 0.00085f,
|
||||
float linear_end = 0.0120f,
|
||||
@@ -132,7 +139,7 @@ StableDiffusionGGML::~StableDiffusionGGML() = default;
|
||||
|
||||
const std::map<StableDiffusionGGML::RunnerGroup, std::set<ModelComponent>>& StableDiffusionGGML::runner_components() {
|
||||
static const std::map<RunnerGroup, std::set<ModelComponent>> components{
|
||||
{RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter}},
|
||||
{RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter, ModelComponent::AudioEncoder}},
|
||||
{RunnerGroup::VAE, {ModelComponent::VAE, ModelComponent::PreviewVAE, ModelComponent::AudioVAE}},
|
||||
{RunnerGroup::ControlNet, {ModelComponent::ControlNet}},
|
||||
{RunnerGroup::Extensions, {ModelComponent::PhotoMaker, ModelComponent::PuLID}},
|
||||
@@ -800,6 +807,13 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
|
||||
}
|
||||
}
|
||||
|
||||
if (strlen(SAFE_STR(sd_ctx_params->audio_encoder_path)) > 0) {
|
||||
LOG_INFO("loading audio encoder (wav2vec2) from '%s'", sd_ctx_params->audio_encoder_path);
|
||||
if (!model_loader.init_from_file(sd_ctx_params->audio_encoder_path, "wav2vec2.")) {
|
||||
LOG_WARN("loading audio encoder weights from '%s' failed", sd_ctx_params->audio_encoder_path);
|
||||
}
|
||||
}
|
||||
|
||||
if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) {
|
||||
LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path);
|
||||
if (!model_loader.init_from_file(sd_ctx_params->motion_module_path,
|
||||
@@ -843,6 +857,12 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
|
||||
}
|
||||
|
||||
bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
|
||||
for (float scale : {sd_ctx_params->linear_scale, sd_ctx_params->attn_scale}) {
|
||||
if (!std::isfinite(scale) || scale < 0.f || (scale > 0.f && !std::isfinite(1.f / scale))) {
|
||||
LOG_ERROR("scale overrides must be finite positive values, or 0 to keep model defaults");
|
||||
return false;
|
||||
}
|
||||
}
|
||||
auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
|
||||
n_threads = sd_ctx_params->n_threads;
|
||||
enable_mmap = sd_ctx_params->enable_mmap;
|
||||
@@ -852,7 +872,7 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
|
||||
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 && backend_spec.empty() && params_backend_spec.empty();
|
||||
auto_fit_enabled = sd_ctx_params->auto_fit && params_backend_spec.empty();
|
||||
max_vram_assignment.reset(0.f);
|
||||
{
|
||||
std::string error;
|
||||
@@ -1000,6 +1020,7 @@ bool StableDiffusionGGML::build_core_runners() {
|
||||
high_noise_diffusion_model = std::move(runners.high_noise_diffusion);
|
||||
clip_vision = std::move(runners.clip_vision);
|
||||
ip_adapter = std::move(runners.ip_adapter);
|
||||
audio_encoder = std::move(runners.audio_encoder);
|
||||
|
||||
cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE));
|
||||
diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
|
||||
@@ -1009,11 +1030,15 @@ bool StableDiffusionGGML::build_core_runners() {
|
||||
if (clip_vision) {
|
||||
clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION));
|
||||
}
|
||||
if (audio_encoder) {
|
||||
audio_encoder->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::AUDIO_ENCODER));
|
||||
}
|
||||
return register_runner_params(ModelComponent::Conditioner, cond_stage_model, SDBackendModule::TE) &&
|
||||
register_runner_params(ModelComponent::Diffusion, diffusion_model, SDBackendModule::DIFFUSION) &&
|
||||
register_runner_params(ModelComponent::HighNoiseDiffusion, high_noise_diffusion_model, SDBackendModule::DIFFUSION) &&
|
||||
register_runner_params(ModelComponent::CLIPVision, clip_vision, SDBackendModule::CLIP_VISION) &&
|
||||
register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION);
|
||||
register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION) &&
|
||||
register_runner_params(ModelComponent::AudioEncoder, audio_encoder, SDBackendModule::AUDIO_ENCODER);
|
||||
}
|
||||
|
||||
bool StableDiffusionGGML::build_vae_runners() {
|
||||
@@ -1111,6 +1136,12 @@ bool StableDiffusionGGML::validate_and_load_runners() {
|
||||
ignore_tensors.insert("model.diffusion_model.__32x32__");
|
||||
ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__");
|
||||
|
||||
if (audio_encoder != nullptr) {
|
||||
// These wav2vec2 tensors are unused during feature extraction.
|
||||
ignore_tensors.insert("wav2vec2.lm_head.");
|
||||
ignore_tensors.insert("wav2vec2.masked_spec_embed");
|
||||
}
|
||||
|
||||
if (audio_vae_model) {
|
||||
if (!sd_version_is_minimax_h3(version)) {
|
||||
ignore_tensors.insert("audio_vae.encoder");
|
||||
@@ -1306,6 +1337,9 @@ bool StableDiffusionGGML::build_denoiser() {
|
||||
pred_type = SEFI_FLOW_PRED;
|
||||
} else if (sd_version_is_minit2i(version)) {
|
||||
pred_type = MINIT2I_FLOW_PRED;
|
||||
} else if (sd_version_is_sensenova_u1(version)) {
|
||||
pred_type = SENSENOVA_U1_FLOW_PRED;
|
||||
default_flow_shift = 3.f;
|
||||
} else {
|
||||
pred_type = EPS_PRED;
|
||||
}
|
||||
@@ -1351,6 +1385,11 @@ bool StableDiffusionGGML::build_denoiser() {
|
||||
denoiser = std::make_shared<MiniT2IFlowDenoiser>();
|
||||
break;
|
||||
}
|
||||
case SENSENOVA_U1_FLOW_PRED: {
|
||||
LOG_INFO("running in SenseNova U1.5 FLOW mode");
|
||||
denoiser = std::make_shared<SenseNovaU1FlowDenoiser>(default_flow_shift);
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
LOG_ERROR("Unknown predition type %i", pred_type);
|
||||
return false;
|
||||
@@ -1737,6 +1776,29 @@ sd::Tensor<float> StableDiffusionGGML::get_clip_vision_output(const sd::Tensor<f
|
||||
return output;
|
||||
}
|
||||
|
||||
// Returns 50 Hz wav2vec2 states in sd::Tensor layout: [dim, frames, layers].
|
||||
sd::Tensor<float> StableDiffusionGGML::get_audio_embedding(const sd_audio_t& audio) {
|
||||
if (audio_encoder == nullptr) {
|
||||
LOG_ERROR("audio encoder model is not loaded");
|
||||
return {};
|
||||
}
|
||||
if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) {
|
||||
LOG_ERROR("invalid driving audio");
|
||||
return {};
|
||||
}
|
||||
auto mono = sd::audio::downmix_to_mono(audio.data, audio.sample_count, audio.channels);
|
||||
if (mono.empty()) {
|
||||
LOG_ERROR("audio mono downmix failed");
|
||||
return {};
|
||||
}
|
||||
mono = sd::audio::resample_audio(mono.data(), mono.size(), audio.sample_rate, 16000);
|
||||
if (mono.empty()) {
|
||||
LOG_ERROR("audio resample to 16 kHz failed");
|
||||
return {};
|
||||
}
|
||||
return audio_encoder->compute(n_threads, mono);
|
||||
}
|
||||
|
||||
void StableDiffusionGGML::compute_ip_adapter_tokens(const sd_image_t& image, float strength) {
|
||||
ip_adapter_tokens = {};
|
||||
ip_adapter_uncond_tokens = {};
|
||||
@@ -1792,6 +1854,10 @@ std::vector<float> StableDiffusionGGML::process_timesteps(const std::vector<floa
|
||||
}
|
||||
}
|
||||
return new_timesteps;
|
||||
}
|
||||
if (diffusion_model->get_desc() == "Wan2.2-S2V-14B") {
|
||||
int64_t frame_count = init_latent.shape()[2];
|
||||
return std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
|
||||
} else {
|
||||
return timesteps;
|
||||
}
|
||||
@@ -2300,7 +2366,8 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights};
|
||||
} else if (sd_version_is_wan(version)) {
|
||||
diffusion_params.extra = WanDiffusionExtra{vace_context.empty() ? nullptr : &vace_context,
|
||||
vace_strength};
|
||||
vace_strength,
|
||||
condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios[0]};
|
||||
} else if (sd_version_is_hunyuan_video(version)) {
|
||||
diffusion_params.extra = HunyuanVideoDiffusionExtra{
|
||||
&guidance_tensor,
|
||||
@@ -2333,6 +2400,9 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
} else if (sd_version_is_minit2i(version)) {
|
||||
diffusion_params.extra = MiniT2IDiffusionExtra{
|
||||
condition.c_vector.empty() ? nullptr : &condition.c_vector};
|
||||
} else if (sd_version_is_sensenova_u1(version)) {
|
||||
diffusion_params.extra = SenseNovaU1DiffusionExtra{
|
||||
condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids};
|
||||
} else {
|
||||
diffusion_params.extra = std::monostate{};
|
||||
}
|
||||
@@ -2489,7 +2559,9 @@ int StableDiffusionGGML::get_vae_scale_factor() {
|
||||
int StableDiffusionGGML::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) || sd_version_is_minimax_h3(version)) {
|
||||
if (sd_version_is_sensenova_u1(version)) {
|
||||
down_factor = 32;
|
||||
} else 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;
|
||||
@@ -2515,6 +2587,8 @@ int StableDiffusionGGML::get_latent_channel() {
|
||||
latent_channel = 3;
|
||||
} else if (sd_version_is_minit2i(version)) {
|
||||
latent_channel = 3;
|
||||
} else if (sd_version_is_sensenova_u1(version)) {
|
||||
latent_channel = 3;
|
||||
} else if (sd_version_is_pid(version)) {
|
||||
latent_channel = 3;
|
||||
} else if (sd_version_is_sefi_image(version)) {
|
||||
@@ -2621,7 +2695,6 @@ sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float
|
||||
auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
|
||||
const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode();
|
||||
while (decoded.empty() &&
|
||||
auto_fit_enabled &&
|
||||
sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) {
|
||||
decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
|
||||
}
|
||||
|
||||
@@ -27,6 +27,9 @@ struct LoraModel;
|
||||
struct ConditionerParams;
|
||||
struct SDCondition;
|
||||
struct RefImageParams;
|
||||
namespace Wav2Vec2 {
|
||||
class Wav2Vec2ModelRunner;
|
||||
}
|
||||
|
||||
extern const char* model_version_to_str[];
|
||||
|
||||
@@ -63,6 +66,7 @@ public:
|
||||
std::shared_ptr<VAE> first_stage_model;
|
||||
std::shared_ptr<VAE> preview_vae;
|
||||
std::shared_ptr<AudioVAERunner> audio_vae_model;
|
||||
std::shared_ptr<Wav2Vec2::Wav2Vec2ModelRunner> audio_encoder;
|
||||
std::shared_ptr<ControlNet> control_net;
|
||||
std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
|
||||
sd::Tensor<float> ip_adapter_tokens;
|
||||
@@ -363,6 +367,8 @@ public:
|
||||
int clip_skip = -1,
|
||||
bool zero_out_masked = false);
|
||||
|
||||
sd::Tensor<float> get_audio_embedding(const sd_audio_t& audio);
|
||||
|
||||
void compute_ip_adapter_tokens(const sd_image_t& image, float strength);
|
||||
|
||||
std::vector<float> process_timesteps(const std::vector<float>& timesteps,
|
||||
|
||||
@@ -28,6 +28,7 @@ namespace sd::pipeline {
|
||||
sd::Tensor<float> denoise_mask;
|
||||
sd::Tensor<float> clip_vision_output;
|
||||
sd::Tensor<float> vace_context;
|
||||
sd::Tensor<float> s2v_audio_embed;
|
||||
int64_t ref_image_num = 0;
|
||||
int64_t video_conditioning_frame_count = 0;
|
||||
int64_t video_target_frame_count = 0;
|
||||
@@ -59,7 +60,8 @@ namespace sd::pipeline {
|
||||
const sd_vid_gen_params_t* sd_vid_gen_params,
|
||||
sd_image_t** frames_out,
|
||||
int* num_frames_out,
|
||||
sd_audio_t** audio_out);
|
||||
sd_audio_t** audio_out,
|
||||
int* fps_out);
|
||||
|
||||
sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd,
|
||||
const char* model_path,
|
||||
|
||||
@@ -438,6 +438,10 @@ namespace sd::pipeline {
|
||||
condition_params.zero_out_masked = false;
|
||||
auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
|
||||
condition_params);
|
||||
if (cond.empty()) {
|
||||
LOG_ERROR("failed to encode prompt");
|
||||
return std::nullopt;
|
||||
}
|
||||
if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
|
||||
cond.c_concat = latents->concat_latent; // TODO: optimize
|
||||
}
|
||||
@@ -455,6 +459,9 @@ namespace sd::pipeline {
|
||||
// states with a zeroed prompt mask, so no extra text encode is needed.
|
||||
uncond.c_crossattn = cond.c_crossattn;
|
||||
uncond.c_vector = sd::Tensor<float>::zeros_like(cond.c_vector);
|
||||
} else if (sd_version_is_sensenova_u1(sd->version)) {
|
||||
auto* sensenova_conditioner = static_cast<SenseNovaU1Conditioner*>(sd->cond_stage_model.get());
|
||||
uncond = sensenova_conditioner->get_unconditional_condition(request->negative_prompt);
|
||||
} else {
|
||||
bool zero_out_masked = false;
|
||||
if (sd_version_is_sdxl(sd->version) &&
|
||||
@@ -466,6 +473,10 @@ namespace sd::pipeline {
|
||||
condition_params.zero_out_masked = zero_out_masked;
|
||||
uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
|
||||
condition_params);
|
||||
if (uncond.empty()) {
|
||||
LOG_ERROR("failed to encode negative prompt");
|
||||
return std::nullopt;
|
||||
}
|
||||
}
|
||||
if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
|
||||
uncond.c_concat = latents->concat_latent; // TODO: optimize
|
||||
@@ -491,6 +502,10 @@ namespace sd::pipeline {
|
||||
}
|
||||
img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
|
||||
condition_params);
|
||||
if (img_uncond.empty()) {
|
||||
LOG_ERROR("failed to encode image guidance prompt");
|
||||
return std::nullopt;
|
||||
}
|
||||
if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
|
||||
img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
#include "core/util.h"
|
||||
#include "extensions/generation_extension.h"
|
||||
#include "model/adapter/ip_adapter.hpp"
|
||||
#include "model/audio/wav2vec2.hpp"
|
||||
#include "model/diffusion/anima.hpp"
|
||||
#include "model/diffusion/boogu.hpp"
|
||||
#include "model/diffusion/control.hpp"
|
||||
@@ -27,6 +28,7 @@
|
||||
#include "model/diffusion/model.hpp"
|
||||
#include "model/diffusion/pid.hpp"
|
||||
#include "model/diffusion/qwen_image.hpp"
|
||||
#include "model/diffusion/sensenova_u1.h"
|
||||
#include "model/diffusion/unet.hpp"
|
||||
#include "model/diffusion/wan.hpp"
|
||||
#include "model/diffusion/z_image.hpp"
|
||||
@@ -233,6 +235,16 @@ namespace sd::model_builders {
|
||||
tensor_storage_map,
|
||||
weight_manager);
|
||||
}
|
||||
if (version == VERSION_WAN2_2_S2V &&
|
||||
tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) {
|
||||
if (!ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) {
|
||||
return false;
|
||||
}
|
||||
result.audio_encoder = std::make_shared<Wav2Vec2::Wav2Vec2ModelRunner>(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER),
|
||||
tensor_storage_map,
|
||||
"wav2vec2.",
|
||||
weight_manager);
|
||||
}
|
||||
} else if (sd_version_is_lingbot_video(version)) {
|
||||
bool enable_vision = false;
|
||||
for (const auto& [name, _] : tensor_storage_map) {
|
||||
@@ -306,6 +318,12 @@ namespace sd::model_builders {
|
||||
tensor_storage_map,
|
||||
"model.diffusion_model.model.net",
|
||||
weight_manager);
|
||||
} else if (sd_version_is_sensenova_u1(version)) {
|
||||
result.conditioner = std::make_shared<SenseNovaU1Conditioner>();
|
||||
result.diffusion = std::make_shared<SenseNovaU1::SenseNovaU1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
|
||||
tensor_storage_map,
|
||||
"",
|
||||
weight_manager);
|
||||
} else if (sd_version_is_anima(version)) {
|
||||
result.conditioner = std::make_shared<AnimaConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
|
||||
tensor_storage_map,
|
||||
@@ -396,6 +414,21 @@ namespace sd::model_builders {
|
||||
"ip_adapter",
|
||||
weight_manager);
|
||||
}
|
||||
if (result.conditioner) {
|
||||
result.conditioner->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
if (result.diffusion) {
|
||||
result.diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
if (result.high_noise_diffusion) {
|
||||
result.high_noise_diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
if (result.clip_vision) {
|
||||
result.clip_vision->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
if (result.ip_adapter) {
|
||||
result.ip_adapter->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
runners = std::move(result);
|
||||
return true;
|
||||
}
|
||||
@@ -493,7 +526,7 @@ namespace sd::model_builders {
|
||||
}
|
||||
};
|
||||
|
||||
if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
|
||||
if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) {
|
||||
LOG_INFO("using FakeVAE");
|
||||
result.vae = std::make_shared<FakeVAE>(version,
|
||||
ctx.backends.runtime_backend(SDBackendModule::VAE),
|
||||
@@ -531,6 +564,15 @@ namespace sd::model_builders {
|
||||
result.preview->set_conv2d_direct_enabled(true);
|
||||
}
|
||||
}
|
||||
if (result.vae) {
|
||||
result.vae->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
if (result.preview) {
|
||||
result.preview->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
if (result.audio) {
|
||||
result.audio->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
}
|
||||
runners = std::move(result);
|
||||
return true;
|
||||
}
|
||||
@@ -552,6 +594,7 @@ namespace sd::model_builders {
|
||||
LOG_INFO("Using Conv2d direct in the control net");
|
||||
control_net->set_conv2d_direct_enabled(true);
|
||||
}
|
||||
control_net->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
||||
runner = std::move(control_net);
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -15,6 +15,9 @@ struct DiffusionModelRunner;
|
||||
struct VAE;
|
||||
struct AudioVAERunner;
|
||||
struct ControlNet;
|
||||
namespace Wav2Vec2 {
|
||||
class Wav2Vec2ModelRunner;
|
||||
}
|
||||
struct GenerationExtension;
|
||||
struct GenerationExtensionInitContext;
|
||||
namespace IPAdapter {
|
||||
@@ -37,6 +40,7 @@ namespace sd::model_builders {
|
||||
std::shared_ptr<DiffusionModelRunner> high_noise_diffusion;
|
||||
std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;
|
||||
std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
|
||||
std::shared_ptr<Wav2Vec2::Wav2Vec2ModelRunner> audio_encoder;
|
||||
};
|
||||
|
||||
struct VAEOptions {
|
||||
|
||||
@@ -161,7 +161,10 @@ namespace sd::pipeline {
|
||||
frames = 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->version) && fps != 24) {
|
||||
if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->fps != 16) {
|
||||
LOG_WARN("Wan2.2 S2V uses 16 fps; overriding requested fps %d", sd_vid_gen_params->fps);
|
||||
fps = 16;
|
||||
} else if (sd_version_is_minimax_h3(sd->version) && fps != 24) {
|
||||
LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps);
|
||||
fps = 24;
|
||||
}
|
||||
|
||||
+149
-1
@@ -5,6 +5,7 @@
|
||||
#include <cstdlib>
|
||||
#include <optional>
|
||||
|
||||
#include "conditioning/wan_audio.h"
|
||||
#include "core/rng.hpp"
|
||||
#include "core/rng_philox.hpp"
|
||||
#include "diffusion_engine.h"
|
||||
@@ -420,6 +421,45 @@ namespace sd::pipeline {
|
||||
return audio;
|
||||
}
|
||||
|
||||
// Build the first 16 fps audio window, zero-padding past the track end.
|
||||
static sd::Tensor<float> build_s2v_audio_window(const sd::Tensor<float>& stacked, int64_t batch_frames) {
|
||||
const int64_t embed_dim = stacked.shape()[0];
|
||||
const int64_t in_frames = stacked.shape()[1];
|
||||
const int64_t num_layers = stacked.shape()[2];
|
||||
if (embed_dim <= 0 || in_frames <= 0 || num_layers <= 0 || batch_frames <= 0) {
|
||||
return {};
|
||||
}
|
||||
std::vector<float> layer_first(static_cast<size_t>(num_layers) * in_frames * embed_dim);
|
||||
for (int64_t l = 0; l < num_layers; ++l) {
|
||||
for (int64_t f = 0; f < in_frames; ++f) {
|
||||
const float* src = stacked.data() + l * embed_dim * in_frames + f * embed_dim;
|
||||
std::copy_n(src,
|
||||
static_cast<size_t>(embed_dim),
|
||||
layer_first.data() + (static_cast<size_t>(l) * in_frames + f) * embed_dim);
|
||||
}
|
||||
}
|
||||
sd::wan_audio::BucketPlan plan;
|
||||
std::vector<float> buckets = sd::wan_audio::build_audio_buckets(layer_first.data(),
|
||||
static_cast<int>(num_layers),
|
||||
static_cast<int>(in_frames),
|
||||
static_cast<int>(embed_dim),
|
||||
static_cast<int>(batch_frames),
|
||||
&plan);
|
||||
if (buckets.empty() || plan.bucket_frames < batch_frames) {
|
||||
return {};
|
||||
}
|
||||
// Reorder frame-major buckets into sd::Tensor's [dim, frame, layer] layout.
|
||||
sd::Tensor<float> window({embed_dim, batch_frames, num_layers});
|
||||
for (int64_t f = 0; f < batch_frames; ++f) {
|
||||
for (int64_t l = 0; l < num_layers; ++l) {
|
||||
const float* src = buckets.data() + (static_cast<size_t>(f) * num_layers + l) * embed_dim;
|
||||
float* dst = window.data() + l * embed_dim * batch_frames + f * embed_dim;
|
||||
std::copy_n(src, static_cast<size_t>(embed_dim), dst);
|
||||
}
|
||||
}
|
||||
return window;
|
||||
}
|
||||
|
||||
static std::optional<ImageGenerationLatents> prepare_video_generation_latents(StableDiffusionGGML* sd,
|
||||
const sd_vid_gen_params_t* sd_vid_gen_params,
|
||||
GenerationRequest* request) {
|
||||
@@ -1033,6 +1073,53 @@ namespace sd::pipeline {
|
||||
latents.vace_context = sd::ops::concat(vace_context, mask_context, 3); // [b, 2*c + vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor]
|
||||
int64_t t2 = ggml_time_ms();
|
||||
LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
|
||||
} else if (sd->diffusion_model->get_desc() == "Wan2.2-S2V-14B") {
|
||||
LOG_INFO("S2V");
|
||||
if (!end_image.empty()) {
|
||||
LOG_WARN("Wan2.2 S2V ignores end_image");
|
||||
}
|
||||
if (sd_vid_gen_params->ref_audios_count > 1) {
|
||||
LOG_ERROR("Wan2.2 S2V supports a single driving audio track");
|
||||
return std::nullopt;
|
||||
}
|
||||
int64_t t1 = ggml_time_ms();
|
||||
if (!start_image.empty()) {
|
||||
auto ref_img = start_image.reshape({start_image.shape()[0],
|
||||
start_image.shape()[1],
|
||||
1,
|
||||
start_image.shape()[2],
|
||||
1});
|
||||
auto encoded_ref = sd->encode_first_stage(ref_img);
|
||||
if (encoded_ref.empty()) {
|
||||
LOG_ERROR("failed to encode S2V reference image");
|
||||
return std::nullopt;
|
||||
}
|
||||
// Wan consumes reference latents in 4D.
|
||||
latents.ref_latents.push_back(encoded_ref.reshape({encoded_ref.shape()[0],
|
||||
encoded_ref.shape()[1],
|
||||
encoded_ref.shape()[2],
|
||||
encoded_ref.shape()[3]}));
|
||||
}
|
||||
if (sd_vid_gen_params->ref_audios_count == 1) {
|
||||
if (sd->audio_encoder == nullptr) {
|
||||
LOG_ERROR("S2V audio conditioning requires --audio-encoder (wav2vec2)");
|
||||
return std::nullopt;
|
||||
}
|
||||
auto stacked = sd->get_audio_embedding(sd_vid_gen_params->ref_audios[0]);
|
||||
if (stacked.empty()) {
|
||||
LOG_ERROR("failed to compute wav2vec2 embedding for driving audio");
|
||||
return std::nullopt;
|
||||
}
|
||||
int64_t latent_t = sd->video_frames_to_latent_frames(request->frames);
|
||||
int64_t batch_frames = latent_t * 4;
|
||||
latents.s2v_audio_embed = build_s2v_audio_window(stacked, batch_frames);
|
||||
if (latents.s2v_audio_embed.empty()) {
|
||||
LOG_ERROR("failed to build S2V audio window");
|
||||
return std::nullopt;
|
||||
}
|
||||
}
|
||||
int64_t t2 = ggml_time_ms();
|
||||
LOG_INFO("s2v conditioning prepared, taking %" PRId64 " ms", t2 - t1);
|
||||
}
|
||||
|
||||
if (latents.init_latent.empty()) {
|
||||
@@ -1084,6 +1171,12 @@ namespace sd::pipeline {
|
||||
latents.keyframe_indices);
|
||||
}
|
||||
}
|
||||
if (sd->version == VERSION_WAN2_2_S2V) {
|
||||
embeds.cond.c_ref_images = latents.ref_latents;
|
||||
if (!latents.s2v_audio_embed.empty()) {
|
||||
embeds.cond.c_ref_audios = {latents.s2v_audio_embed};
|
||||
}
|
||||
}
|
||||
if (request.use_uncond) {
|
||||
condition_params.text = request.negative_prompt;
|
||||
embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
|
||||
@@ -1096,6 +1189,12 @@ namespace sd::pipeline {
|
||||
embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks;
|
||||
embeds.uncond.c_position_ids = embeds.cond.c_position_ids;
|
||||
}
|
||||
if (sd->version == VERSION_WAN2_2_S2V) {
|
||||
embeds.uncond.c_ref_images = latents.ref_latents;
|
||||
if (!latents.s2v_audio_embed.empty()) {
|
||||
embeds.uncond.c_ref_audios = {sd::Tensor<float>::zeros_like(latents.s2v_audio_embed)};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int64_t t1 = ggml_time_ms();
|
||||
@@ -1422,10 +1521,14 @@ namespace sd::pipeline {
|
||||
const sd_vid_gen_params_t* sd_vid_gen_params,
|
||||
sd_image_t** frames_out,
|
||||
int* num_frames_out,
|
||||
sd_audio_t** audio_out) {
|
||||
sd_audio_t** audio_out,
|
||||
int* fps_out) {
|
||||
if (sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) {
|
||||
LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d",
|
||||
sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height);
|
||||
if (fps_out != nullptr) {
|
||||
*fps_out = std::max(1, sd_vid_gen_params->fps);
|
||||
}
|
||||
return generate_animatediff_video(sd, sd_vid_gen_params, frames_out, num_frames_out);
|
||||
}
|
||||
|
||||
@@ -1437,6 +1540,9 @@ namespace sd::pipeline {
|
||||
sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
|
||||
sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
|
||||
GenerationRequest request(sd, sd_vid_gen_params);
|
||||
if (fps_out != nullptr) {
|
||||
*fps_out = request.fps;
|
||||
}
|
||||
bool latent_upscale_enabled = request.hires.enabled;
|
||||
GenerationRequest hires_request = request;
|
||||
if (latent_upscale_enabled) {
|
||||
@@ -1725,6 +1831,33 @@ namespace sd::pipeline {
|
||||
LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000);
|
||||
|
||||
sd_audio_t* generated_audio = nullptr;
|
||||
if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->ref_audios_count > 0) {
|
||||
// Return the driving track for muxing with the generated video.
|
||||
const sd_audio_t& driving = sd_vid_gen_params->ref_audios[0];
|
||||
generated_audio = (sd_audio_t*)malloc(sizeof(sd_audio_t));
|
||||
if (generated_audio != nullptr) {
|
||||
generated_audio->sample_rate = driving.sample_rate;
|
||||
generated_audio->channels = driving.channels;
|
||||
generated_audio->sample_count = driving.sample_count;
|
||||
generated_audio->data = (float*)malloc(sizeof(float) * driving.sample_count * driving.channels);
|
||||
if (generated_audio->data == nullptr) {
|
||||
free(generated_audio);
|
||||
generated_audio = nullptr;
|
||||
} else {
|
||||
memcpy(generated_audio->data,
|
||||
driving.data,
|
||||
sizeof(float) * driving.sample_count * driving.channels);
|
||||
}
|
||||
}
|
||||
if (generated_audio != nullptr) {
|
||||
LOG_DEBUG("s2v output audio: %u Hz, %u channels, %llu samples",
|
||||
generated_audio->sample_rate,
|
||||
generated_audio->channels,
|
||||
(unsigned long long)generated_audio->sample_count);
|
||||
} else {
|
||||
LOG_DEBUG("s2v output audio copy failed (out of memory)");
|
||||
}
|
||||
}
|
||||
if ((sd_version_is_ltxav(sd->version) || sd_version_is_minimax_h3(sd->version)) &&
|
||||
latents.audio_length > 0 &&
|
||||
sd->audio_vae_model != nullptr) {
|
||||
@@ -1774,6 +1907,7 @@ namespace sd::pipeline {
|
||||
return false;
|
||||
}
|
||||
auto result = decode_video_outputs(sd, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out);
|
||||
LOG_DEBUG("decode_video_outputs returned %s", result == nullptr ? "nullptr (failed)" : "frames");
|
||||
if (result == nullptr) {
|
||||
free_sd_audio(generated_audio);
|
||||
return false;
|
||||
@@ -1786,6 +1920,20 @@ namespace sd::pipeline {
|
||||
if (frames_out != nullptr) {
|
||||
*frames_out = result;
|
||||
}
|
||||
if (sd->version == VERSION_WAN2_2_S2V && generated_audio != nullptr) {
|
||||
// Limit the driving track to the generated video's duration.
|
||||
int fps = request.fps;
|
||||
uint64_t video_frames = num_frames_out != nullptr ? (uint64_t)*num_frames_out : 0;
|
||||
uint64_t want_samples = (uint64_t)((double)video_frames / fps * generated_audio->sample_rate);
|
||||
LOG_DEBUG("s2v audio truncate: %llu samples -> %llu (video %llu frames @ %d fps)",
|
||||
(unsigned long long)generated_audio->sample_count,
|
||||
(unsigned long long)want_samples,
|
||||
(unsigned long long)video_frames,
|
||||
fps);
|
||||
if (want_samples > 0 && want_samples < generated_audio->sample_count) {
|
||||
generated_audio->sample_count = want_samples;
|
||||
}
|
||||
}
|
||||
if (audio_out != nullptr) {
|
||||
*audio_out = generated_audio;
|
||||
} else {
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
#include "audio_processing.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <numeric>
|
||||
|
||||
namespace sd::audio {
|
||||
|
||||
// Match torchaudio's Hann-windowed sinc resampler.
|
||||
std::vector<float> resample_audio(const float* samples,
|
||||
uint64_t sample_count,
|
||||
uint32_t orig_sample_rate,
|
||||
uint32_t target_sample_rate) {
|
||||
if (samples == nullptr || sample_count == 0 || orig_sample_rate == 0 || target_sample_rate == 0) {
|
||||
return {};
|
||||
}
|
||||
if (orig_sample_rate == target_sample_rate) {
|
||||
return std::vector<float>(samples, samples + sample_count);
|
||||
}
|
||||
|
||||
constexpr int kLowpassFilterWidth = 6;
|
||||
constexpr double kRolloff = 0.99;
|
||||
constexpr double kPi = 3.14159265358979323846;
|
||||
|
||||
const uint64_t gcd = std::gcd(static_cast<uint64_t>(orig_sample_rate),
|
||||
static_cast<uint64_t>(target_sample_rate));
|
||||
const int64_t orig_freq = static_cast<int64_t>(orig_sample_rate / gcd);
|
||||
const int64_t new_freq = static_cast<int64_t>(target_sample_rate / gcd);
|
||||
const double base_freq = static_cast<double>(std::min(orig_freq, new_freq)) * kRolloff;
|
||||
const int64_t width = static_cast<int64_t>(std::ceil(kLowpassFilterWidth * orig_freq / base_freq));
|
||||
const int64_t kernel_size = 2 * width + orig_freq;
|
||||
|
||||
std::vector<double> kernel(static_cast<size_t>(new_freq) * kernel_size);
|
||||
for (int64_t j = 0; j < new_freq; ++j) {
|
||||
for (int64_t i = 0; i < kernel_size; ++i) {
|
||||
double t = -static_cast<double>(j) / new_freq + static_cast<double>(i - width) / orig_freq;
|
||||
t *= base_freq;
|
||||
t = std::clamp(t, -static_cast<double>(kLowpassFilterWidth), static_cast<double>(kLowpassFilterWidth));
|
||||
const double cos_arg = std::cos(t * kPi / kLowpassFilterWidth / 2);
|
||||
const double window = cos_arg * cos_arg;
|
||||
double s = t * kPi;
|
||||
const double sinc = (s == 0.0) ? 1.0 : std::sin(s) / s;
|
||||
kernel[j * kernel_size + i] = sinc * window * (base_freq / orig_freq);
|
||||
}
|
||||
}
|
||||
|
||||
const uint64_t num_phases = static_cast<uint64_t>(sample_count / orig_freq) + 1;
|
||||
const uint64_t target_length = (static_cast<uint64_t>(new_freq) * sample_count +
|
||||
static_cast<uint64_t>(orig_freq) - 1) /
|
||||
static_cast<uint64_t>(orig_freq);
|
||||
std::vector<float> out(target_length);
|
||||
for (uint64_t phase = 0; phase < num_phases; ++phase) {
|
||||
const int64_t src_base = static_cast<int64_t>(phase * orig_freq) - width;
|
||||
for (int64_t j = 0; j < new_freq; ++j) {
|
||||
const uint64_t out_index = phase * new_freq + j;
|
||||
if (out_index >= target_length) {
|
||||
break;
|
||||
}
|
||||
const double* k = &kernel[j * kernel_size];
|
||||
double acc = 0.0;
|
||||
for (int64_t i = 0; i < kernel_size; ++i) {
|
||||
const int64_t src = src_base + i;
|
||||
if (src >= 0 && src < static_cast<int64_t>(sample_count)) {
|
||||
acc += samples[src] * k[i];
|
||||
}
|
||||
}
|
||||
out[out_index] = static_cast<float>(acc);
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<float> downmix_to_mono(const float* interleaved_samples,
|
||||
uint64_t sample_count,
|
||||
uint32_t channels) {
|
||||
std::vector<float> mono;
|
||||
if (interleaved_samples == nullptr || sample_count == 0 || channels == 0) {
|
||||
return mono;
|
||||
}
|
||||
mono.resize(static_cast<size_t>(sample_count));
|
||||
if (channels == 1) {
|
||||
std::memcpy(mono.data(), interleaved_samples, static_cast<size_t>(sample_count) * sizeof(float));
|
||||
return mono;
|
||||
}
|
||||
const float scale = 1.0f / static_cast<float>(channels);
|
||||
for (uint64_t i = 0; i < sample_count; ++i) {
|
||||
float sum = 0.0f;
|
||||
for (uint32_t c = 0; c < channels; ++c) {
|
||||
sum += interleaved_samples[i * channels + c];
|
||||
}
|
||||
mono[static_cast<size_t>(i)] = sum * scale;
|
||||
}
|
||||
return mono;
|
||||
}
|
||||
|
||||
} // namespace sd::audio
|
||||
@@ -0,0 +1,22 @@
|
||||
#ifndef __SD_RUNTIME_AUDIO_PROCESSING_H__
|
||||
#define __SD_RUNTIME_AUDIO_PROCESSING_H__
|
||||
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
namespace sd::audio {
|
||||
|
||||
// Returns the input unchanged when sample rates are equal, and an empty vector on invalid input.
|
||||
std::vector<float> resample_audio(const float* samples,
|
||||
uint64_t sample_count,
|
||||
uint32_t orig_sample_rate,
|
||||
uint32_t target_sample_rate);
|
||||
|
||||
// Average interleaved channels; return an empty vector on invalid input.
|
||||
std::vector<float> downmix_to_mono(const float* interleaved_samples,
|
||||
uint64_t sample_count,
|
||||
uint32_t channels);
|
||||
|
||||
} // namespace sd::audio
|
||||
|
||||
#endif // __SD_RUNTIME_AUDIO_PROCESSING_H__
|
||||
@@ -1488,6 +1488,82 @@ struct MiniT2IFlowDenoiser : public Denoiser {
|
||||
}
|
||||
};
|
||||
|
||||
// SenseNova U1.5 integrates velocity over t=0..1 while the generic sampler
|
||||
// integrates over descending sigma. With sigma=1-t, returning
|
||||
// denoised=x+sigma*v makes the generic Euler derivative exactly -v, so the
|
||||
// descending-sigma update is identical to the official ascending-time update.
|
||||
struct SenseNovaU1FlowDenoiser : public DiscreteFlowDenoiser {
|
||||
explicit SenseNovaU1FlowDenoiser(float shift = 3.f)
|
||||
: DiscreteFlowDenoiser(shift) {}
|
||||
|
||||
float sigma_min() override {
|
||||
return 0.f;
|
||||
}
|
||||
|
||||
float sigma_max() override {
|
||||
return 1.f;
|
||||
}
|
||||
|
||||
float sigma_to_t(float sigma) override {
|
||||
return 1.f - sigma;
|
||||
}
|
||||
|
||||
float t_to_sigma(float t) override {
|
||||
float sigma = 1.f - t;
|
||||
return shift * sigma / (1.f + (shift - 1.f) * sigma);
|
||||
}
|
||||
|
||||
std::vector<float> get_scalings(float sigma) override {
|
||||
return {1.f, sigma, 1.f};
|
||||
}
|
||||
|
||||
sd::Tensor<float> noise_scaling(float sigma,
|
||||
const sd::Tensor<float>& noise,
|
||||
const sd::Tensor<float>& latent) override {
|
||||
SD_UNUSED(sigma);
|
||||
SD_UNUSED(latent);
|
||||
GGML_ASSERT(noise.dim() >= 2);
|
||||
const float token_w = static_cast<float>(noise.shape()[0]) / 32.f;
|
||||
const float token_h = static_cast<float>(noise.shape()[1]) / 32.f;
|
||||
const float noise_scale = std::min(16.f, std::sqrt((token_w * token_h) / 64.f));
|
||||
return noise * noise_scale;
|
||||
}
|
||||
|
||||
sd::Tensor<float> inverse_noise_scaling(float sigma,
|
||||
const sd::Tensor<float>& latent) override {
|
||||
SD_UNUSED(sigma);
|
||||
return latent;
|
||||
}
|
||||
|
||||
float noise_level_to_sigma(float noise_level) override {
|
||||
SD_UNUSED(noise_level);
|
||||
return 1.f;
|
||||
}
|
||||
|
||||
std::vector<float> get_sigmas(uint32_t n,
|
||||
int image_seq_len,
|
||||
scheduler_t scheduler_type,
|
||||
SDVersion version,
|
||||
const char* extra_sample_args = nullptr) override {
|
||||
SD_UNUSED(image_seq_len);
|
||||
SD_UNUSED(scheduler_type);
|
||||
SD_UNUSED(version);
|
||||
SD_UNUSED(extra_sample_args);
|
||||
std::vector<float> sigmas;
|
||||
sigmas.reserve(n + 1);
|
||||
if (n == 0) {
|
||||
sigmas.push_back(0.f);
|
||||
return sigmas;
|
||||
}
|
||||
for (uint32_t i = 0; i <= n; ++i) {
|
||||
const float t = static_cast<float>(i) / static_cast<float>(n);
|
||||
sigmas.push_back(t_to_sigma(t));
|
||||
}
|
||||
sigmas.back() = 0.f;
|
||||
return sigmas;
|
||||
}
|
||||
};
|
||||
|
||||
typedef std::function<sd::guidance::GuiderOutput(const sd::Tensor<float>&, float, int)> denoise_cb_t;
|
||||
|
||||
static std::pair<float, float> get_ancestral_step(float sigma_from,
|
||||
|
||||
@@ -156,6 +156,7 @@ const char* prediction_to_str[] = {
|
||||
"flux_flow",
|
||||
"sefi_flow",
|
||||
"minit2i_flow",
|
||||
"sensenova_u1_flow",
|
||||
};
|
||||
|
||||
const char* sd_prediction_name(enum prediction_t prediction) {
|
||||
@@ -322,6 +323,8 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
|
||||
sd_ctx_params->eager_load = false;
|
||||
sd_ctx_params->enable_mmap = false;
|
||||
sd_ctx_params->diffusion_flash_attn = false;
|
||||
sd_ctx_params->linear_scale = 0.f;
|
||||
sd_ctx_params->attn_scale = 0.f;
|
||||
sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO;
|
||||
sd_ctx_params->backend = nullptr;
|
||||
sd_ctx_params->params_backend = nullptr;
|
||||
@@ -352,6 +355,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
|
||||
"embeddings_connectors_path: %s\n"
|
||||
"vae_path: %s\n"
|
||||
"audio_vae_path: %s\n"
|
||||
"audio_encoder_path: %s\n"
|
||||
"taesd_path: %s\n"
|
||||
"control_net_path: %s\n"
|
||||
"photo_maker_path: %s\n"
|
||||
@@ -373,6 +377,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
|
||||
"auto_fit: %s\n"
|
||||
"flash_attn: %s\n"
|
||||
"diffusion_flash_attn: %s\n"
|
||||
"linear_scale: %g\n"
|
||||
"attn_scale: %g\n"
|
||||
"vae_format: %s\n",
|
||||
SAFE_STR(sd_ctx_params->model_path),
|
||||
SAFE_STR(sd_ctx_params->clip_l_path),
|
||||
@@ -387,6 +393,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
|
||||
SAFE_STR(sd_ctx_params->embeddings_connectors_path),
|
||||
SAFE_STR(sd_ctx_params->vae_path),
|
||||
SAFE_STR(sd_ctx_params->audio_vae_path),
|
||||
SAFE_STR(sd_ctx_params->audio_encoder_path),
|
||||
SAFE_STR(sd_ctx_params->taesd_path),
|
||||
SAFE_STR(sd_ctx_params->control_net_path),
|
||||
SAFE_STR(sd_ctx_params->photo_maker_path),
|
||||
@@ -408,6 +415,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
|
||||
BOOL_STR(sd_ctx_params->auto_fit),
|
||||
BOOL_STR(sd_ctx_params->flash_attn),
|
||||
BOOL_STR(sd_ctx_params->diffusion_flash_attn),
|
||||
sd_ctx_params->linear_scale,
|
||||
sd_ctx_params->attn_scale,
|
||||
sd_vae_format_name(sd_ctx_params->vae_format));
|
||||
|
||||
return buf;
|
||||
@@ -695,6 +704,13 @@ SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) {
|
||||
return sd_ctx->sd->control_net != nullptr;
|
||||
}
|
||||
|
||||
const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx) {
|
||||
if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_ctx->sd->version >= VERSION_COUNT) {
|
||||
return "Unknown";
|
||||
}
|
||||
return model_version_to_str[sd_ctx->sd->version];
|
||||
}
|
||||
|
||||
enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
|
||||
return sd::pipeline::default_sample_method(sd_ctx != nullptr ? sd_ctx->sd : nullptr);
|
||||
}
|
||||
@@ -722,8 +738,12 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
|
||||
const sd_vid_gen_params_t* sd_vid_gen_params,
|
||||
sd_image_t** frames_out,
|
||||
int* num_frames_out,
|
||||
sd_audio_t** audio_out) {
|
||||
sd_audio_t** audio_out,
|
||||
int* fps_out) {
|
||||
if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) {
|
||||
if (fps_out != nullptr) {
|
||||
*fps_out = 0;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -739,10 +759,13 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
|
||||
|
||||
StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd);
|
||||
if (!execution.ready) {
|
||||
if (fps_out != nullptr) {
|
||||
*fps_out = 0;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out);
|
||||
return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out, fps_out);
|
||||
}
|
||||
|
||||
SD_API void free_sd_images(sd_image_t* result_images, int num_images) {
|
||||
|
||||
@@ -45,16 +45,8 @@ void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) {
|
||||
bpe_len = rank;
|
||||
}
|
||||
|
||||
Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) {
|
||||
UNK_TOKEN = "<|endoftext|>";
|
||||
EOS_TOKEN = "<|endoftext|>";
|
||||
PAD_TOKEN = "<|endoftext|>";
|
||||
|
||||
UNK_TOKEN_ID = 151643;
|
||||
EOS_TOKEN_ID = 151643;
|
||||
PAD_TOKEN_ID = 151643;
|
||||
|
||||
special_tokens = {
|
||||
static const std::vector<std::string>& qwen2_special_tokens() {
|
||||
static const std::vector<std::string> tokens = {
|
||||
"<|endoftext|>",
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
@@ -87,6 +79,24 @@ Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) {
|
||||
"<|bot_token|>",
|
||||
"<|tms_token|>",
|
||||
};
|
||||
return tokens;
|
||||
}
|
||||
|
||||
Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str)
|
||||
: Qwen2Tokenizer(merges_utf8_str, qwen2_special_tokens()) {
|
||||
}
|
||||
|
||||
Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str,
|
||||
const std::vector<std::string>& special_tokens_override) {
|
||||
UNK_TOKEN = "<|endoftext|>";
|
||||
EOS_TOKEN = "<|endoftext|>";
|
||||
PAD_TOKEN = "<|endoftext|>";
|
||||
|
||||
UNK_TOKEN_ID = 151643;
|
||||
EOS_TOKEN_ID = 151643;
|
||||
PAD_TOKEN_ID = 151643;
|
||||
|
||||
special_tokens = special_tokens_override;
|
||||
|
||||
if (merges_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str);
|
||||
|
||||
@@ -2,12 +2,15 @@
|
||||
#define __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "bpe_tokenizer.h"
|
||||
|
||||
class Qwen2Tokenizer : public BPETokenizer {
|
||||
protected:
|
||||
void load_from_merges(const std::string& merges_utf8_str);
|
||||
Qwen2Tokenizer(const std::string& merges_utf8_str,
|
||||
const std::vector<std::string>& special_tokens_override);
|
||||
|
||||
public:
|
||||
explicit Qwen2Tokenizer(const std::string& merges_utf8_str = "");
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
#include "sensenova_u1_tokenizer.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
static const std::vector<std::string>& sensenova_u1_special_tokens() {
|
||||
static const std::vector<std::string> tokens = {
|
||||
"<|endoftext|>",
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
"<|object_ref_start|>",
|
||||
"<|object_ref_end|>",
|
||||
"<|box_start|>",
|
||||
"<|box_end|>",
|
||||
"<|quad_start|>",
|
||||
"<|quad_end|>",
|
||||
"<|vision_start|>",
|
||||
"<|vision_end|>",
|
||||
"<|vision_pad|>",
|
||||
"<|image_pad|>",
|
||||
"<|video_pad|>",
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
"<|fim_prefix|>",
|
||||
"<|fim_middle|>",
|
||||
"<|fim_suffix|>",
|
||||
"<|fim_pad|>",
|
||||
"<|repo_name|>",
|
||||
"<|file_sep|>",
|
||||
"<tool_response>",
|
||||
"</tool_response>",
|
||||
"<think>",
|
||||
"</think>",
|
||||
"<IMG_CONTEXT>",
|
||||
"<img>",
|
||||
"</img>",
|
||||
};
|
||||
return tokens;
|
||||
}
|
||||
|
||||
SenseNovaU1Tokenizer::SenseNovaU1Tokenizer(const std::string& merges_utf8_str)
|
||||
: Qwen2Tokenizer(merges_utf8_str, sensenova_u1_special_tokens()) {
|
||||
EOS_TOKEN = "<|im_end|>";
|
||||
EOS_TOKEN_ID = 151645;
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
#ifndef __SD_TOKENIZERS_SENSENOVA_U1_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_SENSENOVA_U1_TOKENIZER_H__
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "qwen2_tokenizer.h"
|
||||
|
||||
class SenseNovaU1Tokenizer : public Qwen2Tokenizer {
|
||||
public:
|
||||
explicit SenseNovaU1Tokenizer(const std::string& merges_utf8_str = "");
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_SENSENOVA_U1_TOKENIZER_H__
|
||||
Reference in New Issue
Block a user