Compare commits

...
Author SHA1 Message Date
Wagner Bruna 07a85c74cb feat: support Brownian tree noise in all noise injection samplers (#1899) 2026-09-15 02:37:31 +08:00
leejet f9ddc0f388 refactor: require external Gemma 2 and GPT-OSS tokenizers (#1974) 2026-09-15 01:25:13 +08:00
leejet 4964abdfc5 feat: support external Hugging Face tokenizer JSON files (#1973) 2026-09-15 00:27:39 +08:00
Санька Четвёртыйandleejet 42d6c0ab92 feat: Add generation parameters into video metadata (#1901)
Co-authored-by: leejet <leejet714@gmail.com>
2026-09-14 00:01:09 +08:00
leejet 5a5400bf0c fix: resolve MSVC narrowing conversion warnings (#1969) 2026-09-13 23:45:54 +08:00
fszontagh ca37fad89a fix: validate vision projector output dim against LLM hidden size (#1918) 2026-09-13 23:42:46 +08:00
Georgeandleejet 0bd72f075a feat: add Wan2.2 S2V (audio+img-to-video) support (#1925)
Co-authored-by: leejet <leejet714@gmail.com>
2026-09-13 23:33:50 +08:00
fszontagh 4a7da26b73 fix: bound plain-text runs in parse_prompt_attention regex (#1919) 2026-09-13 23:29:19 +08:00
leejet 9a977388a8 fix: guard GPU memory capacity and propagate encoding failures (#1958) 2026-09-13 21:35:46 +08:00
leejet 44dd13716d feat: preserve explicit backend assignments during auto-fit (#1967) 2026-09-13 17:40:38 +08:00
leejet 7f410a3793 feat: add linear and attention scale overrides (#1964) 2026-09-12 01:41:28 +08:00
leejet 5ebce93342 fix: reuse graph plans when scale parameters change (#1963) 2026-09-12 01:04:25 +08:00
Maphist0 7f986a9d73 feat: add SenseNova U1.5 support (#1935) 2026-09-12 00:39:59 +08:00
fszontagh e06b205384 feat: expose the loaded model version name through the public API (#1962) 2026-09-11 23:39:56 +08:00
vmobilis 3191b23d4b fix: handle invalid option numbers (#1961) 2026-09-11 23:16:39 +08:00
146 changed files with 101137 additions and 965 deletions
+3 -1
View File
@@ -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:
+4
View File
@@ -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.
+2 -1
View File
@@ -118,7 +118,8 @@ documentation.
6. Run the narrowest useful build, test, or inspection command available.
Follow `CONTRIBUTING.md` for formatting, naming, PR expectations, dependency
update policy, and security rules.
update policy, and security rules. For tokenizer additions, follow its embedded-data
allowlist and default to an external `tokenizer.json`.
---
+3
View File
@@ -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)
@@ -340,6 +342,7 @@ add_subdirectory(thirdparty)
target_sources(${SD_LIB} PRIVATE $<TARGET_OBJECTS:zip>)
target_link_libraries(${SD_LIB} PUBLIC ggml)
target_link_libraries(${SD_LIB} PRIVATE onig sd-utf8proc)
target_include_directories(${SD_LIB} PUBLIC . src include)
target_include_directories(${SD_LIB} PRIVATE src/core)
target_include_directories(${SD_LIB} PUBLIC . thirdparty)
+15
View File
@@ -46,6 +46,21 @@ Some older code in the project may not fully follow the current conventions. Ple
When adding or modifying model implementations, follow the model config and weight detection conventions in [docs/model_config.md](docs/model_config.md).
## Tokenizer Data
New model integrations must use an external `tokenizer.json` by default. Do not
embed new vocabularies or merge tables solely for less widely used models;
these tables increase the binary size for every user.
The embedded-data allowlist is CLIP, T5/UMT5, Qwen 2/3, Mistral, and Gemma 3/4.
Models may reuse an existing embedded tokenizer when its vocabulary and behavior
match their text encoder. Gemma 2 and GPT-OSS require external JSON files.
Adding to this allowlist requires maintainer approval, supported by the model's
usage, reuse across models, and measured binary-size cost. Document the matching
JSON and CLI option for models that require an external tokenizer, and fail
initialization clearly when it is missing.
## AI-Assisted Contributions
AI tools may be used to assist development, but contributors are responsible for the quality and correctness of the submitted code.
+2
View File
@@ -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)
+30 -12
View File
@@ -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
@@ -292,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.
+6 -2
View File
@@ -12,13 +12,17 @@ Lens uses a Lens diffusion transformer, the FLUX.2 VAE, and GPT-OSS-20B as the L
- safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
- Download GPT-OSS-20B
- gguf: https://huggingface.co/unsloth/gpt-oss-20b-GGUF/tree/main
- Download GPT-OSS-20B tokenizer.json
- https://huggingface.co/openai/gpt-oss-20b/tree/main
Lens and Lens Turbo require an external GPT-OSS `tokenizer.json` matching the text encoder checkpoint. Save it as `tokenizer_gpt_oss.json` and pass it with `--tokenizer`; the tokenizer is not embedded in sd.cpp. See [JSON tokenizers](tokenizers.md) for CLI and C API usage.
## Examples
### Lens
```
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 5.0 -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --tokenizer ..\models\tokenizers\tokenizer_gpt_oss.json --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 5.0 -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v
```
<img width="256" alt="Lens example" src="../assets/lens/example.png" />
@@ -26,7 +30,7 @@ Lens uses a Lens diffusion transformer, the FLUX.2 VAE, and GPT-OSS-20B as the L
### Lens Turbo
```
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_turbo_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 1.0 -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v --steps 4
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_turbo_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --tokenizer ..\models\tokenizers\tokenizer_gpt_oss.json --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 1.0 -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v --steps 4
```
<img width="256" alt="Lens Turbo example" src="../assets/lens/turbo_example.png" />
+1 -1
View File
@@ -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:
+5 -1
View File
@@ -11,6 +11,8 @@ In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a
- safetensors: https://huggingface.co/Comfy-Org/PixelDiT/tree/main/diffusion_models
- Download Gemma 2 2B
- safetensors: https://huggingface.co/Comfy-Org/PixelDiT/tree/main/text_encoders
- Download Gemma 2 2B tokenizer.json
- https://huggingface.co/google/gemma-2-2b/tree/main
- Download the VAE that matches the PiD checkpoint backbone
- safetensors: https://huggingface.co/nvidia/PiD/tree/main/checkpoints
- Flux / Z-Image PiD: use the Flux VAE and pass `--vae-format flux`
@@ -20,10 +22,12 @@ In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a
The official PiD model card should be checked before use. At the time of the initial PiD release, the official weights are under the NSCLv1 non-commercial license.
PiD and PiD 1.5 require an external Gemma 2 `tokenizer.json` matching the text encoder checkpoint. Save it as `tokenizer_gemma2.json` and pass it with `--tokenizer`; the tokenizer is not embedded in sd.cpp. See [JSON tokenizers](tokenizers.md) for CLI and C API usage.
## Examples
```
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\pid_flux1_512_to_2048_4step_bf16.safetensors --llm "..\models\text_encoders\gemma_2_2b_it_elm_bf16.safetensors" --vae ..\models\vae\ae.sft --vae-format flux --cfg-scale 1.0 -p "a lovely cat" -r ..\assets\ernie_image\turbo_example.png --diffusion-fa -v --steps 4 -H 2048 -W 2048 --rng cpu
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\pid_flux1_512_to_2048_4step_bf16.safetensors --llm "..\models\text_encoders\gemma_2_2b_it_elm_bf16.safetensors" --tokenizer ..\models\tokenizers\tokenizer_gemma2.json --vae ..\models\vae\ae.sft --vae-format flux --cfg-scale 1.0 -p "a lovely cat" -r ..\assets\ernie_image\turbo_example.png --diffusion-fa -v --steps 4 -H 2048 -W 2048 --rng cpu
```
Before:
+46
View File
@@ -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.
+107
View File
@@ -0,0 +1,107 @@
# JSON tokenizers
Use a Hugging Face `tokenizer.json` to supply the tokenizer vocabulary, merges,
added tokens, and processing stages. **PiD (including PiD 1.5) and Lens (including
Lens Turbo) require an external JSON**; their Gemma 2 and GPT-OSS tokenizers are
not embedded. Initialization fails if the main tokenizer is missing. Other
models keep their embedded tokenizer when this option is omitted.
```shell
sd-cli --diffusion-model model.gguf --llm text_encoder.gguf \
--tokenizer tokenizer_gemma2.json --vae vae.safetensors -p "a cat"
```
Choose the JSON belonging to the text encoder checkpoint. Checking that IDs fit
the embedding table does not establish that two vocabularies have the same
meaning. The JSON file is loaded when the text encoder is created; its embedded
vocabulary is not loaded in this case.
| Model | Required text encoder tokenizer | Example |
| --- | --- | --- |
| PiD / PiD 1.5 | Gemma 2 matching the text encoder checkpoint | `--tokenizer tokenizer_gemma2.json` |
| Lens / Lens Turbo | GPT-OSS matching the text encoder checkpoint | `--tokenizer tokenizer_gpt_oss.json` |
The Gemma 3/4 tokenizer used by LTX-2 remains embedded.
| Option | Encoder |
| --- | --- |
| `--tokenizer FILE` | Main LLM/BPE encoder: Gemma 2, Gemma 3, Qwen 2/3, Mistral, GPT-OSS; also Anima and HiDream-O1 |
| `--tokenizer FILE` | Shared CLIP tokenizer in SD1/SD2/SDXL, or CLIP-L in Flux |
| `--tokenizer clip-l=FILE` | Separate CLIP-L in SD3 or Flux |
| `--tokenizer clip-g=FILE` | Separate CLIP-G in SD3 |
Use comma-separated assignments to configure multiple slots, for example
`--tokenizer main=main.json,clip-l=clip_l.json,clip-g=clip_g.json`.
A plain file path is equivalent to `main=FILE`. You may also repeat `--tokenizer`
with explicit assignments, such as `--tokenizer main=main.json --tokenizer clip-l=clip.json`.
Empty assignment paths, unknown keys, malformed assignments and
duplicate slots are rejected. Commas separate entries in the assignment form;
quote the complete argument when paths contain spaces.
SD3 overrides must name the `clip-l` or `clip-g` slot. SDXL uses one shared
tokenizer for both CLIP encoders. Do not supply both `main` and `clip-l` for Flux.
A slot targeting an absent or unsupported encoder fails initialization.
T5/SentencePiece Unigram tokenizers are outside this implementation's scope.
For example, SD3 can load the same CLIP JSON into both slots:
```shell
sd-cli --diffusion-model sd3.gguf --clip_l clip_l.safetensors \
--clip_g clip_g.safetensors --t5xxl t5xxl.gguf --vae vae.safetensors \
--tokenizer clip-l=tokenizer_clip.json,clip-g=tokenizer_clip.json \
-p "a cat"
```
The C API accepts the same string in `sd_ctx_params_t::tokenizer`. A null or
empty value keeps an embedded tokenizer where available; PiD and Lens require
a nonempty main tokenizer path. The CLI passes the string through;
`TokenizerConfig` parses and validates it when text encoders are initialized.
```c
sd_ctx_params_t params;
sd_ctx_params_init(&params);
params.tokenizer = "clip-l=tokenizer_clip.json,clip-g=tokenizer_clip.json";
```
Rebuild applications against the updated public header when using the updated
library.
## Supported components
| Stage | Supported configurations |
| --- | --- |
| Normalizer | `Sequence`, `NFC`, `Lowercase`, `Replace` with String/Regex patterns |
| PreTokenizer | `Sequence`, `Split` with String/Regex patterns, all five delimiter behaviors and `invert`; `ByteLevel` with `add_prefix_space` and `use_regex` |
| Model | Deterministic `BPE`, string or array-pair merges, `unk_token`, `fuse_unk`, `byte_fallback`, `ignore_merges`, `end_of_word_suffix` |
| PostProcessor | Single-sequence `TemplateProcessing` with at most one prefix and one suffix token, `RobertaProcessing`, `ByteLevel` |
| Decoder | `Sequence`, `Replace`, `ByteLevel`, `ByteFallback`, `Fuse` |
| AddedToken | Special and ordinary added tokens, original IDs, raw or normalized matching, leftmost-longest matching |
`ByteLevel.use_regex` defaults to true when omitted. ByteLevel postprocessing
changes offsets only and adds no tokens. Added tokens with `single_word`,
`lstrip`, or `rstrip` enabled, nonzero BPE dropout, nonempty
`continuing_subword_prefix`, and unsupported component types fail loading.
New added-token IDs must follow the model vocabulary consecutively; configurations
whose IDs Hugging Face would reassign are rejected.
JSON `padding` and `truncation` must be null. This API returns IDs, not offsets,
type IDs, or paired-input encodings; the pair template is not used.
The pipeline covers the CLIP, Gemma 2, Gemma 3, GPT-OSS, Mistral 3, Qwen 2 and
Qwen 3 JSON configurations used by the differential test. It does not imply
support for every tokenizer published under those model names.
## Prompt integration
Prompt attention parsing and model-specific chat/image templates remain in the
conditioner. Raw `encode()` does not add BOS/EOS. The conditioner concatenates
weighted prompt fragments, then the existing padding/chunking step applies the
JSON single-sequence template once per sequence or CLIP chunk. Padding ID,
direction, length limits and attention masks remain text encoder policies.
CLIP requires both BOS and EOS because its chunking reserves those positions.
The internal `encode()`, `tokenize()`, and `decode()` interfaces return a success
flag and write to an output parameter. A successful result may be empty; a failed
call clears its output. JSON tokenizer input, normalization, and regex failures
return `false` with diagnostic information instead of throwing. Invalid
JSON, unsupported stages, conflicting IDs and IDs outside the encoder embedding
table fail initialization instead of falling back to the embedded tokenizer.
+45
View File
@@ -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
View File
@@ -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
```
+3
View File
@@ -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).
+9 -5
View File
@@ -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());
}
}
@@ -540,12 +541,16 @@ bool save_results(const SDCliParams& cli_params,
if (cli_params.mode == VID_GEN && num_results > 1) {
if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
ext = ".avi";
std::string params = gen_params.embed_image_metadata
? get_image_params(ctx_params, gen_params, gen_params.seed, cli_params.mode)
: "";
fs::path video_path = base_path;
video_path += ext;
std::string final_ext_lower = ext.string();
std::transform(final_ext_lower.begin(), final_ext_lower.end(), final_ext_lower.begin(), ::tolower);
const bool mux_audio = generated_audio != nullptr && (final_ext_lower == ".avi" || final_ext_lower == ".webm");
if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr) == 0) {
if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr, params) == 0) {
LOG_INFO("save result video to '%s'", video_path.string().c_str());
if (generated_audio != nullptr && !mux_audio) {
fs::path wav_path = video_path;
@@ -687,8 +692,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 +954,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);
}
+78 -8
View File
@@ -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 = {
@@ -382,6 +410,11 @@ ArgOptions SDContextParams::get_options() {
"path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, mistral-small3.2 for flux2, ...)",
0,
&llm_path},
{"",
"--tokenizer",
"tokenizer.json path, or comma-separated main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens",
(int)',',
&tokenizer},
{"",
"--llm_vision",
"path to the llm vit",
@@ -432,6 +465,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 +716,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",
@@ -851,6 +901,7 @@ std::string SDContextParams::to_string() const {
<< " t5xxl_path: \"" << t5xxl_path << "\",\n"
<< " llm_path: \"" << llm_path << "\",\n"
<< " llm_vision_path: \"" << llm_vision_path << "\",\n"
<< " tokenizer: \"" << tokenizer << "\",\n"
<< " diffusion_model_path: \"" << diffusion_model_path << "\",\n"
<< " high_noise_diffusion_model_path: \"" << high_noise_diffusion_model_path << "\",\n"
<< " uncond_diffusion_model_path: \"" << uncond_diffusion_model_path << "\",\n"
@@ -858,6 +909,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 +938,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"
@@ -915,12 +969,14 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.t5xxl_path = t5xxl_path.c_str();
sd_ctx_params.llm_path = llm_path.c_str();
sd_ctx_params.llm_vision_path = llm_vision_path.c_str();
sd_ctx_params.tokenizer = tokenizer.c_str();
sd_ctx_params.diffusion_model_path = diffusion_model_path.c_str();
sd_ctx_params.high_noise_diffusion_model_path = high_noise_diffusion_model_path.c_str();
sd_ctx_params.uncond_diffusion_model_path = uncond_diffusion_model_path.c_str();
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 +995,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;
@@ -1051,7 +1109,7 @@ ArgOptions SDGenerationParams::get_options() {
&hires_upscaler},
{"",
"--extra-sample-args",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
(int)',',
&extra_sample_args},
{"",
@@ -1471,6 +1529,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 +1726,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)",
+4
View File
@@ -124,6 +124,7 @@ struct SDContextParams {
std::string t5xxl_path;
std::string llm_path;
std::string llm_vision_path;
std::string tokenizer;
std::string diffusion_model_path;
std::string high_noise_diffusion_model_path;
std::string uncond_diffusion_model_path;
@@ -131,6 +132,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 +177,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();
+53 -11
View File
@@ -810,7 +810,31 @@ uint8_t* load_image_from_memory(const char* image_bytes,
return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel);
}
std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
static void append_avi_metadata(std::vector<uint8_t>& data, const std::string& parameters) {
if (parameters.empty()) {
return;
}
std::vector<uint8_t> info_content;
write_fourcc(info_content, "INFO");
const size_t comment_size = parameters.size() + 1;
write_fourcc(info_content, "ICMT");
write_u32_le(info_content, static_cast<uint32_t>(comment_size));
info_content.insert(info_content.end(), parameters.begin(), parameters.end());
info_content.push_back(0);
if (comment_size & 1u) {
info_content.push_back(0);
}
write_fourcc(data, "LIST");
write_u32_le(data, static_cast<uint32_t>(info_content.size()));
data.insert(data.end(), info_content.begin(), info_content.end());
size_t start_pos = data.size();
}
std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
if (num_images == 0) {
fprintf(stderr, "Error: Image array is empty.\n");
return {};
@@ -1000,6 +1024,8 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
const size_t movi_size = avi_data.size() - movi_size_pos - 4;
patch_u32_le(avi_data, movi_size_pos, static_cast<uint32_t>(movi_size));
append_avi_metadata(avi_data, parameters);
write_fourcc(avi_data, "idx1");
write_u32_le(avi_data, static_cast<uint32_t>(index.size() * 16));
for (const auto& entry : index) {
@@ -1015,8 +1041,8 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
return avi_data;
}
int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio);
int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
if (avi_data.empty()) {
return -1;
}
@@ -1146,7 +1172,7 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
#endif
#ifdef SD_USE_WEBM
std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
if (num_images == 0) {
fprintf(stderr, "Error: Image array is empty.\n");
return {};
@@ -1213,6 +1239,21 @@ std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, in
segment.GetSegmentInfo()->set_writing_app("stable-diffusion.cpp");
segment.GetSegmentInfo()->set_muxing_app("stable-diffusion.cpp");
LOG_DEBUG("Embedding parameters to metadata: %s", parameters.c_str());
if (!parameters.empty()) {
mkvmuxer::Tag* tag = segment.AddTag();
if (tag) {
if (!tag->add_simple_tag("COMMENT", parameters.c_str())) {
LOG_WARN("Failed to add COMMENT simple tag.");
}
} else {
LOG_WARN("Failed to add tag to segment.");
}
} else {
LOG_INFO("Paramaters is empty, COMMENT tag not embedded.\n");
}
const uint64_t frame_duration_ns = std::max<uint64_t>(
1, static_cast<uint64_t>(std::llround(1000000000.0 / static_cast<double>(fps))));
uint64_t timestamp_ns = 0;
@@ -1271,8 +1312,8 @@ std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, in
return writer.data();
}
int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio);
int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
if (webm_data.empty()) {
return -1;
}
@@ -1289,7 +1330,8 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
int num_images,
int fps,
int quality,
const sd_audio_t* audio) {
const sd_audio_t* audio,
const std::string& parameters) {
std::string format = output_format;
std::transform(format.begin(), format.end(), format.begin(),
[](unsigned char c) { return static_cast<char>(tolower(c)); });
@@ -1299,7 +1341,7 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
#ifdef SD_USE_WEBM
if (format == "webm") {
return create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio);
return create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
}
#endif
@@ -1309,14 +1351,14 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
}
#endif
return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio);
return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
}
int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
std::string path = filename ? filename : "";
auto pos = path.find_last_of('.');
std::string ext = pos == std::string::npos ? "" : path.substr(pos);
std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality, audio);
std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality, audio, parameters);
if (video_data.empty()) {
return -1;
}
+18 -12
View File
@@ -57,13 +57,15 @@ int create_mjpg_avi_from_sd_images(const char* filename,
sd_image_t* images,
int num_images,
int fps,
int quality = 90,
const sd_audio_t* audio = nullptr);
int quality = 90,
const sd_audio_t* audio = nullptr,
const std::string& parameters = "");
std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images,
int num_images,
int fps,
int quality = 90,
const sd_audio_t* audio = nullptr);
int quality = 90,
const sd_audio_t* audio = nullptr,
const std::string& parameters = "");
#ifdef SD_USE_WEBP
int create_animated_webp_from_sd_images(const char* filename,
@@ -82,27 +84,31 @@ int create_webm_from_sd_images(const char* filename,
sd_image_t* images,
int num_images,
int fps,
int quality = 90,
const sd_audio_t* audio = nullptr);
int quality = 90,
const sd_audio_t* audio = nullptr,
const std::string& parameters = "");
std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images,
int num_images,
int fps,
int quality = 90,
const sd_audio_t* audio = nullptr);
int quality = 90,
const sd_audio_t* audio = nullptr,
const std::string& parameters = "");
#endif
int create_video_from_sd_images(const char* filename,
sd_image_t* images,
int num_images,
int fps,
int quality = 90,
const sd_audio_t* audio = nullptr);
int quality = 90,
const sd_audio_t* audio = nullptr,
const std::string& parameters = "");
std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& output_format,
sd_image_t* images,
int num_images,
int fps,
int quality = 90,
const sd_audio_t* audio = nullptr);
int quality = 90,
const sd_audio_t* audio = nullptr,
const std::string& parameters = "");
bool write_wav_to_file(const std::string& path,
const float* interleaved_samples,
+3
View File
@@ -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).
+7 -4
View File
@@ -237,6 +237,9 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
int& output_fps,
std::string& error_message) {
sd_vid_gen_params_t params = job.vid_gen.to_sd_vid_gen_params_t();
std::string str_params = job.vid_gen.gen_params.embed_image_metadata
? get_image_params(*runtime.ctx_params, job.vid_gen.gen_params, job.vid_gen.gen_params.seed, VID_GEN)
: "";
SDImageVec results;
int num_results = 0;
@@ -245,7 +248,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, &params, &raw_results, &num_results, &generated_audio)) {
if (!generate_video(runtime.sd_ctx, &params, &raw_results, &num_results, &generated_audio, &output_fps)) {
raw_results = nullptr;
}
results.adopt(raw_results, num_results);
@@ -261,9 +264,10 @@ 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);
generated_audio,
str_params);
free_sd_audio(generated_audio);
if (video_bytes.empty()) {
error_message = "failed to encode generated video container";
@@ -273,7 +277,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;
}
+11 -1
View File
@@ -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,9 @@ 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
const char* tokenizer; // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens
} sd_ctx_params_t;
typedef struct {
@@ -493,6 +498,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 +523,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;
+357 -81
View File
@@ -7,6 +7,7 @@
#include <limits>
#include <optional>
#include <sstream>
#include <stdexcept>
#include "core/ggml_tensor_utils.h"
#include "core/tensor_ggml.hpp"
@@ -16,6 +17,8 @@
#include "model/te/llm.hpp"
#include "model/te/t5.hpp"
#include "model_loader.h"
#include "tokenizers/sensenova_u1_tokenizer.h"
#include "tokenizers/tokenizer_config.h"
struct SDCondition {
sd::Tensor<float> c_crossattn;
@@ -149,6 +152,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() {}
};
@@ -157,7 +161,7 @@ public:
// Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/sd_hijack_clip.py#L283
struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
SDVersion version = VERSION_SD1;
CLIPTokenizer tokenizer;
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<CLIPTextModelRunner> text_model;
std::shared_ptr<CLIPTextModelRunner> text_model2;
@@ -171,12 +175,18 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
const String2TensorStorage& tensor_storage_map,
const std::map<std::string, std::string>& orig_embedding_map,
SDVersion version = VERSION_SD1,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: version(version), tokenizer(sd_version_is_sd2(version) ? 0 : 49407) {
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: version(version) {
const int pad_id = sd_version_is_sd2(version) ? 0 : 49407;
tokenizer = tokenizers.create(TokenizerConfig::MAIN, 49408, pad_id, false, true);
if (!tokenizer) {
tokenizer = std::make_shared<CLIPTokenizer>(pad_id);
}
for (const auto& kv : orig_embedding_map) {
std::string name = normalize_embedding_name(kv.first);
embedding_map[name] = kv.second;
tokenizer.add_special_token(name);
tokenizer->add_special_token(name);
}
bool force_clip_f32 = !embedding_map.empty();
if (sd_version_is_sd1(version)) {
@@ -231,6 +241,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)) {
@@ -356,16 +373,15 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
return load_embedding(name, iter->second, bpe_tokens);
}
std::vector<int> convert_token_to_id(std::string text) {
bool convert_token_to_id(const std::string& text, std::vector<int>& tokens) {
auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
return append_embedding_tokens(str, bpe_tokens);
};
std::vector<int> curr_tokens = tokenizer.encode(text, on_new_token_cb);
return curr_tokens;
return tokenizer->encode(text, tokens, on_new_token_cb);
}
std::string decode(const std::vector<int>& tokens) {
return tokenizer.decode(tokens);
bool decode(const std::vector<int>& tokens, std::string& text) {
return tokenizer->decode(tokens, text);
}
std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
@@ -403,18 +419,21 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
if (padding_size > 0) {
LOG_VERBOSE("BREAK token encountered, padding current chunk by %zu tokens.", padding_size);
tokens.insert(tokens.end(), padding_size, tokenizer.EOS_TOKEN_ID);
tokens.insert(tokens.end(), padding_size, tokenizer->EOS_TOKEN_ID);
weights.insert(weights.end(), padding_size, 1.0f);
}
continue; // Skip to the next item after handling BREAK
}
std::vector<int> curr_tokens = tokenizer.encode(curr_text, on_new_token_cb);
std::vector<int> curr_tokens;
if (!tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
tokenizer.pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
tokenizer->pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
// for (int i = 0; i < tokens.size(); i++) {
// std::cout << tokens[i] << ":" << weights[i] << ", ";
@@ -451,7 +470,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
sd::Tensor<int32_t> input_ids2;
size_t max_token_idx = 0;
if (sd_version_is_sdxl(version)) {
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer.EOS_TOKEN_ID);
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer->EOS_TOKEN_ID);
if (it != chunk_tokens.end()) {
std::fill(std::next(it), chunk_tokens.end(), 0);
}
@@ -552,7 +571,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
auto tokens_and_weights = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
if (tokens_and_weights.first.empty()) {
return {};
}
std::vector<int>& tokens = tokens_and_weights.first;
std::vector<float>& weights = tokens_and_weights.second;
return get_learned_condition_common(n_threads,
@@ -620,8 +642,8 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
};
struct SD3CLIPEmbedder : public Conditioner {
CLIPTokenizer clip_l_tokenizer;
CLIPTokenizer clip_g_tokenizer;
std::shared_ptr<Tokenizer> clip_l_tokenizer;
std::shared_ptr<Tokenizer> clip_g_tokenizer;
T5UniGramTokenizer t5_tokenizer;
std::shared_ptr<CLIPTextModelRunner> clip_l;
std::shared_ptr<CLIPTextModelRunner> clip_g;
@@ -629,8 +651,8 @@ struct SD3CLIPEmbedder : public Conditioner {
SD3CLIPEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: clip_g_tokenizer(0) {
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {}) {
bool use_clip_l = false;
bool use_clip_g = false;
bool use_t5 = false;
@@ -648,9 +670,17 @@ struct SD3CLIPEmbedder : public Conditioner {
return;
}
if (use_clip_l) {
clip_l_tokenizer = tokenizers.create(TokenizerConfig::CLIP_L, 49408, 49407, false, true);
if (!clip_l_tokenizer) {
clip_l_tokenizer = std::make_shared<CLIPTokenizer>();
}
clip_l = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, false, weight_manager);
}
if (use_clip_g) {
clip_g_tokenizer = tokenizers.create(TokenizerConfig::CLIP_G, 49408, 0, false, true);
if (!clip_g_tokenizer) {
clip_g_tokenizer = std::make_shared<CLIPTokenizer>(0);
}
clip_g = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, false, weight_manager);
}
if (use_t5) {
@@ -736,6 +766,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);
@@ -790,27 +832,36 @@ struct SD3CLIPEmbedder : public Conditioner {
const std::string& curr_text = item.first;
float curr_weight = item.second;
if (clip_l) {
std::vector<int> curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb);
std::vector<int> curr_tokens;
if (!clip_l_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end());
clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
}
if (clip_g) {
std::vector<int> curr_tokens = clip_g_tokenizer.encode(curr_text, on_new_token_cb);
std::vector<int> curr_tokens;
if (!clip_g_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
clip_g_tokens.insert(clip_g_tokens.end(), curr_tokens.begin(), curr_tokens.end());
clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight);
}
if (t5) {
std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
}
if (clip_l) {
clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
clip_l_tokenizer->pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
}
if (clip_g) {
clip_g_tokenizer.pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
clip_g_tokenizer->pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
}
if (t5) {
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
@@ -881,7 +932,7 @@ struct SD3CLIPEmbedder : public Conditioner {
chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights);
if (chunk_idx == 0) {
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer.EOS_TOKEN_ID);
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer->EOS_TOKEN_ID);
max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
pooled_l = clip_l->compute(n_threads,
input_ids,
@@ -924,7 +975,7 @@ struct SD3CLIPEmbedder : public Conditioner {
chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights);
if (chunk_idx == 0) {
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_g_tokenizer.EOS_TOKEN_ID);
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_g_tokenizer->EOS_TOKEN_ID);
max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
pooled_g = clip_g->compute(n_threads,
input_ids,
@@ -1002,6 +1053,9 @@ struct SD3CLIPEmbedder : public Conditioner {
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, 77, 77, true);
if (tokens_and_weights.empty()) {
return {};
}
return get_learned_condition_common(n_threads,
tokens_and_weights,
conditioner_params.clip_skip,
@@ -1010,7 +1064,7 @@ struct SD3CLIPEmbedder : public Conditioner {
};
struct FluxCLIPEmbedder : public Conditioner {
CLIPTokenizer clip_l_tokenizer;
std::shared_ptr<Tokenizer> clip_l_tokenizer;
T5UniGramTokenizer t5_tokenizer;
std::shared_ptr<CLIPTextModelRunner> clip_l;
std::shared_ptr<T5Runner> t5;
@@ -1018,7 +1072,8 @@ struct FluxCLIPEmbedder : public Conditioner {
FluxCLIPEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {}) {
bool use_clip_l = false;
bool use_t5 = false;
for (auto pair : tensor_storage_map) {
@@ -1035,6 +1090,11 @@ struct FluxCLIPEmbedder : public Conditioner {
}
if (use_clip_l) {
auto slot = tokenizers.has(TokenizerConfig::CLIP_L) ? TokenizerConfig::CLIP_L : TokenizerConfig::MAIN;
clip_l_tokenizer = tokenizers.create(slot, 49408, 49407, false, true);
if (!clip_l_tokenizer) {
clip_l_tokenizer = std::make_shared<CLIPTokenizer>();
}
clip_l = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, false, weight_manager);
} else {
LOG_WARN("clip_l text encoder not found! Prompt adherence might be degraded.");
@@ -1106,6 +1166,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);
@@ -1151,19 +1220,25 @@ struct FluxCLIPEmbedder : public Conditioner {
const std::string& curr_text = item.first;
float curr_weight = item.second;
if (clip_l) {
std::vector<int> curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb);
std::vector<int> curr_tokens;
if (!clip_l_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end());
clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
}
if (t5) {
std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
}
if (clip_l) {
clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
clip_l_tokenizer->pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
}
if (t5) {
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
@@ -1213,7 +1288,7 @@ struct FluxCLIPEmbedder : public Conditioner {
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
size_t max_token_idx = 0;
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer.EOS_TOKEN_ID);
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer->EOS_TOKEN_ID);
max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
pooled = clip_l->compute(n_threads,
@@ -1224,7 +1299,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 +1321,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);
@@ -1270,6 +1351,9 @@ struct FluxCLIPEmbedder : public Conditioner {
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
if (tokens_and_weights.empty()) {
return {};
}
return get_learned_condition_common(n_threads,
tokens_and_weights,
conditioner_params.clip_skip,
@@ -1368,6 +1452,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);
@@ -1407,7 +1497,10 @@ struct T5CLIPEmbedder : public Conditioner {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
@@ -1505,6 +1598,9 @@ struct T5CLIPEmbedder : public Conditioner {
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
if (std::get<0>(tokens_and_weights).empty()) {
return {};
}
return get_learned_condition_common(n_threads,
tokens_and_weights,
conditioner_params.clip_skip,
@@ -1576,6 +1672,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);
@@ -1597,7 +1699,10 @@ struct MiniT2IConditioner : public Conditioner {
return result;
}
std::vector<int> tokens = tokenizer.encode(conditioner_params.text);
std::vector<int> tokens;
if (!tokenizer.encode(conditioner_params.text, tokens)) {
return {};
}
if (tokens.size() > prompt_length) {
tokens.resize(prompt_length);
}
@@ -1623,21 +1728,93 @@ 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) {
std::vector<int> tokens;
if (!tokenizer.encode(build_query(text, is_negative), tokens)) {
return {};
}
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;
std::shared_ptr<Tokenizer> qwen_tokenizer;
T5UniGramTokenizer t5_tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
AnimaConditioner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {}) {
llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
backend,
tensor_storage_map,
"text_encoders.llm",
false,
weight_manager);
qwen_tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 151643);
if (!qwen_tokenizer) {
qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
@@ -1672,6 +1849,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);
}
@@ -1700,7 +1881,10 @@ struct AnimaConditioner : public Conditioner {
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
std::vector<int> curr_tokens = qwen_tokenizer->tokenize(curr_text, nullptr);
std::vector<int> curr_tokens;
if (!qwen_tokenizer->tokenize(curr_text, curr_tokens, nullptr)) {
return {};
}
qwen_tokens.insert(qwen_tokens.end(), curr_tokens.begin(), curr_tokens.end());
// Anima uses uniform Qwen token weights.
qwen_weights.insert(qwen_weights.end(), curr_tokens.size(), 1.f);
@@ -1713,7 +1897,10 @@ struct AnimaConditioner : public Conditioner {
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
@@ -1732,6 +1919,10 @@ struct AnimaConditioner : public Conditioner {
auto& t5_tokens = std::get<2>(tokenized);
auto& t5_weights = std::get<3>(tokenized);
if (qwen_tokens.empty()) {
return {};
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(qwen_tokens.size()), 1}, qwen_tokens);
auto hidden_states = llm->compute(n_threads,
input_ids,
@@ -1758,7 +1949,7 @@ struct AnimaConditioner : public Conditioner {
struct LLMEmbedder : public Conditioner {
SDVersion version;
std::shared_ptr<BPETokenizer> tokenizer;
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<T5Runner> byt5;
@@ -1767,8 +1958,17 @@ struct LLMEmbedder : public Conditioner {
SDVersion version = VERSION_QWEN_IMAGE,
const std::string prefix = "",
bool enable_vision = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: version(version) {
if (!tokenizers.has(TokenizerConfig::MAIN)) {
if (sd_version_is_lens(version)) {
throw std::runtime_error("Lens requires an external GPT-OSS tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
}
if (sd_version_is_pid(version)) {
throw std::runtime_error("PiD requires an external Gemma 2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
}
}
LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
if (version == VERSION_FLUX2) {
arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
@@ -1789,21 +1989,28 @@ struct LLMEmbedder : public Conditioner {
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
arch = LLM::LLMArch::QWEN3;
}
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
tokenizer = std::make_shared<MistralTokenizer>();
} else if (arch == LLM::LLMArch::GPT_OSS_20B) {
tokenizer = std::make_shared<GPTOSSTokenizer>();
} else if (arch == LLM::LLMArch::GEMMA2_2B) {
tokenizer = std::make_shared<Gemma2Tokenizer>();
} else {
tokenizer = std::make_shared<Qwen2Tokenizer>();
}
llm = std::make_shared<LLM::LLMRunner>(arch,
llm = std::make_shared<LLM::LLMRunner>(arch,
backend,
tensor_storage_map,
"text_encoders.llm",
enable_vision,
weight_manager);
int pad_id = 151643;
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
pad_id = 11;
} else if (arch == LLM::LLMArch::GPT_OSS_20B) {
pad_id = 199999;
} else if (arch == LLM::LLMArch::GEMMA2_2B) {
pad_id = 0;
}
tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, pad_id);
if (!tokenizer) {
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
tokenizer = std::make_shared<MistralTokenizer>();
} else {
tokenizer = std::make_shared<Qwen2Tokenizer>();
}
}
if (sd_version_is_hunyuan_video(version)) {
const std::string byt5_prefix = "text_encoders.t5xxl.transformer";
for (const auto& [name, _] : tensor_storage_map) {
@@ -1876,6 +2083,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);
@@ -1935,7 +2149,10 @@ struct LLMEmbedder : public Conditioner {
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens = tokenizer->encode(curr_text, nullptr);
std::vector<int> curr_tokens;
if (!tokenizer->encode(curr_text, curr_tokens, nullptr)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
@@ -1968,6 +2185,10 @@ struct LLMEmbedder : public Conditioner {
auto& weights = std::get<1>(tokens_weights_mask);
auto& mask = std::get<2>(tokens_weights_mask);
if (tokens.empty()) {
return {};
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
sd::Tensor<float> attention_mask;
if (!mask.empty()) {
@@ -2127,7 +2348,11 @@ struct LLMEmbedder : public Conditioner {
GGML_ASSERT(image_outputs.size() == 4);
auto image_embed = std::move(image_outputs[0]);
prompt += "<|vision_start|>";
int image_embed_idx = static_cast<int>(tokenizer->encode(prompt, nullptr).size());
std::vector<int> prefix_tokens;
if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) {
return false;
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
if (deepstack_image_embeds.empty()) {
deepstack_image_embeds.resize(image_outputs.size() - 1);
@@ -2143,6 +2368,7 @@ struct LLMEmbedder : public Conditioner {
prompt += placeholder;
}
prompt += "<|vision_end|>";
return true;
};
const auto* references = conditioner_params.minimax_h3_references;
@@ -2159,11 +2385,13 @@ struct LLMEmbedder : public Conditioner {
GGML_ASSERT(item.frames.size() == 1);
auto resized = resize_for_vision(item.frames[0]);
prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
if (!add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size)) {
return {};
}
continue;
}
@@ -2187,22 +2415,26 @@ struct LLMEmbedder : public Conditioner {
second.shape()[3]});
}
auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
add_vision_outputs(llm->encode_video_block_outputs(n_threads,
pair,
false),
static_cast<int>(first.shape()[1]) / patch_size,
static_cast<int>(first.shape()[0]) / patch_size);
if (!add_vision_outputs(llm->encode_video_block_outputs(n_threads,
pair,
false),
static_cast<int>(first.shape()[1]) / patch_size,
static_cast<int>(first.shape()[0]) / patch_size)) {
return {};
}
}
}
} else if (conditioner_params.ref_images != nullptr) {
for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
auto resized = resize_for_vision((*conditioner_params.ref_images)[i]);
prompt += "<Picture " + std::to_string(i + 1) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
if (!add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size)) {
return {};
}
}
}
}
@@ -2238,7 +2470,10 @@ struct LLMEmbedder : public Conditioner {
"enhanced description for the prompt below and avoid including any additional "
"commentary or evaluations:<|im_end|>\n<|im_start|>user\n";
auto prefix_tokens = tokenizer->encode(prompt_prefix, nullptr);
std::vector<int> prefix_tokens;
if (!tokenizer->encode(prompt_prefix, prefix_tokens, nullptr)) {
return {};
}
prompt_template_encode_start_idx = 0;
for (int token : prefix_tokens) {
if (token != pad_token) {
@@ -2291,7 +2526,11 @@ struct LLMEmbedder : public Conditioner {
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "<|vision_start|>";
@@ -2433,7 +2672,11 @@ struct LLMEmbedder : public Conditioner {
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>";
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "<|vision_start|>";
@@ -2501,7 +2744,11 @@ struct LLMEmbedder : public Conditioner {
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
@@ -2706,7 +2953,10 @@ struct LLMEmbedder : public Conditioner {
"- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n"
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:\n"
"User Prompt: ";
auto chi_tokens = std::get<0>(tokenize(chi_prompt, {0, 0}));
auto chi_tokens = std::get<0>(tokenize(chi_prompt, {0, 0}));
if (chi_tokens.empty()) {
return {};
}
size_t num_chi_tokens = chi_tokens.size();
max_length = (int)num_chi_tokens + pixeldit_max_length - 2;
min_length = max_length;
@@ -2725,7 +2975,9 @@ struct LLMEmbedder : public Conditioner {
0,
false,
max_length);
GGML_ASSERT(!hidden_states.empty());
if (hidden_states.empty()) {
return {};
}
if (hidden_states.shape()[1] > pixeldit_max_length) {
auto bos = sd::ops::slice(hidden_states, 1, 0, 1);
@@ -2758,6 +3010,9 @@ struct LLMEmbedder : public Conditioner {
max_length,
deepstack_image_embeds,
image_grids);
if (hidden_states.empty()) {
return {};
}
std::vector<sd::Tensor<float>> extra_hidden_states_vec;
if (sd_version_is_hunyuan_video(version) && byt5) {
std::vector<std::string> quoted_texts;
@@ -2808,6 +3063,9 @@ struct LLMEmbedder : public Conditioner {
prompt_template_encode_start_idx,
spell_quotes,
max_length);
if (extra_hidden_states.empty()) {
return {};
}
extra_hidden_states_vec.push_back(std::move(extra_hidden_states));
}
@@ -2906,7 +3164,7 @@ struct LTXAVEmbedder : public Conditioner {
static constexpr int64_t kNumStates = 49;
static constexpr int64_t kMinLength = 1024;
std::shared_ptr<GemmaTokenizer> tokenizer;
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<LTXAVTextProjectionRunner> projector;
std::string projector_prefix;
@@ -2933,17 +3191,21 @@ struct LTXAVEmbedder : public Conditioner {
const String2TensorStorage& tensor_storage_map = {},
const std::string& llm_prefix = "text_encoders.llm",
const std::string& projector_prefix = "text_embedding_projection",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: projector_prefix(projector_prefix) {
LLM::LLMArch arch = detect_gemma_arch(tensor_storage_map, llm_prefix);
LOG_INFO("ltxav text encoder: %s", arch == LLM::LLMArch::GEMMA4_12B ? "gemma 4" : "gemma 3");
tokenizer = std::make_shared<GemmaTokenizer>();
llm = std::make_shared<LLM::LLMRunner>(arch,
llm = std::make_shared<LLM::LLMRunner>(arch,
backend,
tensor_storage_map,
llm_prefix,
false,
weight_manager);
tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 0, true);
if (!tokenizer) {
tokenizer = std::make_shared<GemmaTokenizer>();
}
dual_projection = tensor_storage_map.find(projector_prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end();
projector = std::make_shared<LTXAVTextProjectionRunner>(backend,
tensor_storage_map,
@@ -2965,6 +3227,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);
@@ -3017,7 +3284,10 @@ struct LTXAVEmbedder : public Conditioner {
std::vector<int> tokens;
std::vector<float> weights;
for (const auto& item : parsed_attention) {
auto curr_tokens = tokenizer->encode(item.first, nullptr);
std::vector<int> curr_tokens;
if (!tokenizer->encode(item.first, curr_tokens, nullptr)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), item.second);
}
@@ -3035,6 +3305,10 @@ struct LTXAVEmbedder : public Conditioner {
auto& weights = std::get<1>(tokens_weights_mask);
auto& mask = std::get<2>(tokens_weights_mask);
if (tokens.empty()) {
return {};
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, std::vector<int32_t>(tokens.begin(), tokens.end()));
sd::Tensor<float> attention_mask;
if (!mask.empty()) {
@@ -3133,7 +3407,9 @@ struct LTXAVEmbedder : public Conditioner {
prompt_attn_range.second = static_cast<int>(prompt.size());
auto hidden_states = encode_prompt(n_threads, prompt, prompt_attn_range);
GGML_ASSERT(!hidden_states.empty());
if (hidden_states.empty()) {
return {};
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0);
+103
View File
@@ -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
+32
View File
@@ -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__
+134 -39
View File
@@ -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,7 @@ 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) {
+18 -4
View File
@@ -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() {
+1 -1
View File
@@ -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_; }
+73 -4
View File
@@ -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
View File
@@ -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: [OCIC, 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,
+14 -2
View File
@@ -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";
}
+2
View File
@@ -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);
+14 -2
View File
@@ -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;
+31 -2
View File
@@ -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);
}
+19
View File
@@ -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;
}
+171
View File
@@ -0,0 +1,171 @@
#include "regex.h"
#include <cstdint>
#include <limits>
#include <mutex>
#define ONIG_ESCAPE_UCHAR_COLLISION
#define ONIG_ESCAPE_REGEX_T_COLLISION
#include <oniguruma.h>
namespace sd {
struct Regex::Impl {
OnigRegex regex = nullptr;
~Impl() {
if (regex) {
onig_free(regex);
}
}
};
struct RegexRegionDeleter {
void operator()(OnigRegion* region) const {
onig_region_free(region, 1);
}
};
static bool regex_error(std::string* error, const std::string& message) {
if (error) {
*error = message;
}
return false;
}
static bool regex_onig_error(std::string* error, int code, OnigErrorInfo* info = nullptr) {
OnigUChar buffer[ONIG_MAX_ERROR_MESSAGE_LEN];
onig_error_code_to_str(buffer, code, info);
return regex_error(error, reinterpret_cast<const char*>(buffer));
}
static int regex_initialize() {
static std::once_flag once;
static int result = ONIG_NORMAL;
std::call_once(once, [] {
OnigEncoding encodings[] = {ONIG_ENCODING_UTF8};
result = onig_initialize(encodings, 1);
});
// onig_end() would invalidate expressions held by other Regex instances.
return result;
}
// Rust str excludes overlong encodings, surrogates and extended UTF-8 accepted by Oniguruma.
static bool regex_valid_utf8(const std::string& text) {
size_t position = 0;
while (position < text.size()) {
const auto lead = static_cast<unsigned char>(text[position++]);
if (lead < 0x80) {
continue;
}
int count = 0;
if (lead >= 0xC2 && lead <= 0xDF) {
count = 1;
} else if (lead >= 0xE0 && lead <= 0xEF) {
count = 2;
} else if (lead >= 0xF0 && lead <= 0xF4) {
count = 3;
}
if (count == 0 || text.size() - position < static_cast<size_t>(count)) {
return false;
}
uint32_t codepoint = lead & (0x7F >> count);
for (int i = 0; i < count; ++i) {
const auto byte = static_cast<unsigned char>(text[position++]);
if ((byte & 0xC0) != 0x80) {
return false;
}
codepoint = (codepoint << 6) | (byte & 0x3F);
}
constexpr uint32_t minimum[] = {0, 0x80, 0x800, 0x10000};
if (codepoint < minimum[count] || codepoint > 0x10FFFF || (codepoint >= 0xD800 && codepoint <= 0xDFFF)) {
return false;
}
}
return true;
}
Regex::Regex() = default;
Regex::~Regex() = default;
Regex::Regex(Regex&&) noexcept = default;
Regex& Regex::operator=(Regex&&) noexcept = default;
bool Regex::compile(const std::string& pattern, std::string* error) {
if (error) {
error->clear();
}
const int initialized = regex_initialize();
if (initialized != ONIG_NORMAL) {
return regex_onig_error(error, initialized);
}
if (pattern.size() > static_cast<size_t>(std::numeric_limits<int>::max())) {
return regex_error(error, "regex pattern exceeds Oniguruma's offset range");
}
const auto* begin = reinterpret_cast<const OnigUChar*>(pattern.data());
const auto* end = begin + pattern.size();
if (!regex_valid_utf8(pattern)) {
return regex_error(error, "regex pattern is not valid UTF-8");
}
auto next = std::make_unique<Impl>();
OnigErrorInfo info{};
static std::mutex compile_mutex;
std::lock_guard<std::mutex> lock(compile_mutex);
const int result = onig_new(&next->regex, begin, end, ONIG_OPTION_NONE,
ONIG_ENCODING_UTF8, ONIG_SYNTAX_ONIGURUMA, &info);
if (result != ONIG_NORMAL) {
return regex_onig_error(error, result, &info);
}
impl_ = std::move(next);
return true;
}
bool Regex::find_matches(const std::string& text, std::vector<Match>& matches, std::string* error) const {
matches.clear();
if (error) {
error->clear();
}
if (!impl_) {
return regex_error(error, "regex has not been compiled");
}
if (text.size() > static_cast<size_t>(std::numeric_limits<int>::max())) {
return regex_error(error, "regex input exceeds Oniguruma's offset range");
}
const auto* begin = reinterpret_cast<const OnigUChar*>(text.data());
const auto* end = begin + text.size();
if (!regex_valid_utf8(text)) {
return regex_error(error, "regex input is not valid UTF-8");
}
std::unique_ptr<OnigRegion, RegexRegionDeleter> region(onig_region_new());
if (!region) {
return regex_error(error, "failed to allocate regex match region");
}
size_t position = 0;
while (position <= text.size()) {
const int result = onig_search(impl_->regex, begin, end, begin + position, end,
region.get(), ONIG_OPTION_NONE);
if (result == ONIG_MISMATCH) {
break;
}
if (result < 0) {
matches.clear();
return regex_onig_error(error, result);
}
const size_t match_begin = static_cast<size_t>(region->beg[0]);
const size_t match_end = static_cast<size_t>(region->end[0]);
// Match rust-onig's find_iter: suppress an empty match at the previous match's end.
if (match_begin == match_end && !matches.empty() && matches.back().second == match_end) {
if (position == text.size()) {
break;
}
position += static_cast<size_t>(ONIGENC_MBC_ENC_LEN(ONIG_ENCODING_UTF8, begin + position));
continue;
}
matches.emplace_back(match_begin, match_end);
position = match_end;
}
return true;
}
} // namespace sd
+32
View File
@@ -0,0 +1,32 @@
#ifndef __SD_CORE_REGEX_H__
#define __SD_CORE_REGEX_H__
#include <cstddef>
#include <memory>
#include <string>
#include <utility>
#include <vector>
namespace sd {
class Regex {
struct Impl;
std::unique_ptr<Impl> impl_;
public:
using Match = std::pair<size_t, size_t>;
Regex();
~Regex();
Regex(Regex&&) noexcept;
Regex& operator=(Regex&&) noexcept;
// Failed compilation leaves the previous expression intact.
bool compile(const std::string& pattern, std::string* error = nullptr);
// Matches are non-overlapping UTF-8 byte ranges; each call owns its search state.
bool find_matches(const std::string& text, std::vector<Match>& matches, std::string* error = nullptr) const;
};
} // namespace sd
#endif // __SD_CORE_REGEX_H__
+7
View File
@@ -1,6 +1,8 @@
#ifndef __SD_CORE_RNG_HPP__
#define __SD_CORE_RNG_HPP__
#include <cstdint>
#include <memory>
#include <random>
#include <vector>
@@ -8,6 +10,7 @@ class RNG {
public:
virtual void manual_seed(uint64_t seed) = 0;
virtual std::vector<float> randn(uint32_t n) = 0;
virtual std::shared_ptr<RNG> clone() const = 0;
};
class STDDefaultRNG : public RNG {
@@ -15,6 +18,10 @@ private:
std::default_random_engine generator;
public:
std::shared_ptr<RNG> clone() const override {
return std::make_shared<STDDefaultRNG>(*this);
}
void manual_seed(uint64_t seed) override {
generator.seed((unsigned int)seed);
}
+6
View File
@@ -1,7 +1,9 @@
#ifndef __SD_CORE_RNG_MT19937_HPP__
#define __SD_CORE_RNG_MT19937_HPP__
#include <array>
#include <cmath>
#include <limits>
#include <vector>
#include "core/rng.hpp"
@@ -123,6 +125,10 @@ class MT19937RNG : public RNG {
public:
MT19937RNG(uint64_t seed = 0) { manual_seed(seed); }
std::shared_ptr<RNG> clone() const override {
return std::make_shared<MT19937RNG>(*this);
}
void manual_seed(uint64_t seed) override {
s.seed_ = seed;
s.seeded_ = true;
+4
View File
@@ -93,6 +93,10 @@ public:
this->offset = 0;
}
std::shared_ptr<RNG> clone() const override {
return std::make_shared<PhiloxRNG>(*this);
}
void manual_seed(uint64_t seed) override {
this->seed = seed;
this->offset = 0;
+45 -8
View File
@@ -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) {
+65 -30
View File
@@ -18,12 +18,15 @@ tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
auto tokens_and_weights = clip_conditioner.tokenize(text);
std::vector<int> source_tokens = std::move(tokens_and_weights.first);
std::vector<float> source_weights = std::move(tokens_and_weights.second);
if (source_tokens.empty()) {
return {};
}
if (!source_tokens.empty() && source_tokens.front() == clip_conditioner.tokenizer.BOS_TOKEN_ID) {
if (!source_tokens.empty() && source_tokens.front() == clip_conditioner.tokenizer->BOS_TOKEN_ID) {
source_tokens.erase(source_tokens.begin());
source_weights.erase(source_weights.begin());
}
if (!source_tokens.empty() && source_tokens.back() == clip_conditioner.tokenizer.EOS_TOKEN_ID) {
if (!source_tokens.empty() && source_tokens.back() == clip_conditioner.tokenizer->EOS_TOKEN_ID) {
source_tokens.pop_back();
source_weights.pop_back();
}
@@ -49,12 +52,12 @@ tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
weights.push_back(source_weights[i]);
}
clip_conditioner.tokenizer.pad_tokens(tokens,
&weights,
nullptr,
clip_conditioner.text_model->model.n_token,
clip_conditioner.text_model->model.n_token,
true);
clip_conditioner.tokenizer->pad_tokens(tokens,
&weights,
nullptr,
clip_conditioner.text_model->model.n_token,
clip_conditioner.text_model->model.n_token,
true);
std::vector<bool> class_token_mask;
for (int i = 0; i < tokens.size(); i++) {
class_token_mask.push_back(class_idx >= 0 && class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
@@ -69,8 +72,14 @@ get_photomaker_condition_with_trigger(FrozenCLIPEmbedderWithCustomWords& clip_co
const ConditionerParams& conditioner_params,
const std::string& trigger_word,
int trigger_token_count) {
auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word);
GGML_ASSERT(image_tokens.size() == 1);
std::vector<int> image_tokens;
if (!clip_conditioner.convert_token_to_id(trigger_word, image_tokens)) {
return {};
}
if (image_tokens.size() != 1) {
LOG_ERROR("PhotoMaker trigger word must encode to one token");
return {};
}
auto tokens_and_weights = tokenize_photomaker_trigger(clip_conditioner,
conditioner_params.text,
trigger_token_count,
@@ -78,27 +87,43 @@ get_photomaker_condition_with_trigger(FrozenCLIPEmbedderWithCustomWords& clip_co
std::vector<int>& tokens = std::get<0>(tokens_and_weights);
std::vector<float>& weights = std::get<1>(tokens_and_weights);
std::vector<bool>& trigger_mask = std::get<2>(tokens_and_weights);
auto cond = clip_conditioner.get_learned_condition_common(n_threads,
tokens,
weights,
conditioner_params.clip_skip,
conditioner_params.width,
conditioner_params.height,
conditioner_params.zero_out_masked);
if (tokens.empty()) {
return {};
}
auto cond = clip_conditioner.get_learned_condition_common(n_threads,
tokens,
weights,
conditioner_params.clip_skip,
conditioner_params.width,
conditioner_params.height,
conditioner_params.zero_out_masked);
return std::make_tuple(std::move(cond), trigger_mask);
}
static std::string remove_photomaker_trigger_from_prompt(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
const std::string& prompt,
const std::string& trigger_word) {
auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word);
GGML_ASSERT(image_tokens.size() == 1);
static bool remove_photomaker_trigger_from_prompt(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
const std::string& prompt,
const std::string& trigger_word,
std::string& result) {
std::vector<int> image_tokens;
if (!clip_conditioner.convert_token_to_id(trigger_word, image_tokens)) {
return false;
}
if (image_tokens.size() != 1) {
LOG_ERROR("PhotoMaker trigger word must encode to one token");
return false;
}
auto tokens_and_weights = clip_conditioner.tokenize(prompt);
std::vector<int>& tokens = tokens_and_weights.first;
auto it = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
GGML_ASSERT(it != tokens.end());
if (tokens.empty()) {
return false;
}
auto it = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
if (it == tokens.end()) {
LOG_ERROR("PhotoMaker trigger word was not found in tokenized prompt");
return false;
}
tokens.erase(it);
return clip_conditioner.decode(tokens);
return clip_conditioner.decode(tokens, result);
}
struct PhotoMakerExtension : public GenerationExtension {
@@ -135,6 +160,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");
}
@@ -222,6 +248,10 @@ struct PhotoMakerExtension : public GenerationExtension {
trigger_token_count);
SDCondition prepared_id_condition = std::get<0>(cond_tup);
auto class_tokens_mask = std::get<1>(cond_tup);
if (prepared_id_condition.empty()) {
LOG_ERROR("failed to encode PhotoMaker prompt");
return false;
}
if (std::find(class_tokens_mask.begin(), class_tokens_mask.end(), true) == class_tokens_mask.end()) {
LOG_WARN("PhotoMaker trigger word '%s' was not found in prompt", trigger_word.c_str());
LOG_WARN("Turn off PhotoMaker for this request");
@@ -262,11 +292,16 @@ struct PhotoMakerExtension : public GenerationExtension {
prepared_id_condition.c_crossattn = std::move(res);
int64_t t1 = ggml_time_ms();
id_condition = std::move(prepared_id_condition);
start_merge_step = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps);
ctx.condition_params.text = remove_photomaker_trigger_from_prompt(*clip_conditioner,
ctx.condition_params.text,
trigger_word);
std::string prompt;
if (!remove_photomaker_trigger_from_prompt(*clip_conditioner,
ctx.condition_params.text,
trigger_word,
prompt)) {
return false;
}
id_condition = std::move(prepared_id_condition);
start_merge_step = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps);
ctx.condition_params.text = std::move(prompt);
LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0);
LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step);
+9 -2
View File
@@ -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;
+1 -1
View File
@@ -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);
+5 -6
View File
@@ -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;
+413
View File
@@ -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__
+2 -2
View File
@@ -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));
}
+58 -2
View File
@@ -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;
@@ -367,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;
@@ -765,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);
}
};
@@ -870,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;
+4 -3
View File
@@ -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
+2 -4
View File
@@ -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,
+1 -1
View File
@@ -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);
}
};
+1 -1
View File
@@ -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;
}
+26 -6
View File
@@ -484,13 +484,19 @@ namespace HiDreamO1 {
};
struct HiDreamO1Conditioner : public Conditioner {
Qwen2Tokenizer tokenizer;
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<HiDreamO1VisionRunner> vision_runner;
HiDreamO1Conditioner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, tensor_storage_map, "model.visual", weight_manager)) {}
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, tensor_storage_map, "model.visual", weight_manager)) {
tokenizer = tokenizers.create(TokenizerConfig::MAIN, HiDreamO1Config::detect_from_weights(tensor_storage_map, "").llm.vocab_size, 151643);
if (!tokenizer) {
tokenizer = std::make_shared<Qwen2Tokenizer>();
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
vision_runner->get_param_tensors(tensors);
@@ -504,6 +510,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);
}
@@ -534,7 +544,10 @@ namespace HiDreamO1 {
if (ref_images.empty()) {
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
auto input_ids = tokenizer.encode(prompt, nullptr);
std::vector<int> input_ids;
if (!tokenizer->encode(prompt, input_ids, nullptr)) {
return {};
}
std::vector<int32_t> input_ids_pad = input_ids;
input_ids_pad.push_back(VISION_START_TOKEN_ID);
@@ -608,7 +621,11 @@ namespace HiDreamO1 {
auto patch_img = resized_ref * 2.0f - 1.0f;
result.c_ref_images.push_back(std::move(patch_img));
int64_t prompt_start = static_cast<int64_t>(tokenizer.encode(prompt + "<|vision_start|>", nullptr).size());
std::vector<int> prefix_tokens;
if (!tokenizer->encode(prompt + "<|vision_start|>", prefix_tokens, nullptr)) {
return {};
}
int64_t prompt_start = static_cast<int64_t>(prefix_tokens.size());
prompt += "<|vision_start|>";
prompt += repeat_special_token("<|image_pad|>", image_tokens);
prompt += "<|vision_end|>";
@@ -619,7 +636,10 @@ namespace HiDreamO1 {
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
auto input_ids = tokenizer.encode(prompt, nullptr);
std::vector<int> input_ids;
if (!tokenizer->encode(prompt, input_ids, nullptr)) {
return {};
}
std::vector<int32_t> input_ids_pad = input_ids;
input_ids_pad.push_back(VISION_START_TOKEN_ID);
+1 -1
View File
@@ -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
+1 -2
View File
@@ -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,
+1 -2
View File
@@ -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,
+1 -2
View File
@@ -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,
+7 -7
View File
@@ -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,
+7
View File
@@ -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 {
+846
View File
@@ -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,
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__
+175 -49
View File
@@ -1,6 +1,7 @@
#ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
#define __SD_MODEL_DIFFUSION_WAN_HPP__
#include <algorithm>
#include <cinttypes>
#include <map>
#include <memory>
@@ -19,25 +20,30 @@ namespace WAN {
constexpr int WAN_GRAPH_SIZE = 10240;
struct WanConfig {
std::string model_type = "t2v";
std::tuple<int, int, int> patch_size = {1, 2, 2};
int64_t text_len = 512;
int64_t in_dim = 16;
int64_t dim = 2048;
int64_t ffn_dim = 8192;
int freq_dim = 256;
int64_t text_dim = 4096;
int64_t out_dim = 16;
int64_t num_heads = 16;
int num_layers = 32;
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;
std::string model_type = "t2v";
std::tuple<int, int, int> patch_size = {1, 2, 2};
int64_t text_len = 512;
int64_t in_dim = 16;
int64_t dim = 2048;
int64_t ffn_dim = 8192;
int freq_dim = 256;
int64_t text_dim = 4096;
int64_t out_dim = 16;
int64_t num_heads = 16;
int num_layers = 32;
int vace_layers = 0;
int64_t vace_in_dim = 96;
std::map<int, int> vace_layers_mapping = {};
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() {
+215
View File
@@ -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__
+47 -17
View File
@@ -15,6 +15,7 @@
#include <regex>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
@@ -31,9 +32,9 @@
#include "model_manager.h"
#include "tokenizers/bpe_tokenizer.h"
#include "tokenizers/gemma_tokenizer.h"
#include "tokenizers/gpt_oss_tokenizer.h"
#include "tokenizers/mistral_tokenizer.h"
#include "tokenizers/qwen2_tokenizer.h"
#include "tokenizers/tokenizer_config.h"
namespace LLM {
constexpr int LLM_GRAPH_SIZE = 65536;
@@ -139,7 +140,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 +232,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 +280,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 +334,20 @@ 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 +1377,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 +1904,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) {
@@ -2338,7 +2352,7 @@ namespace LLM {
};
struct LLMEmbedder {
std::shared_ptr<BPETokenizer> tokenizer;
std::shared_ptr<Tokenizer> tokenizer;
LLMRunner model;
LLMEmbedder(LLMArch arch,
@@ -2346,14 +2360,27 @@ namespace LLM {
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
bool enable_vision = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: model(arch, backend, tensor_storage_map, prefix, enable_vision, weight_manager) {
int pad_id = 151643;
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
tokenizer = std::make_shared<MistralTokenizer>();
pad_id = 11;
} else if (arch == LLMArch::GPT_OSS_20B) {
tokenizer = std::make_shared<GPTOSSTokenizer>();
} else {
tokenizer = std::make_shared<Qwen2Tokenizer>();
pad_id = 199999;
} else if (arch == LLMArch::GEMMA2_2B) {
pad_id = 0;
}
tokenizer = tokenizers.create(TokenizerConfig::MAIN, model.config.vocab_size, pad_id);
if (!tokenizer) {
if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B) {
throw std::runtime_error("GPT-OSS and Gemma 2 require an external tokenizer.json in the main tokenizer slot");
}
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
tokenizer = std::make_shared<MistralTokenizer>();
} else {
tokenizer = std::make_shared<Qwen2Tokenizer>();
}
}
}
@@ -2389,7 +2416,10 @@ namespace LLM {
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens = tokenizer->tokenize(curr_text, nullptr);
std::vector<int> curr_tokens;
if (!tokenizer->tokenize(curr_text, curr_tokens, nullptr)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
+5 -2
View File
@@ -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};
@@ -567,7 +567,10 @@ struct T5Embedder {
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens = tokenizer.encode(curr_text);
std::vector<int> curr_tokens;
if (!tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
+1 -1
View File
@@ -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]
+1 -1
View File
@@ -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]
+1 -1
View File
@@ -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);
+1 -2
View File
@@ -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,
+1 -2
View File
@@ -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,
+1 -1
View File
@@ -564,7 +564,7 @@ public:
int64_t chunk_frames = 5 * decoder->t_upscale;
int64_t pad = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames;
result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, pad, false, false);
result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, static_cast<int>(pad), false, false);
int64_t num_chunks = (num_frames + pad) / chunk_frames;
auto to_trim = decoder->t_upscale - 1;
+1 -1
View File
@@ -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;
+2 -2
View File
@@ -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]
+3
View File
@@ -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:
+5
View File
@@ -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) {
+23
View File
@@ -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
View File
@@ -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
View File
@@ -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;
+79 -5
View File
@@ -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)) {
+7
View File
@@ -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;
@@ -124,6 +128,7 @@ public:
&sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path,
&sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path,
&sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path,
&sd_ctx_params_t::tokenizer,
&sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path,
&sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path,
&sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path,
@@ -363,6 +368,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,
+3 -1
View File
@@ -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,
+15
View File
@@ -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
}
+92 -23
View File
@@ -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"
@@ -71,12 +73,13 @@ namespace sd::model_builders {
}
}
bool build_core_runners(const Context& ctx, CoreRunners& runners) {
bool build_core_runners(const Context& ctx, CoreRunners& runners) try {
const auto* sd_ctx_params = &ctx.params;
const auto& tensor_storage_map = ctx.tensor_storage_map;
const auto version = ctx.version;
const auto& weight_manager = ctx.weight_manager;
CoreRunners result;
TokenizerConfig tokenizers(sd_ctx_params->tokenizer);
if (!ensure_backend_pair(ctx.backends, SDBackendModule::TE) ||
!ensure_backend_pair(ctx.backends, SDBackendModule::DIFFUSION)) {
return false;
@@ -85,7 +88,8 @@ namespace sd::model_builders {
if (sd_version_is_sd3(version)) {
result.conditioner = std::make_shared<SD3CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<MMDiTRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -96,7 +100,8 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Pid::PiDRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model.net",
@@ -107,7 +112,8 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Ideogram4::Ideogram4Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -118,7 +124,8 @@ namespace sd::model_builders {
version,
"",
true,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Krea2::Krea2Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -145,11 +152,13 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
} else {
result.conditioner = std::make_shared<FluxCLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
weight_manager,
tokenizers);
}
result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
@@ -164,7 +173,8 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -176,7 +186,8 @@ namespace sd::model_builders {
tensor_storage_map,
"text_encoders.llm",
"text_embedding_projection",
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<LTXV::LTXAVRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -187,7 +198,8 @@ namespace sd::model_builders {
version,
"",
true,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -198,7 +210,8 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Hunyuan::HunyuanVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -233,6 +246,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) {
@@ -246,7 +269,8 @@ namespace sd::model_builders {
version,
"",
enable_vision,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<LingBotVideo::LingBotVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -259,7 +283,8 @@ namespace sd::model_builders {
version,
"",
enable_vision,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -272,7 +297,8 @@ namespace sd::model_builders {
version,
"",
true,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<MageFlow::MageFlowRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -283,7 +309,8 @@ namespace sd::model_builders {
version,
"",
true,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -293,7 +320,8 @@ namespace sd::model_builders {
} else if (version == VERSION_HIDREAM_O1) {
result.conditioner = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<HiDreamO1::HiDreamO1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model",
@@ -306,10 +334,17 @@ 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,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Anima::AnimaRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -320,7 +355,8 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<ZImage::ZImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -332,7 +368,8 @@ namespace sd::model_builders {
version,
"",
true,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Boogu::BooguImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -344,7 +381,8 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<ErnieImage::ErnieImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -355,7 +393,8 @@ namespace sd::model_builders {
version,
"",
false,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<Lens::LensRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -369,7 +408,8 @@ namespace sd::model_builders {
tensor_storage_map,
embbeding_map,
version,
weight_manager);
weight_manager,
tokenizers);
result.diffusion = std::make_shared<UNetModelRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
@@ -396,8 +436,27 @@ 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);
}
tokenizers.validate_usage();
runners = std::move(result);
return true;
} catch (const std::exception& error) {
LOG_ERROR("failed to build model runners: %s", error.what());
return false;
}
bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) {
@@ -493,7 +552,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 +590,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 +620,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;
}
+4
View File
@@ -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 {
+4 -1
View File
@@ -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;
}
+174 -14
View File
@@ -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()) {
@@ -1052,10 +1139,10 @@ namespace sd::pipeline {
return latents;
}
static ImageGenerationEmbeds prepare_video_generation_embeds(StableDiffusionGGML* sd,
const sd_vid_gen_params_t* sd_vid_gen_params,
const GenerationRequest& request,
const ImageGenerationLatents& latents) {
static std::optional<ImageGenerationEmbeds> prepare_video_generation_embeds(StableDiffusionGGML* sd,
const sd_vid_gen_params_t* sd_vid_gen_params,
const GenerationRequest& request,
const ImageGenerationLatents& latents) {
ConditionerRunnerEndOnExit conditioner_runner_end{sd->cond_stage_model.get()};
ImageGenerationEmbeds embeds;
@@ -1072,8 +1159,12 @@ namespace sd::pipeline {
int64_t prepare_start_ms = ggml_time_ms();
embeds.cond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params);
embeds.cond.c_concat = latents.concat_latent;
embeds.cond.c_vector = latents.clip_vision_output;
if (embeds.cond.empty()) {
LOG_ERROR("failed to encode video prompt");
return std::nullopt;
}
embeds.cond.c_concat = latents.concat_latent;
embeds.cond.c_vector = latents.clip_vision_output;
if (sd_version_is_minimax_h3(sd->version)) {
embeds.cond.c_ref_images = latents.ref_latents;
embeds.cond.c_ref_audios = latents.reference_audio_latents;
@@ -1084,10 +1175,20 @@ 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,
condition_params);
condition_params.text = request.negative_prompt;
embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params);
if (embeds.uncond.empty()) {
LOG_ERROR("failed to encode negative video prompt");
return std::nullopt;
}
embeds.uncond.c_concat = latents.concat_latent;
embeds.uncond.c_vector = latents.clip_vision_output;
if (sd_version_is_minimax_h3(sd->version)) {
@@ -1096,6 +1197,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 +1529,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 +1548,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) {
@@ -1468,10 +1582,14 @@ namespace sd::pipeline {
}
ImageGenerationLatents latents = std::move(*latent_inputs_opt);
ImageGenerationEmbeds embeds = prepare_video_generation_embeds(sd,
sd_vid_gen_params,
request,
latents);
auto embeds_opt = prepare_video_generation_embeds(sd,
sd_vid_gen_params,
request,
latents);
if (!embeds_opt) {
return false;
}
ImageGenerationEmbeds embeds = std::move(*embeds_opt);
if (latent_upscale_enabled) {
LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale",
request.width,
@@ -1725,6 +1843,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 +1919,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 +1932,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 {
+97
View File
@@ -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
+22
View File
@@ -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__
+293 -192
View File
@@ -12,6 +12,8 @@
#include <utility>
#include "core/rng.hpp"
#include "core/rng_mt19937.hpp"
#include "core/rng_philox.hpp"
#include "core/tensor.hpp"
#include "core/util.h"
#include "model.h"
@@ -1488,6 +1490,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,
@@ -1556,12 +1634,18 @@ static std::tuple<float, float, float> get_ancestral_step(float sigma_from,
}
}
class NoiseSampler {
public:
virtual sd::Tensor<float> operator()(double sigma_from, double sigma_to) = 0;
virtual ~NoiseSampler() = default;
};
static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng = nullptr,
bool is_flow_denoiser = false,
float eta = 0.f) {
NoiseSampler& noise_sampler,
bool is_flow_denoiser = false,
float eta = 0.f) {
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
float sigma = sigmas[i];
@@ -1584,7 +1668,7 @@ static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
if (is_flow_denoiser) {
x *= alpha_scale;
}
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
x += noise_sampler(sigma, sigma_to) * sigma_up;
}
}
}
@@ -1702,7 +1786,7 @@ static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
float eta) {
auto t_fn = [](float sigma) -> float { return -log(sigma); };
auto sigma_fn = [](float t) -> float { return exp(-t); };
@@ -1734,7 +1818,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
}
if (sigmas[i + 1] > 0) {
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
x += noise_sampler(sigmas[i], sigmas[i + 1]) * sigma_up;
}
}
return x;
@@ -1743,7 +1827,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
float eta = 1.0f) {
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
@@ -1826,7 +1910,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
x = (x * sigma_down_i_ratio) + (D_i * (1.0f - sigma_down_i_ratio));
if (sigma_to > 0.0f && eta > 0.0f) {
x = alpha_scale * x + sd::Tensor<float>::randn_like(x, rng) * sigma_up;
x = alpha_scale * x + noise_sampler(sigma, sigma_to) * sigma_up;
}
}
}
@@ -1902,169 +1986,17 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
return x;
}
// DPM-Solver++(2M) SDE, midpoint variant. Ref: Lu et al. arXiv:2211.01095;
// k-diffusion sample_dpmpp_2m_sde.
// DPM-Solver++(2M) SDE, midpoint variant.
// Ref: Lu et al. arXiv:2211.01095; k-diffusion sample_dpmpp_2m_sde
static sd::Tensor<float> sample_dpmpp_2m_sde(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
float eta) {
sd::Tensor<float> old_denoised;
bool have_old_denoised = false;
float h_last = 0.f;
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1);
if (denoised_opt.pred.empty()) {
return {};
}
sd::Tensor<float> denoised = std::move(denoised_opt.pred);
if (sigmas[i + 1] == 0.f) {
x = denoised;
} else {
float t = -std::log(sigmas[i]);
float s = -std::log(sigmas[i + 1]);
float h = s - t;
float eta_h = eta * h;
float a = sigmas[i + 1] / sigmas[i] * std::exp(-eta_h);
float b = -std::expm1(-h - eta_h);
x = a * x + b * denoised;
if (have_old_denoised) {
float r = h_last / h;
x += (0.5f * b / r) * (denoised - old_denoised);
}
if (eta > 0.f) {
x += sd::Tensor<float>::randn_like(x, rng) * (sigmas[i + 1] * std::sqrt(-std::expm1(-2.f * eta_h)));
}
h_last = h;
}
old_denoised = denoised;
have_old_denoised = true;
}
return x;
}
// Seeded Brownian tree providing deterministic, step-count-stable Gaussian
// increments for stochastic samplers. Constructed once per generation; each
// call returns unit-variance noise for interval [sigma_a, sigma_b].
// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree.
class BrownianTreeNoiseSampler {
public:
BrownianTreeNoiseSampler(const sd::Tensor<float>& x_template,
double sigma_min,
double sigma_max,
uint64_t seed)
: t_min_(sigma_min),
t_max_(sigma_max),
shape_(x_template.shape()),
root_seed_(mix64(seed, 0x9E3779B97F4A7C15ULL)) {
auto rng = std::make_shared<STDDefaultRNG>();
rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL));
w_at_tmax_ = sd::Tensor<float>::randn(shape_, rng) * std::sqrt(static_cast<float>(t_max_ - t_min_));
}
sd::Tensor<float> operator()(double sigma_a, double sigma_b) {
double a = clamp(std::min(sigma_a, sigma_b));
double b = clamp(std::max(sigma_a, sigma_b));
auto dW = w(b) - w(a);
float span = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
return dW * (1.0f / std::sqrt(span));
}
private:
static constexpr int kMaxDepth = 24;
static uint64_t mix64(uint64_t v, uint64_t salt) {
uint64_t z = v + salt;
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
return z ^ (z >> 31);
}
double clamp(double t) const {
return std::min(std::max(t, t_min_), t_max_);
}
sd::Tensor<float> w(double t) {
auto it = cache_.find(t);
if (it != cache_.end()) {
return it->second;
}
sd::Tensor<float> zero = sd::Tensor<float>::zeros(shape_);
sd::Tensor<float> out = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth);
cache_.emplace(t, out);
return out;
}
sd::Tensor<float> bridge(double a,
double c,
const sd::Tensor<float>& w_a,
const sd::Tensor<float>& w_c,
double t,
uint64_t node_seed,
int depth) {
if (depth <= 0 || c - a < 1e-9) {
float alpha = (c > a) ? static_cast<float>((t - a) / (c - a)) : 0.5f;
return (1.0f - alpha) * w_a + alpha * w_c;
}
double m = 0.5 * (a + c);
double std_dev = std::sqrt((c - m) * (m - a) / (c - a));
auto rng = std::make_shared<STDDefaultRNG>();
rng->manual_seed(node_seed);
auto z = sd::Tensor<float>::randn(shape_, rng);
auto w_m = 0.5f * (w_a + w_c) + static_cast<float>(std_dev) * z;
if (t == m) {
return w_m;
}
if (t < m) {
return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1);
}
return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1);
}
double t_min_;
double t_max_;
std::vector<int64_t> shape_;
uint64_t root_seed_;
sd::Tensor<float> w_at_tmax_;
std::map<double, sd::Tensor<float>> cache_;
};
// DPM-Solver++(2M) SDE, midpoint variant, with step-count-stable Brownian-tree
// noise. Same trajectory shape at any step count for a given seed. Aliased in
// k-diffusion / ComfyUI as sample_dpmpp_2m_sde_gpu.
// Ref: Lu et al. arXiv:2211.01095; torchsde BrownianTree.
static sd::Tensor<float> sample_dpmpp_2m_sde_bt(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
double sigma_max = 0.0;
double sigma_min = std::numeric_limits<double>::infinity();
for (float s : sigmas) {
if (s > 0.0f) {
sigma_max = std::max(sigma_max, static_cast<double>(s));
sigma_min = std::min(sigma_min, static_cast<double>(s));
}
}
if (sigma_max <= sigma_min) {
return x;
}
uint64_t tree_seed = 0;
{
auto draw = rng->randn(2);
std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed));
}
BrownianTreeNoiseSampler noise_sampler(x, sigma_min, sigma_max, tree_seed);
sd::Tensor<float> old_denoised;
bool have_old_denoised = false;
float h_last = 0.f;
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1);
@@ -2105,7 +2037,7 @@ using SamplerExtraArgs = KeyValueArgs;
static sd::Tensor<float> sample_lcm(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
bool is_flow_denoiser,
const SamplerExtraArgs& extra_sample_args) {
struct LCMSampleArgs {
@@ -2158,7 +2090,7 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
if (is_flow_denoiser) {
x *= (1 - sigmas[i + 1]);
}
auto noise = sd::Tensor<float>::randn_like(x, rng);
auto noise = noise_sampler(sigmas[i], sigmas[i + 1]);
if (args.noise_clip_std > 0.0f && noise.numel() > 0) {
double mean = 0.0;
for (int64_t j = 0; j < noise.numel(); ++j) {
@@ -2276,7 +2208,7 @@ static sd::Tensor<float> sample_ipndm_v(denoise_cb_t model,
static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
bool is_flow_denoiser,
float eta) {
sd::Tensor<float> old_denoised = x;
@@ -2341,7 +2273,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
if (is_flow_denoiser) {
x *= alpha_scale;
}
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
x += noise_sampler(sigma_from, sigma_to) * sigma_up;
}
old_denoised = denoised;
@@ -2354,7 +2286,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
bool is_flow_denoiser,
float eta) {
const float c2 = 0.5f;
@@ -2417,7 +2349,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
if (is_flow_denoiser) {
x *= alpha_scale;
}
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
x += noise_sampler(sigma_from, sigma_to) * sigma_up;
}
}
return x;
@@ -2426,7 +2358,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
sd::Tensor<float> x,
std::vector<float> sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
bool is_flow_denoiser,
float eta) {
constexpr int max_stage = 3;
@@ -2548,7 +2480,7 @@ static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
float noise_scale_sq = er_lambda_t * er_lambda_t - er_lambda_s * er_lambda_s * r * r;
if (s_noise > 0.0f && noise_scale_sq > 0.0f) {
float noise_scale = alpha_t * std::sqrt(std::max(noise_scale_sq, 0.0f));
x += sd::Tensor<float>::randn_like(x, rng) * noise_scale;
x += noise_sampler(sigmas[i], sigmas[i + 1]) * noise_scale;
}
}
@@ -2561,7 +2493,7 @@ static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
static sd::Tensor<float> sample_tcd(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
float eta) {
float beta_start = 0.00085f;
float beta_end = 0.0120f;
@@ -2618,7 +2550,7 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
if (eta > 0 && sigma_to > 0.0f) {
x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * noise_sampler(sigma, sigma_to);
}
}
return x;
@@ -2713,9 +2645,9 @@ static sd::Tensor<float> sample_lms(denoise_cb_t model,
sd::Tensor<float> d_cur = (x - denoised) / sigma;
x += d_cur * lms_coeff[0];
if (max_order > 1) { // if max_order == 1, the history is not used (order always < 2)
int hist_size_p1 = hist.size() + 1;
int hist_size_p1 = static_cast<int>(hist.size()) + 1;
if (i) { // history does not exist at 1st step
int hist_max = hist.size() - 1;
int hist_max = static_cast<int>(hist.size()) - 1;
for (int c = 2; c <= order; c++)
x += hist[std::min(hist_max, hist_size_p1 - c + shift)] * lms_coeff[c - 1];
// max_order == 4 => hist[] index = 2, 1, 0
@@ -2753,7 +2685,7 @@ static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
float eta) {
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
@@ -2772,7 +2704,7 @@ static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
x = denoised + d * sigma_down;
if (sigmas[i + 1] > 0) {
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
x += noise_sampler(sigmas[i], sigmas[i + 1]) * sigma_up;
}
}
return x;
@@ -2782,7 +2714,7 @@ static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
NoiseSampler& noise_sampler,
bool is_flow_denoiser,
float eta,
const SamplerExtraArgs& extra_sample_args) {
@@ -2829,13 +2761,180 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
if (is_flow_denoiser) {
x *= alpha_scale;
}
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
x += noise_sampler(sigma, sigma_to) * sigma_up;
}
}
}
return x;
}
// independent and identically distributed Gaussian noise (default for most samplers)
class IIDGaussianNoiseSampler : public NoiseSampler {
public:
IIDGaussianNoiseSampler(const sd::Tensor<float>& x_template, std::shared_ptr<RNG> r)
: rng(std::move(r)), shape(x_template.shape()) {}
sd::Tensor<float> operator()(double sigma_from, double sigma_to) override {
(void)sigma_from;
(void)sigma_to;
return sd::Tensor<float>::randn(shape, rng);
}
private:
std::shared_ptr<RNG> rng;
std::vector<int64_t> shape;
};
// A fixed tree seed, shape and sigma range give consistent increments across
// interval subdivisions. Each query returns normalized Gaussian noise.
// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree.
class BrownianTreeNoiseSampler : public NoiseSampler {
public:
BrownianTreeNoiseSampler(const sd::Tensor<float>& x_template,
double sigma_min,
double sigma_max,
std::shared_ptr<RNG> seed_rng,
std::shared_ptr<RNG> node_rng)
: t_min_(sigma_min),
t_max_(sigma_max),
shape_(x_template.shape()),
seed_rng_(std::move(seed_rng)),
node_rng_(std::move(node_rng)) {}
sd::Tensor<float> operator()(double sigma_a, double sigma_b) override {
if (!initialized_) {
uint64_t seed = 0;
auto draw = seed_rng_->randn(2);
std::memcpy(&seed, draw.data(), sizeof(seed));
root_seed_ = mix64(seed, 0x9E3779B97F4A7C15ULL);
node_rng_->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL));
w_at_tmax_ = sd::Tensor<float>::randn(shape_, node_rng_) * std::sqrt(static_cast<float>(t_max_ - t_min_));
initialized_ = true;
}
double a = clamp(std::min(sigma_a, sigma_b));
double b = clamp(std::max(sigma_a, sigma_b));
auto dW = w(b) - w(a);
float span = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
return dW * (1.0f / std::sqrt(span));
}
private:
static constexpr int kMaxDepth = 24;
static uint64_t mix64(uint64_t v, uint64_t salt) {
uint64_t z = v + salt;
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
return z ^ (z >> 31);
}
double clamp(double t) const {
return std::min(std::max(t, t_min_), t_max_);
}
sd::Tensor<float> w(double t) {
auto it = cache_.find(t);
if (it != cache_.end()) {
return it->second;
}
sd::Tensor<float> zero = sd::Tensor<float>::zeros(shape_);
sd::Tensor<float> out = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth);
cache_.emplace(t, out);
return out;
}
sd::Tensor<float> bridge(double a,
double c,
const sd::Tensor<float>& w_a,
const sd::Tensor<float>& w_c,
double t,
uint64_t node_seed,
int depth) {
if (depth <= 0 || c - a < 1e-9) {
float alpha = (c > a) ? static_cast<float>((t - a) / (c - a)) : 0.5f;
return (1.0f - alpha) * w_a + alpha * w_c;
}
double m = 0.5 * (a + c);
double std_dev = std::sqrt((c - m) * (m - a) / (c - a));
node_rng_->manual_seed(node_seed);
auto z = sd::Tensor<float>::randn(shape_, node_rng_);
auto w_m = 0.5f * (w_a + w_c) + static_cast<float>(std_dev) * z;
if (t == m) {
return w_m;
}
if (t < m) {
return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1);
}
return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1);
}
double t_min_;
double t_max_;
std::vector<int64_t> shape_;
std::shared_ptr<RNG> seed_rng_;
std::shared_ptr<RNG> node_rng_;
uint64_t root_seed_ = 0;
bool initialized_ = false;
sd::Tensor<float> w_at_tmax_;
std::map<double, sd::Tensor<float>> cache_;
};
static std::unique_ptr<NoiseSampler> make_noise_sampler(const sd::Tensor<float>& x, std::shared_ptr<RNG> rng, sample_method_t method, const std::vector<float>& sigmas, const SamplerExtraArgs& extra_args) {
bool brownian_tree = (method == DPMPP2M_SDE_BT_SAMPLE_METHOD);
bool def_brownian_tree = brownian_tree;
std::string brownian_tree_rng = "cpu";
for (const auto& [key, value] : extra_args) {
if (key == "noise_sampler") {
if (value == "iid") {
brownian_tree = false;
} else if (value == "brownian_tree") {
brownian_tree = true;
} else {
LOG_WARN("unknown noise_sampler value '%s'; using default", value.c_str());
}
} else if (key == "brownian_tree_rng") {
if (value == "cpu" || value == "cuda" || value == "std_default" || value == "sampler_rng") {
brownian_tree_rng = value;
} else {
LOG_WARN("ignoring invalid brownian_tree_rng value '%s'; expected cpu, cuda, std_default or sampler_rng", value.c_str());
}
}
}
if (brownian_tree) {
double sigma_max = 0.0;
double sigma_min = std::numeric_limits<double>::infinity();
for (float s : sigmas) {
if (s > 0.0f) {
sigma_max = std::max(sigma_max, static_cast<double>(s));
sigma_min = std::min(sigma_min, static_cast<double>(s));
}
}
if (sigma_max > sigma_min) {
std::shared_ptr<RNG> node_rng;
if (brownian_tree_rng == "sampler_rng") {
node_rng = rng->clone();
} else if (brownian_tree_rng == "std_default") {
node_rng = std::make_shared<STDDefaultRNG>();
} else if (brownian_tree_rng == "cuda") {
node_rng = std::make_shared<PhiloxRNG>();
} else {
node_rng = std::make_shared<MT19937RNG>();
}
if (!def_brownian_tree) {
LOG_INFO("setting noise sampler to Brownian tree (%s RNG)", brownian_tree_rng.c_str());
}
return std::make_unique<BrownianTreeNoiseSampler>(x, sigma_min, sigma_max, rng, std::move(node_rng));
}
}
if (def_brownian_tree) {
LOG_INFO("setting noise sampler to independent and identically distributed (iid)");
}
return std::make_unique<IIDGaussianNoiseSampler>(x, rng);
}
// k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
denoise_cb_t model,
@@ -2852,9 +2951,12 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
}
}
SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg");
std::unique_ptr<NoiseSampler> noise_sampler = make_noise_sampler(x, rng, method, sigmas, extra_args);
switch (method) {
case EULER_A_SAMPLE_METHOD:
return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
return sample_euler_ancestral(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
case EULER_SAMPLE_METHOD:
return sample_euler(model, std::move(x), sigmas);
case HEUN_SAMPLE_METHOD:
@@ -2863,42 +2965,41 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
return sample_dpm2(model, std::move(x), sigmas);
case DPMPP2S_A_SAMPLE_METHOD:
if (is_flow_denoiser)
return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, rng, eta);
return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, *noise_sampler, eta);
else
return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng, eta);
return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, *noise_sampler, eta);
case DPMPP2M_SAMPLE_METHOD:
return sample_dpmpp_2m(model, std::move(x), sigmas);
case DPMPP2Mv2_SAMPLE_METHOD:
return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
case LCM_SAMPLE_METHOD:
return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser, extra_args);
return sample_lcm(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, extra_args);
case IPNDM_SAMPLE_METHOD:
return sample_ipndm(model, std::move(x), sigmas);
case IPNDM_V_SAMPLE_METHOD:
return sample_ipndm_v(model, std::move(x), sigmas);
case RES_MULTISTEP_SAMPLE_METHOD:
return sample_res_multistep(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
return sample_res_multistep(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
case RES_2S_SAMPLE_METHOD:
return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
return sample_res_2s(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
case ER_SDE_SAMPLE_METHOD:
return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
return sample_er_sde(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
case DPMPP2M_SDE_SAMPLE_METHOD:
return sample_dpmpp_2m_sde(model, std::move(x), sigmas, rng, eta);
case DPMPP2M_SDE_BT_SAMPLE_METHOD:
return sample_dpmpp_2m_sde_bt(model, std::move(x), sigmas, rng, eta);
return sample_dpmpp_2m_sde(model, std::move(x), sigmas, *noise_sampler, eta);
case DDIM_TRAILING_SAMPLE_METHOD:
// DDIM is equivalent to Euler Ancestral with the Simple scheduler
return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
return sample_euler_ancestral(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
case TCD_SAMPLE_METHOD:
return sample_tcd(model, std::move(x), sigmas, rng, eta);
return sample_tcd(model, std::move(x), sigmas, *noise_sampler, eta);
case LMS_SAMPLE_METHOD:
return sample_lms(model, std::move(x), sigmas, extra_args);
case EULER_CFG_PP_SAMPLE_METHOD:
return sample_euler_cfg_pp(model, std::move(x), sigmas);
case EULER_A_CFG_PP_SAMPLE_METHOD:
return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, rng, eta);
return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, *noise_sampler, eta);
case EULER_GE_SAMPLE_METHOD:
return sample_gradient_estimation(model, std::move(x), sigmas, rng, is_flow_denoiser, eta, extra_args);
return sample_gradient_estimation(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta, extra_args);
default:
return {};
}
+27 -2
View File
@@ -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;
@@ -346,12 +349,14 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"t5xxl_path: %s\n"
"llm_path: %s\n"
"llm_vision_path: %s\n"
"tokenizer: %s\n"
"diffusion_model_path: %s\n"
"high_noise_diffusion_model_path: %s\n"
"uncond_diffusion_model_path: %s\n"
"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 +378,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),
@@ -381,12 +388,14 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
SAFE_STR(sd_ctx_params->t5xxl_path),
SAFE_STR(sd_ctx_params->llm_path),
SAFE_STR(sd_ctx_params->llm_vision_path),
SAFE_STR(sd_ctx_params->tokenizer),
SAFE_STR(sd_ctx_params->diffusion_model_path),
SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path),
SAFE_STR(sd_ctx_params->uncond_diffusion_model_path),
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 +417,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 +706,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 +740,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 +761,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) {
+7 -2
View File
@@ -130,7 +130,11 @@ std::vector<std::u32string> BPETokenizer::bpe(const std::u32string& token) const
return word;
}
std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t on_new_token_cb) {
bool BPETokenizer::encode(const std::string& text, std::vector<int>& result, on_new_token_cb_t on_new_token_cb, std::string* error) {
result.clear();
if (error) {
error->clear();
}
std::vector<int32_t> bpe_tokens;
std::vector<std::string> token_strs;
@@ -206,7 +210,8 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
}
ss << "]";
LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str());
return bpe_tokens;
result = std::move(bpe_tokens);
return true;
}
std::string BPETokenizer::decode_token(int token_id) const {
+1 -1
View File
@@ -36,7 +36,7 @@ public:
BPETokenizer() = default;
virtual ~BPETokenizer() = default;
std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) override;
bool encode(const std::string& text, std::vector<int>& tokens, on_new_token_cb_t on_new_token_cb = nullptr, std::string* error = nullptr) override;
};
#endif // __SD_TOKENIZERS_BPE_TOKENIZER_H__
-161
View File
@@ -189,164 +189,3 @@ GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::st
load_from_merges(load_gemma_merges(), load_gemma_vocab_json());
}
}
std::string Gemma2Tokenizer::normalize(const std::string& text) const {
std::string normalized = text;
size_t pos = 0;
while ((pos = normalized.find(' ', pos)) != std::string::npos) {
normalized.replace(pos, 1, "\xE2\x96\x81");
pos += 3;
}
return normalized;
}
void Gemma2Tokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
nlohmann::json vocab;
try {
vocab = nlohmann::json::parse(vocab_utf8_str);
} catch (const nlohmann::json::parse_error&) {
GGML_ABORT("invalid vocab json str");
}
for (const auto& [key, value] : vocab.items()) {
std::u32string token = utf8_to_utf32(key);
int i = value;
encoder[token] = i;
decoder[i] = token;
}
encoder_len = static_cast<int>(vocab.size());
LOG_VERBOSE("vocab size: %d", encoder_len);
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
for (const auto& merge : merges) {
size_t space_pos = merge.find(' ');
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
}
LOG_VERBOSE("merges size %zu", merge_pairs.size());
int rank = 0;
for (const auto& merge : merge_pairs) {
bpe_ranks[merge] = rank++;
}
bpe_len = rank;
}
Gemma2Tokenizer::Gemma2Tokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
byte_level_bpe = false;
byte_fallback = true;
add_bos_token = true;
PAD_TOKEN = "<pad>";
EOS_TOKEN = "<eos>";
BOS_TOKEN = "<bos>";
UNK_TOKEN = "<unk>";
PAD_TOKEN_ID = 0;
EOS_TOKEN_ID = 1;
BOS_TOKEN_ID = 2;
UNK_TOKEN_ID = 3;
std::vector<std::string> special_tokens_before_merge = {
PAD_TOKEN,
EOS_TOKEN,
BOS_TOKEN,
UNK_TOKEN,
"<mask>",
"<2mass>",
"[@BOS@]",
};
for (int i = 0; i <= 98; i++) {
special_tokens_before_merge.push_back("<unused" + std::to_string(i) + ">");
}
special_tokens_before_merge.push_back("<start_of_turn>");
special_tokens_before_merge.push_back("<end_of_turn>");
for (int i = 1; i <= 31; i++) {
special_tokens_before_merge.push_back(std::string(i, '\n'));
}
for (int i = 2; i <= 31; i++) {
std::string whitespace_token;
for (int j = 0; j < i; j++) {
whitespace_token += "\xE2\x96\x81";
}
special_tokens_before_merge.push_back(whitespace_token);
}
std::vector<std::string> html_tokens = {
"<table>",
"<caption>",
"<thead>",
"<tbody>",
"<tfoot>",
"<tr>",
"<th>",
"<td>",
"</table>",
"</caption>",
"</thead>",
"</tbody>",
"</tfoot>",
"</tr>",
"</th>",
"</td>",
"<h1>",
"<h2>",
"<h3>",
"<h4>",
"<h5>",
"<h6>",
"<blockquote>",
"</h1>",
"</h2>",
"</h3>",
"</h4>",
"</h5>",
"</h6>",
"</blockquote>",
"<strong>",
"<em>",
"<b>",
"<i>",
"<u>",
"<s>",
"<sub>",
"<sup>",
"<code>",
"</strong>",
"</em>",
"</b>",
"</i>",
"</u>",
"</s>",
"</sub>",
"</sup>",
"</code>",
};
special_tokens_before_merge.insert(special_tokens_before_merge.end(),
html_tokens.begin(),
html_tokens.end());
for (int i = 0; i <= 0xFF; i++) {
char hex_buf[16];
snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", i);
special_tokens_before_merge.push_back(hex_buf);
}
std::vector<std::string> special_tokens_after_merge = {
"[toxicity=0]",
};
for (int i = 1; i <= 31; i++) {
special_tokens_after_merge.insert(special_tokens_after_merge.begin() + i - 1,
std::string(i, '\t'));
}
for (int i = 99; i <= 99; i++) {
special_tokens_after_merge.push_back("<unused" + std::to_string(i) + ">");
}
special_tokens = special_tokens_before_merge;
special_tokens.insert(special_tokens.end(),
special_tokens_after_merge.begin(),
special_tokens_after_merge.end());
if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
load_from_merges(merges_utf8_str, vocab_utf8_str);
} else {
load_from_merges(load_gemma2_merges(), load_gemma2_vocab_json());
}
}
-9
View File
@@ -14,13 +14,4 @@ public:
explicit GemmaTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
};
class Gemma2Tokenizer : public BPETokenizer {
protected:
void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
std::string normalize(const std::string& text) const override;
public:
explicit Gemma2Tokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
};
#endif // __SD_TOKENIZERS_GEMMA_TOKENIZER_H__
-91
View File
@@ -1,91 +0,0 @@
#include "gpt_oss_tokenizer.h"
#include "core/util.h"
#include "json.hpp"
#include "vocab/vocab.h"
void GPTOSSTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
auto byte_unicode_pairs = bytes_to_unicode();
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
for (auto& pair : byte_unicode_pairs) {
byte_decoder[pair.second] = pair.first;
}
nlohmann::json vocab;
try {
vocab = nlohmann::json::parse(vocab_utf8_str);
} catch (const nlohmann::json::parse_error&) {
GGML_ABORT("invalid vocab json str");
}
for (const auto& [key, value] : vocab.items()) {
std::u32string token = utf8_to_utf32(key);
int i = value;
encoder[token] = i;
decoder[i] = token;
}
encoder_len = static_cast<int>(encoder.size());
for (auto& special_token : special_tokens) {
auto token = utf8_to_utf32(special_token);
encoder[token] = encoder_len;
decoder[encoder_len] = token;
encoder_len++;
}
encoder_len = static_cast<int>(encoder.size());
LOG_VERBOSE("vocab size: %d", encoder_len);
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
for (const auto& merge : merges) {
size_t space_pos = merge.find(' ');
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
}
LOG_VERBOSE("merges size %zu", merge_pairs.size());
int rank = 0;
for (const auto& merge : merge_pairs) {
bpe_ranks[merge] = rank++;
}
bpe_len = rank;
}
GPTOSSTokenizer::GPTOSSTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
BOS_TOKEN = "<|startoftext|>";
UNK_TOKEN = "<|endoftext|>";
EOS_TOKEN = "<|endoftext|>";
PAD_TOKEN = "<|endoftext|>";
BOS_TOKEN_ID = 199998;
EOS_TOKEN_ID = 199999;
UNK_TOKEN_ID = 199999;
PAD_TOKEN_ID = 199999;
special_tokens = {
"<|startoftext|>",
"<|endoftext|>",
"<|reserved_200000|>",
"<|reserved_200001|>",
"<|return|>",
"<|constrain|>",
"<|reserved_200004|>",
"<|channel|>",
"<|start|>",
"<|end|>",
"<|message|>",
"<|reserved_200009|>",
"<|reserved_200010|>",
"<|reserved_200011|>",
"<|call|>",
"<|reserved_200013|>",
"<|reserved_200014|>",
"<|reserved_200015|>",
"<|reserved_200016|>",
"<|reserved_200017|>",
"<|endofprompt|>",
};
if (merges_utf8_str.size() > 0) {
load_from_merges(merges_utf8_str, vocab_utf8_str);
} else {
load_from_merges(load_gpt_oss_merges(), load_gpt_oss_vocab_json());
}
}
-16
View File
@@ -1,16 +0,0 @@
#ifndef __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
#define __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
#include <string>
#include "bpe_tokenizer.h"
class GPTOSSTokenizer : public BPETokenizer {
protected:
void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
public:
explicit GPTOSSTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
};
#endif // __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
+856
View File
@@ -0,0 +1,856 @@
#include "hf_tokenizer.h"
#include <algorithm>
#include <array>
#include <climits>
#include <cstdlib>
#include <fstream>
#include <queue>
#include <sstream>
#include <stdexcept>
#include <unordered_map>
#include <unordered_set>
#include "core/regex.h"
#include "core/util.h"
#include "json.hpp"
#include "utf8proc.h"
using TokenizerJSON = nlohmann::json;
static void tokenizer_require(bool condition, const std::string& message) {
if (!condition) {
throw std::runtime_error("tokenizer.json: " + message);
}
}
static std::string tokenizer_utf8(int32_t codepoint) {
utf8proc_uint8_t bytes[4];
auto count = utf8proc_encode_char(codepoint, bytes);
return std::string(reinterpret_cast<const char*>(bytes), count);
}
static bool tokenizer_error(std::string* error, const std::string& message) {
if (error) {
*error = "tokenizer.json: " + message;
} else {
LOG_ERROR("tokenizer.json: %s", message.c_str());
}
return false;
}
static bool tokenizer_next(const std::string& text, size_t& offset, int32_t& codepoint, std::string* error) {
auto size = utf8proc_iterate(reinterpret_cast<const utf8proc_uint8_t*>(text.data() + offset), text.size() - offset, &codepoint);
if (size <= 0) {
return tokenizer_error(error, "invalid UTF-8 input");
}
offset += size;
return true;
}
static int tokenizer_id(const TokenizerJSON& value) {
tokenizer_require(value.is_number_integer(), "token ID must be an integer");
auto id = value.get<int64_t>();
tokenizer_require(id >= 0 && id <= INT_MAX, "token ID outside int32 range");
return static_cast<int>(id);
}
static uint64_t tokenizer_pair(int left, int right) {
return (static_cast<uint64_t>(left) << 32) | static_cast<uint32_t>(right);
}
struct HFTokenizer::Impl {
struct Pattern {
std::string literal;
std::shared_ptr<sd::Regex> regex;
explicit Pattern(const TokenizerJSON& config) {
tokenizer_require(config.is_object() && config.size() == 1, "invalid String/Regex pattern");
if (config.contains("String")) {
literal = config.at("String").get<std::string>();
} else {
tokenizer_require(config.contains("Regex"), "unsupported pattern");
regex = std::make_shared<sd::Regex>();
std::string error;
bool ok = regex->compile(config.at("Regex").get<std::string>(), &error);
tokenizer_require(ok, "invalid regex: " + error);
}
}
bool matches(const std::string& text, std::vector<sd::Regex::Match>& result, std::string* error) const {
result.clear();
if (regex) {
std::string regex_error;
if (!regex->find_matches(text, result, &regex_error)) {
return tokenizer_error(error, "regex search failed: " + regex_error);
}
} else if (literal.empty()) {
size_t offset = 0;
for (;;) {
result.emplace_back(offset, offset);
if (offset == text.size()) {
break;
}
int32_t cp;
if (!tokenizer_next(text, offset, cp, error)) {
result.clear();
return false;
}
}
} else {
size_t offset = 0;
while ((offset = text.find(literal, offset)) != std::string::npos) {
result.emplace_back(offset, offset + literal.size());
offset += literal.size();
}
}
return true;
}
bool replace(const std::string& text, const std::string& replacement, std::string& result, std::string* error) const {
result.clear();
size_t offset = 0;
std::vector<sd::Regex::Match> found;
if (!matches(text, found, error)) {
return false;
}
for (const auto& match : found) {
result.append(text, offset, match.first - offset);
result += replacement;
offset = match.second;
}
result.append(text, offset, std::string::npos);
return true;
}
bool split(const std::string& text, const std::string& behavior, bool invert, std::vector<std::string>& result, std::string* error) const {
result.clear();
struct Part {
size_t start, end;
bool matched;
};
std::vector<Part> parts;
size_t offset = 0;
std::vector<sd::Regex::Match> found;
if (!matches(text, found, error)) {
return false;
}
for (const auto& match : found) {
if (match.first > offset) {
parts.push_back({offset, match.first, invert});
}
parts.push_back({match.first, match.second, !invert});
offset = match.second;
}
if (offset < text.size()) {
parts.push_back({offset, text.size(), invert});
}
if (behavior == "MergedWithNext") {
std::reverse(parts.begin(), parts.end());
}
std::vector<Part> merged;
bool previous = false;
for (const auto& part : parts) {
bool join = (behavior == "Contiguous" && part.matched == previous) ||
((behavior == "MergedWithPrevious" || behavior == "MergedWithNext") && part.matched && !previous);
if (join && !merged.empty()) {
merged.back().start = std::min(merged.back().start, part.start);
merged.back().end = std::max(merged.back().end, part.end);
} else if (behavior != "Removed" || !part.matched) {
merged.push_back(part);
}
previous = part.matched;
}
if (behavior == "MergedWithNext") {
std::reverse(merged.begin(), merged.end());
}
for (const auto& part : merged) {
if (part.start != part.end) {
result.push_back(text.substr(part.start, part.end - part.start));
}
}
return true;
}
};
struct Step {
std::string type, content, behavior;
std::shared_ptr<Pattern> pattern;
bool invert = false, prefix_space = false;
};
struct Trie {
struct Node {
std::unordered_map<unsigned char, size_t> children;
int id = -1;
};
std::vector<Node> nodes{1};
void add(const std::string& text, int id) {
size_t index = 0;
for (unsigned char c : text) {
auto found = nodes[index].children.find(c);
if (found == nodes[index].children.end()) {
size_t next = nodes.size();
nodes[index].children.emplace(c, next);
nodes.emplace_back();
index = next;
} else {
index = found->second;
}
}
nodes[index].id = id;
}
std::pair<size_t, int> match(const std::string& text, size_t start) const {
size_t index = 0;
std::pair<size_t, int> result{start, -1};
for (size_t end = start; end < text.size(); ++end) {
auto found = nodes[index].children.find(static_cast<unsigned char>(text[end]));
if (found == nodes[index].children.end()) {
break;
}
index = found->second;
if (nodes[index].id >= 0) {
result = {end + 1, nodes[index].id};
}
}
return result;
}
};
struct Merge {
size_t rank;
int id;
};
std::unordered_map<std::string, int> vocab;
std::unordered_map<std::string, int> added_vocab;
std::unordered_map<int, std::string> tokens;
std::unordered_map<uint64_t, Merge> merges;
std::unordered_set<int> special_ids;
std::vector<std::string> custom_tokens;
std::vector<Step> normalizers, pre_tokenizers, decoders;
Trie raw_added, normalized_added;
std::array<std::string, 256> byte_encoder;
std::unordered_map<int32_t, unsigned char> byte_decoder;
std::string suffix;
int unk = -1;
bool fuse_unk = false, byte_fallback = false, ignore_merges = false, has_decoder = false;
Impl() {
int extra = 256;
for (int byte = 0; byte < 256; ++byte) {
int cp = ((byte >= 33 && byte <= 126) || (byte >= 161 && byte <= 172) || byte >= 174) ? byte : extra++;
byte_encoder[byte] = tokenizer_utf8(cp);
byte_decoder[cp] = static_cast<unsigned char>(byte);
}
}
static void parse_steps(const TokenizerJSON& config, const std::string& stage, std::vector<Step>& out, int depth = 0) {
tokenizer_require(depth < 32, stage + " nesting is too deep");
if (config.is_null()) {
return;
}
Step step;
step.type = config.at("type").get<std::string>();
if (step.type == "Sequence") {
const char* key = stage == "normalizer" ? "normalizers" : stage == "pre_tokenizer" ? "pretokenizers"
: "decoders";
for (const auto& child : config.at(key)) {
parse_steps(child, stage, out, depth + 1);
}
return;
}
if ((stage == "normalizer" || stage == "decoder") && step.type == "Replace") {
step.pattern = std::make_shared<Pattern>(config.at("pattern"));
step.content = config.at("content").get<std::string>();
} else if (stage == "normalizer" && (step.type == "NFC" || step.type == "Lowercase")) {
} else if (stage == "pre_tokenizer" && step.type == "Split") {
step.pattern = std::make_shared<Pattern>(config.at("pattern"));
step.behavior = config.at("behavior").get<std::string>();
tokenizer_require(step.behavior == "Removed" || step.behavior == "Isolated" || step.behavior == "Contiguous" || step.behavior == "MergedWithPrevious" || step.behavior == "MergedWithNext", "unsupported Split behavior: " + step.behavior);
step.invert = config.value("invert", false);
} else if ((stage == "pre_tokenizer" || stage == "decoder") && step.type == "ByteLevel") {
step.prefix_space = config.value("add_prefix_space", true);
if (stage == "pre_tokenizer" && config.value("use_regex", true)) {
step.pattern = std::make_shared<Pattern>(TokenizerJSON{{"Regex", R"('s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+)"}});
}
} else if (stage == "decoder" && (step.type == "ByteFallback" || step.type == "Fuse")) {
} else {
tokenizer_require(false, "unsupported " + stage + ": " + step.type);
}
out.push_back(std::move(step));
}
bool normalize(std::string text, std::string& result, std::string* error) const {
result.clear();
for (const auto& step : normalizers) {
if (step.type == "Replace") {
std::string replaced;
if (!step.pattern->replace(text, step.content, replaced, error)) {
return false;
}
text = std::move(replaced);
} else if (step.type == "NFC") {
utf8proc_uint8_t* output = nullptr;
auto size = utf8proc_map(reinterpret_cast<const utf8proc_uint8_t*>(text.data()), text.size(), &output, static_cast<utf8proc_option_t>(UTF8PROC_STABLE | UTF8PROC_COMPOSE));
std::unique_ptr<utf8proc_uint8_t, decltype(&std::free)> buffer(output, &std::free);
if (size < 0) {
return tokenizer_error(error, std::string("NFC normalization failed: ") + utf8proc_errmsg(size));
}
text.assign(reinterpret_cast<const char*>(output), size);
} else {
std::string lower;
for (size_t i = 0; i < text.size();) {
int32_t cp;
if (!tokenizer_next(text, i, cp, error)) {
return false;
}
// Rust char::to_lowercase uses full, context-free lowercase. U+0130 expands.
lower += cp == 0x130 ? "i\xcc\x87" : tokenizer_utf8(utf8proc_tolower(cp));
}
text = std::move(lower);
}
}
result = std::move(text);
return true;
}
bool pre_tokenize(const std::string& text, std::vector<std::string>& result, std::string* error) const {
result.clear();
std::vector<std::string> pieces{text};
for (const auto& step : pre_tokenizers) {
std::vector<std::string> next;
for (auto piece : pieces) {
if (piece.empty()) {
continue;
}
std::vector<std::string> split;
if (step.type == "Split") {
if (!step.pattern->split(piece, step.behavior, step.invert, split, error)) {
return false;
}
next.insert(next.end(), split.begin(), split.end());
} else {
if (step.prefix_space && piece.front() != ' ') {
piece.insert(piece.begin(), ' ');
}
if (step.pattern) {
if (!step.pattern->split(piece, "Isolated", false, split, error)) {
return false;
}
} else {
split.push_back(piece);
}
for (const auto& part : split) {
std::string encoded;
for (unsigned char byte : part) {
encoded += byte_encoder[byte];
}
next.push_back(std::move(encoded));
}
}
}
pieces = std::move(next);
}
result = std::move(pieces);
return true;
}
bool bpe(const std::string& text, std::vector<int>& ids, std::string* error) const {
ids.clear();
if (ignore_merges) {
auto found = vocab.find(text);
if (found != vocab.end()) {
ids.push_back(found->second);
return true;
}
}
bool pending_unk = false;
for (size_t i = 0; i < text.size();) {
int32_t cp;
size_t end = i;
if (!tokenizer_next(text, end, cp, error)) {
ids.clear();
return false;
}
std::string symbol = text.substr(i, end - i);
if (end == text.size()) {
symbol += suffix;
}
i = end;
auto found = vocab.find(symbol);
if (found != vocab.end()) {
if (pending_unk) {
ids.push_back(unk);
pending_unk = false;
}
ids.push_back(found->second);
continue;
}
if (byte_fallback) {
std::vector<int> bytes;
for (unsigned char byte : symbol) {
const char* hex = "0123456789ABCDEF";
std::string token = "<0x00>";
token[3] = hex[byte >> 4];
token[4] = hex[byte & 15];
auto fallback = vocab.find(token);
if (fallback == vocab.end()) {
break;
}
bytes.push_back(fallback->second);
}
if (bytes.size() == symbol.size()) {
ids.insert(ids.end(), bytes.begin(), bytes.end());
continue;
}
}
if (unk >= 0) {
if (pending_unk && !fuse_unk) {
ids.push_back(unk);
}
pending_unk = true;
}
}
if (pending_unk) {
ids.push_back(unk);
}
struct Symbol {
int id;
size_t prev, next, generation = 0;
bool alive = true;
};
struct Candidate {
size_t rank, left, right, left_generation, right_generation;
int id;
bool operator<(const Candidate& other) const {
return rank != other.rank ? rank > other.rank : left > other.left;
}
};
const size_t none = ids.size();
std::vector<Symbol> symbols;
for (size_t i = 0; i < ids.size(); ++i) {
symbols.push_back({ids[i], i == 0 ? none : i - 1, i + 1});
}
std::priority_queue<Candidate> queue;
auto push = [&](size_t left) {
if (left == none || symbols[left].next == none) {
return;
}
size_t right = symbols[left].next;
auto found = merges.find(tokenizer_pair(symbols[left].id, symbols[right].id));
if (found != merges.end()) {
queue.push({found->second.rank, left, right, symbols[left].generation, symbols[right].generation, found->second.id});
}
};
for (size_t i = 0; i < symbols.size(); ++i) {
push(i);
}
while (!queue.empty()) {
Candidate item = queue.top();
queue.pop();
auto& left = symbols[item.left];
auto& right = symbols[item.right];
if (!left.alive || !right.alive || left.next != item.right || left.generation != item.left_generation || right.generation != item.right_generation) {
continue;
}
left.id = item.id;
left.next = right.next;
++left.generation;
right.alive = false;
if (left.next != none) {
symbols[left.next].prev = item.left;
}
push(left.prev);
push(item.left);
}
ids.clear();
for (const auto& symbol : symbols) {
if (symbol.alive) {
ids.push_back(symbol.id);
}
}
return true;
}
int lookup(const std::string& token) const {
auto added = added_vocab.find(token);
if (added != added_vocab.end()) {
return added->second;
}
auto found = vocab.find(token);
return found == vocab.end() ? -1 : found->second;
}
void add_token(const std::string& token, int id, bool added = false) {
auto old = tokens.find(id);
tokenizer_require(old == tokens.end() || old->second == token, "conflicting token ID " + std::to_string(id));
int old_id = lookup(token);
tokenizer_require(old_id < 0 || old_id == id, "conflicting ID for token " + token);
tokens[id] = token;
(added ? added_vocab : vocab)[token] = id;
}
};
HFTokenizer::HFTokenizer(const std::string& path)
: impl_(new Impl) {
std::ifstream stream(path, std::ios::binary);
tokenizer_require(stream.good(), "cannot open " + path);
TokenizerJSON config;
stream >> config;
tokenizer_require(config.value("version", std::string("1.0")) == "1.0", "unsupported version");
tokenizer_require(config.value("padding", TokenizerJSON()).is_null(), "JSON padding is unsupported; padding is controlled by the text encoder");
tokenizer_require(config.value("truncation", TokenizerJSON()).is_null(), "JSON truncation is unsupported; truncation is controlled by the text encoder");
const auto& model = config.at("model");
tokenizer_require(model.at("type") == "BPE", "only BPE models are supported");
tokenizer_require(model.value("dropout", TokenizerJSON()).is_null() || model.at("dropout") == 0, "BPE dropout is unsupported");
const auto& prefix = model.value("continuing_subword_prefix", TokenizerJSON());
tokenizer_require(prefix.is_null() || prefix == "", "nonempty continuing_subword_prefix is unsupported");
const auto& suffix = model.value("end_of_word_suffix", TokenizerJSON());
impl_->suffix = suffix.is_null() ? "" : suffix.get<std::string>();
impl_->fuse_unk = model.value("fuse_unk", false);
impl_->byte_fallback = model.value("byte_fallback", false);
impl_->ignore_merges = model.value("ignore_merges", false);
tokenizer_require(model.at("vocab").is_object(), "BPE vocab must be an object");
impl_->vocab.reserve(model.at("vocab").size());
impl_->tokens.reserve(model.at("vocab").size());
for (const auto& entry : model.at("vocab").items()) {
impl_->add_token(entry.key(), tokenizer_id(entry.value()));
}
if (!model.value("unk_token", TokenizerJSON()).is_null()) {
UNK_TOKEN = model.at("unk_token").get<std::string>();
impl_->unk = impl_->lookup(UNK_TOKEN);
tokenizer_require(impl_->unk >= 0, "unk_token is absent from vocab");
}
UNK_TOKEN_ID = impl_->unk;
tokenizer_require(model.at("merges").is_array(), "BPE merges must be an array");
impl_->merges.reserve(model.at("merges").size());
size_t rank = 0;
for (const auto& merge : model.at("merges")) {
std::string left, right;
if (merge.is_string()) {
auto value = merge.get<std::string>();
auto space = value.find(' ');
tokenizer_require(space != std::string::npos && value.find(' ', space + 1) == std::string::npos, "invalid legacy BPE merge");
left = value.substr(0, space);
right = value.substr(space + 1);
} else {
tokenizer_require(merge.is_array() && merge.size() == 2, "BPE merge must contain two tokens");
left = merge.at(0).get<std::string>();
right = merge.at(1).get<std::string>();
}
int a = impl_->lookup(left), b = impl_->lookup(right), id = impl_->lookup(left + right);
tokenizer_require(a >= 0 && b >= 0 && id >= 0, "BPE merge references a missing vocab token");
impl_->merges[tokenizer_pair(a, b)] = {rank++, id};
}
Impl::parse_steps(config.value("normalizer", TokenizerJSON()), "normalizer", impl_->normalizers);
Impl::parse_steps(config.value("pre_tokenizer", TokenizerJSON()), "pre_tokenizer", impl_->pre_tokenizers);
impl_->has_decoder = !config.value("decoder", TokenizerJSON()).is_null();
Impl::parse_steps(config.value("decoder", TokenizerJSON()), "decoder", impl_->decoders);
size_t next_added_id = impl_->vocab.size();
for (const auto& token : config.value("added_tokens", TokenizerJSON::array())) {
for (const char* flag : {"single_word", "lstrip", "rstrip"}) {
tokenizer_require(!token.value(flag, false), std::string("added_tokens.") + flag + "=true is unsupported");
}
auto content = token.at("content").get<std::string>();
tokenizer_require(!content.empty(), "empty added token is unsupported");
int id = tokenizer_id(token.at("id"));
if (impl_->lookup(content) < 0) {
tokenizer_require(static_cast<size_t>(id) == next_added_id++, "nonconsecutive added token IDs would be reassigned by Hugging Face tokenizers");
}
impl_->add_token(content, id, true);
if (token.value("special", false)) {
special_tokens.push_back(content);
impl_->special_ids.insert(id);
}
bool normalized = token.value("normalized", true);
std::string pattern = content;
if (normalized) {
std::string error;
bool ok = impl_->normalize(content, pattern, &error);
tokenizer_require(ok, error);
}
tokenizer_require(!pattern.empty(), "added token normalizes to an empty string");
(normalized ? impl_->normalized_added : impl_->raw_added).add(pattern, id);
}
const auto& processor = config.value("post_processor", TokenizerJSON());
BOS_TOKEN_ID = EOS_TOKEN_ID = -1;
if (!processor.is_null()) {
auto type = processor.at("type").get<std::string>();
auto special = [&](const TokenizerJSON& pair) {
tokenizer_require(pair.is_array() && pair.size() == 2, "invalid postprocessor special token");
int id = tokenizer_id(pair.at(1));
tokenizer_require(impl_->lookup(pair.at(0).get<std::string>()) == id, "postprocessor token/ID does not match vocab");
return id;
};
if (type == "RobertaProcessing") {
BOS_TOKEN_ID = special(processor.at("cls"));
EOS_TOKEN_ID = special(processor.at("sep"));
} else if (type == "TemplateProcessing") {
bool seen_sequence = false;
for (const auto& item : processor.at("single")) {
if (item.contains("Sequence")) {
tokenizer_require(!seen_sequence && item.at("Sequence").at("id") == "A", "single template must contain exactly one sequence A");
seen_sequence = true;
} else {
auto name = item.at("SpecialToken").at("id").get<std::string>();
const auto& token = processor.at("special_tokens").at(name);
tokenizer_require(token.at("ids").size() == 1 && token.at("tokens").size() == 1, "multi-ID template special tokens are unsupported");
int id = special(TokenizerJSON::array({token.at("tokens").at(0), token.at("ids").at(0)}));
int& target = seen_sequence ? EOS_TOKEN_ID : BOS_TOKEN_ID;
tokenizer_require(target < 0, "single template supports at most one prefix and one suffix token");
target = id;
}
}
tokenizer_require(seen_sequence, "single template has no sequence A");
} else {
tokenizer_require(type == "ByteLevel", "unsupported post_processor: " + type);
}
}
add_bos_token = BOS_TOKEN_ID >= 0;
add_eos_token = EOS_TOKEN_ID >= 0;
BOS_TOKEN = decode_token(BOS_TOKEN_ID);
EOS_TOKEN = decode_token(EOS_TOKEN_ID);
set_padding(0, false);
}
HFTokenizer::~HFTokenizer() = default;
void HFTokenizer::set_padding(int token_id, bool left) {
PAD_TOKEN_ID = token_id;
PAD_TOKEN = decode_token(token_id);
pad_left = left;
}
void HFTokenizer::validate_vocab_size(int64_t embedding_rows) const {
tokenizer_require(embedding_rows > 0, "text encoder has no token embedding rows");
for (const auto& token : impl_->tokens) {
tokenizer_require(token.first < embedding_rows, "token ID " + std::to_string(token.first) + " exceeds text encoder vocabulary (" + std::to_string(embedding_rows) + ")");
}
tokenizer_require(PAD_TOKEN_ID >= 0 && PAD_TOKEN_ID < embedding_rows, "padding ID exceeds text encoder vocabulary");
}
int HFTokenizer::token_to_id(const std::string& token) const {
return impl_->lookup(token);
}
void HFTokenizer::add_special_token(const std::string& token) {
Tokenizer::add_special_token(token);
if (!token.empty()) {
impl_->custom_tokens.push_back(token);
}
}
bool HFTokenizer::encode(const std::string& text, std::vector<int>& tokens, on_new_token_cb_t callback, std::string* error) {
tokens.clear();
if (error) {
error->clear();
}
for (size_t i = 0; i < text.size();) {
int32_t cp;
if (!tokenizer_next(text, i, cp, error)) {
return false;
}
}
std::vector<int> result;
Impl::Trie raw_custom, normalized_custom;
if (callback) {
for (size_t index = 0; index < impl_->custom_tokens.size(); ++index) {
const auto& token = impl_->custom_tokens[index];
raw_custom.add(token, static_cast<int>(index));
std::string normalized;
if (!impl_->normalize(token, normalized, error)) {
return false;
}
if (!normalized.empty()) {
normalized_custom.add(normalized, static_cast<int>(index));
}
}
}
auto encode_plain = [&](const std::string& value) {
std::vector<std::string> pieces;
if (!impl_->pre_tokenize(value, pieces, error)) {
return false;
}
for (auto& piece : pieces) {
if (callback && callback(piece, result)) {
continue;
}
std::vector<int> ids;
if (!impl_->bpe(piece, ids, error)) {
return false;
}
result.insert(result.end(), ids.begin(), ids.end());
}
return true;
};
auto extract = [&](const std::string& value, const Impl::Trie& added, const Impl::Trie& custom_tokens, const auto& encode_gap) {
size_t start = 0, i = 0;
while (i < value.size()) {
auto match = added.match(value, i);
auto custom = custom_tokens.match(value, i);
bool use_custom = custom.second >= 0 && custom.first >= match.first;
if (match.second < 0 && !use_custom) {
++i;
continue;
}
if (!encode_gap(value.substr(start, i - start))) {
return false;
}
if (use_custom) {
auto token = impl_->custom_tokens[custom.second];
if (!callback(token, result)) {
if (match.second >= 0 && match.first == custom.first) {
result.push_back(match.second);
} else if (!encode_gap(value.substr(i, custom.first - i))) {
return false;
}
}
} else {
result.push_back(match.second);
}
start = i = use_custom ? custom.first : match.first;
}
return encode_gap(value.substr(start));
};
auto encode_normalized = [&](const std::string& value) {
std::string normalized;
if (!impl_->normalize(value, normalized, error)) {
return false;
}
return extract(normalized, impl_->normalized_added, normalized_custom, encode_plain);
};
if (!extract(text, impl_->raw_added, raw_custom, encode_normalized)) {
return false;
}
std::stringstream ss;
ss << "[";
for (int id : result) {
auto token = impl_->tokens.find(id);
if (token != impl_->tokens.end()) {
ss << "\"" << token->second << "\", ";
} else {
ss << "\"<id:" << id << ">\", ";
}
}
ss << "]";
LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), result.size(), ss.str().c_str());
tokens = std::move(result);
return true;
}
std::string HFTokenizer::decode_token(int id) const {
auto found = impl_->tokens.find(id);
return found == impl_->tokens.end() ? "" : found->second;
}
static std::string tokenizer_lossy_utf8(const std::string& bytes, bool fallback) {
std::string result;
for (size_t i = 0; i < bytes.size();) {
int32_t cp;
auto count = utf8proc_iterate(reinterpret_cast<const utf8proc_uint8_t*>(bytes.data() + i), bytes.size() - i, &cp);
if (count > 0) {
result.append(bytes, i, count);
i += count;
} else if (fallback) {
result.clear();
for (size_t j = 0; j < bytes.size(); ++j) {
result += "\xef\xbf\xbd";
}
return result;
} else {
result += "\xef\xbf\xbd";
unsigned char lead = bytes[i++];
size_t expected = lead >= 0xc2 && lead <= 0xdf ? 2 : lead >= 0xe0 && lead <= 0xef ? 3
: lead >= 0xf0 && lead <= 0xf4 ? 4
: 1;
for (size_t j = 1; j < expected && i < bytes.size(); ++j) {
unsigned char c = bytes[i];
if (c < 0x80 || c > 0xbf || (j == 1 && ((lead == 0xe0 && c < 0xa0) || (lead == 0xed && c > 0x9f) || (lead == 0xf0 && c < 0x90) || (lead == 0xf4 && c > 0x8f)))) {
break;
}
++i;
}
}
}
return result;
}
bool HFTokenizer::decode(const std::vector<int>& ids, std::string& text, std::string* error) const {
text.clear();
if (error) {
error->clear();
}
std::vector<std::string> pieces;
for (int id : ids) {
if (!impl_->special_ids.count(id) && impl_->tokens.count(id)) {
pieces.push_back(decode_token(id));
}
}
for (const auto& step : impl_->decoders) {
if (step.type == "Replace") {
for (auto& piece : pieces) {
std::string replaced;
if (!step.pattern->replace(piece, step.content, replaced, error)) {
return false;
}
piece = std::move(replaced);
}
} else if (step.type == "ByteLevel" || step.type == "Fuse") {
std::string joined;
for (const auto& piece : pieces) {
std::string bytes;
if (step.type == "ByteLevel") {
for (size_t i = 0; i < piece.size();) {
int32_t cp;
if (!tokenizer_next(piece, i, cp, error)) {
return false;
}
auto found = impl_->byte_decoder.find(cp);
if (found == impl_->byte_decoder.end()) {
bytes = piece;
break;
}
bytes += static_cast<char>(found->second);
}
} else {
bytes = piece;
}
joined += bytes;
}
pieces = {step.type == "ByteLevel" ? tokenizer_lossy_utf8(joined, false) : joined};
} else {
std::vector<std::string> decoded;
std::string bytes;
auto flush = [&] {
if (!bytes.empty()) {
decoded.push_back(tokenizer_lossy_utf8(bytes, true));
bytes.clear();
}
};
for (const auto& piece : pieces) {
auto hex = [](char c) { return c >= '0' && c <= '9' ? c - '0' : c >= 'A' && c <= 'F' ? c - 'A' + 10
: c >= 'a' && c <= 'f' ? c - 'a' + 10
: -1; };
if (piece.size() == 6 && piece.compare(0, 3, "<0x") == 0 && piece[5] == '>' && hex(piece[3]) >= 0 && hex(piece[4]) >= 0) {
bytes += static_cast<char>((hex(piece[3]) << 4) | hex(piece[4]));
} else {
flush();
decoded.push_back(piece);
}
}
flush();
pieces = std::move(decoded);
}
}
std::string result;
for (size_t i = 0; i < pieces.size(); ++i) {
if (i && !impl_->has_decoder) {
result += ' ';
}
result += pieces[i];
}
text = std::move(result);
return true;
}
+26
View File
@@ -0,0 +1,26 @@
#ifndef __SD_TOKENIZERS_HF_TOKENIZER_H__
#define __SD_TOKENIZERS_HF_TOKENIZER_H__
#include <memory>
#include "tokenizer.h"
class HFTokenizer : public Tokenizer {
struct Impl;
std::unique_ptr<Impl> impl_;
std::string decode_token(int token_id) const override;
public:
explicit HFTokenizer(const std::string& path);
~HFTokenizer() override;
// Padding belongs to the encoder; tokenizer.json supplies the single-sequence template.
void set_padding(int token_id, bool left);
void validate_vocab_size(int64_t embedding_rows) const;
int token_to_id(const std::string& token) const;
void add_special_token(const std::string& token) override;
bool encode(const std::string& text, std::vector<int>& tokens, on_new_token_cb_t on_new_token_cb = nullptr, std::string* error = nullptr) override;
bool decode(const std::vector<int>& tokens, std::string& text, std::string* error = nullptr) const override;
};
#endif // __SD_TOKENIZERS_HF_TOKENIZER_H__
+20 -10
View File
@@ -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);
+3
View File
@@ -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 = "");

Some files were not shown because too many files have changed in this diff Show More