Compare commits

...

17 Commits

Author SHA1 Message Date
leejet
b167b942f7
fix: align Qwen Image 2.1 flow schedule with official defaults (#2048) 2026-09-25 02:43:10 +08:00
leejet
1a2330de68
fix: reserve 128 MiB headroom when selecting monolithic execution (#2046) 2026-09-25 01:53:11 +08:00
leejet
740c7ae193
feat: add configurable Qwen cache types and early cache scheduling (#2045) 2026-09-25 01:51:01 +08:00
Nick Beerbower
4dfe8f5d45
feat: add a stand-alone upscale endpoint to the server (#2026) 2026-09-25 01:13:23 +08:00
stduhpf
caa111adf3
feat: optimize cfg special cases with guidance schdeule (#2033) 2026-09-25 01:10:22 +08:00
leejet
88411ef1e0
refactor: centralize circular RoPE and extend image model support (#2039) 2026-09-24 01:45:14 +08:00
Dong Wang
500ef5fa7c
fix: map mmapped weights through Metal buffers instead of CPU buffers (#2037) 2026-09-24 00:53:49 +08:00
leejet
2a4ebba818 ci: automatically close PRs from organization-owned forks 2026-09-24 00:40:42 +08:00
Dong Wang
70c1dbc01e
perf: run one-frame Wan VAE convolutions as 2D convolutions (#2038) 2026-09-24 00:23:31 +08:00
assouan
36746936c0
fix: add graph cuts for MiniMax-H3 text conditioning (#1900) 2026-09-24 00:03:29 +08:00
leejet
2dc7f5408a
feat: add Qwen Image 2.1 prefix KV cache (#2035) 2026-09-23 23:08:21 +08:00
leejet
e6281b6318
feat: add configurable conditioning cache for all models (#2034) 2026-09-23 22:25:42 +08:00
stduhpf
241518b35d
feat: add latent2rgba preview for Qwen-Image 2.1 (#2032) 2026-09-23 22:16:55 +08:00
leejet
c92d73c408
fix: preserve alpha when upscaling RGBA images with ESRGAN (#2029) 2026-09-23 02:23:10 +08:00
leejet
28b454bda1
feat: add configurable image input preprocessing (#2028) 2026-09-23 02:12:49 +08:00
LED-M
2bb72947cb
perf: cache MiniMax H3 text conditioning (#1966) 2026-09-22 22:19:39 +08:00
leejet
ac45422a05
fix: honor reference image resize settings in OpenAI edits (#2025) 2026-09-22 22:02:21 +08:00
72 changed files with 2921 additions and 1007 deletions

View File

@ -0,0 +1,61 @@
name: Close PRs from organization forks
on:
pull_request_target:
types: [opened, reopened]
permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
cancel-in-progress: false
jobs:
close-organization-fork-pr:
if: >-
github.event.pull_request.head.repo.owner.type == 'Organization' &&
github.event.pull_request.head.repo.id != github.event.pull_request.base.repo.id
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- name: Explain the contribution policy and close the PR
uses: actions/github-script@v9
with:
script: |
const { data: pr } = await github.rest.pulls.get({
...context.repo,
pull_number: context.issue.number,
});
const headRepo = pr.head.repo;
if (pr.state !== 'open' || !headRepo ||
headRepo.id === pr.base.repo.id || headRepo.owner.type !== 'Organization') {
return;
}
const marker = '<!-- organization-fork-policy -->';
const comments = await github.paginate(github.rest.issues.listComments, {
...context.repo,
issue_number: pr.number,
per_page: 100,
});
const alreadyExplained = comments.some(comment =>
comment.user?.login === 'github-actions[bot]' && comment.body?.includes(marker));
if (!alreadyExplained) {
await github.rest.issues.createComment({
...context.repo,
issue_number: pr.number,
body: [
marker,
'This repository requires contributions from forks to use a personal fork with **Allow edits from maintainers** enabled.',
'GitHub does not support this option for organization-owned forks, so this PR is being closed automatically.',
'Please open a new PR from a fork in your personal GitHub account and enable **Allow edits from maintainers** so maintainers can help update the branch.',
'See [the GitHub documentation](https://docs.github.com/en/pull-requests/how-tos/work-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork).',
].join('\n\n'),
});
}
await github.rest.pulls.update({
...context.repo,
pull_number: pr.number,
state: 'closed',
});

View File

@ -12,6 +12,10 @@ If you want to update a third-party dependency, please open an issue first inste
## Pull Requests ## Pull Requests
When contributing from a fork, use a fork under your personal GitHub account and enable **Allow edits from maintainers**. This lets maintainers make follow-up fixes directly on the PR branch.
PRs from organization-owned forks are automatically closed when opened or reopened because GitHub does not support this maintainer-edit option for those forks. Submit the changes from a personal fork instead. See [GitHub's documentation](https://docs.github.com/en/pull-requests/how-tos/work-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork).
Keep each PR focused on one clear change. Large or overly complex PRs are harder to review and may not be merged. Keep each PR focused on one clear change. Large or overly complex PRs are harder to review and may not be merged.
Do not include test code or test scripts in commits or PRs. Keep them local and report verification results in the PR description. Do not include test code or test scripts in commits or PRs. Keep them local and report verification results in the PR description.

View File

@ -161,6 +161,9 @@ resident allocations. Vulkan reports exceeding total memory are rejected because
its heap-budget subtraction can underflow. Other backends use the cap instead of its heap-budget subtraction can underflow. Other backends use the cap instead of
treating such reports as zero free memory. Failed checks log the reported free and treating such reports as zero free memory. Failed checks log the reported free and
total memory alongside tracked weight and runtime allocations. total memory alongside tracked weight and runtime allocations.
With `--mmap`, device-backed mappings count toward these budgets at their full
mapped-file size, once per device buffer even when multiple parameter blocks
share it. Mappings retained in the loader cache continue to count.
Components are considered in `diffusion`, `te`, `vae` order so that repeatedly Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
used diffusion weights have priority. Each component's weights use the first used diffusion weights have priority. Each component's weights use the first

View File

@ -2,6 +2,16 @@
Caching methods accelerate diffusion inference by reusing intermediate computations when changes between steps are small. Caching methods accelerate diffusion inference by reusing intermediate computations when changes between steps are small.
### Conditioning Cache
Conditioning results are cached per model context using an LRU cache. The default
capacity is **0 (disabled) for `sd-cli`** and **4 entries for `sd-server` and the C
API**. Set `--conditioning-cache-size N` to change the limit; `0` disables caching.
For example, `sd-cli -m model.safetensors -p "a cat" --conditioning-cache-size 4`
enables the cache in the CLI. The C API option is
`sd_ctx_params_t::conditioning_cache_size`, initialized by `sd_ctx_params_init()`.
This cache is independent of the diffusion-step `--cache-mode` options below.
### Cache Modes ### Cache Modes
| Mode | Target | Description | | Mode | Target | Description |

View File

@ -26,6 +26,9 @@ Stable-diffusion.spp also supports basic Unet-based editing models like instruct
## Configuring Reference Modes (`--ref-image-args`) ## Configuring Reference Modes (`--ref-image-args`)
For a one-time input transform before reference presets and model processing,
including cropping, padding, and resizing algorithms, see [Image preprocessing](./image_preprocessing.md).
Different DiT-based editing models require different configurations to process reference images correctly (e.g., whether to use a Vision Language Model (VLM) encoder or pass VAE-encoded images directly to the DiT). Different DiT-based editing models require different configurations to process reference images correctly (e.g., whether to use a Vision Language Model (VLM) encoder or pass VAE-encoded images directly to the DiT).
To simplify this, we provide **Presets**. By default, the system automatically selects the best preset based on the model architecture. However, you can override this using the `--ref-image-args` argument. To simplify this, we provide **Presets**. By default, the system automatically selects the best preset based on the model architecture. However, you can override this using the `--ref-image-args` argument.

View File

@ -2,6 +2,8 @@
You can use ESRGAN—such as the model [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)—to upscale the generated images and improve their overall resolution and clarity. You can use ESRGAN—such as the model [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)—to upscale the generated images and improve their overall resolution and clarity.
RGBA images, including Qwen Image 2.1 output, keep their alpha channel during model upscaling and hires fix. ESRGAN processes the RGB channels; the alpha channel is resized with bilinear interpolation and recombined with the upscaled image.
- Specify the model path using the `--upscale-model PATH` parameter. example: - Specify the model path using the `--upscale-model PATH` parameter. example:
```bash ```bash

173
docs/image_preprocessing.md Normal file
View File

@ -0,0 +1,173 @@
# Image preprocessing
Use `--image-preprocess` to transform each image input once, before generation:
```sh
sd-cli ... \
--image-preprocess "target=init,mode=crop-resize,filter=lanczos,antialias=true" \
--image-preprocess "target=mask,filter=nearest-exact" \
--image-preprocess "target=ref,index=0,mode=fit-pad,width=768,height=768,filter=bicubic"
```
CLI and server image loaders decode at the original resolution. The generation
entry point merges input defaults with user rules and prepares one transformed
image per input. The original pipeline then consumes those images, including
its mandatory canvas adaptation, reference resizing, and encoder preprocessing.
```text
native-resolution image
-> input defaults + user overrides
-> one input transform
-> original generation pipeline and model-specific processing
```
These rules do not override internal VAE, CLIP/VLM, ControlNet, or pixel-patch preprocessing.
`--ref-image-args` retains its existing meaning and runs after this input transform.
## Inputs and defaults
| `target` | Input | Default geometry | Indexed? |
| --- | --- | --- | --- |
| `init` | img2img image or video first frame | Center crop to the generation aspect ratio, then resize | No |
| `end` | Video last frame | Center crop, then resize | No |
| `mask` | Inpainting mask | Inherit init geometry; otherwise center crop, then resize | No |
| `control` | Control image | Center crop, then resize | No |
| `ref` | Reference images | Preserve source dimensions | Yes |
| `ip-adapter` | IP-Adapter image | Preserve source dimensions | No |
| `id` | PhotoMaker identity images | Preserve source dimensions | Yes |
| `control-frame` | Control video frames | Center crop, then resize | Yes |
Canvas defaults use the aligned generation dimensions. Reference, IP-Adapter,
and identity inputs use their original dimensions unless overridden. Default
resampling is nearest for images and nearest-exact for masks.
These defaults are shared by CLI, server, and C API. Moving geometry out of
the loaders replaces the previous CLI/server BOX/sRGB resizing, so default
pixels are not guaranteed to match earlier builds.
Reference video and audio preprocessing are outside these image rules.
Preprocessing options apply to `img_gen` and `vid_gen`, not standalone upscale
or ADetailer mode. ADetailer clears the user's rules for its internal crops.
## Rules
Rules are comma-separated `key=value` lists. Repeat the CLI option or separate
rules with semicolons. Every rule requires a `target` and at least one option.
Rule syntax and input compatibility are checked when image/video generation
starts. Unknown keys, invalid values, duplicate keys in a rule, missing images,
and out-of-range indices cause generation to fail with an error log.
Omit `index` to configure every image of that type; otherwise use a zero-based
index. CLI directory inputs follow filename order. Indexed rules override
type-wide rules field by field, regardless of order. At equal specificity,
the last value for a field wins. `auto` selects the input preset.
| `mode` | Input transform |
| --- | --- |
| `auto` | Use the input's default geometry |
| `none` | Keep source dimensions without resizing, cropping, or padding |
| `stretch` | Resize to the target dimensions |
| `crop` | Crop a target-sized rectangle without resizing; fail if the source is too small |
| `crop-resize` | Crop to the target aspect ratio, then resize |
| `fit-pad` | Fit the entire image inside the target dimensions, preserving aspect ratio, then pad |
`width` and `height` must be specified together as positive integers. They
override the input transform's dimensions, not the generation or encoder size.
For a native-size preset, specifying dimensions without a mode selects stretch.
`mode=none` with explicit dimensions different from the source is contradictory
and is rejected.
`anchor=center|top|bottom|left|right` selects crop/padding placement.
`pad_color=#RRGGBB` or `#RRGGBBAA` selects padding, defaulting to opaque black.
A grayscale mask uses the first color component.
`filter=auto|nearest|nearest-exact|bilinear|bicubic|lanczos` selects resampling.
`antialias=auto|true|false` enables antialiasing automatically for filtered
downscaling; explicit true requires bilinear, bicubic, or Lanczos.
Filtered RGBA resizing uses premultiplied alpha.
`canny=true|false` enables edge detection for any supported image target,
defaulting to `false`. It runs once after geometry, before the original
generation pipeline, including with `mode=none`. Grayscale, grayscale-alpha,
RGB, and RGBA inputs are supported; alpha is preserved.
Each input has its own Canny setting. Indexed rules can enable or disable it
for individual references, identity images, or video control frames.
```sh
--image-preprocess "target=init,mode=fit-pad,canny=true"
--image-preprocess "target=ref,index=0,mode=none,canny=true"
--image-preprocess "target=control-frame,index=2,canny=true"
```
Init and mask sources must have the same dimensions. The mask inherits the
init crop, resize, and padding coordinates, while retaining its own filter,
padding value, and Canny setting. Conflicting mask geometry is rejected. An
omitted mask remains absent until the original pipeline creates its default mask.
## Downstream behavior
`mode=none` only skips the input geometry transform. For example:
```sh
--image-preprocess "target=init,mode=none" \
--image-preprocess "target=ref,mode=none"
```
The init image is still adapted to the generation canvas by the original
pipeline. Reference images still follow `--ref-image-args` and model-specific
resizing. CLIP retains its fixed input dimensions and normalization. HiDream-O1
retains its original pixel-reference and visual preprocessing.
Existing sharing between consumers is preserved: for example, Wan img2video
uses the same adapted first frame for VAE conditioning and CLIP. High-resolution
passes reuse the prepared images and apply their original size adaptation;
they do not apply the user's crop a second time.
To disable reference resizing before VAE encoding, use
`--ref-image-args "resize_before_vae=false"` or the server field
`"ref_image_args": "resize_before_vae=false"`. This is separate from
`target=ref,mode=none`, which only skips input geometry. Model constraints
still apply.
## Server requests
Native image/video requests and SDAPI accept `image_preprocess` as a string or
an array of rule strings:
```json
{
"image_preprocess": [
"target=init,mode=fit-pad,filter=bicubic",
"target=mask,filter=nearest-exact",
"target=ref,index=0,mode=none"
]
}
```
OpenAI-compatible requests accept it through
`<sd_cpp_extra_args>{...}</sd_cpp_extra_args>` in the prompt.
Request rules replace server-default rules. Generation metadata records the
user rules; image encodings and channel conventions are unchanged.
## C API
Set `image_preprocess` on the existing image/video generation parameters.
The `generate_image()` and `generate_video()` signatures are unchanged:
```c
sd_img_gen_params_t params;
sd_img_gen_params_init(&params);
/* Set prompt, original-resolution input images, and generation options. */
params.image_preprocess.rules = "target=init,mode=crop-resize,filter=lanczos;"
"target=mask,filter=nearest-exact";
bool ok = generate_image(ctx, &params, &images, &count);
```
Both generation parameter initializers set `image_preprocess.rules` to `NULL`,
selecting input presets. Rule strings are borrowed for the synchronous call.
The library owns temporary transformed pixels; caller images and arrays are
not modified. Add `canny=true` to the desired target's rule in
`image_preprocess.rules` to enable Canny.
The parameter structs have grown; applications and bindings must be rebuilt.

View File

@ -63,6 +63,12 @@ See [backend selection](./backend.md) for full syntax.
When a graph has cut markers and its missing weights plus incremental compute workspace exceed the available device headroom, it runs its fixed segment list in order. A reusable monolithic compute buffer is not counted as a new allocation. An explicit `--max-vram` budget deducts already-resident managed weights and compute/cache buffers registered by every runner sharing the device, so later graph runs remain segmented when the full graph exceeds the budget. The current segment's weights are pinned during compute, and the next parameter-bearing segment is prefetched when the device supports asynchronous transfer. No opt-in streaming flag is required. When a graph has cut markers and its missing weights plus incremental compute workspace exceed the available device headroom, it runs its fixed segment list in order. A reusable monolithic compute buffer is not counted as a new allocation. An explicit `--max-vram` budget deducts already-resident managed weights and compute/cache buffers registered by every runner sharing the device, so later graph runs remain segmented when the full graph exceeds the budget. The current segment's weights are pinned during compute, and the next parameter-bearing segment is prefetched when the device supports asynchronous transfer. No opt-in streaming flag is required.
When choosing between monolithic and segmented execution, the runner requires
an additional 128 MiB of headroom in both available device memory and any explicit
managed budget. This planning headroom absorbs small allocation estimate changes;
subsequent capacity checks can consume it while still preserving the 512 MiB device
scratch reserve and respecting the managed budget.
- `--max-vram <GiB>` optionally lowers the live-memory limit. A positive value is a managed per-device budget, `0` uses the device's current free memory without an explicit budget, and a negative value snapshots free memory at startup while reserving that many GiB (`--max-vram -1` reserves about 1 GiB). Driver contexts and unrelated external allocations remain outside the managed budget. - `--max-vram <GiB>` optionally lowers the live-memory limit. A positive value is a managed per-device budget, `0` uses the device's current free memory without an explicit budget, and a negative value snapshots free memory at startup while reserving that many GiB (`--max-vram -1` reserves about 1 GiB). Driver contexts and unrelated external allocations remain outside the managed budget.
- `--disable-prefetch` disables asynchronous next-segment prefetch while retaining synchronous loading, eviction, and segmented execution. - `--disable-prefetch` disables asynchronous next-segment prefetch while retaining synchronous loading, eviction, and segmented execution.
- `--disable-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments. - `--disable-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments.

View File

@ -23,7 +23,7 @@ Run the following commands from the build directory. Use image dimensions divisi
### Text to image ### Text to image
```powershell ```powershell
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1.png .\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1.png
``` ```
<img alt="Qwen Image 2.1 example" src="../assets/qwen/qwen_image_2.1.png" /> <img alt="Qwen Image 2.1 example" src="../assets/qwen/qwen_image_2.1.png" />
@ -35,11 +35,39 @@ To use GGUF diffusion weights, set `--diffusion-model` to the path of a file suc
Pass the reference image with `-r` and describe the edit in `-p`. Vision weights are required; the example below loads them separately with `--llm_vision`. Pass the reference image with `-r` and describe the edit in `-p`. Vision weights are required; the example below loads them separately with `--llm_vision`.
```powershell ```powershell
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1_edit.png .\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1_edit.png
``` ```
For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`. For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`.
### Prefix cache
By default, the first denoising call for each fixed condition saves the text and reference-image keys and values from every transformer layer. Later calls only compute the target-image tokens. Positive and negative conditions use separate caches, which are released when sampling ends.
Set `qwen_image_2_1_prefix_cache_type` in `--model-args` to `auto` or a type name using the same parser and case-sensitive names as `--type`:
- `auto` (default): use FP16 only when Flash Attention is enabled, Sage Attention is disabled, the attention scale is unchanged, and every attention operation in the cache-writing or cache-reading graph selects Flash Attention after backend support checks. If an operation falls back, rebuild the prefix in FP32 before executing and keep FP32 for the rest of that sampling run.
- `f32`: always store FP32 keys and values.
- `f16`: always store FP16 keys and values, including with ordinary attention or custom attention scaling. This saves cache memory but can introduce additional rounding error.
- Other types, such as `bf16`, `q4_1`, `q5_0`, `q5_1`, `q8_0`, `q4_K`, `q6_K`, `iq4_nl`, and `iq4_xs`: use the requested storage type if the ggml build provides runtime conversion to and from FP32. Quantization is lossy and must be selected explicitly; `auto` never selects a quantized type.
Cache data is packed into contiguous rows of `hidden_size` elements before conversion, so 256-element quantization blocks work with the model's 128-element attention heads without padding. The type's block size must divide `hidden_size`. Unknown types, types lacking runtime conversion (for example `q8_1` and several IQ formats), and incompatible block sizes are ignored with a warning, leaving the previous setting or the default `auto` unchanged.
For example, use `--model-args qwen_image_2_1_prefix_cache_type=q8_0` to enable 8-bit cache storage. Cached keys and values are converted back to the attention input type before concatenating with the current target tokens. This reduces persistent cache memory; attention working buffers still use floating-point values, and conversion adds work on each step. Backends without the required conversion operations use the existing CPU fallback.
For the default 32-layer model, a prefix of 4096 tokens takes approximately the following memory per condition, excluding weights, working buffers, and allocation overhead:
| Cache type | Memory |
| --- | ---: |
| `f32` | 4 GiB |
| `f16` | 2 GiB |
| `q8_0` | 1.0625 GiB |
| `q4_0` | 0.5625 GiB |
The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path.
Disable this optimization with `--model-args qwen_image_2_1_prefix_cache=false`. It reuses step-independent activations; numerical results can still differ slightly because the matrix sizes change.
### Alpha channel ### Alpha channel
This model supports alpha channel output. As the model determines whether to output a regular image or with transparency through the prompt, according to [official recommendation](https://github.com/QwenLM/Qwen-Image-2.1#transparent-image-generation-rgba), use the following prompt format for better results: This model supports alpha channel output. As the model determines whether to output a regular image or with transparency through the prompt, according to [official recommendation](https://github.com/QwenLM/Qwen-Image-2.1#transparent-image-generation-rgba), use the following prompt format for better results:

View File

@ -14,6 +14,12 @@ equivalent to `--log-level verbose`. If repeated, the last logging option wins.
For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see
[ADetailer](../../docs/adetailer.md). [ADetailer](../../docs/adetailer.md).
Use repeatable `--image-preprocess` rules to select resizing, cropping, padding,
and resampling separately for each image input. Add `canny=true` to any input
rule for edge detection. See
[Image preprocessing](../../docs/image_preprocessing.md) for input selectors,
input defaults, downstream model processing, mask alignment, and examples.
Metadata mode inspects PNG/JPEG container metadata without loading any model: Metadata mode inspects PNG/JPEG container metadata without loading any model:
```bash ```bash

View File

@ -41,7 +41,6 @@ struct SDCliParams {
std::string metadata_format = "text"; std::string metadata_format = "text";
sd_log_level_t log_level = SD_LOG_INFO; sd_log_level_t log_level = SD_LOG_INFO;
bool canny_preprocess = false;
bool convert_name = false; bool convert_name = false;
preview_t preview_method = PREVIEW_NONE; preview_t preview_method = PREVIEW_NONE;
@ -107,10 +106,6 @@ struct SDCliParams {
}; };
options.bool_options = { options.bool_options = {
{"",
"--canny",
"apply canny preprocessor (edge detection)",
true, &canny_preprocess},
{"", {"",
"--convert-name", "--convert-name",
"convert tensor name (for convert mode)", "convert tensor name (for convert mode)",
@ -268,7 +263,6 @@ struct SDCliParams {
<< " metadata_format: \"" << metadata_format << "\",\n" << " metadata_format: \"" << metadata_format << "\",\n"
<< " log_level: " << log_level_name(log_level) << ",\n" << " log_level: " << log_level_name(log_level) << ",\n"
<< " color: " << (color ? "true" : "false") << ",\n" << " color: " << (color ? "true" : "false") << ",\n"
<< " canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
<< " convert_name: " << (convert_name ? "true" : "false") << ",\n" << " convert_name: " << (convert_name ? "true" : "false") << ",\n"
<< " preview_method: " << previews_str[preview_method] << ",\n" << " preview_method: " << previews_str[preview_method] << ",\n"
<< " preview_interval: " << preview_interval << ",\n" << " preview_interval: " << preview_interval << ",\n"
@ -328,9 +322,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
bool load_images_from_dir(const std::string dir, bool load_images_from_dir(const std::string dir,
std::vector<SDImageOwner>& images, std::vector<SDImageOwner>& images,
int expected_width = 0, int max_image_num = 0) {
int expected_height = 0,
int max_image_num = 0) {
if (!fs::exists(dir) || !fs::is_directory(dir)) { if (!fs::exists(dir) || !fs::is_directory(dir)) {
LOG_ERROR("'%s' is not a valid directory\n", dir.c_str()); LOG_ERROR("'%s' is not a valid directory\n", dir.c_str());
return false; return false;
@ -358,7 +350,7 @@ bool load_images_from_dir(const std::string dir,
int width = 0; int width = 0;
int height = 0; int height = 0;
int loaded_channel = 0; int loaded_channel = 0;
uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, expected_width, expected_height); uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, 0, 0);
if (image_buffer == nullptr) { if (image_buffer == nullptr) {
LOG_ERROR("load image from '%s' failed", path.c_str()); LOG_ERROR("load image from '%s' failed", path.c_str());
return false; return false;
@ -652,10 +644,11 @@ int main(int argc, const char* argv[]) {
SDCliParams cli_params; SDCliParams cli_params;
SDContextParams ctx_params; SDContextParams ctx_params;
ctx_params.conditioning_cache_size = 0;
SDGenerationParams gen_params; SDGenerationParams gen_params;
parse_args(argc, argv, cli_params, ctx_params, gen_params);
sd_set_log_callback(sd_log_cb, (void*)&cli_params); sd_set_log_callback(sd_log_cb, (void*)&cli_params);
parse_args(argc, argv, cli_params, ctx_params, gen_params);
if (cli_params.mode == METADATA) { if (cli_params.mode == METADATA) {
MetadataReadOptions options; MetadataReadOptions options;
@ -751,16 +744,8 @@ int main(int argc, const char* argv[]) {
auto load_image_and_update_size = [&](const std::string& path, auto load_image_and_update_size = [&](const std::string& path,
SDImageOwner& image, SDImageOwner& image,
bool resize_image = true,
int expected_channel = 3) -> bool { int expected_channel = 3) -> bool {
int expected_width = 0; if (!load_sd_image_from_file(image.put(), path.c_str(), 0, 0, expected_channel)) {
int expected_height = 0;
if (resize_image && gen_params.width_and_height_are_set()) {
expected_width = gen_params.width;
expected_height = gen_params.height;
}
if (!load_sd_image_from_file(image.put(), path.c_str(), expected_width, expected_height, expected_channel)) {
LOG_ERROR("load image from '%s' failed", path.c_str()); LOG_ERROR("load image from '%s' failed", path.c_str());
return false; return false;
} }
@ -783,7 +768,7 @@ int main(int argc, const char* argv[]) {
if (gen_params.init_image_path.size() > 0) { if (gen_params.init_image_path.size() > 0) {
const bool native_init = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER; const bool native_init = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER;
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image, true, native_init ? 0 : 3)) { if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image, native_init ? 0 : 3)) {
return 1; return 1;
} }
} }
@ -798,7 +783,7 @@ int main(int argc, const char* argv[]) {
gen_params.ref_images.clear(); gen_params.ref_images.clear();
for (auto& path : gen_params.ref_image_paths) { for (auto& path : gen_params.ref_image_paths) {
SDImageOwner ref_image({0, 0, 0, nullptr}); SDImageOwner ref_image({0, 0, 0, nullptr});
if (!load_image_and_update_size(path, ref_image, false, 0)) { if (!load_image_and_update_size(path, ref_image, 0)) {
return 1; return 1;
} }
gen_params.ref_images.push_back(std::move(ref_image)); gen_params.ref_images.push_back(std::move(ref_image));
@ -839,41 +824,22 @@ int main(int argc, const char* argv[]) {
if (gen_params.mask_image_path.size() > 0) { if (gen_params.mask_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.mask_image.put(), if (!load_sd_image_from_file(gen_params.mask_image.put(),
gen_params.mask_image_path.c_str(), gen_params.mask_image_path.c_str(),
gen_params.get_resolved_width(), 0,
gen_params.get_resolved_height(), 0,
1)) { 1)) {
LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str()); LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
return 1; return 1;
} }
} else {
sd_image_t generated_mask = {0, 0, 1, nullptr};
generated_mask.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
if (generated_mask.data == nullptr) {
LOG_ERROR("malloc mask image failed");
return 1;
}
generated_mask.width = gen_params.get_resolved_width();
generated_mask.height = gen_params.get_resolved_height();
memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
gen_params.mask_image.reset(generated_mask);
} }
if (gen_params.control_image_path.size() > 0) { if (gen_params.control_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.control_image.put(), if (!load_sd_image_from_file(gen_params.control_image.put(),
gen_params.control_image_path.c_str(), gen_params.control_image_path.c_str(),
gen_params.get_resolved_width(), 0,
gen_params.get_resolved_height())) { 0)) {
LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str()); LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
return 1; return 1;
} }
if (cli_params.canny_preprocess) { // apply preprocessor
preprocess_canny(gen_params.control_image.get(),
0.08f,
0.08f,
0.8f,
1.0f,
false);
}
} }
if (gen_params.ip_adapter_image_path.size() > 0) { if (gen_params.ip_adapter_image_path.size() > 0) {
@ -890,8 +856,6 @@ int main(int argc, const char* argv[]) {
gen_params.control_frames.clear(); gen_params.control_frames.clear();
if (!load_images_from_dir(gen_params.control_video_path, if (!load_images_from_dir(gen_params.control_video_path,
gen_params.control_frames, gen_params.control_frames,
gen_params.get_resolved_width(),
gen_params.get_resolved_height(),
gen_params.video_frames)) { gen_params.video_frames)) {
return 1; return 1;
} }
@ -900,10 +864,7 @@ int main(int argc, const char* argv[]) {
if (!gen_params.pm_id_images_dir.empty()) { if (!gen_params.pm_id_images_dir.empty()) {
gen_params.pm_id_images.clear(); gen_params.pm_id_images.clear();
if (!load_images_from_dir(gen_params.pm_id_images_dir, if (!load_images_from_dir(gen_params.pm_id_images_dir,
gen_params.pm_id_images, gen_params.pm_id_images)) {
0,
0,
0)) {
return 1; return 1;
} }
} }

View File

@ -518,7 +518,8 @@ ArgOptions SDContextParams::get_options() {
{"", {"",
"--model-args", "--model-args",
"extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, " "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
"chroma_t5_mask_pad, qwen_image_zero_cond_t", "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache, "
"qwen_image_2_1_prefix_cache_type (auto or a type name from --type)",
(int)',', (int)',',
&model_args}, &model_args},
{"", {"",
@ -571,6 +572,10 @@ ArgOptions SDContextParams::get_options() {
"number of threads to use during computation (default: -1). " "number of threads to use during computation (default: -1). "
"If threads <= 0, then threads will be set to the number of CPU physical cores", "If threads <= 0, then threads will be set to the number of CPU physical cores",
&n_threads}, &n_threads},
{"",
"--conditioning-cache-size",
"maximum number of conditioning results cached per model context (default: " + std::to_string(conditioning_cache_size) + ", 0 disables caching)",
&conditioning_cache_size},
}; };
options.bool_options = { options.bool_options = {
@ -822,6 +827,10 @@ bool SDContextParams::resolve(SDMode mode) {
} }
bool SDContextParams::validate(SDMode mode) { bool SDContextParams::validate(SDMode mode) {
if (conditioning_cache_size < 0) {
LOG_ERROR("error: conditioning-cache-size must be non-negative");
return false;
}
if (mode == CONVERT) { if (mode == CONVERT) {
const bool has_convert_input = model_path.length() != 0 || const bool has_convert_input = model_path.length() != 0 ||
clip_l_path.length() != 0 || clip_l_path.length() != 0 ||
@ -898,6 +907,7 @@ std::string SDContextParams::to_string() const {
std::ostringstream oss; std::ostringstream oss;
oss << "SDContextParams {\n" oss << "SDContextParams {\n"
<< " n_threads: " << n_threads << ",\n" << " n_threads: " << n_threads << ",\n"
<< " conditioning_cache_size: " << conditioning_cache_size << ",\n"
<< " model_path: \"" << model_path << "\",\n" << " model_path: \"" << model_path << "\",\n"
<< " clip_l_path: \"" << clip_l_path << "\",\n" << " clip_l_path: \"" << clip_l_path << "\",\n"
<< " clip_g_path: \"" << clip_g_path << "\",\n" << " clip_g_path: \"" << clip_g_path << "\",\n"
@ -992,6 +1002,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.pulid_weights_path = pulid_weights_path.c_str(); sd_ctx_params.pulid_weights_path = pulid_weights_path.c_str();
sd_ctx_params.tensor_type_rules = tensor_type_rules.c_str(); sd_ctx_params.tensor_type_rules = tensor_type_rules.c_str();
sd_ctx_params.n_threads = n_threads; sd_ctx_params.n_threads = n_threads;
sd_ctx_params.conditioning_cache_size = conditioning_cache_size;
sd_ctx_params.wtype = wtype; sd_ctx_params.wtype = wtype;
sd_ctx_params.rng_type = rng_type; sd_ctx_params.rng_type = rng_type;
sd_ctx_params.sampler_rng_type = sampler_rng_type; sd_ctx_params.sampler_rng_type = sampler_rng_type;
@ -1128,6 +1139,9 @@ ArgOptions SDGenerationParams::get_options() {
"Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)", "Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)",
(int)',', (int)',',
&ref_image_args}, &ref_image_args},
{"", "--image-preprocess",
"Image preprocessing rule: target=init|end|mask|control|ref|ip-adapter|id|control-frame,index=N,mode=auto|none|stretch|crop|crop-resize|fit-pad,filter=auto|nearest|nearest-exact|bilinear|bicubic|lanczos,antialias=auto|true|false,width=W,height=H,anchor=center|top|bottom|left|right,pad_color=#RRGGBB[AA],canny=true|false. Repeat for multiple rules.",
(int)';', &image_preprocess},
}; };
options.int_options = { options.int_options = {
@ -1308,11 +1322,6 @@ ArgOptions SDGenerationParams::get_options() {
"automatically increase the indices of references images based on the order they are listed (starting with 1).", "automatically increase the indices of references images based on the order they are listed (starting with 1).",
true, true,
&increase_ref_index}, &increase_ref_index},
{"",
"--disable-auto-resize-ref-image",
"disable auto resize of ref images",
false,
&auto_resize_ref_image},
{"", {"",
"--circular", "--circular",
"enable circular padding on both axes for tileable output", "enable circular padding on both axes for tileable output",
@ -1870,8 +1879,6 @@ bool decode_base64_image(const std::string& encoded_input,
static bool parse_image_json_field(const json& parent, static bool parse_image_json_field(const json& parent,
const char* key, const char* key,
int channels, int channels,
int expected_width,
int expected_height,
SDImageOwner& out_image) { SDImageOwner& out_image) {
if (!parent.contains(key)) { if (!parent.contains(key)) {
return true; return true;
@ -1883,14 +1890,12 @@ static bool parse_image_json_field(const json& parent,
if (!parent.at(key).is_string()) { if (!parent.at(key).is_string()) {
return false; return false;
} }
return decode_base64_image(parent.at(key).get<std::string>(), channels, expected_width, expected_height, out_image); return decode_base64_image(parent.at(key).get<std::string>(), channels, 0, 0, out_image);
} }
static bool parse_image_array_json_field(const json& parent, static bool parse_image_array_json_field(const json& parent,
const char* key, const char* key,
int channels, int channels,
int expected_width,
int expected_height,
std::vector<SDImageOwner>& out_images) { std::vector<SDImageOwner>& out_images) {
if (!parent.contains(key)) { if (!parent.contains(key)) {
return true; return true;
@ -1909,7 +1914,7 @@ static bool parse_image_array_json_field(const json& parent,
return false; return false;
} }
SDImageOwner image; SDImageOwner image;
if (!decode_base64_image(item.get<std::string>(), channels, expected_width, expected_height, image)) { if (!decode_base64_image(item.get<std::string>(), channels, 0, 0, image)) {
return false; return false;
} }
out_images.push_back(std::move(image)); out_images.push_back(std::move(image));
@ -2008,6 +2013,29 @@ static bool resolve_model_file_from_dir(const std::string& model_name,
return false; return false;
} }
bool SDGenerationParams::parse_image_preprocess_json(const std::string& json_str) {
const auto value = json::parse(json_str, nullptr, false);
std::string rules;
if (value.is_string()) {
rules = value.get<std::string>();
} else if (value.is_array()) {
for (const auto& item : value) {
if (!item.is_string()) {
LOG_ERROR("image_preprocess must contain rule strings");
return false;
}
if (!rules.empty())
rules += ";";
rules += item.get<std::string>();
}
} else {
LOG_ERROR("image_preprocess must be a string or array of strings");
return false;
}
image_preprocess = std::move(rules);
return true;
}
bool SDGenerationParams::from_json_str( bool SDGenerationParams::from_json_str(
const std::string& json_str, const std::string& json_str,
const std::function<std::string(const std::string&)>& lora_path_resolver) { const std::function<std::string(const std::string&)>& lora_path_resolver) {
@ -2019,6 +2047,9 @@ bool SDGenerationParams::from_json_str(
return false; return false;
} }
if (j.contains("image_preprocess") && !parse_image_preprocess_json(j["image_preprocess"].dump()))
return false;
auto load_if_exists = [&](const char* key, auto& out) { auto load_if_exists = [&](const char* key, auto& out) {
if (j.contains(key)) { if (j.contains(key)) {
using T = std::decay_t<decltype(out)>; using T = std::decay_t<decltype(out)>;
@ -2056,6 +2087,7 @@ bool SDGenerationParams::from_json_str(
load_if_exists("cache_mode", cache_mode); load_if_exists("cache_mode", cache_mode);
load_if_exists("cache_option", cache_option); load_if_exists("cache_option", cache_option);
load_if_exists("scm_mask", scm_mask); load_if_exists("scm_mask", scm_mask);
load_if_exists("ref_image_args", ref_image_args);
load_if_exists("clip_skip", clip_skip); load_if_exists("clip_skip", clip_skip);
load_if_exists("width", width); load_if_exists("width", width);
@ -2073,7 +2105,6 @@ bool SDGenerationParams::from_json_str(
load_if_exists("moe_boundary", moe_boundary); load_if_exists("moe_boundary", moe_boundary);
load_if_exists("vace_strength", vace_strength); load_if_exists("vace_strength", vace_strength);
load_if_exists("auto_resize_ref_image", auto_resize_ref_image);
load_if_exists("increase_ref_index", increase_ref_index); load_if_exists("increase_ref_index", increase_ref_index);
load_if_exists("embed_image_metadata", embed_image_metadata); load_if_exists("embed_image_metadata", embed_image_metadata);
@ -2217,37 +2248,23 @@ bool SDGenerationParams::from_json_str(
LOG_ERROR("invalid lora"); LOG_ERROR("invalid lora");
return false; return false;
} }
if (!parse_image_json_field(j, "init_image", 0, width, height, init_image)) { auto load_image = [&](const char* key, int channels, SDImageOwner& image) {
LOG_ERROR("invalid init_image"); if (!parse_image_json_field(j, key, channels, image)) {
LOG_ERROR("invalid %s", key);
return false;
}
return true;
};
if (!load_image("init_image", 0, init_image) ||
!load_image("end_image", 3, end_image) ||
!load_image("mask_image", 1, mask_image) ||
!load_image("control_image", 3, control_image) ||
!load_image("ip_adapter_image", 3, ip_adapter_image)) {
return false; return false;
} }
if (!parse_image_json_field(j, "end_image", 3, width, height, end_image)) { if (!parse_image_array_json_field(j, "ref_images", 0, ref_images) ||
LOG_ERROR("invalid end_image"); !parse_image_array_json_field(j, "control_frames", 3, control_frames)) {
return false; LOG_ERROR("invalid input image array");
}
if (!parse_image_array_json_field(j,
"ref_images",
0,
auto_resize_ref_image ? width : 0,
auto_resize_ref_image ? height : 0,
ref_images)) {
LOG_ERROR("invalid ref_images");
return false;
}
if (!parse_image_array_json_field(j, "control_frames", 3, width, height, control_frames)) {
LOG_ERROR("invalid control_frames");
return false;
}
if (!parse_image_json_field(j, "mask_image", 1, width, height, mask_image)) {
LOG_ERROR("invalid mask_image");
return false;
}
if (!parse_image_json_field(j, "control_image", 3, width, height, control_image)) {
LOG_ERROR("invalid control_image");
return false;
}
if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) {
LOG_ERROR("invalid ip_adapter_image");
return false; return false;
} }
@ -2491,6 +2508,10 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s
} }
bool SDGenerationParams::validate(SDMode mode) { bool SDGenerationParams::validate(SDMode mode) {
if (!image_preprocess.empty() && mode != IMG_GEN && mode != VID_GEN) {
LOG_ERROR("--image-preprocess requires img_gen or vid_gen mode");
return false;
}
if (batch_count <= 0) { if (batch_count <= 0) {
LOG_ERROR("error: batch_count must be greater than 0"); LOG_ERROR("error: batch_count must be greater than 0");
return false; return false;
@ -2666,14 +2687,6 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
pulid_id_weight, pulid_id_weight,
}; };
if (!auto_resize_ref_image) {
if (!ref_image_args.empty()) {
ref_image_args += ",";
}
ref_image_args += "resize_before_vae=0";
LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
}
if (increase_ref_index) { if (increase_ref_index) {
if (!ref_image_args.empty()) { if (!ref_image_args.empty()) {
ref_image_args += ","; ref_image_args += ",";
@ -2721,6 +2734,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size()); params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
params.circular_x = circular || circular_x; params.circular_x = circular || circular_x;
params.circular_y = circular || circular_y; params.circular_y = circular || circular_y;
params.image_preprocess = {image_preprocess.c_str()};
return params; return params;
} }
@ -2823,6 +2837,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size()); params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
params.circular_x = circular || circular_x; params.circular_x = circular || circular_x;
params.circular_y = circular || circular_y; params.circular_y = circular || circular_y;
params.image_preprocess = {image_preprocess.c_str()};
return params; return params;
} }
@ -2879,7 +2894,8 @@ std::string SDGenerationParams::to_string() const {
<< " ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n" << " ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n"
<< " ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n" << " ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n"
<< " control_video_path: \"" << control_video_path << "\",\n" << " control_video_path: \"" << control_video_path << "\",\n"
<< " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n" << " image_preprocess: " << image_preprocess << ",\n"
<< " ref_image_args: " << ref_image_args << ",\n"
<< " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n" << " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n"
<< " pm_id_images_dir: \"" << pm_id_images_dir << "\",\n" << " pm_id_images_dir: \"" << pm_id_images_dir << "\",\n"
<< " pm_id_embed_path: \"" << pm_id_embed_path << "\",\n" << " pm_id_embed_path: \"" << pm_id_embed_path << "\",\n"
@ -3030,12 +3046,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path); set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path);
root["models"] = std::move(models); root["models"] = std::move(models);
root["clip_skip"] = gen_params.clip_skip; root["clip_skip"] = gen_params.clip_skip;
root["strength"] = gen_params.strength; root["strength"] = gen_params.strength;
root["control_strength"] = gen_params.control_strength; root["control_strength"] = gen_params.control_strength;
root["ip_adapter_strength"] = gen_params.ip_adapter_strength; root["ip_adapter_strength"] = gen_params.ip_adapter_strength;
root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image; root["ref_image_args"] = gen_params.ref_image_args;
root["increase_ref_index"] = gen_params.increase_ref_index; root["image_preprocess"] = gen_params.image_preprocess;
root["increase_ref_index"] = gen_params.increase_ref_index;
if (mode == VID_GEN) { if (mode == VID_GEN) {
root["video"] = { root["video"] = {
{"frame_count", gen_params.video_frames}, {"frame_count", gen_params.video_frames},

View File

@ -116,7 +116,8 @@ bool decode_base64_image(const std::string& encoded_input,
SDImageOwner& out_image); SDImageOwner& out_image);
struct SDContextParams { struct SDContextParams {
int n_threads = -1; int n_threads = -1;
int conditioning_cache_size = 4;
std::string model_path; std::string model_path;
std::string clip_l_path; std::string clip_l_path;
std::string clip_g_path; std::string clip_g_path;
@ -200,18 +201,17 @@ struct SDGenerationParams {
std::string ad_prompt; std::string ad_prompt;
std::string ad_negative_prompt; std::string ad_negative_prompt;
std::string extra_ad_args; std::string extra_ad_args;
int clip_skip = -1; // <= 0 represents unspecified int clip_skip = -1; // <= 0 represents unspecified
int width = -1; int width = -1;
int height = -1; int height = -1;
int batch_count = 1; int batch_count = 1;
int qwen_image_layers = 3; int qwen_image_layers = 3;
int64_t seed = 42; int64_t seed = 42;
float strength = 0.75f; float strength = 0.75f;
float control_strength = 0.9f; float control_strength = 0.9f;
float ip_adapter_strength = 1.0f; float ip_adapter_strength = 1.0f;
bool auto_resize_ref_image = true; bool increase_ref_index = false;
bool increase_ref_index = false; bool embed_image_metadata = true;
bool embed_image_metadata = true;
std::string init_image_path; std::string init_image_path;
std::string end_image_path; std::string end_image_path;
@ -247,6 +247,7 @@ struct SDGenerationParams {
std::string extra_tiling_args; std::string extra_tiling_args;
std::string ref_image_args; std::string ref_image_args;
std::string image_preprocess;
std::string pm_id_images_dir; std::string pm_id_images_dir;
std::string pm_id_embed_path; std::string pm_id_embed_path;
@ -310,6 +311,7 @@ struct SDGenerationParams {
ArgOptions get_options(); ArgOptions get_options();
bool from_json_str(const std::string& json_str, bool from_json_str(const std::string& json_str,
const std::function<std::string(const std::string&)>& lora_path_resolver = {}); const std::function<std::string(const std::string&)>& lora_path_resolver = {});
bool parse_image_preprocess_json(const std::string& json_str);
bool initialize_cache_params(); bool initialize_cache_params();
void extract_and_remove_lora(const std::string& lora_model_dir); void extract_and_remove_lora(const std::string& lora_model_dir);
bool width_and_height_are_set() const; bool width_and_height_are_set() const;

View File

@ -56,6 +56,7 @@ Current endpoints include:
- `GET /sdcpp/v1/jobs/{id}` - `GET /sdcpp/v1/jobs/{id}`
- `POST /sdcpp/v1/jobs/{id}/cancel` - `POST /sdcpp/v1/jobs/{id}/cancel`
- `POST /sdcpp/v1/vid_gen` - `POST /sdcpp/v1/vid_gen`
- `POST /sdcpp/v1/upscale`
## `sd_cpp_extra_args` ## `sd_cpp_extra_args`
@ -148,6 +149,19 @@ Native extension fields:
- any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt` - any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
Uploaded images are decoded at their original dimensions. The first decoded
image establishes the generation dimensions if `size` is omitted. Input
geometry follows `image_preprocess`: references preserve their dimensions by
default, while init and mask use the generation canvas preset.
Reference encoding then follows model presets and `ref_image_args`. To skip
input geometry for references and disable resizing before VAE encoding, include
this in `prompt`:
```text
edit this image <sd_cpp_extra_args>{"image_preprocess":"target=ref,mode=none","ref_image_args":"resize_before_vae=false"}</sd_cpp_extra_args>
```
Response fields: Response fields:
| Field | Type | Notes | | Field | Type | Notes |
@ -421,7 +435,8 @@ Top-level fields:
| `samplers` | `array<string>` | Available sampling methods | | `samplers` | `array<string>` | Available sampling methods |
| `schedulers` | `array<string>` | Available schedulers | | `schedulers` | `array<string>` | Available schedulers |
| `loras` | `array<object>` | Available LoRA entries | | `loras` | `array<object>` | Available LoRA entries |
| `upscalers` | `array<object>` | Available model-backed highres upscalers | | `upscalers` | `array<object>` | Available highres upscalers, built-in and model-backed |
| `upscale` | `boolean` | Whether a compatible RGB ESRGAN model is available for `POST /sdcpp/v1/upscale` |
| `limits` | `object` | Shared queue and size limits | | `limits` | `object` | Shared queue and size limits |
`model` `model`
@ -463,6 +478,8 @@ Shared nested fields:
| Field | Type | Notes | | Field | Type | Notes |
| --- | --- | --- | | --- | --- | --- |
| `upscalers[].name` | `string` | Built-in name or model stem; use this value in `hires.upscaler` | | `upscalers[].name` | `string` | Built-in name or model stem; use this value in `hires.upscaler` |
| `upscalers[].model` | `boolean` | True for a model-backed upscaler, false for a built-in scaling filter |
| `upscalers[].image_upscale` | `boolean` | Whether this model can be selected by `POST /sdcpp/v1/upscale`; false for latent upscalers and built-in filters |
Built-in entries include `None`, `Lanczos`, `Nearest`, `Latent`, `Latent (nearest)`, `Latent (nearest-exact)`, `Latent (antialiased)`, `Latent (bicubic)`, and `Latent (bicubic antialiased)`. Model-backed entries are scanned from the top level of `--hires-upscalers-dir`; subdirectories are not scanned. Built-in entries include `None`, `Lanczos`, `Nearest`, `Latent`, `Latent (nearest)`, `Latent (nearest-exact)`, `Latent (antialiased)`, `Latent (bicubic)`, and `Latent (bicubic antialiased)`. Model-backed entries are scanned from the top level of `--hires-upscalers-dir`; subdirectories are not scanned.
@ -476,6 +493,8 @@ Built-in entries include `None`, `Lanczos`, `Nearest`, `Latent`, `Latent (neares
| `limits.max_height` | `integer` | | `limits.max_height` | `integer` |
| `limits.max_batch_count` | `integer` | | `limits.max_batch_count` | `integer` |
| `limits.max_queue_size` | `integer` | | `limits.max_queue_size` | `integer` |
| `limits.max_upscale_width` | `integer` |
| `limits.max_upscale_height` | `integer` |
Shared default fields used by both `img_gen` and `vid_gen`: Shared default fields used by both `img_gen` and `vid_gen`:
@ -526,7 +545,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH
| Field | Type | | Field | Type |
| --- | --- | | --- | --- |
| `batch_count` | `integer` | | `batch_count` | `integer` |
| `auto_resize_ref_image` | `boolean` | | `ref_image_args` | `string` |
| `increase_ref_index` | `boolean` | | `increase_ref_index` | `boolean` |
| `control_strength` | `number` | | `control_strength` | `number` |
| `ip_adapter_strength` | `number` | | `ip_adapter_strength` | `number` |
@ -628,6 +647,52 @@ Typical status codes:
- `404 Not Found` - `404 Not Found`
- `410 Gone` - `410 Gone`
#### `POST /sdcpp/v1/upscale`
Runs one RGB ESRGAN upscaler over an image, with no generation involved. Latent upscaler models remain available for hires generation but cannot be used here.
This is the HTTP equivalent of `sd-cli -M upscale`: no diffusion model, text
encoder or sampling is used, so it is fast enough to answer synchronously and
does not create a job.
Request fields:
| Field | Type | Notes |
| --- | --- | --- |
| `image` | `string` | Required. Base64 or data URL image |
| `upscaler` | `string` | A name from `upscalers` with `image_upscale: true`; the first compatible entry when omitted |
| `repeats` | `integer` | Run the upscaler this many times, 1 to 4 (default `1`) |
| `tile_size` | `integer` | Tile size, defaulting to the server's `--upscale-tile-size` |
| `output_format` | `string` | `png`, `jpeg`, or `webp` when built with WebP support (default `png`); unsupported formats return 400 |
| `output_compression` | `integer` | Range is clamped to `0..100` |
Response fields:
| Field | Type | Notes |
| --- | --- | --- |
| `images` | `array<object>` | One image |
| `images[].index` | `integer` | |
| `images[].b64_json` | `string` | Base64-encoded image bytes |
| `upscaler` | `string` | The upscaler actually used |
| `scale` | `integer` | The model's scale factor |
| `repeats` | `integer` | How many times it was run |
| `width` | `integer` | Result width |
| `height` | `integer` | Result height |
| `output_format` | `string` | Final encoded image format |
Typical status codes:
- `200 OK`
- `400 Bad Request` (invalid request, unsupported output format, unreadable image, incompatible upscaler, or output dimensions exceeding the limit)
- `500 Internal Server Error`
Notes:
- Final output dimensions, including all repeats, must not exceed 8192 pixels on either axis (`limits.max_upscale_width` and `limits.max_upscale_height`). Requests exceeding this bound are rejected before upscaling.
- The upscaler models are three-channel; alpha is not preserved.
- The request holds the generation context lock, so an upscale and a
generation never run on the device at the same time.
#### `POST /sdcpp/v1/jobs/{id}/cancel` #### `POST /sdcpp/v1/jobs/{id}/cancel`
Attempts to cancel an accepted job. Attempts to cancel an accepted job.
@ -653,7 +718,7 @@ Example:
"strength": 0.75, "strength": 0.75,
"seed": -1, "seed": -1,
"batch_count": 1, "batch_count": 1,
"auto_resize_ref_image": true, "ref_image_args": "",
"increase_ref_index": false, "increase_ref_index": false,
"control_strength": 0.9, "control_strength": 0.9,
"ip_adapter_strength": 1.0, "ip_adapter_strength": 1.0,
@ -728,6 +793,17 @@ Example:
### Image Encoding Rules ### Image Encoding Rules
Native image/video requests and SDAPI accept `image_preprocess` as a rule string
or array of rule strings. OpenAI-compatible requests can supply it in
`sd_cpp_extra_args`. See [Image preprocessing](../../docs/image_preprocessing.md)
for one-time input geometry, native-resolution decoding, mask alignment, and
`canny=true` for edge detection on any supported image input.
Image generation also accepts `ref_image_args` as a string (for example,
`"resize_before_vae=false"`) in native and SDAPI requests, or through
`sd_cpp_extra_args` in OpenAI-compatible requests. It controls downstream
reference encoding and is independent of input geometry rules.
Any image field accepts: Any image field accepts:
- a raw base64 string, or - a raw base64 string, or
@ -763,7 +839,8 @@ Top-level scalar fields:
| `strength` | `number` | | `strength` | `number` |
| `seed` | `integer` | | `seed` | `integer` |
| `batch_count` | `integer` | | `batch_count` | `integer` |
| `auto_resize_ref_image` | `boolean` | | `ref_image_args` | `string` |
| `image_preprocess` | `string \| array<string>` |
| `increase_ref_index` | `boolean` | | `increase_ref_index` | `boolean` |
| `control_strength` | `number` | | `control_strength` | `number` |
| `ip_adapter_strength` | `number` | | `ip_adapter_strength` | `number` |

View File

@ -76,9 +76,9 @@ int main(int argc, const char** argv) {
SDSvrParams svr_params; SDSvrParams svr_params;
SDContextParams ctx_params; SDContextParams ctx_params;
SDGenerationParams default_gen_params; SDGenerationParams default_gen_params;
parse_args(argc, argv, svr_params, ctx_params, default_gen_params);
sd_set_log_callback(sd_log_cb, (void*)&svr_params); sd_set_log_callback(sd_log_cb, (void*)&svr_params);
parse_args(argc, argv, svr_params, ctx_params, default_gen_params);
LOG_VERBOSE("version: %s", version_string().c_str()); LOG_VERBOSE("version: %s", version_string().c_str());
LOG_VERBOSE("%s", sd_get_system_info()); LOG_VERBOSE("%s", sd_get_system_info());

View File

@ -157,6 +157,7 @@ static bool build_openai_edit_request(const httplib::Request& req,
request.gen_params.height = height; request.gen_params.height = height;
request.gen_params.batch_count = n; request.gen_params.batch_count = n;
std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
for (auto& bytes : images_bytes) { for (auto& bytes : images_bytes) {
int img_w = 0; int img_w = 0;
int img_h = 0; int img_h = 0;
@ -165,7 +166,8 @@ static bool build_openai_edit_request(const httplib::Request& req,
reinterpret_cast<const char*>(bytes.data()), reinterpret_cast<const char*>(bytes.data()),
static_cast<int>(bytes.size()), static_cast<int>(bytes.size()),
img_w, img_h, resolved_channel, img_w, img_h, resolved_channel,
0, 0, 0); 0, 0,
0);
if (raw_pixels == nullptr) { if (raw_pixels == nullptr) {
continue; continue;
} }
@ -175,23 +177,10 @@ static bool build_openai_edit_request(const httplib::Request& req,
request.gen_params.set_width_and_height_if_unset(image_owner.get().width, image_owner.get().height); request.gen_params.set_width_and_height_if_unset(image_owner.get().width, image_owner.get().height);
if (is_first_ref_image) { if (is_first_ref_image) {
int init_w = 0; request.gen_params.init_image = image_owner;
int init_h = 0; if (request.gen_params.init_image.get().data == nullptr) {
if (request.gen_params.width_and_height_are_set()) { error_message = "could not allocate init image";
init_w = request.gen_params.width; return false;
init_h = request.gen_params.height;
}
int init_img_w = 0;
int init_img_h = 0;
int init_resolved_channel = 0;
uint8_t* init_pixels = load_image_from_memory(
reinterpret_cast<const char*>(bytes.data()),
static_cast<int>(bytes.size()),
init_img_w, init_img_h, init_resolved_channel,
init_w, init_h, 0);
if (init_pixels != nullptr) {
request.gen_params.init_image.reset({(uint32_t)init_img_w, (uint32_t)init_img_h, (uint32_t)init_resolved_channel, init_pixels});
} }
} }
@ -199,12 +188,6 @@ static bool build_openai_edit_request(const httplib::Request& req,
} }
if (!mask_bytes.empty()) { if (!mask_bytes.empty()) {
int expected_width = 0;
int expected_height = 0;
if (request.gen_params.width_and_height_are_set()) {
expected_width = request.gen_params.width;
expected_height = request.gen_params.height;
}
int mask_w = 0; int mask_w = 0;
int mask_h = 0; int mask_h = 0;
int mask_channel = 0; int mask_channel = 0;
@ -213,7 +196,7 @@ static bool build_openai_edit_request(const httplib::Request& req,
reinterpret_cast<const char*>(mask_bytes.data()), reinterpret_cast<const char*>(mask_bytes.data()),
static_cast<int>(mask_bytes.size()), static_cast<int>(mask_bytes.size()),
mask_w, mask_h, mask_channel, mask_w, mask_h, mask_channel,
expected_width, expected_height, 1); 0, 0, 1);
request.gen_params.mask_image.reset({(uint32_t)mask_w, (uint32_t)mask_h, 1, mask_raw}); request.gen_params.mask_image.reset({(uint32_t)mask_w, (uint32_t)mask_h, 1, mask_raw});
const sd_image_t& mask_image = request.gen_params.mask_image.get(); const sd_image_t& mask_image = request.gen_params.mask_image.get();
request.gen_params.set_width_and_height_if_unset(mask_image.width, mask_image.height); request.gen_params.set_width_and_height_if_unset(mask_image.width, mask_image.height);
@ -226,7 +209,6 @@ static bool build_openai_edit_request(const httplib::Request& req,
}); });
} }
std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) { if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) {
error_message = "invalid sd_cpp_extra_args"; error_message = "invalid sd_cpp_extra_args";
return false; return false;

View File

@ -80,17 +80,6 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) {
return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT; return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT;
} }
static void assign_solid_mask(SDImageOwner& mask_owner, int width, int height) {
const size_t pixel_count = static_cast<size_t>(width) * static_cast<size_t>(height);
uint8_t* raw_mask = static_cast<uint8_t*>(malloc(pixel_count));
if (raw_mask == nullptr) {
mask_owner.reset({0, 0, 1, nullptr});
return;
}
std::memset(raw_mask, 255, pixel_count);
mask_owner.reset({(uint32_t)width, (uint32_t)height, 1, raw_mask});
}
static bool build_sdapi_img_gen_request(const json& j, static bool build_sdapi_img_gen_request(const json& j,
ServerRuntime& runtime, ServerRuntime& runtime,
bool img2img, bool img2img,
@ -193,15 +182,25 @@ static bool build_sdapi_img_gen_request(const json& j,
} }
} }
if (img2img) { if (j.contains("ref_image_args")) {
const int expected_width = request.gen_params.width_and_height_are_set() ? request.gen_params.width : 0; if (!j["ref_image_args"].is_string()) {
const int expected_height = request.gen_params.width_and_height_are_set() ? request.gen_params.height : 0; error_message = "ref_image_args must be a string";
return false;
}
request.gen_params.ref_image_args = j["ref_image_args"].get<std::string>();
}
if (j.contains("image_preprocess") && !request.gen_params.parse_image_preprocess_json(j["image_preprocess"].dump())) {
error_message = "invalid image_preprocess";
return false;
}
if (img2img) {
if (j.contains("init_images") && j["init_images"].is_array() && !j["init_images"].empty()) { if (j.contains("init_images") && j["init_images"].is_array() && !j["init_images"].empty()) {
if (decode_base64_image(j["init_images"][0].get<std::string>(), if (decode_base64_image(j["init_images"][0].get<std::string>(),
0, 0,
expected_width, 0,
expected_height, 0,
request.gen_params.init_image)) { request.gen_params.init_image)) {
const sd_image_t& image = request.gen_params.init_image.get(); const sd_image_t& image = request.gen_params.init_image.get();
request.gen_params.set_width_and_height_if_unset(image.width, image.height); request.gen_params.set_width_and_height_if_unset(image.width, image.height);
@ -211,8 +210,8 @@ static bool build_sdapi_img_gen_request(const json& j,
if (j.contains("mask") && j["mask"].is_string()) { if (j.contains("mask") && j["mask"].is_string()) {
if (decode_base64_image(j["mask"].get<std::string>(), if (decode_base64_image(j["mask"].get<std::string>(),
1, 1,
expected_width, 0,
expected_height, 0,
request.gen_params.mask_image)) { request.gen_params.mask_image)) {
const sd_image_t& image = request.gen_params.mask_image.get(); const sd_image_t& image = request.gen_params.mask_image.get();
request.gen_params.set_width_and_height_if_unset(image.width, image.height); request.gen_params.set_width_and_height_if_unset(image.width, image.height);
@ -225,9 +224,7 @@ static bool build_sdapi_img_gen_request(const json& j,
} }
} }
} else { } else {
const int resolved_width = request.gen_params.get_resolved_width(); request.gen_params.mask_image.reset({0, 0, 1, nullptr});
const int resolved_height = request.gen_params.get_resolved_height();
assign_solid_mask(request.gen_params.mask_image, resolved_width, resolved_height);
} }
float denoising_strength = j.value("denoising_strength", -1.f); float denoising_strength = j.value("denoising_strength", -1.f);
@ -244,12 +241,7 @@ static bool build_sdapi_img_gen_request(const json& j,
SDImageOwner image_owner; SDImageOwner image_owner;
if (decode_base64_image(extra_image.get<std::string>(), if (decode_base64_image(extra_image.get<std::string>(),
0, 0,
request.gen_params.auto_resize_ref_image && request.gen_params.width_and_height_are_set() 0, 0,
? request.gen_params.width
: 0,
request.gen_params.auto_resize_ref_image && request.gen_params.width_and_height_are_set()
? request.gen_params.height
: 0,
image_owner)) { image_owner)) {
const sd_image_t& image = image_owner.get(); const sd_image_t& image = image_owner.get();
request.gen_params.set_width_and_height_if_unset(image.width, image.height); request.gen_params.set_width_and_height_if_unset(image.width, image.height);

View File

@ -3,12 +3,33 @@
#include <algorithm> #include <algorithm>
#include <cmath> #include <cmath>
#include <filesystem> #include <filesystem>
#include <limits>
#include "async_jobs.h" #include "async_jobs.h"
#include "common/common.h" #include "common/common.h"
#include "common/media_io.h"
#include "common/resource_owners.hpp"
namespace fs = std::filesystem; namespace fs = std::filesystem;
static constexpr uint32_t k_max_upscale_dimension = 8192;
static bool valid_upscale_dimensions(const sd_image_t& image, int factor, int repeats) {
if (image.width == 0 || image.height == 0 || factor < 1 || repeats < 1 || repeats > 4) {
return false;
}
uint32_t width = image.width;
uint32_t height = image.height;
for (int i = 0; i < repeats; ++i) {
if (width > k_max_upscale_dimension / factor || height > k_max_upscale_dimension / factor) {
return false;
}
width *= factor;
height *= factor;
}
return true;
}
static bool parse_cache_mode(const std::string& mode_str, sd_cache_mode_t& mode_out) { static bool parse_cache_mode(const std::string& mode_str, sd_cache_mode_t& mode_out) {
if (mode_str == "disabled") { if (mode_str == "disabled") {
mode_out = SD_CACHE_DISABLED; mode_out = SD_CACHE_DISABLED;
@ -127,7 +148,8 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
{"seed", defaults.seed}, {"seed", defaults.seed},
{"batch_count", defaults.batch_count}, {"batch_count", defaults.batch_count},
{"qwen_image_layers", defaults.qwen_image_layers}, {"qwen_image_layers", defaults.qwen_image_layers},
{"auto_resize_ref_image", defaults.auto_resize_ref_image}, {"ref_image_args", defaults.ref_image_args},
{"image_preprocess", defaults.image_preprocess},
{"increase_ref_index", defaults.increase_ref_index}, {"increase_ref_index", defaults.increase_ref_index},
{"control_strength", defaults.control_strength}, {"control_strength", defaults.control_strength},
{"ip_adapter_strength", defaults.ip_adapter_strength}, {"ip_adapter_strength", defaults.ip_adapter_strength},
@ -153,6 +175,7 @@ static json make_vid_gen_defaults_json(const SDGenerationParams& defaults, const
{"strength", defaults.strength}, {"strength", defaults.strength},
{"seed", defaults.seed}, {"seed", defaults.seed},
{"video_frames", defaults.video_frames}, {"video_frames", defaults.video_frames},
{"image_preprocess", defaults.image_preprocess},
{"fps", defaults.fps}, {"fps", defaults.fps},
{"moe_boundary", defaults.moe_boundary}, {"moe_boundary", defaults.moe_boundary},
{"vace_strength", defaults.vace_strength}, {"vace_strength", defaults.vace_strength},
@ -239,37 +262,59 @@ static json make_capabilities_json(ServerRuntime& runtime) {
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "None"}, {"name", "None"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Lanczos"}, {"name", "Lanczos"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Nearest"}, {"name", "Nearest"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Latent"}, {"name", "Latent"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Latent (nearest)"}, {"name", "Latent (nearest)"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Latent (nearest-exact)"}, {"name", "Latent (nearest-exact)"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Latent (antialiased)"}, {"name", "Latent (antialiased)"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Latent (bicubic)"}, {"name", "Latent (bicubic)"},
{"model", false},
{"image_upscale", false},
}); });
available_upscalers.push_back({ available_upscalers.push_back({
{"name", "Latent (bicubic antialiased)"}, {"name", "Latent (bicubic antialiased)"},
{"model", false},
{"image_upscale", false},
}); });
bool have_upscaler_models = false;
{ {
std::lock_guard<std::mutex> lock(*runtime.upscaler_mutex); std::lock_guard<std::mutex> lock(*runtime.upscaler_mutex);
for (const auto& entry : *runtime.upscaler_cache) { for (const auto& entry : *runtime.upscaler_cache) {
available_upscalers.push_back({ available_upscalers.push_back({
{"name", entry.name}, {"name", entry.name},
{"model", true},
{"image_upscale", entry.image_upscale_factor > 0},
}); });
have_upscaler_models = have_upscaler_models || entry.image_upscale_factor > 0;
} }
} }
@ -339,6 +384,8 @@ static json make_capabilities_json(ServerRuntime& runtime) {
{"max_height", 4096}, {"max_height", 4096},
{"max_batch_count", 8}, {"max_batch_count", 8},
{"max_queue_size", manager.max_pending_jobs}, {"max_queue_size", manager.max_pending_jobs},
{"max_upscale_width", k_max_upscale_dimension},
{"max_upscale_height", k_max_upscale_dimension},
}; };
result["samplers"] = samplers; result["samplers"] = samplers;
result["schedulers"] = schedulers; result["schedulers"] = schedulers;
@ -348,6 +395,7 @@ static json make_capabilities_json(ServerRuntime& runtime) {
result["features_by_mode"] = features_by_mode; result["features_by_mode"] = features_by_mode;
result["loras"] = available_loras; result["loras"] = available_loras;
result["upscalers"] = available_upscalers; result["upscalers"] = available_upscalers;
result["upscale"] = have_upscaler_models;
return result; return result;
} }
@ -413,6 +461,171 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
res.set_content(make_capabilities_json(*runtime).dump(), "application/json"); res.set_content(make_capabilities_json(*runtime).dump(), "application/json");
}); });
svr.Post("/sdcpp/v1/upscale", [runtime](const httplib::Request& req, httplib::Response& res) {
try {
if (req.body.empty()) {
res.status = 400;
res.set_content(R"({"error":"empty body"})", "application/json");
return;
}
json body = json::parse(req.body);
if (!body.is_object()) {
res.status = 400;
res.set_content(R"({"error":"body must be an object"})", "application/json");
return;
}
for (const char* key : {"repeats", "tile_size", "output_compression"}) {
if (!body.contains(key)) {
continue;
}
const auto& value = body[key];
const bool valid = value.is_number_unsigned()
? value.get<uint64_t>() <= static_cast<uint64_t>(std::numeric_limits<int>::max())
: value.is_number_integer() && value.get<int64_t>() >= std::numeric_limits<int>::min() &&
value.get<int64_t>() <= std::numeric_limits<int>::max();
if (!valid) {
res.status = 400;
res.set_content(json({{"error", std::string(key) + " must be a 32-bit integer"}}).dump(), "application/json");
return;
}
}
ImgGenJobRequest output_options;
std::string error_message;
if (!assign_output_options(output_options,
body.value("output_format", std::string("png")),
body.value("output_compression", 100),
true,
error_message)) {
res.status = 400;
res.set_content(json({{"error", error_message}}).dump(), "application/json");
return;
}
const int tile_size = std::max(32, body.value("tile_size", runtime->default_gen_params->upscale_tile_size));
const int repeats = std::clamp(body.value("repeats", 1), 1, 4);
const std::string wanted = body.value("upscaler", std::string());
const std::string encoded = body.value("image", std::string());
if (encoded.empty()) {
res.status = 400;
res.set_content(R"({"error":"image is required"})", "application/json");
return;
}
SDImageOwner input;
if (!decode_base64_image(encoded, 3, 0, 0, input) || input.get().data == nullptr) {
res.status = 400;
res.set_content(R"({"error":"image could not be read"})", "application/json");
return;
}
refresh_upscaler_cache(*runtime);
int model_scale = 0;
std::string model_path;
std::string used_name;
{
std::lock_guard<std::mutex> lock(*runtime->upscaler_mutex);
for (const auto& entry : *runtime->upscaler_cache) {
if (entry.image_upscale_factor > 0 && (wanted.empty() || entry.name == wanted)) {
model_path = entry.fullpath;
used_name = entry.name;
model_scale = entry.image_upscale_factor;
break;
}
}
}
if (model_path.empty()) {
res.status = 400;
res.set_content(json({{"error", wanted.empty()
? std::string("no RGB ESRGAN upscaler models are available; "
"start the server with --hires-upscalers-dir")
: "no compatible image upscaler called " + wanted}})
.dump(),
"application/json");
return;
}
if (!valid_upscale_dimensions(input.get(), model_scale, repeats)) {
res.status = 400;
res.set_content(R"({"error":"upscaled dimensions must not exceed 8192 x 8192"})", "application/json");
return;
}
// One GPU: an upscale must not run while a generation is using it.
std::lock_guard<std::mutex> ctx_lock(*runtime->sd_ctx_mutex);
UpscalerCtxPtr upscaler_ctx(new_upscaler_ctx(model_path.c_str(),
runtime->ctx_params->diffusion_conv_direct,
runtime->ctx_params->n_threads,
tile_size,
runtime->ctx_params->backend.c_str(),
runtime->ctx_params->params_backend.c_str()));
if (upscaler_ctx == nullptr) {
res.status = 500;
res.set_content(R"({"error":"the upscaler model could not be loaded"})", "application/json");
return;
}
const int factor = get_upscale_factor(upscaler_ctx.get());
// The model file may have changed since its metadata was cached.
if (!valid_upscale_dimensions(input.get(), factor, repeats)) {
res.status = 400;
res.set_content(R"({"error":"upscaled dimensions must not exceed 8192 x 8192"})", "application/json");
return;
}
SDImageOwner current(input.release());
for (int i = 0; i < repeats; ++i) {
sd_image_t* out_images = nullptr;
int out_count = 0;
if (!upscale(upscaler_ctx.get(), current.get(), (uint32_t)factor, &out_images, &out_count) ||
out_count <= 0 || out_images[0].data == nullptr) {
free_sd_images(out_images, out_count);
res.status = 500;
res.set_content(R"({"error":"upscale failed"})", "application/json");
return;
}
sd_image_t produced = out_images[0];
out_images[0] = {0, 0, 0, nullptr};
free_sd_images(out_images, out_count);
current.reset(produced);
}
const std::string& format = output_options.output_format;
const int compression = output_options.output_compression;
const sd_image_t result = current.get();
auto image_bytes = encode_image_to_vector(format == "jpeg" ? EncodedImageFormat::JPEG
: format == "webp" ? EncodedImageFormat::WEBP
: EncodedImageFormat::PNG,
result.data,
result.width,
result.height,
result.channel,
"",
compression);
if (image_bytes.empty()) {
res.status = 500;
res.set_content(R"({"error":"the result could not be encoded"})", "application/json");
return;
}
json out;
out["upscaler"] = used_name;
out["scale"] = factor;
out["repeats"] = repeats;
out["width"] = result.width;
out["height"] = result.height;
out["output_format"] = format;
json images = json::array();
images.push_back({{"index", 0}, {"b64_json", base64_encode(image_bytes)}});
out["images"] = std::move(images);
res.set_content(out.dump(), "application/json");
res.status = 200;
} catch (const json::exception& e) {
res.status = 400;
res.set_content(json({{"error", "invalid request"}, {"message", e.what()}}).dump(), "application/json");
} catch (const std::exception& e) {
res.status = 500;
res.set_content(json({{"error", std::string("server_error: ") + e.what()}}).dump(), "application/json");
}
});
svr.Post("/sdcpp/v1/img_gen", [runtime](const httplib::Request& req, httplib::Response& res) { svr.Post("/sdcpp/v1/img_gen", [runtime](const httplib::Request& req, httplib::Response& res) {
try { try {
if (req.body.empty()) { if (req.body.empty()) {

View File

@ -295,6 +295,11 @@ std::string get_lora_full_path(ServerRuntime& rt, const std::string& path) {
void refresh_upscaler_cache(ServerRuntime& rt) { void refresh_upscaler_cache(ServerRuntime& rt) {
std::vector<UpscalerEntry> new_cache; std::vector<UpscalerEntry> new_cache;
std::vector<UpscalerEntry> previous_cache;
{
std::lock_guard<std::mutex> lock(*rt.upscaler_mutex);
previous_cache = *rt.upscaler_cache;
}
fs::path upscaler_dir = rt.ctx_params->hires_upscalers_dir; fs::path upscaler_dir = rt.ctx_params->hires_upscalers_dir;
if (fs::exists(upscaler_dir) && fs::is_directory(upscaler_dir)) { if (fs::exists(upscaler_dir) && fs::is_directory(upscaler_dir)) {
@ -308,10 +313,24 @@ void refresh_upscaler_cache(ServerRuntime& rt) {
} }
UpscalerEntry upscaler_entry; UpscalerEntry upscaler_entry;
upscaler_entry.name = p.stem().u8string(); upscaler_entry.name = p.stem().u8string();
upscaler_entry.fullpath = fs::absolute(p).lexically_normal().u8string(); upscaler_entry.fullpath = fs::absolute(p).lexically_normal().u8string();
upscaler_entry.model_name = "ESRGAN_4x"; upscaler_entry.model_name = "ESRGAN_4x";
upscaler_entry.path = p.filename().u8string(); upscaler_entry.path = p.filename().u8string();
upscaler_entry.file_size = entry.file_size();
upscaler_entry.last_modified = entry.last_write_time();
auto previous = std::find_if(previous_cache.begin(), previous_cache.end(), [&](const UpscalerEntry& cached) {
return cached.fullpath == upscaler_entry.fullpath &&
cached.file_size == upscaler_entry.file_size &&
cached.last_modified == upscaler_entry.last_modified;
});
upscaler_entry.image_upscale_factor = previous != previous_cache.end()
? previous->image_upscale_factor
: get_upscaler_model_scale(upscaler_entry.fullpath.c_str());
if (upscaler_entry.image_upscale_factor > 0) {
upscaler_entry.scale = upscaler_entry.image_upscale_factor;
upscaler_entry.model_name = "ESRGAN_" + std::to_string(upscaler_entry.scale) + "x";
}
new_cache.push_back(std::move(upscaler_entry)); new_cache.push_back(std::move(upscaler_entry));
} }

View File

@ -2,6 +2,7 @@
#include <algorithm> #include <algorithm>
#include <cstdint> #include <cstdint>
#include <filesystem>
#include <mutex> #include <mutex>
#include <string> #include <string>
#include <vector> #include <vector>
@ -43,6 +44,9 @@ struct UpscalerEntry {
std::string fullpath; std::string fullpath;
std::string model_name; std::string model_name;
int scale = 4; int scale = 4;
int image_upscale_factor = 0;
uintmax_t file_size = 0;
std::filesystem::file_time_type last_modified;
}; };
struct ServerRuntime { struct ServerRuntime {

View File

@ -247,6 +247,7 @@ typedef struct {
float attn_scale; // Override flash-attention K/V 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 const char* tokenizer; // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens
bool sage_attn; bool sage_attn;
int conditioning_cache_size; // Maximum cached conditioning entries per context; 0 disables caching (default: 4)
} sd_ctx_params_t; } sd_ctx_params_t;
typedef struct { typedef struct {
@ -263,6 +264,11 @@ typedef struct {
uint8_t* data; uint8_t* data;
} sd_image_t; } sd_image_t;
typedef struct {
// Semicolon-separated target=...,key=value rules. NULL preserves defaults.
const char* rules;
} sd_image_preprocess_params_t;
typedef struct { typedef struct {
sd_image_t* frames; sd_image_t* frames;
int frame_count; int frame_count;
@ -410,6 +416,7 @@ typedef struct {
int qwen_image_layers; int qwen_image_layers;
bool circular_x; bool circular_x;
bool circular_y; bool circular_y;
sd_image_preprocess_params_t image_preprocess;
} sd_img_gen_params_t; } sd_img_gen_params_t;
typedef struct { typedef struct {
@ -443,6 +450,7 @@ typedef struct {
sd_hires_params_t hires; sd_hires_params_t hires;
bool circular_x; bool circular_x;
bool circular_y; bool circular_y;
sd_image_preprocess_params_t image_preprocess;
} sd_vid_gen_params_t; } sd_vid_gen_params_t;
typedef struct sd_ctx_t sd_ctx_t; typedef struct sd_ctx_t sd_ctx_t;
@ -550,6 +558,8 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
int* num_images_out); int* num_images_out);
SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
// Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model.
SD_API int get_upscaler_model_scale(const char* model_path);
typedef struct adetailer_ctx_t adetailer_ctx_t; typedef struct adetailer_ctx_t adetailer_ctx_t;

View File

@ -3166,6 +3166,7 @@ struct LLMEmbedder : public Conditioner {
int64_t tag_count = static_cast<int64_t>(tags.size()); int64_t tag_count = static_cast<int64_t>(tags.size());
result.c_token_types = sd::Tensor<int32_t>({tag_count}, std::move(tags)); result.c_token_types = sd::Tensor<int32_t>({tag_count}, std::move(tags));
} }
return result; return result;
} }
}; };

View File

@ -0,0 +1,107 @@
#ifndef __SD_CONDITIONING_CONDITIONING_CACHE_H__
#define __SD_CONDITIONING_CONDITIONING_CACHE_H__
#include <algorithm>
#include <list>
#include <tuple>
#include "conditioning/conditioner.hpp"
class ConditioningCache {
struct Entry {
ConditionerParams params;
std::vector<sd::Tensor<float>> ref_images;
std::vector<MiniMaxH3PresentationItem> references;
SDCondition condition;
Entry(const ConditionerParams& input, const SDCondition& output)
: params(input), condition(output) {
// Request-owned reference pointers must not outlive the request.
if (input.ref_images != nullptr) {
ref_images = *input.ref_images;
params.ref_images = &ref_images;
}
if (input.minimax_h3_references != nullptr) {
references = *input.minimax_h3_references;
params.minimax_h3_references = &references;
}
}
Entry(const Entry&) = delete;
Entry& operator=(const Entry&) = delete;
};
size_t capacity_ = 4;
std::list<Entry> entries_;
static bool same_images(const std::vector<sd::Tensor<float>>& a,
const std::vector<sd::Tensor<float>>& b) {
return std::equal(a.begin(), a.end(), b.begin(), b.end(),
[](const sd::Tensor<float>& x, const sd::Tensor<float>& y) {
return x.shape() == y.shape() && x.values() == y.values();
});
}
static bool same_params(const ConditionerParams& a, const ConditionerParams& b) {
const auto fields = [](const ConditionerParams& p) {
const auto& r = p.ref_image_params;
return std::tie(p.text, p.clip_skip, p.width, p.height, p.zero_out_masked,
r.pass_to_vlm, r.pass_to_dit, r.ref_index_mode,
r.force_ref_timestep_zero, r.resize_before_vae, r.vae_input_max_pixels,
r.vlm_resize_mode, r.vlm_min_size, r.vlm_max_size, r.resize_vae_to_target);
};
if (fields(a) != fields(b) ||
(a.ref_images == nullptr) != (b.ref_images == nullptr) ||
(a.minimax_h3_references == nullptr) != (b.minimax_h3_references == nullptr)) {
return false;
}
if (a.ref_images != nullptr && !same_images(*a.ref_images, *b.ref_images)) {
return false;
}
if (a.minimax_h3_references != nullptr &&
!std::equal(a.minimax_h3_references->begin(), a.minimax_h3_references->end(),
b.minimax_h3_references->begin(), b.minimax_h3_references->end(),
[](const MiniMaxH3PresentationItem& x, const MiniMaxH3PresentationItem& y) {
return x.kind == y.kind && x.timestamps == y.timestamps && same_images(x.frames, y.frames);
})) {
return false;
}
return true;
}
public:
void set_capacity(size_t capacity) {
capacity_ = capacity;
while (entries_.size() > capacity_) {
entries_.pop_back();
}
}
void clear() {
entries_.clear();
}
SDCondition get(Conditioner& conditioner, int n_threads, const ConditionerParams& params) {
if (capacity_ == 0) {
return conditioner.get_learned_condition(n_threads, params);
}
for (auto it = entries_.begin(); it != entries_.end(); ++it) {
if (same_params(it->params, params)) {
entries_.splice(entries_.begin(), entries_, it);
LOG_INFO("conditioning cache hit");
return entries_.front().condition;
}
}
auto condition = conditioner.get_learned_condition(n_threads, params);
if (!condition.empty()) {
if (entries_.size() == capacity_) {
entries_.pop_back();
}
entries_.emplace_front(params, condition);
LOG_VERBOSE("conditioning cache stored (%zu/%zu)", entries_.size(), capacity_);
}
return condition;
}
};
#endif // __SD_CONDITIONING_CONDITIONING_CACHE_H__

View File

@ -623,7 +623,11 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
bool skip_reshape, bool skip_reshape,
bool flash_attn, bool flash_attn,
float kv_scale, float kv_scale,
bool sage_attn) { // avoid overflow bool sage_attn,
bool* used_flash_attn) { // avoid overflow
if (used_flash_attn != nullptr) {
*used_flash_attn = false;
}
int64_t L_q; int64_t L_q;
int64_t L_k; int64_t L_k;
int64_t C; int64_t C;
@ -755,6 +759,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
if (can_use_flash_attn) { if (can_use_flash_attn) {
kqv = build_kqv(q, k, v, mask); kqv = build_kqv(q, k, v, mask);
if (kqv != nullptr) { if (kqv != nullptr) {
if (used_flash_attn != nullptr) {
*used_flash_attn = true;
}
kqv = ggml_view_4d(ctx, kqv = ggml_view_4d(ctx,
kqv, kqv,
d_head, d_head,

View File

@ -217,11 +217,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
ggml_tensor* k, ggml_tensor* k,
ggml_tensor* v, ggml_tensor* v,
int64_t n_head, int64_t n_head,
ggml_tensor* mask = nullptr, ggml_tensor* mask = nullptr,
bool skip_reshape = false, bool skip_reshape = false,
bool flash_attn = false, bool flash_attn = false,
float kv_scale = 1.0f, float kv_scale = 1.0f,
bool sage_attn = false); bool sage_attn = false,
bool* used_flash_attn = nullptr);
ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
ggml_tensor* x, ggml_tensor* x,

View File

@ -13,6 +13,7 @@
#endif #endif
#include "core/util.h" #include "core/util.h"
#include "ggml-backend-impl.h"
#include "ggml-impl.h" #include "ggml-impl.h"
#include "stable-diffusion.h" #include "stable-diffusion.h"
@ -433,6 +434,24 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU; return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
} }
ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
void* ptr,
size_t size,
size_t max_tensor_size) {
ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
if (buffer != nullptr && buffer->context == nullptr) {
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device);
if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
// Metal can wrap a failed mapping in a non-null buffer. Its free callback also
// dereferences the missing context, so only release the outer buffer.
buffer->iface.free_buffer = nullptr;
ggml_backend_buffer_free(buffer);
return nullptr;
}
}
return buffer;
}
bool sd_backend_supports_cuda_mma(ggml_backend_t backend) { bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
#ifdef SD_USE_CUDA #ifdef SD_USE_CUDA
if (!sd_backend_is(backend, "CUDA")) { if (!sd_backend_is(backend, "CUDA")) {

View File

@ -88,6 +88,10 @@ private:
bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is(ggml_backend_t backend, const std::string& name);
bool sd_backend_is_cpu(ggml_backend_t backend); bool sd_backend_is_cpu(ggml_backend_t backend);
bool sd_backend_supports_cuda_mma(ggml_backend_t backend); bool sd_backend_supports_cuda_mma(ggml_backend_t backend);
ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
void* ptr,
size_t size,
size_t max_tensor_size);
ggml_backend_t sd_backend_cpu_init(); ggml_backend_t sd_backend_cpu_init();
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,

View File

@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* mask, ggml_tensor* mask,
bool skip_reshape, bool skip_reshape,
bool flash_attn, bool flash_attn,
float kv_scale) { float kv_scale,
bool* used_flash_attn) {
if (ctx->attn_scale > 0.f) { if (ctx->attn_scale > 0.f) {
kv_scale = ctx->attn_scale; 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, ctx->sage_attn_enabled); return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
} }
void GGMLRunner::alloc_params_ctx() { void GGMLRunner::alloc_params_ctx() {
@ -515,9 +516,10 @@ GGMLRunner::~GGMLRunner() {
free_params_ctx(); free_params_ctx();
} }
GGMLRunnerContext GGMLRunner::get_context() { GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) {
GGMLRunnerContext runner_ctx; GGMLRunnerContext runner_ctx;
runner_ctx.ggml_ctx = compute_ctx; runner_ctx.ggml_ctx = compute_ctx;
runner_ctx.graph = graph;
runner_ctx.backend = runtime_backend; runner_ctx.backend = runtime_backend;
runner_ctx.flash_attn_enabled = flash_attn_enabled; runner_ctx.flash_attn_enabled = flash_attn_enabled;
runner_ctx.sage_attn_enabled = sage_attn_enabled; runner_ctx.sage_attn_enabled = sage_attn_enabled;
@ -532,8 +534,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
runner_ctx.get_cache_tensor = [this](const std::string& name) { runner_ctx.get_cache_tensor = [this](const std::string& name) {
return this->get_cache_tensor_by_name(name); return this->get_cache_tensor_by_name(name);
}; };
runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) { runner_ctx.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) {
this->cache(name, tensor); this->cache(name, tensor, graph);
}; };
runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) { runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
this->set_backend_tensor_data(tensor, data); this->set_backend_tensor_data(tensor, data);
@ -575,7 +577,7 @@ ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
} }
} }
void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) { void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph) {
if (tensor != nullptr && tensor->view_src != nullptr) { if (tensor != nullptr && tensor->view_src != nullptr) {
tensor = ggml_cont(compute_ctx, tensor); tensor = ggml_cont(compute_ctx, tensor);
} }
@ -583,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
ggml_set_output(tensor); ggml_set_output(tensor);
} }
cache_.stage(name, tensor); cache_.stage(name, tensor);
if (graph != nullptr && tensor != nullptr) {
// Schedule the cache output here so its source can be reused before graph end.
ggml_build_forward_expand(graph, tensor);
}
} }
std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph, std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
@ -644,6 +650,10 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
std::optional<sd::Tensor<float>> output; std::optional<sd::Tensor<float>> output;
try { try {
output = execute_graph(graph, n_threads, no_return, read_outputs); output = execute_graph(graph, n_threads, no_return, read_outputs);
} catch (const std::bad_alloc&) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
LOG_ERROR("%s graph allocation failed", get_desc().c_str());
return std::nullopt;
} catch (const std::exception& error) { } catch (const std::exception& error) {
last_compute_status_ = GGML_STATUS_FAILED; last_compute_status_ = GGML_STATUS_FAILED;
LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
@ -827,11 +837,20 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
last_compute_status_ = GGML_STATUS_ALLOC_FAILED; last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return std::nullopt; return std::nullopt;
} }
auto fits_monolithic = [&]() {
// Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
for (auto& request : requests) {
request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
}
return fits(requests, params);
};
auto manager = residency_manager.lock(); auto manager = residency_manager.lock();
const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
manager != nullptr && manager->segmented_compute_enabled() && manager != nullptr && manager->segmented_compute_enabled() &&
cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 && cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
!fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params); !fits_monolithic();
ggml_graph_cut::Plan monolithic_plan; ggml_graph_cut::Plan monolithic_plan;
if (!segmented) { if (!segmented) {
monolithic_plan.segments.emplace_back(); monolithic_plan.segments.emplace_back();
@ -964,10 +983,16 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
} }
LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
index + 1, plan.segments.size(), segment.group_name.c_str()); index + 1, plan.segments.size(), segment.group_name.c_str());
if (!execute_segment(segment_graph, n_threads) || if (!execute_segment(segment_graph, n_threads)) {
!cache_.capture(segment_graph) || return fail_segment("execution");
!cut_cache_.capture(graph, segment, get_desc().c_str())) { }
return fail_segment("execution or output caching"); auto cache_status = cache_.capture(segment_graph);
if (cache_status == GGML_STATUS_SUCCESS) {
cache_status = cut_cache_.capture(graph, segment, get_desc().c_str());
}
if (cache_status != GGML_STATUS_SUCCESS) {
last_compute_status_ = cache_status;
return fail_segment("output caching");
} }
sync_runtime_residency(); sync_runtime_residency();
if (last) { if (last) {

View File

@ -67,6 +67,7 @@ struct WeightAdapter {
struct GGMLRunnerContext { struct GGMLRunnerContext {
ggml_backend_t backend = nullptr; ggml_backend_t backend = nullptr;
ggml_context* ggml_ctx = nullptr; ggml_context* ggml_ctx = nullptr;
ggml_cgraph* graph = nullptr;
bool flash_attn_enabled = false; bool flash_attn_enabled = false;
bool sage_attn_enabled = false; bool sage_attn_enabled = false;
float linear_scale = 0.f; float linear_scale = 0.f;
@ -102,6 +103,12 @@ struct GGMLRunnerContext {
return get_cache_tensor(name); return get_cache_tensor(name);
} }
void expand_graph(ggml_tensor* tensor) const {
if (graph != nullptr && tensor != nullptr) {
ggml_build_forward_expand(graph, tensor);
}
}
void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const { void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
if (!cache_tensor || tensor == nullptr) { if (!cache_tensor || tensor == nullptr) {
return; return;
@ -122,10 +129,11 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* k, ggml_tensor* k,
ggml_tensor* v, ggml_tensor* v,
int64_t n_head, int64_t n_head,
ggml_tensor* mask = nullptr, ggml_tensor* mask = nullptr,
bool skip_reshape = false, bool skip_reshape = false,
bool flash_attn = false, bool flash_attn = false,
float kv_scale = 1.f); float kv_scale = 1.f,
bool* used_flash_attn = nullptr);
struct GGMLRunner { struct GGMLRunner {
private: private:
@ -289,7 +297,8 @@ public:
virtual ~GGMLRunner(); virtual ~GGMLRunner();
virtual GGMLRunnerContext get_context(); // Binding a graph schedules cache outputs at registration instead of graph end.
virtual GGMLRunnerContext get_context(ggml_cgraph* graph = nullptr);
void reset_compute_ctx(); void reset_compute_ctx();
@ -324,7 +333,7 @@ public:
ggml_tensor* to_backend(ggml_tensor* tensor); ggml_tensor* to_backend(ggml_tensor* tensor);
void cache(const std::string name, ggml_tensor* tensor); void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr);
ggml_tensor* get_cache_tensor_by_name(const std::string& name) { ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
return cache_.get(name); return cache_.get(name);

View File

@ -26,10 +26,13 @@ namespace sd {
std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend, std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
const std::string& name, const std::string& name,
ggml_tensor* source) { ggml_tensor* source,
ggml_status& status) {
status = GGML_STATUS_FAILED;
if (ggml_graph_cut::tensor_buffer(source) == nullptr) { if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
return nullptr; return nullptr;
} }
status = GGML_STATUS_ALLOC_FAILED;
auto entry = std::make_unique<CachedTensor>(); auto entry = std::make_unique<CachedTensor>();
entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
if (entry->context == nullptr) { if (entry->context == nullptr) {
@ -50,6 +53,7 @@ namespace sd {
} else { } else {
ggml_backend_tensor_copy(source, entry->tensor); ggml_backend_tensor_copy(source, entry->tensor);
} }
status = GGML_STATUS_SUCCESS;
return entry; return entry;
} }
@ -106,9 +110,9 @@ namespace sd {
return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending; return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending;
} }
bool RunnerCache::capture(ggml_cgraph* graph) { ggml_status RunnerCache::capture(ggml_cgraph* graph) {
if (outputs_.empty()) { if (outputs_.empty()) {
return true; return GGML_STATUS_SUCCESS;
} }
const auto tensors = cache_graph_tensors(graph); const auto tensors = cache_graph_tensors(graph);
for (const auto& output : outputs_) { for (const auto& output : outputs_) {
@ -116,14 +120,15 @@ namespace sd {
continue; continue;
} }
GGML_ASSERT(ggml_is_contiguous(output.second)); GGML_ASSERT(ggml_is_contiguous(output.second));
auto entry = CachedTensor::copy(backend_, output.first, output.second); ggml_status status;
auto entry = CachedTensor::copy(backend_, output.first, output.second, status);
if (entry == nullptr) { if (entry == nullptr) {
return false; return status;
} }
pending_[output.first] = std::move(entry); pending_[output.first] = std::move(entry);
} }
ggml_backend_synchronize(backend_); ggml_backend_synchronize(backend_);
return true; return GGML_STATUS_SUCCESS;
} }
void RunnerCache::graph_end(bool success) { void RunnerCache::graph_end(bool success) {
@ -180,9 +185,9 @@ namespace sd {
} }
} }
bool GraphCutTensorCache::capture(ggml_cgraph* graph, ggml_status GraphCutTensorCache::capture(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment, const ggml_graph_cut::Segment& segment,
const char* log_desc) { const char* log_desc) {
size_t copied_bytes = 0; size_t copied_bytes = 0;
size_t copied_count = 0; size_t copied_count = 0;
for (int index : segment.output_node_indices) { for (int index : segment.output_node_indices) {
@ -191,10 +196,11 @@ namespace sd {
!segment.future_cut_names.count(output->name)) { !segment.future_cut_names.count(output->name)) {
continue; continue;
} }
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output)); ggml_status status;
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status);
if (entry == nullptr) { if (entry == nullptr) {
LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name); LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name);
return false; return status;
} }
const size_t size = ggml_backend_buffer_get_size(entry->buffer); const size_t size = ggml_backend_buffer_get_size(entry->buffer);
copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size; copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size;
@ -206,6 +212,6 @@ namespace sd {
LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)", LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)",
log_desc, copied_bytes / (1024.f * 1024.f), copied_count); log_desc, copied_bytes / (1024.f * 1024.f), copied_count);
} }
return true; return GGML_STATUS_SUCCESS;
} }
} }

View File

@ -20,7 +20,8 @@ namespace sd {
~CachedTensor(); ~CachedTensor();
static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend, static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
const std::string& name, const std::string& name,
ggml_tensor* source); ggml_tensor* source,
ggml_status& status);
}; };
using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>; using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>;
@ -41,7 +42,8 @@ namespace sd {
const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; } const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; }
size_t pending_bytes(ggml_cgraph* graph) const; size_t pending_bytes(ggml_cgraph* graph) const;
size_t resident_bytes(ggml_backend_dev_t device) const; size_t resident_bytes(ggml_backend_dev_t device) const;
bool capture(ggml_cgraph* graph); bool empty() const { return committed_.empty(); }
ggml_status capture(ggml_cgraph* graph);
void graph_end(bool success); void graph_end(bool success);
void clear(); void clear();
}; };
@ -57,7 +59,7 @@ namespace sd {
size_t resident_bytes(ggml_backend_dev_t device) const; size_t resident_bytes(ggml_backend_dev_t device) const;
size_t estimate_output_bytes(ggml_cgraph* graph, size_t estimate_output_bytes(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const; const ggml_graph_cut::Segment& segment) const;
bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
void prune(const std::unordered_set<std::string>& keep_names); void prune(const std::unordered_set<std::string>& keep_names);
void clear() { tensors_.clear(); } void clear() { tensors_.clear(); }
}; };

View File

@ -970,6 +970,7 @@ bool adetail_image(adetailer_ctx_t* context,
generation.pm_params = {}; generation.pm_params = {};
generation.pulid_params = {}; generation.pulid_params = {};
generation.hires.enabled = false; generation.hires.enabled = false;
generation.image_preprocess = {};
if (params.steps > 0) { if (params.steps > 0) {
generation.sample_params.sample_steps = params.steps; generation.sample_params.sample_steps = params.steps;
generation.sample_params.custom_sigmas = nullptr; generation.sample_params.custom_sigmas = nullptr;

View File

@ -5,6 +5,7 @@
#include <cassert> #include <cassert>
#include <cmath> #include <cmath>
#include <set> #include <set>
#include <utility>
#include <vector> #include <vector>
#include "core/ggml_extend.h" #include "core/ggml_extend.h"
#include "core/ggml_runner.h" #include "core/ggml_runner.h"
@ -16,6 +17,45 @@ namespace Rope {
ErnieImage, ErnieImage,
}; };
struct SpatialRegion {
size_t begin;
size_t count;
float height_period;
float width_period;
int height_axis = 1;
int width_axis = 2;
};
struct PositionLayout {
// Token ranges are relative to one batch item.
std::vector<SpatialRegion> images;
size_t token_count = 0;
void append_tokens(size_t count) {
token_count += count;
}
void append_image(int height, int width, int frames = 1, float height_step = 1.f, float width_step = 1.f) {
size_t count = static_cast<size_t>(height) * width * frames;
images.push_back({token_count, count, height * height_step, width * width_step});
append_tokens(count);
}
};
struct Frequency {
size_t axis;
float omega;
};
struct Embedding {
std::vector<float> values;
std::vector<std::vector<float>> ids;
PositionLayout positions;
std::vector<Frequency> frequencies;
EmbedNDLayout layout = EmbedNDLayout::Matrix;
int batch_size = 1;
};
enum class RefIndexMode { enum class RefIndexMode {
FIXED, FIXED,
INCREASE, INCREASE,
@ -56,40 +96,25 @@ namespace Rope {
return flat_vec; return flat_vec;
} }
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos, __STATIC_INLINE__ std::vector<float> rope_frequencies(int dim, float theta) {
int dim,
float theta,
const std::vector<int>& axis_wrap_dims = {}) {
assert(dim % 2 == 0); assert(dim % 2 == 0);
int half_dim = dim / 2; int half_dim = dim / 2;
std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim); std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
std::vector<float> omega(half_dim); std::vector<float> omega(half_dim);
for (int i = 0; i < half_dim; ++i) { for (int i = 0; i < half_dim; ++i) {
omega[i] = 1.0f / ::powf(1.f * theta, scale[i]); omega[i] = 1.0f / ::powf(1.f * theta, scale[i]);
} }
return omega;
}
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
const std::vector<float>& omega) {
int half_dim = static_cast<int>(omega.size());
size_t pos_size = pos.size(); size_t pos_size = pos.size();
std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim)); std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
for (size_t i = 0; i < pos_size; ++i) { for (size_t i = 0; i < pos_size; ++i) {
for (size_t j = 0; j < half_dim; ++j) { for (size_t j = 0; j < half_dim; ++j) {
float angle = pos[i] * omega[j]; float angle = pos[i] * omega[j];
if (!axis_wrap_dims.empty()) {
size_t wrap_size = axis_wrap_dims.size();
// mod batch size since we only store this for one item in the batch
size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0;
int wrap_dim = axis_wrap_dims[wrap_idx];
if (wrap_dim > 0) {
constexpr float TWO_PI = 6.28318530717958647692f;
float cycles = omega[j] * wrap_dim / TWO_PI;
// closest periodic harmonic, necessary to ensure things neatly tile
// without this round, things don't tile at the boundaries and you end up
// with the model knowing what is "center"
float rounded = std::round(cycles);
angle = pos[i] * TWO_PI * rounded / wrap_dim;
}
}
out[i][j] = angle; out[i][j] = angle;
} }
@ -108,6 +133,12 @@ namespace Rope {
return result; return result;
} }
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
int dim,
float theta) {
return rope(pos, rope_frequencies(dim, theta));
}
// Generate IDs for image patches and text // Generate IDs for image patches and text
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) { __STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) {
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f)); auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
@ -136,12 +167,16 @@ namespace Rope {
int patch_size, int patch_size,
int bs, int bs,
int axes_dim_num, int axes_dim_num,
int index = 0, int index = 0,
int h_offset = 0, int h_offset = 0,
int w_offset = 0, int w_offset = 0,
bool scale_rope = false) { bool scale_rope = false,
PositionLayout* layout = nullptr) {
int h_len = (h + (patch_size / 2)) / patch_size; int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size; int w_len = (w + (patch_size / 2)) / patch_size;
if (layout) {
layout->append_image(h_len, w_len);
}
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0)); std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
int h_start = h_offset; int h_start = h_offset;
@ -192,8 +227,8 @@ namespace Rope {
int bs, int bs,
const std::vector<float>& axis_thetas, const std::vector<float>& axis_thetas,
const std::vector<int>& axes_dim, const std::vector<int>& axes_dim,
const std::vector<std::vector<int>>& wrap_dims = {}, EmbedNDLayout layout = EmbedNDLayout::Matrix,
EmbedNDLayout layout = EmbedNDLayout::Matrix) { std::vector<Frequency>* frequencies = nullptr) {
std::vector<std::vector<float>> trans_ids = transpose(ids); std::vector<std::vector<float>> trans_ids = transpose(ids);
size_t pos_len = ids.size() / bs; size_t pos_len = ids.size() / bs;
size_t num_axes = axes_dim.size(); size_t num_axes = axes_dim.size();
@ -205,19 +240,25 @@ namespace Rope {
for (int d : axes_dim) for (int d : axes_dim)
emb_dim += d / 2; emb_dim += d / 2;
if (frequencies) {
frequencies->clear();
frequencies->reserve(emb_dim);
}
std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0)); std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
size_t offset = 0; size_t offset = 0;
for (size_t i = 0; i < num_axes; ++i) { for (size_t i = 0; i < num_axes; ++i) {
std::vector<int> axis_wrap_dims;
if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
axis_wrap_dims = wrap_dims[i];
}
float axis_theta = 10000.0f; float axis_theta = 10000.0f;
if (!axis_thetas.empty()) { if (!axis_thetas.empty()) {
axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)]; axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
} }
auto omega = rope_frequencies(axes_dim[i], axis_theta);
if (frequencies) {
for (float frequency : omega) {
frequencies->push_back({i, frequency});
}
}
std::vector<std::vector<float>> rope_emb = std::vector<std::vector<float>> rope_emb =
rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2] rope(trans_ids[i], omega); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
for (int b = 0; b < bs; ++b) { for (int b = 0; b < bs; ++b) {
for (int j = 0; j < pos_len; ++j) { for (int j = 0; j < pos_len; ++j) {
for (int k = 0; k < rope_emb[0].size(); ++k) { for (int k = 0; k < rope_emb[0].size(); ++k) {
@ -253,10 +294,10 @@ namespace Rope {
int bs, int bs,
float theta, float theta,
const std::vector<int>& axes_dim, const std::vector<int>& axes_dim,
const std::vector<std::vector<int>>& wrap_dims = {}, EmbedNDLayout layout = EmbedNDLayout::Matrix,
EmbedNDLayout layout = EmbedNDLayout::Matrix) { std::vector<Frequency>* frequencies = nullptr) {
std::vector<float> axis_thetas(axes_dim.size(), theta); std::vector<float> axis_thetas(axes_dim.size(), theta);
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout); return embed_nd(ids, bs, axis_thetas, axes_dim, layout, frequencies);
} }
__STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids, __STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
@ -264,7 +305,7 @@ namespace Rope {
float theta, float theta,
int head_dim, int head_dim,
const std::vector<int>& mrope_section, const std::vector<int>& mrope_section,
const std::vector<std::vector<int>>& axis_wrap_dims = {}) { std::vector<Frequency>* frequencies = nullptr) {
GGML_ASSERT(bs > 0); GGML_ASSERT(bs > 0);
GGML_ASSERT(head_dim % 2 == 0); GGML_ASSERT(head_dim % 2 == 0);
GGML_ASSERT(mrope_section.size() >= 3); GGML_ASSERT(mrope_section.size() >= 3);
@ -273,20 +314,26 @@ namespace Rope {
size_t pos_len = ids.size() / bs; size_t pos_len = ids.size() / bs;
int half_dim = head_dim / 2; int half_dim = head_dim / 2;
auto omega = rope_frequencies(head_dim, theta);
if (frequencies) {
frequencies->clear();
for (float frequency : omega) {
frequencies->push_back({0, frequency});
}
}
std::vector<std::vector<std::vector<float>>> axis_embs; std::vector<std::vector<std::vector<float>>> axis_embs;
axis_embs.reserve(3); axis_embs.reserve(3);
for (int axis = 0; axis < 3; ++axis) { for (int axis = 0; axis < 3; ++axis) {
std::vector<int> axis_wrap; axis_embs.push_back(rope(trans_ids[axis], omega));
if (axis < static_cast<int>(axis_wrap_dims.size())) {
axis_wrap = axis_wrap_dims[axis];
}
axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
} }
std::vector<std::vector<float>> emb = axis_embs[0]; std::vector<std::vector<float>> emb = axis_embs[0];
for (int axis = 1; axis < 3; ++axis) { for (int axis = 1; axis < 3; ++axis) {
int length = std::min<int>(mrope_section[axis] * 3, half_dim); int length = std::min<int>(mrope_section[axis] * 3, half_dim);
for (int freq_idx = axis; freq_idx < length; freq_idx += 3) { for (int freq_idx = axis; freq_idx < length; freq_idx += 3) {
if (frequencies) {
(*frequencies)[freq_idx].axis = axis;
}
for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) { for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
for (int k = 0; k < 4; ++k) { for (int k = 0; k < 4; ++k) {
emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k]; emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
@ -298,13 +345,13 @@ namespace Rope {
return flatten(emb); return flatten(emb);
} }
__STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height, __STATIC_INLINE__ Embedding embed_2d_interleaved(int height,
int width, int width,
int dim, int dim,
float theta = 10000.f, float theta = 10000.f,
float scale = 16.f, float scale = 16.f,
int ref_grid_h = 0, int ref_grid_h = 0,
int ref_grid_w = 0) { int ref_grid_w = 0) {
assert(dim % 4 == 0); assert(dim % 4 == 0);
int half_dim = dim / 2; int half_dim = dim / 2;
int dim_axis = dim / 2; int dim_axis = dim / 2;
@ -318,6 +365,10 @@ namespace Rope {
w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power); w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
} }
Embedding result;
result.positions.append_image(height, width, 1,
height > 1 ? scale / (height - 1) : 1.f,
width > 1 ? scale / (width - 1) : 1.f);
std::vector<float> x_pos; std::vector<float> x_pos;
std::vector<float> y_pos; std::vector<float> y_pos;
x_pos.reserve(static_cast<size_t>(height) * width); x_pos.reserve(static_cast<size_t>(height) * width);
@ -326,13 +377,20 @@ namespace Rope {
float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1); float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
for (int ix = 0; ix < width; ++ix) { for (int ix = 0; ix < width; ++ix) {
float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1); float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
result.ids.push_back({0.f, y, x});
x_pos.push_back(x); x_pos.push_back(x);
y_pos.push_back(y); y_pos.push_back(y);
} }
} }
auto x_emb = rope(x_pos, dim_axis, theta * w_ntk); auto x_freq = rope_frequencies(dim_axis, theta * w_ntk);
auto y_emb = rope(y_pos, dim_axis, theta * h_ntk); auto y_freq = rope_frequencies(dim_axis, theta * h_ntk);
auto x_emb = rope(x_pos, x_freq);
auto y_emb = rope(y_pos, y_freq);
for (int i = 0; i < axis_half_dim; ++i) {
result.frequencies.push_back({2, x_freq[i]});
result.frequencies.push_back({1, y_freq[i]});
}
std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4); std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
for (int pos = 0; pos < height * width; ++pos) { for (int pos = 0; pos < height * width; ++pos) {
@ -348,7 +406,8 @@ namespace Rope {
} }
} }
} }
return out; result.values = std::move(out);
return result;
} }
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size, __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
@ -359,7 +418,8 @@ namespace Rope {
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
bool scale_rope, bool scale_rope,
int base_offset = 0) { int base_offset = 0,
PositionLayout* layout = nullptr) {
std::vector<std::vector<float>> ids; std::vector<std::vector<float>> ids;
int curr_h_offset = 0; int curr_h_offset = 0;
int curr_w_offset = 0; int curr_w_offset = 0;
@ -386,7 +446,8 @@ namespace Rope {
static_cast<int>(index * ref_index_scale), static_cast<int>(index * ref_index_scale),
h_offset + base_offset, h_offset + base_offset,
w_offset + base_offset, w_offset + base_offset,
scale_rope); scale_rope,
layout);
ids = concat_ids(ids, ref_ids, bs); ids = concat_ids(ids, ref_ids, bs);
if (ref_index_mode == RefIndexMode::INCREASE) { if (ref_index_mode == RefIndexMode::INCREASE) {
@ -409,88 +470,53 @@ namespace Rope {
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
bool is_longcat) { bool is_longcat,
PositionLayout* layout = nullptr) {
if (layout) {
layout->append_tokens(context_len);
}
int x_index = is_longcat ? 1 : 0; int x_index = is_longcat ? 1 : 0;
auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims); auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
int offset = is_longcat ? context_len : 0; int offset = is_longcat ? context_len : 0;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset); auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset, false, layout);
auto ids = concat_ids(txt_ids, img_ids, bs); auto ids = concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) { if (ref_latents.size() > 0) {
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset); auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset, layout);
ids = concat_ids(ids, refs_ids, bs); ids = concat_ids(ids, refs_ids, bs);
} }
return ids; return ids;
} }
// Generate flux positional embeddings // Generate flux positional embeddings
__STATIC_INLINE__ std::vector<float> gen_flux_pe(int h, __STATIC_INLINE__ Embedding gen_flux_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
std::set<int> txt_arange_dims, std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
float ref_index_scale, float ref_index_scale,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim,
bool circular_w, bool is_longcat) {
const std::vector<int>& axes_dim, Embedding result;
bool is_longcat) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_flux_ids(h, result.ids = gen_flux_ids(h,
w, w,
patch_size, patch_size,
bs, bs,
static_cast<int>(axes_dim.size()), static_cast<int>(axes_dim.size()),
context_len, context_len,
txt_arange_dims, txt_arange_dims,
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
ref_index_scale, ref_index_scale,
is_longcat); is_longcat, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len; // text first
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
cursor += img_tokens;
// reference latents
for (ggml_tensor* ref : ref_latents) {
if (ref == nullptr) {
continue;
}
int ref_h = static_cast<int>(ref->ne[1]);
int ref_w = static_cast<int>(ref->ne[0]);
int ref_h_l = (ref_h + (patch_size / 2)) / patch_size;
int ref_w_l = (ref_w + (patch_size / 2)) / patch_size;
size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = ref_h_l;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = ref_w_l;
}
}
cursor += ref_tokens;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
} }
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t, __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
@ -500,14 +526,18 @@ namespace Rope {
int ph, int ph,
int pw, int pw,
int bs, int bs,
int t_offset = 0, int t_offset = 0,
int h_offset = 0, int h_offset = 0,
int w_offset = 0, int w_offset = 0,
bool scale_rope = false) { bool scale_rope = false,
PositionLayout* layout = nullptr) {
int t_len = (t + (pt / 2)) / pt; int t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph; int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw; int w_len = (w + (pw / 2)) / pw;
if (layout) {
layout->append_image(h_len, w_len, t_len);
}
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0)); std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
if (scale_rope) { if (scale_rope) {
@ -573,7 +603,11 @@ namespace Rope {
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode) { RefIndexMode ref_index_mode,
PositionLayout* layout = nullptr) {
if (layout) {
layout->append_tokens(context_len);
}
int h_len = (h + (patch_size / 2)) / patch_size; int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size; int w_len = (w + (patch_size / 2)) / patch_size;
int txt_id_start = std::max(h_len, w_len) / 2; int txt_id_start = std::max(h_len, w_len) / 2;
@ -585,90 +619,49 @@ namespace Rope {
} }
} }
int axes_dim_num = 3; int axes_dim_num = 3;
auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true); auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true, layout);
auto ids = concat_ids(txt_ids_repeated, img_ids, bs); auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
if (ref_latents.size() > 0) { if (ref_latents.size() > 0) {
int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1; int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true); auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true, 0, layout);
ids = concat_ids(ids, refs_ids, bs); ids = concat_ids(ids, refs_ids, bs);
} }
return ids; return ids;
} }
// Generate qwen_image positional embeddings // Generate qwen_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t, __STATIC_INLINE__ Embedding gen_qwen_image_pe(int t,
int h, int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode); result.ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles return result;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
const size_t total_tokens = ids.size();
// Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
size_t cursor = context_len; // ignore text tokens
const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
cursor += img_tokens;
// For each reference image, store wrap sizes as well
for (ggml_tensor* ref : ref_latents) {
if (ref == nullptr) {
continue;
}
int ref_h = static_cast<int>(ref->ne[1]);
int ref_w = static_cast<int>(ref->ne[0]);
int ref_pad_h = (patch_size - (ref_h % patch_size)) % patch_size;
int ref_pad_w = (patch_size - (ref_w % patch_size)) % patch_size;
int ref_h_len = (ref_h + ref_pad_h) / patch_size;
int ref_w_len = (ref_w + ref_pad_w) / patch_size;
size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = ref_h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = ref_w_len;
}
}
cursor += ref_n_tokens;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
} }
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h, __STATIC_INLINE__ Embedding gen_mage_flow_pe(int h,
int w, int w,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
int theta, int theta,
const std::vector<int>& axes_dim) { const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
const int axes_dim_num = static_cast<int>(axes_dim.size()); const int axes_dim_num = static_cast<int>(axes_dim.size());
auto make_image_ids = [=](int image_h, int image_w, int image_index) { auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) {
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w, std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
std::vector<float>(axes_dim_num, 0.f)); std::vector<float>(axes_dim_num, 0.f));
result.positions.append_image(image_h, image_w);
int h_start = -(image_h - image_h / 2); int h_start = -(image_h - image_h / 2);
int w_start = -(image_w - image_w / 2); int w_start = -(image_w - image_w / 2);
for (int b = 0; b < bs; ++b) { for (int b = 0; b < bs; ++b) {
@ -692,15 +685,18 @@ namespace Rope {
static_cast<int>(i + 1)); static_cast<int>(i + 1));
ids = concat_ids(ids, ref_ids, bs); ids = concat_ids(ids, ref_ids, bs);
} }
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim); result.ids = std::move(ids);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
} }
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h, __STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
int w, int w,
int bs, int bs,
int context_len, int context_len,
bool scale_rope = true) { bool scale_rope = true,
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope); PositionLayout* layout = nullptr) {
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope, layout);
int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0; int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len); auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
@ -711,44 +707,37 @@ namespace Rope {
} }
} }
if (layout) {
layout->append_tokens(context_len);
}
return concat_ids(img_ids_repeated, txt_ids_repeated, bs); return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
} }
__STATIC_INLINE__ std::vector<float> gen_lens_pe(int h, __STATIC_INLINE__ Embedding gen_lens_pe(int h,
int w, int w,
int bs, int bs,
int context_len, int context_len,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true); result.ids = gen_lens_ids(h, w, bs, context_len, true, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][token_i] = h;
}
if (circular_w) {
wrap_dims[2][token_i] = w;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
} }
__STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h, __STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len) { int context_len,
PositionLayout* layout = nullptr) {
int h_len = h / patch_size; int h_len = h / patch_size;
int w_len = w / patch_size; int w_len = w / patch_size;
if (layout) {
layout->append_image(h_len, w_len);
}
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f)); std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len); std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len); std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
@ -774,39 +763,25 @@ namespace Rope {
} }
} }
if (layout) {
layout->append_tokens(context_len);
}
return concat_ids(img_ids_repeated, txt_ids, bs); return concat_ids(img_ids_repeated, txt_ids, bs);
} }
__STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h, __STATIC_INLINE__ Embedding gen_ernie_image_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len); result.layout = EmbedNDLayout::ErnieImage;
std::vector<std::vector<int>> wrap_dims; result.ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len, &result.positions);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
int h_len = h / patch_size; return result;
int w_len = w / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
} }
// Generate wan positional embeddings // Generate wan positional embeddings
@ -905,7 +880,8 @@ namespace Rope {
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode) { RefIndexMode ref_index_mode,
PositionLayout* layout = nullptr) {
SD_UNUSED(ref_index_mode); SD_UNUSED(ref_index_mode);
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of); int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f)); auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
@ -913,11 +889,17 @@ namespace Rope {
txt_ids[i][0] = (i % padded_context_len) + 1.f; txt_ids[i][0] = (i % padded_context_len) + 1.f;
} }
if (layout) {
layout->append_tokens(padded_context_len);
}
int axes_dim_num = 3; int axes_dim_num = 3;
int index = padded_context_len + 1; int index = padded_context_len + 1;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index); auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of); int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
if (layout) {
layout->append_tokens(img_pad_len);
}
if (img_pad_len > 0) { if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f)); std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
img_ids = concat_ids(img_ids, img_pad_ids, bs); img_ids = concat_ids(img_ids, img_pad_ids, bs);
@ -936,7 +918,8 @@ namespace Rope {
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
int seq_multi_of) { int seq_multi_of,
PositionLayout* layout = nullptr) {
int context_pad_len = bound_mod(context_len, seq_multi_of); int context_pad_len = bound_mod(context_len, seq_multi_of);
int padded_context_len = context_len + context_pad_len; int padded_context_len = context_len + context_pad_len;
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f)); auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
@ -947,11 +930,17 @@ namespace Rope {
} }
} }
if (layout) {
layout->append_tokens(padded_context_len);
}
int axes_dim_num = 3; int axes_dim_num = 3;
int index = padded_context_len + 1; int index = padded_context_len + 1;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index); auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of); int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
if (layout) {
layout->append_tokens(img_pad_len);
}
if (img_pad_len > 0) { if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f)); std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
img_ids = concat_ids(img_ids, img_pad_ids, bs); img_ids = concat_ids(img_ids, img_pad_ids, bs);
@ -968,7 +957,8 @@ namespace Rope {
int patch_size, int patch_size,
int context_len, int context_len,
int sigvq_len, int sigvq_len,
int seq_multi_of) { int seq_multi_of,
PositionLayout* layout = nullptr) {
const int context_pad = bound_mod(context_len, seq_multi_of); const int context_pad = bound_mod(context_len, seq_multi_of);
const int padded_context = context_len + context_pad; const int padded_context = context_len + context_pad;
const int h_len = (h + (patch_size / 2)) / patch_size; const int h_len = (h + (patch_size / 2)) / patch_size;
@ -994,11 +984,17 @@ namespace Rope {
cursor += 2; cursor += 2;
} }
if (layout) {
layout->append_tokens(cap_ids.size());
}
std::vector<std::vector<float>> img_ids; std::vector<std::vector<float>> img_ids;
for (int copy = 0; copy < 2; ++copy) { for (int copy = 0; copy < 2; ++copy) {
auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy]); auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy], 0, 0, false, layout);
img_ids.insert(img_ids.end(), ids.begin(), ids.end()); img_ids.insert(img_ids.end(), ids.begin(), ids.end());
img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f)); img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f));
if (layout) {
layout->append_tokens(image_pad);
}
} }
const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1; const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
@ -1016,95 +1012,59 @@ namespace Rope {
ids.insert(ids.end(), cap_ids.begin(), cap_ids.end()); ids.insert(ids.end(), cap_ids.begin(), cap_ids.end());
ids.insert(ids.end(), img_ids.begin(), img_ids.end()); ids.insert(ids.end(), img_ids.begin(), img_ids.end());
ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end()); ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end());
if (layout) {
layout->append_tokens(sigvq_ids.size());
}
SD_UNUSED(padded_image); SD_UNUSED(padded_image);
return ids; return ids;
} }
__STATIC_INLINE__ std::vector<float> gen_llada_image_edit_pe(int h, __STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h,
int w, int w,
int patch_size, int patch_size,
int context_len, int context_len,
int sigvq_len, int sigvq_len,
int seq_multi_of, int seq_multi_of,
int theta, int theta,
const std::vector<int>& axes_dim) { const std::vector<int>& axes_dim) {
auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of); Embedding result;
return embed_nd(ids, 1, static_cast<float>(theta), axes_dim, {}); result.batch_size = 1;
result.ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of, &result.positions);
result.values = embed_nd(result.ids, 1, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
} }
__STATIC_INLINE__ std::vector<float> gen_llada_image_pe(int h, __STATIC_INLINE__ Embedding gen_llada_image_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of); result.ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len + bound_mod(context_len, seq_multi_of);
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
} }
// Generate z_image positional embeddings // Generate z_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h, __STATIC_INLINE__ Embedding gen_z_image_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
int seq_multi_of, int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode, RefIndexMode ref_index_mode,
int theta, int theta,
bool circular_h, const std::vector<int>& axes_dim) {
bool circular_w, Embedding result;
const std::vector<int>& axes_dim) { result.batch_size = bs;
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode); result.ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode, &result.positions);
std::vector<std::vector<int>> wrap_dims; result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { return result;
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len + bound_mod(context_len, seq_multi_of); // skip text (and its padding)
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
} }
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx, __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,

View File

@ -0,0 +1,65 @@
#ifndef __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
#define __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
#include "model/common/rope.hpp"
namespace Rope {
__STATIC_INLINE__ void apply_circular(Embedding& embedding, bool circular_x, bool circular_y) {
if (!circular_x && !circular_y) {
return;
}
GGML_ASSERT(embedding.batch_size > 0);
GGML_ASSERT(embedding.ids.size() % embedding.batch_size == 0);
size_t pos_len = embedding.ids.size() / embedding.batch_size;
size_t half_dim = embedding.frequencies.size();
GGML_ASSERT(embedding.positions.token_count == pos_len);
GGML_ASSERT(embedding.values.size() == embedding.ids.size() * half_dim * 4);
constexpr float TWO_PI = 6.28318530717958647692f;
for (const auto& region : embedding.positions.images) {
GGML_ASSERT(region.begin <= pos_len && region.count <= pos_len - region.begin);
for (size_t j = 0; j < half_dim; ++j) {
const auto& frequency = embedding.frequencies[j];
float period = 0.f;
if (circular_y && frequency.axis == static_cast<size_t>(region.height_axis)) {
period = region.height_period;
} else if (circular_x && frequency.axis == static_cast<size_t>(region.width_axis)) {
period = region.width_period;
}
if (period <= 0) {
continue;
}
// Quantize to periodic harmonics while preserving the original coordinate offsets.
float rounded = std::round(frequency.omega * period / TWO_PI);
for (int b = 0; b < embedding.batch_size; ++b) {
size_t begin = b * pos_len + region.begin;
for (size_t i = begin; i < begin + region.count; ++i) {
GGML_ASSERT(frequency.axis < embedding.ids[i].size());
float angle = embedding.ids[i][frequency.axis] * TWO_PI * rounded / period;
float cos_val = std::cos(angle);
float sin_val = std::sin(angle);
if (embedding.layout == EmbedNDLayout::ErnieImage) {
size_t cos_offset = (i * half_dim + j) * 2;
size_t sin_offset = embedding.ids.size() * half_dim * 2 + cos_offset;
embedding.values[cos_offset] = cos_val;
embedding.values[cos_offset + 1] = cos_val;
embedding.values[sin_offset] = sin_val;
embedding.values[sin_offset + 1] = sin_val;
} else {
size_t offset = (i * half_dim + j) * 4;
embedding.values[offset] = cos_val;
embedding.values[offset + 1] = -sin_val;
embedding.values[offset + 2] = sin_val;
embedding.values[offset + 3] = cos_val;
}
}
}
}
}
}
} // namespace Rope
#endif // __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__

View File

@ -603,34 +603,37 @@ namespace Anima {
return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2)); return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2));
} }
static std::vector<float> gen_anima_image_pe_vec(int bs, static Rope::Embedding gen_anima_image_pe_vec(int bs,
int h, int h,
int w, int w,
int patch_size, int patch_size,
int theta, int theta,
const std::vector<int>& axes_dim, const std::vector<int>& axes_dim,
float h_extrapolation_ratio, float h_extrapolation_ratio,
float w_extrapolation_ratio, float w_extrapolation_ratio,
float t_extrapolation_ratio, float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) { const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h, Rope::Embedding result;
w, result.batch_size = bs;
patch_size, result.ids = Rope::gen_flux_ids(h,
bs, w,
static_cast<int>(axes_dim.size()), patch_size,
0, bs,
{}, static_cast<int>(axes_dim.size()),
ref_latents, 0,
Rope::RefIndexMode::FIXED, {},
1.0f, ref_latents,
false); Rope::RefIndexMode::FIXED,
1.0f,
false, &result.positions);
std::vector<float> axis_thetas = { std::vector<float> axis_thetas = {
static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]), static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]), static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]),
static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]), static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]),
}; };
return Rope::embed_nd(ids, bs, axis_thetas, axes_dim); result.values = Rope::embed_nd(result.ids, bs, axis_thetas, axes_dim, result.layout, &result.frequencies);
return result;
} }
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor, ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@ -657,16 +660,16 @@ namespace Anima {
int64_t h_pad = x->ne[1] + pad_h; int64_t h_pad = x->ne[1] + pad_h;
int64_t w_pad = x->ne[0] + pad_w; int64_t w_pad = x->ne[0] + pad_w;
image_pe_vec = gen_anima_image_pe_vec(1, image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1,
static_cast<int>(h_pad), static_cast<int>(h_pad),
static_cast<int>(w_pad), static_cast<int>(w_pad),
static_cast<int>(config.patch_size), static_cast<int>(config.patch_size),
config.theta, config.theta,
config.axes_dim, config.axes_dim,
4.0f, 4.0f,
4.0f, 4.0f,
1.0f, 1.0f,
ref_latents); ref_latents));
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len); auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
set_backend_tensor_data(image_pe, image_pe_vec.data()); set_backend_tensor_data(image_pe, image_pe_vec.data());

View File

@ -720,15 +720,18 @@ namespace Boogu {
} }
} }
__STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h, __STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
int theta, int theta,
const std::vector<int>& axes_dim) { const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids; Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
auto& ids = result.ids;
ids.reserve(static_cast<size_t>(bs) * context_len); ids.reserve(static_cast<size_t>(bs) * context_len);
for (int b = 0; b < bs; b++) { for (int b = 0; b < bs; b++) {
for (int i = 0; i < context_len; i++) { for (int i = 0; i < context_len; i++) {
@ -741,15 +744,18 @@ namespace Boogu {
for (ggml_tensor* ref : ref_latents) { for (ggml_tensor* ref : ref_latents) {
int ref_h_tokens = patched_token_count(ref->ne[1], patch_size); int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
int ref_w_tokens = patched_token_count(ref->ne[0], patch_size); int ref_w_tokens = patched_token_count(ref->ne[0], patch_size);
result.positions.append_image(ref_h_tokens, ref_w_tokens);
append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens); append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
pe_shift += std::max(ref_h_tokens, ref_w_tokens); pe_shift += std::max(ref_h_tokens, ref_w_tokens);
} }
int h_tokens = patched_token_count(h, patch_size); int h_tokens = patched_token_count(h, patch_size);
int w_tokens = patched_token_count(w, patch_size); int w_tokens = patched_token_count(w, patch_size);
result.positions.append_image(h_tokens, w_tokens);
append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens); append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens);
return Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim); result.values = Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
} }
struct BooguImageRunner : public DiffusionModelRunner { struct BooguImageRunner : public DiffusionModelRunner {
@ -793,14 +799,14 @@ namespace Boogu {
ref_latents.push_back(make_input(ref_latent_tensor)); ref_latents.push_back(make_input(ref_latent_tensor));
} }
pe_vec = gen_boogu_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ref_latents, ref_latents,
config.theta, config.theta,
config.axes_dim); config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());

View File

@ -376,7 +376,7 @@ struct ControlNet : public GGMLRunner {
hint = make_input(hint_tensor); hint = make_input(hint_tensor);
} }
auto runner_ctx = get_context(); auto runner_ctx = get_context(gf);
auto outs = control_net.forward(&runner_ctx, auto outs = control_net.forward(&runner_ctx,
x, x,
@ -389,8 +389,7 @@ struct ControlNet : public GGMLRunner {
if (guided_hint_input == nullptr && !outs.empty()) { if (guided_hint_input == nullptr && !outs.empty()) {
guided_hint_output_ggml = outs[0]; guided_hint_output_ggml = outs[0];
ggml_set_output(guided_hint_output_ggml); ggml_set_output(guided_hint_output_ggml);
cache(guided_hint_cache_name(), guided_hint_output_ggml); runner_ctx.persist_cache_tensor(guided_hint_cache_name(), guided_hint_output_ggml);
ggml_build_forward_expand(gf, guided_hint_output_ggml);
} }
control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0); control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0);

View File

@ -415,15 +415,13 @@ namespace ErnieImage {
GGML_ASSERT(!context_tensor.empty()); GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor); ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());

View File

@ -1548,20 +1548,18 @@ namespace Flux {
} else if (version == VERSION_OVIS_IMAGE) { } else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2}; txt_arange_dims = {1, 2};
} }
pe_vec = Rope::gen_flux_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
txt_arange_dims, txt_arange_dims,
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
config.ref_index_scale, config.ref_index_scale,
config.theta, config.theta,
circular_y_enabled, config.axes_dim,
circular_x_enabled, sd_version_is_longcat(version)));
config.axes_dim,
sd_version_is_longcat(version));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len); // 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); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

View File

@ -149,18 +149,21 @@ namespace Ideogram4 {
return std::make_shared<Linear>(in_features, out_features, bias); return std::make_shared<Linear>(in_features, out_features, bias);
} }
__STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h, __STATIC_INLINE__ Rope::Embedding gen_ideogram4_pe(int grid_h,
int grid_w, int grid_w,
int bs, int bs,
int context_len, int context_len,
int head_dim, int head_dim,
int rope_theta, int rope_theta,
const std::vector<int>& mrope_section, const std::vector<int>& mrope_section) {
bool circular_x = false,
bool circular_y = false) {
GGML_ASSERT(bs == 1); GGML_ASSERT(bs == 1);
std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w), Rope::Embedding result;
std::vector<float>(3, 0.f)); result.batch_size = bs;
result.positions.append_tokens(context_len);
result.positions.append_image(grid_h, grid_w);
result.ids.assign(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
std::vector<float>(3, 0.f));
auto& ids = result.ids;
for (int i = 0; i < context_len; ++i) { for (int i = 0; i < context_len; ++i) {
ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)}; ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
@ -175,29 +178,13 @@ namespace Ideogram4 {
} }
} }
std::vector<std::vector<int>> axis_wrap_dims(3); result.values = Rope::embed_interleaved_mrope(ids,
if (circular_y || circular_x) { bs,
size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w); static_cast<float>(rope_theta),
axis_wrap_dims[1].assign(total_len, 0); head_dim,
axis_wrap_dims[2].assign(total_len, 0); mrope_section,
if (circular_y) { &result.frequencies);
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) { return result;
axis_wrap_dims[1][idx] = grid_h;
}
}
if (circular_x) {
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
axis_wrap_dims[2][idx] = grid_w;
}
}
}
return Rope::embed_interleaved_mrope(ids,
bs,
static_cast<float>(rope_theta),
head_dim,
mrope_section,
axis_wrap_dims);
} }
class Ideogram4Attention : public GGMLBlock { class Ideogram4Attention : public GGMLBlock {
@ -509,15 +496,13 @@ namespace Ideogram4 {
int64_t head_dim = config.emb_dim / config.num_heads; int64_t head_dim = config.emb_dim / config.num_heads;
auto runner_ctx = get_context(); auto runner_ctx = get_context();
pe_vec = gen_ideogram4_pe(static_cast<int>(grid_h), pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast<int>(grid_h),
static_cast<int>(grid_w), static_cast<int>(grid_w),
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context_len), static_cast<int>(context_len),
static_cast<int>(head_dim), static_cast<int>(head_dim),
static_cast<int>(config.rope_theta), static_cast<int>(config.rope_theta),
config.mrope_section, config.mrope_section));
runner_ctx.circular_x_enabled,
runner_ctx.circular_y_enabled);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());

View File

@ -689,23 +689,28 @@ namespace Krea2 {
} }
}; };
__STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h, __STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h,
int w, int w,
int patch_size, int patch_size,
int bs, int bs,
int context_len, int context_len,
float theta, float theta,
const std::vector<int>& axes_dim, const std::vector<int>& axes_dim,
const std::vector<ggml_tensor*>& ref_latents, const std::vector<ggml_tensor*>& ref_latents,
Rope::RefIndexMode ref_index_mode) { Rope::RefIndexMode ref_index_mode) {
Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {}); auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false); auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false, &result.positions);
auto ids = Rope::concat_ids(txt_ids, img_ids, bs); auto ids = Rope::concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) { if (ref_latents.size() > 0) {
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0); auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0, &result.positions);
ids = Rope::concat_ids(ids, refs_ids, bs); ids = Rope::concat_ids(ids, refs_ids, bs);
} }
return Rope::embed_nd(ids, bs, theta, axes_dim); result.ids = std::move(ids);
result.values = Rope::embed_nd(result.ids, bs, theta, axes_dim, result.layout, &result.frequencies);
return result;
} }
struct Krea2Runner : public DiffusionModelRunner { struct Krea2Runner : public DiffusionModelRunner {
@ -749,15 +754,15 @@ namespace Krea2 {
ref_latents.push_back(make_input(ref_latent_tensor)); ref_latents.push_back(make_input(ref_latent_tensor));
} }
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
config.theta, config.theta,
config.axes_dim, config.axes_dim,
ref_latents, ref_latents,
ref_image_params.ref_index_mode); ref_image_params.ref_index_mode));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());

View File

@ -384,14 +384,12 @@ namespace Lens {
GGML_ASSERT(!context_tensor.empty()); GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor); ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_lens_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());

View File

@ -412,16 +412,14 @@ namespace LLaDAImage {
GGML_ASSERT(!context_tensor.empty()); GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor); ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ZImage::SEQ_MULTI_OF, ZImage::SEQ_MULTI_OF,
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());
@ -461,14 +459,14 @@ namespace LLaDAImage {
ggml_tensor* source = make_input(source_tensor); ggml_tensor* source = make_input(source_tensor);
GGML_ASSERT(x->ne[3] == 1); GGML_ASSERT(x->ne[3] == 1);
pe_vec = Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0, semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
ZImage::SEQ_MULTI_OF, ZImage::SEQ_MULTI_OF,
config.theta, config.theta,
config.axes_dim); config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());

View File

@ -110,13 +110,13 @@ namespace MageFlow {
} }
int batch_size = static_cast<int>(x->ne[3]); int batch_size = static_cast<int>(x->ne[3]);
pe_vec = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
batch_size, batch_size,
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ref_latents, ref_latents,
config.theta, config.theta,
config.axes_dim); config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data()); set_backend_tensor_data(pe, pe_vec.data());

View File

@ -264,6 +264,9 @@ namespace MiniMaxH3 {
for (int64_t i = 0; i < num_layers; ++i) { for (int64_t i = 0; i < num_layers; ++i) {
auto block = std::dynamic_pointer_cast<TokenRefinerBlock>(blocks["blocks." + std::to_string(i)]); auto block = std::dynamic_pointer_cast<TokenRefinerBlock>(blocks["blocks." + std::to_string(i)]);
x = block->forward(ctx, x); x = block->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(x,
"minimax_h3.token_refiner.blocks." + std::to_string(i),
"hidden_states");
} }
return std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm"])->forward(ctx, x); return std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm"])->forward(ctx, x);
} }
@ -527,7 +530,11 @@ namespace MiniMaxH3 {
GGML_ASSERT(context->ne[0] == config.text_dim); GGML_ASSERT(context->ne[0] == config.text_dim);
auto condition_proj = std::dynamic_pointer_cast<Linear>(blocks["condition_proj"]); auto condition_proj = std::dynamic_pointer_cast<Linear>(blocks["condition_proj"]);
auto token_refiner = std::dynamic_pointer_cast<TokenRefiner>(blocks["token_refiner"]); auto token_refiner = std::dynamic_pointer_cast<TokenRefiner>(blocks["token_refiner"]);
return token_refiner->forward(ctx, condition_proj->forward(ctx, context)); auto projected = condition_proj->forward(ctx, context);
sd::ggml_graph_cut::mark_graph_cut(projected,
"minimax_h3.condition_proj",
"hidden_states");
return token_refiner->forward(ctx, projected);
} }
ggml_tensor* time_embedding(GGMLRunnerContext* ctx, ggml_tensor* time_embedding(GGMLRunnerContext* ctx,

View File

@ -154,18 +154,26 @@ namespace MiniT2I {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f)); return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
} }
inline std::vector<float> make_vision_rope(int side, int head_dim) { inline Rope::Embedding make_vision_rope(int side, int head_dim) {
GGML_ASSERT(head_dim % 4 == 0); GGML_ASSERT(head_dim % 4 == 0);
int dim = head_dim / 2; int dim = head_dim / 2;
int quarter = dim / 2; int quarter = dim / 2;
int length = side * side; int length = side * side;
Rope::Embedding result;
result.positions.append_image(side, side);
std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4); std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
std::vector<float> freqs(quarter); std::vector<float> freqs(quarter);
for (int i = 0; i < quarter; ++i) { for (int i = 0; i < quarter; ++i) {
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim)); freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
} }
for (int axis : {1, 2}) {
for (float frequency : freqs) {
result.frequencies.push_back({static_cast<size_t>(axis), frequency});
}
}
for (int y = 0; y < side; ++y) { for (int y = 0; y < side; ++y) {
for (int x = 0; x < side; ++x) { for (int x = 0; x < side; ++x) {
result.ids.push_back({0.f, static_cast<float>(y), static_cast<float>(x)});
int pos = y * side + x; int pos = y * side + x;
size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4; size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
for (int i = 0; i < quarter; ++i) { for (int i = 0; i < quarter; ++i) {
@ -182,7 +190,8 @@ namespace MiniT2I {
} }
} }
} }
return out; result.values = std::move(out);
return result;
} }
struct SwiGLUMlp : public GGMLBlock { struct SwiGLUMlp : public GGMLBlock {
@ -475,6 +484,8 @@ namespace MiniT2I {
int64_t cached_txt_len = -1; int64_t cached_txt_len = -1;
int64_t cached_hidden_size = -1; int64_t cached_hidden_size = -1;
int64_t cached_head_dim = -1; int64_t cached_head_dim = -1;
bool cached_circular_x = false;
bool cached_circular_y = false;
MiniT2IRunner(ggml_backend_t backend, MiniT2IRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {}, const String2TensorStorage& tensor_storage_map = {},
@ -521,6 +532,8 @@ namespace MiniT2I {
cached_txt_len == txt_len && cached_txt_len == txt_len &&
cached_hidden_size == config.hidden_size && cached_hidden_size == config.hidden_size &&
cached_head_dim == config.head_dim && cached_head_dim == config.head_dim &&
cached_circular_x == circular_x_enabled &&
cached_circular_y == circular_y_enabled &&
cached_pos_embed != nullptr && cached_pos_embed != nullptr &&
cached_txt_pe != nullptr && cached_txt_pe != nullptr &&
cached_joint_pe != nullptr) { cached_joint_pe != nullptr) {
@ -531,7 +544,7 @@ namespace MiniT2I {
auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size)); auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim)); auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
auto img_pe_vec = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)); auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
auto joint_pe_vec = txt_pe_vec; auto joint_pe_vec = txt_pe_vec;
joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end()); joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
@ -561,6 +574,8 @@ namespace MiniT2I {
cached_txt_len = txt_len; cached_txt_len = txt_len;
cached_hidden_size = config.hidden_size; cached_hidden_size = config.hidden_size;
cached_head_dim = config.head_dim; cached_head_dim = config.head_dim;
cached_circular_x = circular_x_enabled;
cached_circular_y = circular_y_enabled;
} }
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor, ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,

View File

@ -7,7 +7,7 @@
#include "core/ggml_runner.h" #include "core/ggml_runner.h"
#include "core/tensor_ggml.hpp" #include "core/tensor_ggml.hpp"
#include "model/common/rope.hpp" #include "model/common/rope_circular.hpp"
#include "model_manager.h" #include "model_manager.h"
enum class RefImageResizeMode { enum class RefImageResizeMode {
@ -71,6 +71,8 @@ struct AnimaDiffusionExtra {
struct QwenImage21DiffusionExtra { struct QwenImage21DiffusionExtra {
const sd::Tensor<int32_t>* image_slots = nullptr; const sd::Tensor<int32_t>* image_slots = nullptr;
// Nonzero IDs identify immutable prefix inputs within one sampling run.
uint64_t prefix_id = 0;
}; };
struct WanDiffusionExtra { struct WanDiffusionExtra {
@ -182,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner {
protected: protected:
std::string prefix; std::string prefix;
std::vector<float> finish_rope_pe(Rope::Embedding embedding) {
Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled);
return std::move(embedding.values);
}
public: public:
DiffusionModelRunner(ggml_backend_t backend, DiffusionModelRunner(ggml_backend_t backend,
const std::string& prefix, const std::string& prefix,

View File

@ -135,13 +135,13 @@ namespace Pid {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta)); return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
} }
inline std::vector<float> make_rope_2d(int height, inline Rope::Embedding make_rope_2d(int height,
int width, int width,
int dim, int dim,
float theta = 10000.f, float theta = 10000.f,
float scale = 16.f, float scale = 16.f,
int ref_grid_h = 0, int ref_grid_h = 0,
int ref_grid_w = 0) { int ref_grid_w = 0) {
GGML_ASSERT(dim % 4 == 0); GGML_ASSERT(dim % 4 == 0);
return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w); return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w);
} }
@ -867,13 +867,13 @@ namespace Pid {
int64_t Hs = Hp / config.patch_size; int64_t Hs = Hp / config.patch_size;
int64_t Ws = Wp / config.patch_size; int64_t Ws = Wp / config.patch_size;
pos_img_vec = make_rope_2d(static_cast<int>(Hs), pos_img_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
static_cast<int>(Ws), static_cast<int>(Ws),
static_cast<int>(config.hidden_size / config.num_groups), static_cast<int>(config.hidden_size / config.num_groups),
10000.f, 10000.f,
16.f, 16.f,
static_cast<int>(config.rope_ref_grid_h), static_cast<int>(config.rope_ref_grid_h),
static_cast<int>(config.rope_ref_grid_w)); static_cast<int>(config.rope_ref_grid_w)));
auto pos_img = ggml_new_tensor_4d(compute_ctx, auto pos_img = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32, GGML_TYPE_F32,
2, 2,
@ -904,13 +904,13 @@ namespace Pid {
1); 1);
set_backend_tensor_data(pixel_pos, pixel_pos_vec.data()); set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
pixel_pos_comp_vec = make_rope_2d(static_cast<int>(Hs), pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
static_cast<int>(Ws), static_cast<int>(Ws),
static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups), static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
10000.f, 10000.f,
16.f, 16.f,
static_cast<int>(config.rope_ref_grid_h), static_cast<int>(config.rope_ref_grid_h),
static_cast<int>(config.rope_ref_grid_w)); static_cast<int>(config.rope_ref_grid_w)));
auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx, auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32, GGML_TYPE_F32,
2, 2,

View File

@ -635,18 +635,16 @@ namespace Qwen {
ref_index_mode = Rope::RefIndexMode::DECREASE; ref_index_mode = Rope::RefIndexMode::DECREASE;
} }
pe_vec = Rope::gen_qwen_image_pe(time_len, pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len,
static_cast<int>(x->ne[1]), static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
batch_size, batch_size,
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len); // 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); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

View File

@ -68,6 +68,7 @@ namespace Qwen {
std::vector<QwenImage21Segment> segments; std::vector<QwenImage21Segment> segments;
std::vector<std::vector<float>> positions; std::vector<std::vector<float>> positions;
int64_t prefix_length = 0; int64_t prefix_length = 0;
Rope::PositionLayout rope_layout;
static QwenImage21Layout build(int64_t text_length, static QwenImage21Layout build(int64_t text_length,
const sd::Tensor<int32_t>& image_slots, const sd::Tensor<int32_t>& image_slots,
@ -82,6 +83,7 @@ namespace Qwen {
auto [height, width] = image_shapes[index]; auto [height, width] = image_shapes[index];
int64_t start = static_cast<int64_t>(layout.positions.size()); int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + height * width, context_start, index}); layout.segments.push_back({start, start + height * width, context_start, index});
layout.rope_layout.append_image(static_cast<int>(height), static_cast<int>(width));
for (int64_t h = 0; h < height; ++h) { for (int64_t h = 0; h < height; ++h) {
for (int64_t w = 0; w < width; ++w) { for (int64_t w = 0; w < width; ++w) {
layout.positions.push_back({static_cast<float>(position), layout.positions.push_back({static_cast<float>(position),
@ -106,6 +108,7 @@ namespace Qwen {
} else { } else {
int64_t start = static_cast<int64_t>(layout.positions.size()); int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + i - begin, begin, -1}); layout.segments.push_back({start, start + i - begin, begin, -1});
layout.rope_layout.append_tokens(i - begin);
for (int64_t j = begin; j < i; ++j, ++position) { for (int64_t j = begin; j < i; ++j, ++position) {
float p = static_cast<float>(position); float p = static_cast<float>(position);
layout.positions.push_back({p, p, p}); layout.positions.push_back({p, p, p});
@ -121,6 +124,20 @@ namespace Qwen {
} }
}; };
struct QwenImage21PrefixCache {
enum class Mode {
NONE,
STORE,
REUSE
};
Mode mode = Mode::NONE;
std::string name;
std::string cut_group;
int64_t prefix_length = 0;
ggml_type type = GGML_TYPE_F32;
bool* flash_attn_used = nullptr;
};
class QwenImage21ZeroCenterRMSNorm : public RMSNorm { class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
public: public:
using RMSNorm::RMSNorm; using RMSNorm::RMSNorm;
@ -160,27 +177,70 @@ namespace Qwen {
} }
} }
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks) { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
int64_t heads = x->ne[0] / dim_head; int64_t heads = x->ne[0] / dim_head;
auto project = [&](const char* name) { auto project = [&](const char* name) {
auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x); auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]); return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]);
}; };
auto q = project("to_q"); auto q = project("to_q");
auto k = project("to_k"); auto k = project("to_k");
auto v = project("to_v"); auto v = project("to_v");
q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q); q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k); k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
q = Rope::apply_rope(ctx->ggml_ctx, q, pe); q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
k = Rope::apply_rope(ctx->ggml_ctx, k, pe); k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
// Preserve query-first attention evaluation while writing each layer's
// prefix before its full-sequence K/V can accumulate across layers.
ctx->expand_graph(q);
auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
// Pack the contiguous data into wider rows so quantization blocks
// can exceed head_dim without padding or changing element order.
part = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
auto copy = ggml_cast(ctx->ggml_ctx, part, cache.type);
// Keep the copy in this layer's segment so graph cuts do not
// retain or recompute the full-sequence K/V in the final segment.
sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name);
ctx->persist_cache_tensor(cache.name + "." + name, copy);
};
persist(k, 1, "k");
persist(v, 2, "v");
}
auto attend = [&](ggml_tensor* aq, ggml_tensor* ak, ggml_tensor* av, ggml_tensor* mask) {
bool used_flash_attn = false;
auto out = ggml_ext_attention_ext(ctx, aq, ak, av, heads, mask, true, ctx->flash_attn_enabled, 1.f, &used_flash_attn);
if (cache.flash_attn_used != nullptr) {
*cache.flash_attn_used &= used_flash_attn;
}
return out;
};
ggml_tensor* result = nullptr; ggml_tensor* result = nullptr;
for (size_t i = 0; i < segments.size(); ++i) { if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
const auto& segment = segments[i]; auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end); auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end); GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end); if (prefix_k->type != k->type) {
auto out = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled); prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type);
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); }
if (prefix_v->type != v->type) {
prefix_v = ggml_cast(ctx->ggml_ctx, prefix_v, v->type);
}
prefix_k = ggml_reshape_4d(ctx->ggml_ctx, prefix_k, dim_head, cache.prefix_length, heads, k->ne[3]);
prefix_v = ggml_reshape_4d(ctx->ggml_ctx, prefix_v, dim_head, heads, cache.prefix_length, v->ne[3]);
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
result = attend(q, k, v, nullptr);
} else {
for (size_t i = 0; i < segments.size(); ++i) {
const auto& segment = segments[i];
auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
auto out = attend(sq, sk, sv, masks[i]);
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
}
} }
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]); auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
@ -219,13 +279,14 @@ namespace Qwen {
return ggml_concat(ctx, prefix, target, 1); return ggml_concat(ctx, prefix, target, 1);
} }
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x); const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length;
h = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length); auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
h = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks); h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length);
x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true)); h = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache);
h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x); x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true));
h = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length); h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
h = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length);
ggml_tensor* gate; ggml_tensor* gate;
auto fused = blocks.find("img_mlp.gate_up"); auto fused = blocks.find("img_mlp.gate_up");
if (fused != blocks.end()) { if (fused != blocks.end()) {
@ -239,7 +300,7 @@ namespace Qwen {
} }
h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate)); h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h); h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], layout.prefix_length, true)); return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true));
} }
}; };
@ -261,7 +322,7 @@ namespace Qwen {
} }
} }
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0); auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
// Runtime flow timesteps already use the [0, 1000] scale. // Runtime flow timesteps already use the [0, 1000] scale.
time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f); time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
@ -269,27 +330,37 @@ namespace Qwen {
time = ggml_silu(ctx->ggml_ctx, time); time = ggml_silu(ctx->ggml_ctx, time);
auto modulation = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time); auto modulation = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0); auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]); auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
ggml_tensor* joint = nullptr; ggml_tensor* joint = nullptr;
for (const auto& segment : layout.segments) { if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
ggml_tensor* h; joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1));
if (segment.image_index < 0) { } else {
h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start, auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
segment.context_start + segment.end - segment.start); for (const auto& segment : layout.segments) {
} else { ggml_tensor* h;
auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index]; if (segment.image_index < 0) {
h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1)); h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
segment.context_start + segment.end - segment.start);
} else {
auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
}
joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
} }
joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
} }
sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint"); sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint");
for (int i = 0; i < config.num_layers; ++i) { for (int i = 0; i < config.num_layers; ++i) {
auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]); const std::string layer = "transformer_blocks." + std::to_string(i);
joint = block->forward(ctx, joint, mod, pe, layout, masks); auto layer_cache = cache;
sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint"); layer_cache.name = cache.name + "." + std::to_string(i);
layer_cache.cut_group = "qwen_image_2_1." + layer;
auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks[layer]);
joint = block->forward(ctx, joint, mod, pe, layout, masks, layer_cache);
sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint");
}
if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) {
joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
} }
joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]); auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
joint = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint); joint = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f)); joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
@ -303,11 +374,61 @@ namespace Qwen {
QwenImage21Model model; QwenImage21Model model;
std::vector<float> pe_data; std::vector<float> pe_data;
std::vector<sd::Tensor<float>> mask_data; std::vector<sd::Tensor<float>> mask_data;
ggml_type prefix_cache_type = GGML_TYPE_COUNT;
bool prefix_cache_enabled = true;
bool prefix_cache_disabled = false;
bool prefix_cache_auto_f32 = false;
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) static bool supports_prefix_cache_type(ggml_type type) {
if (type == GGML_TYPE_F32) {
return true;
}
const auto* traits = ggml_get_type_traits(type);
if (traits->from_float_ref == nullptr || traits->to_float == nullptr) {
return false;
}
auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
if (cpu == nullptr) {
return false;
}
auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free);
if (ctx == nullptr) {
return false;
}
// Some reference quantizers have no runtime copy support. Query the
// device through the registry so dynamically loaded CPU backends work.
auto source = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type));
auto encoded = ggml_cast(ctx.get(), source, type);
auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32);
return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded);
}
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager), : DiffusionModelRunner(backend, prefix, weight_manager),
config(QwenImage21Config::detect_from_weights(weights, prefix)), config(QwenImage21Config::detect_from_weights(weights, prefix)),
model(config) { model(config) {
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) {
LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str());
} else if (key == "qwen_image_2_1_prefix_cache_type") {
if (value == "auto") {
prefix_cache_type = GGML_TYPE_COUNT;
continue;
}
const auto type = sd_type_to_ggml_type(str_to_sd_type(value.c_str()));
if (type == GGML_TYPE_COUNT) {
LOG_WARN("ignoring unknown Qwen Image 2.1 cache type '%s'", value.c_str());
} else if (!supports_prefix_cache_type(type)) {
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': runtime conversion to and from F32 is unavailable", value.c_str());
} else if (config.hidden_size % ggml_blck_size(type) != 0) {
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': block size %" PRId64 " does not divide hidden size %" PRId64,
value.c_str(), ggml_blck_size(type), config.hidden_size);
} else {
prefix_cache_type = type;
}
}
}
model.init(params_ctx, weights, prefix); model.init(params_ctx, weights, prefix);
} }
@ -317,6 +438,20 @@ namespace Qwen {
model.get_param_tensors(tensors, prefix); model.get_param_tensors(tensors, prefix);
} }
bool has_prefix_cache(const QwenImage21PrefixCache& cache) {
for (int i = 0; i < config.num_layers; ++i) {
const auto name = cache.name + "." + std::to_string(i);
auto k = get_cache_tensor_by_name(name + ".k");
auto v = get_cache_tensor_by_name(name + ".v");
if (k == nullptr || v == nullptr || k->type != cache.type || v->type != cache.type ||
k->ne[0] != config.hidden_size || k->ne[1] != cache.prefix_length || k->ne[2] != 1 || k->ne[3] != 1 ||
v->ne[0] != config.hidden_size || v->ne[1] != cache.prefix_length || v->ne[2] != 1 || v->ne[3] != 1) {
return false;
}
}
return true;
}
sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override { sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override {
const auto& x = tensor_or_empty(inputs.x); const auto& x = tensor_or_empty(inputs.x);
const auto& context = tensor_or_empty(inputs.context); const auto& context = tensor_or_empty(inputs.context);
@ -345,38 +480,113 @@ namespace Qwen {
LOG_ERROR("%s", error.what()); LOG_ERROR("%s", error.what());
return {}; return {};
} }
pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim); if (!runner_started()) {
mask_data.clear(); prefix_cache_disabled = false;
for (const auto& segment : layout.segments) { prefix_cache_auto_f32 = false;
sd::Tensor<float> mask; }
if (segment.image_index < 0) { QwenImage21PrefixCache cache;
mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start}); if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
for (int64_t q = segment.start; q < segment.end; ++q) { cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
for (int64_t k = q + 1; k < segment.end; ++k) { ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
mask[k + segment.end * (q - segment.start)] = -INFINITY; cache.prefix_length = layout.prefix_length;
} if (prefix_cache_type != GGML_TYPE_COUNT) {
cache.type = prefix_cache_type;
} else if (!prefix_cache_auto_f32 && flash_attn_enabled && !sage_attn_enabled &&
(attn_scale <= 0.f || attn_scale == 1.f)) {
cache.type = GGML_TYPE_F16;
}
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
}
bool flash_attn_used = true;
auto run = [&](const QwenImage21PrefixCache& active_cache) {
flash_attn_used = true;
auto checked_cache = active_cache;
if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
checked_cache.flash_attn_used = &flash_attn_used;
}
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
Rope::Embedding embedding;
embedding.ids.assign(first_position, layout.positions.end());
const size_t offset = cached ? static_cast<size_t>(layout.prefix_length) : 0;
embedding.positions.token_count = embedding.ids.size();
for (auto region : layout.rope_layout.images) {
if (region.begin >= offset) {
region.begin -= offset;
embedding.positions.images.push_back(region);
} }
} }
mask_data.push_back(std::move(mask)); embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
} pe_data = finish_rope_pe(std::move(embedding));
auto build = [&]() { mask_data.clear();
auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2); if (!cached) {
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, layout.positions.size()); for (const auto& segment : layout.segments) {
set_backend_tensor_data(pe, pe_data.data()); sd::Tensor<float> mask;
std::vector<ggml_tensor*> masks, ref_inputs; if (segment.image_index < 0) {
for (const auto& mask : mask_data) { mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
masks.push_back(mask.empty() ? nullptr : make_input(mask)); for (int64_t q = segment.start; q < segment.end; ++q) {
for (int64_t k = q + 1; k < segment.end; ++k) {
mask[k + segment.end * (q - segment.start)] = -INFINITY;
}
}
}
mask_data.push_back(std::move(mask));
}
} }
for (const auto& ref : refs) { auto build = [&]() {
ref_inputs.push_back(make_input(ref)); auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
} auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
auto ctx = get_context(); layout.positions.size() - (cached ? layout.prefix_length : 0));
auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context), set_backend_tensor_data(pe, pe_data.data());
ref_inputs, pe, layout, masks); std::vector<ggml_tensor*> masks, ref_inputs;
ggml_build_forward_expand(graph, out); for (const auto& mask : mask_data) {
return graph; masks.push_back(mask.empty() ? nullptr : make_input(mask));
}
if (!cached) {
for (const auto& ref : refs) {
ref_inputs.push_back(make_input(ref));
}
}
auto ctx = get_context(graph);
auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
ref_inputs, pe, layout, masks, checked_cache);
if (!flash_attn_used) {
return static_cast<ggml_cgraph*>(nullptr);
}
ggml_build_forward_expand(graph, out);
return graph;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
}; };
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim()); auto result = run(cache);
if (result.empty() && !flash_attn_used) {
// Casting an F16 cache back to F32 cannot recover its original values.
// Recompute the prefix before executing a graph that falls back from FA.
free_cache_ctx_and_buffer();
prefix_cache_auto_f32 = true;
cache.type = GGML_TYPE_F32;
cache.mode = QwenImage21PrefixCache::Mode::STORE;
LOG_DEBUG("Qwen Image 2.1: Flash Attention unavailable; using F32 prefix caching for this sampling run");
result = run(cache);
}
if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
(cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
// The failed graph has ended before persistent inputs are released.
free_cache_ctx_and_buffer();
prefix_cache_disabled = true;
LOG_WARN("Qwen Image 2.1: insufficient memory for prefix caching; retrying without it for this sampling run");
return run(QwenImage21PrefixCache{});
}
if (!result.empty() && cache.mode == QwenImage21PrefixCache::Mode::STORE) {
if (!has_prefix_cache(cache)) {
free_cache_ctx_and_buffer();
prefix_cache_disabled = true;
LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run");
} else {
LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens, %s)", extra->prefix_id, layout.prefix_length, ggml_type_name(cache.type));
}
}
return result;
} }
}; };
} }

View File

@ -442,16 +442,9 @@ namespace SenseNovaU1 {
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2); k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2);
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2); v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
} else { } else {
// Keep dedicated graph outputs alive until the runner copies them ctx->expand_graph(q);
// into its persistent cache buffer after graph execution. ctx->persist_cache_tensor(layer_cache + ".k", k);
auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k); ctx->persist_cache_tensor(layer_cache + ".v", v);
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, q = ggml_cont(ctx->ggml_ctx,
@ -687,7 +680,7 @@ namespace SenseNovaU1 {
ggml_set_name(attention_mask, "snu15.prefix.attention_mask"); ggml_set_name(attention_mask, "snu15.prefix.attention_mask");
set_backend_tensor_data(attention_mask, attention_mask_vec.data()); set_backend_tensor_data(attention_mask, attention_mask_vec.data());
auto runner_ctx = get_context(); auto runner_ctx = get_context(graph);
auto text_model = model.text_model(); auto text_model = model.text_model();
auto hidden = text_model->embed(&runner_ctx, ids); auto hidden = text_model->embed(&runner_ctx, ids);
hidden = text_model->forward(&runner_ctx, hidden = text_model->forward(&runner_ctx,

View File

@ -642,18 +642,16 @@ namespace ZImage {
ref_latents.push_back(make_input(ref_latent_tensor)); ref_latents.push_back(make_input(ref_latent_tensor));
} }
pe_vec = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]), pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]), static_cast<int>(x->ne[0]),
config.patch_size, config.patch_size,
static_cast<int>(x->ne[3]), static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]), static_cast<int>(context->ne[1]),
SEQ_MULTI_OF, SEQ_MULTI_OF,
ref_latents, ref_latents,
ref_index_mode, ref_index_mode,
config.theta, config.theta,
circular_y_enabled, config.axes_dim));
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2); int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len); // 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); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

View File

@ -1300,7 +1300,7 @@ struct LTXVideoVAE : public VAE {
feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx)); feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx));
} }
auto runner_ctx = get_context(); auto runner_ctx = get_context(gf);
int feat_count = 0; int feat_count = 0;
ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx, ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx,
z, z,
@ -1313,8 +1313,7 @@ struct LTXVideoVAE : public VAE {
for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) { for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) {
ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)]; ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
if (feat_cache != nullptr) { if (feat_cache != nullptr) {
cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache); runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
ggml_build_forward_expand(gf, feat_cache);
} }
} }

View File

@ -82,6 +82,20 @@ namespace WAN {
} }
x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
if (w->ne[2] == 1 && x->ne[2] == 1 && x->ne[3] == in_channels) {
// One frame through a one-frame-deep kernel is a 2D conv; backends without
// im2col_3d (Metal) otherwise fall back to a much slower direct conv_3d.
if (!ggml_is_contiguous(x)) {
x = ggml_cont(ctx->ggml_ctx, x);
}
ggml_tensor* x2 = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], in_channels, 1);
ggml_tensor* w2 = ggml_reshape_4d(ctx->ggml_ctx, w, w->ne[0], w->ne[1], in_channels, out_channels);
x2 = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b,
std::get<2>(stride), std::get<1>(stride), 0, 0,
std::get<2>(dilation), std::get<1>(dilation),
ctx->conv2d_direct_enabled);
return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels);
}
return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
0, 0, 0, 0, 0, 0,
@ -1429,15 +1443,14 @@ namespace WAN {
ggml_tensor* z = make_input(z_tensor); ggml_tensor* z = make_input(z_tensor);
auto runner_ctx = get_context(); auto runner_ctx = get_context(gf);
ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx); ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx);
for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) { for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
ggml_tensor* feat_cache = ae._feat_map[feat_idx]; ggml_tensor* feat_cache = ae._feat_map[feat_idx];
if (feat_cache != nullptr) { if (feat_cache != nullptr) {
cache("feat_idx:" + std::to_string(feat_idx), feat_cache); runner_ctx.persist_cache_tensor("feat_idx:" + std::to_string(feat_idx), feat_cache);
ggml_build_forward_expand(gf, feat_cache);
} }
} }

View File

@ -874,7 +874,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors, std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
std::set<std::string> ignore_tensors, std::set<std::string> ignore_tensors,
bool writable_mmap) { bool writable_mmap,
ggml_backend_dev_t device) {
std::set<std::string> names; std::set<std::string> names;
for (const auto& entry : tensors) { for (const auto& entry : tensors) {
names.insert(entry.first); names.insert(entry.first);
@ -896,6 +897,39 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
if (!fdata.mmbuffer) if (!fdata.mmbuffer)
continue; continue;
// Wrapped on first use: a device buffer makes the whole file resident on that device.
std::shared_ptr<struct ggml_backend_buffer> file_buffer = device == nullptr ? fdata.mmbuffer : nullptr;
bool file_unmappable = false;
auto buffer_for_file = [&]() -> ggml_backend_buffer_t {
if (file_buffer || file_unmappable) {
return file_buffer.get();
}
auto cached = fdata.device_mmbuffers.find(device);
if (cached != fdata.device_mmbuffers.end()) {
file_buffer = cached->second;
return file_buffer.get();
}
size_t max_tensor_size = 0;
for (const auto& ts : fdata.tensors) {
max_tensor_size = std::max(max_tensor_size, static_cast<size_t>(ts.nbytes()));
}
ggml_backend_buffer_t buf = sd_backend_dev_buffer_from_host_ptr(device,
fdata.mmapped->writable_data(),
fdata.mmapped->size(),
max_tensor_size);
if (buf == nullptr) {
LOG_WARN("mmap: %s cannot map '%s', loading it instead",
ggml_backend_dev_name(device), fdata.path.c_str());
file_unmappable = true;
return nullptr;
}
LOG_INFO("mmap: mapped '%s' for %s", fdata.path.c_str(), ggml_backend_dev_name(device));
file_buffer = std::shared_ptr<struct ggml_backend_buffer>(buf, ggml_backend_buffer_free);
fdata.device_mmbuffers[device] = file_buffer;
return file_buffer.get();
};
const std::vector<TensorStorage>& file_tensors = fdata.tensors; const std::vector<TensorStorage>& file_tensors = fdata.tensors;
size_t file_mapped_bytes = 0; size_t file_mapped_bytes = 0;
@ -944,10 +978,13 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
continue; continue;
} }
ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get(); ggml_backend_buffer_t buf_mmap = buffer_for_file();
uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap)); if (buf_mmap == nullptr) {
dst_tensor->buffer = buf_mmap; break;
dst_tensor->data = mmap_data + tensor_offset; }
uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
dst_tensor->buffer = buf_mmap;
dst_tensor->data = mmap_data + tensor_offset;
file_mapped_bytes += tensor_size; file_mapped_bytes += tensor_size;
file_mapped_tensors++; file_mapped_tensors++;
@ -956,7 +993,7 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
if (file_mapped_bytes > 0) { if (file_mapped_bytes > 0) {
mapped_tensors += file_mapped_tensors; mapped_tensors += file_mapped_tensors;
mapped_bytes += file_mapped_bytes; mapped_bytes += file_mapped_bytes;
result.push_back({fdata.mmapped, fdata.mmbuffer}); result.push_back({fdata.mmapped, file_buffer});
} }
} }
@ -972,6 +1009,16 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
return result; return result;
} }
std::vector<ggml_backend_buffer_t> ModelLoader::get_device_mmap_buffers() const {
std::vector<ggml_backend_buffer_t> buffers;
for (const auto& fdata : file_data) {
for (const auto& entry : fdata.device_mmbuffers) {
buffers.push_back(entry.second.get());
}
}
return buffers;
}
bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool enable_mmap, bool enable_mmap,
const std::set<std::string>* target_tensor_names, const std::set<std::string>* target_tensor_names,
@ -1115,6 +1162,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) { if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
continue; continue;
} }
if (dst_tensor->buffer != nullptr &&
std::any_of(fdata.device_mmbuffers.begin(), fdata.device_mmbuffers.end(),
[&](const auto& entry) { return entry.second.get() == dst_tensor->buffer; })) {
continue;
}
size_t nbytes_to_read = tensor_storage.nbytes_to_read(); size_t nbytes_to_read = tensor_storage.nbytes_to_read();

View File

@ -20,6 +20,8 @@ struct ModelFileData {
std::vector<TensorStorage> tensors; std::vector<TensorStorage> tensors;
std::shared_ptr<MmapWrapper> mmapped; std::shared_ptr<MmapWrapper> mmapped;
std::shared_ptr<struct ggml_backend_buffer> mmbuffer; std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
// mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr)
std::map<ggml_backend_dev_t, std::shared_ptr<struct ggml_backend_buffer>> device_mmbuffers;
bool is_zip; bool is_zip;
}; };
@ -120,7 +122,9 @@ public:
void process_model_files(bool enable_mmap = false, bool writable_mmap = true); void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors, std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
std::set<std::string> ignore_tensors = {}, std::set<std::string> ignore_tensors = {},
bool writable = true); bool writable = true,
ggml_backend_dev_t device = nullptr);
std::vector<ggml_backend_buffer_t> get_device_mmap_buffers() const;
bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool use_mmap = false, bool use_mmap = false,
const std::set<std::string>* target_tensor_names = nullptr, const std::set<std::string>* target_tensor_names = nullptr,

View File

@ -780,38 +780,52 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
bool ModelManager::mmap_params(const std::vector<TensorState*>& states, bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
std::vector<ParamsStorageBlock*>& created_storage_blocks) { std::vector<ParamsStorageBlock*>& created_storage_blocks) {
std::map<std::string, ggml_tensor*> mmap_candidates; // A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing
std::map<std::string, TensorState*> mmap_states; // stages them for it, so they are mapped through a buffer of that GPU's device.
struct MmapGroup {
std::map<std::string, ggml_tensor*> candidates;
std::map<std::string, TensorState*> states;
};
std::map<ggml_backend_dev_t, MmapGroup> groups;
for (TensorState* state : states) { for (TensorState* state : states) {
if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr || if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr ||
state->tensor->data != nullptr || state->tensor->view_src != nullptr) { state->tensor->data != nullptr || state->tensor->view_src != nullptr) {
continue; continue;
} }
mmap_candidates[state->name] = state->tensor; ggml_backend_dev_t device = nullptr;
mmap_states[state->name] = state; if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) {
} device = ggml_backend_get_device(state->compute_backend);
if (mmap_candidates.empty()) {
return true;
}
auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_);
if (mmap_store.empty()) {
return true;
}
auto block = std::make_unique<ParamsStorageBlock>();
block->mmap_tensor_stores = std::move(mmap_store);
ParamsStorageBlock* raw = block.get();
for (const auto& pair : mmap_states) {
TensorState* state = pair.second;
if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
block->states.push_back(state);
} }
MmapGroup& group = groups[device];
group.candidates[state->name] = state->tensor;
group.states[state->name] = state;
} }
if (!block->states.empty()) { for (auto& [device, group] : groups) {
params_storage_blocks_.push_back(std::move(block)); // Device buffers wrap read-only mappings only; params that LoRAs are merged into in place
created_storage_blocks.push_back(raw); // are loaded instead.
if (device != nullptr && writable_mmap_) {
continue;
}
auto mmap_store = model_loader_.mmap_tensors(group.candidates, {}, writable_mmap_, device);
if (mmap_store.empty()) {
continue;
}
auto block = std::make_unique<ParamsStorageBlock>();
block->mmap_tensor_stores = std::move(mmap_store);
ParamsStorageBlock* raw = block.get();
for (const auto& pair : group.states) {
TensorState* state = pair.second;
if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
block->states.push_back(state);
}
}
if (!block->states.empty()) {
params_storage_blocks_.push_back(std::move(block));
created_storage_blocks.push_back(raw);
}
} }
return true; return true;
} }
@ -1353,15 +1367,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
} }
size_t total_size = 0; size_t total_size = 0;
auto add_buffer = [&](ggml_backend_buffer_t buffer) { std::unordered_set<ggml_backend_buffer_t> seen;
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) { auto add_buffer = [&](ggml_backend_buffer_t buffer) {
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) {
return; return;
} }
ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer); ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer);
auto split_devices = split_buffer_devices_.find(buffer_type); auto split_devices = split_buffer_devices_.find(buffer_type);
const bool on_device = split_devices == split_buffer_devices_.end() const bool on_device = split_devices == split_buffer_devices_.end()
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
return ggml_backend_get_device(entry.first) == compute_device; return ggml_backend_get_device(entry.first) == compute_device;
}); });
if (!on_device) { if (!on_device) {
@ -1371,9 +1386,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size; total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size;
}; };
// The loader may retain device mappings after their parameter blocks are released.
for (ggml_backend_buffer_t buffer : model_loader_.get_device_mmap_buffers()) {
add_buffer(buffer);
}
for (const auto& block : params_storage_blocks_) { for (const auto& block : params_storage_blocks_) {
if (block != nullptr) { if (block != nullptr) {
add_buffer(block->buffer); add_buffer(block->buffer);
for (const auto& store : block->mmap_tensor_stores) {
add_buffer(store.mmbuffer.get());
}
} }
} }
for (const auto& block : compute_staging_blocks_) { for (const auto& block : compute_staging_blocks_) {

View File

@ -1349,9 +1349,8 @@ bool is_first_stage_model_name(const std::string& name) {
} }
static std::string convert_esrgan_tensor_name(std::string name) { static std::string convert_esrgan_tensor_name(std::string name) {
static std::unordered_map<std::string, std::string> esrgan_name_map; static const auto esrgan_name_map = [] {
std::unordered_map<std::string, std::string> esrgan_name_map;
if (esrgan_name_map.empty()) {
esrgan_name_map["model.0."] = "conv_first."; esrgan_name_map["model.0."] = "conv_first.";
constexpr int max_num_blocks = 64; constexpr int max_num_blocks = 64;
@ -1377,7 +1376,8 @@ static std::string convert_esrgan_tensor_name(std::string name) {
esrgan_name_map["model.7."] = "conv_last."; esrgan_name_map["model.7."] = "conv_last.";
esrgan_name_map["model.8."] = "conv_hr."; esrgan_name_map["model.8."] = "conv_hr.";
esrgan_name_map["model.10."] = "conv_last."; esrgan_name_map["model.10."] = "conv_last.";
} return esrgan_name_map;
}();
replace_with_prefix_map(name, esrgan_name_map); replace_with_prefix_map(name, esrgan_name_map);
return name; return name;

View File

@ -7,6 +7,7 @@
#include <list> #include <list>
#include <mutex> #include <mutex>
#include <set> #include <set>
#include <tuple>
#include <type_traits> #include <type_traits>
#include <unordered_set> #include <unordered_set>
#include <utility> #include <utility>
@ -29,6 +30,7 @@
#include "stable-diffusion.h" #include "stable-diffusion.h"
#include "conditioning/conditioner.hpp" #include "conditioning/conditioner.hpp"
#include "conditioning/conditioning_cache.h"
#include "core/backend_fit.h" #include "core/backend_fit.h"
#include "extensions/generation_extension.h" #include "extensions/generation_extension.h"
#include "model/adapter/ip_adapter.hpp" #include "model/adapter/ip_adapter.hpp"
@ -135,6 +137,7 @@ static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
StableDiffusionGGML::StableDiffusionGGML() StableDiffusionGGML::StableDiffusionGGML()
: rng(std::make_shared<PhiloxRNG>()), : rng(std::make_shared<PhiloxRNG>()),
conditioning_cache_(std::make_unique<ConditioningCache>()),
denoiser(std::make_shared<CompVisDenoiser>()) {} denoiser(std::make_shared<CompVisDenoiser>()) {}
StableDiffusionGGML::~StableDiffusionGGML() = default; StableDiffusionGGML::~StableDiffusionGGML() = default;
@ -204,6 +207,8 @@ void StableDiffusionGGML::end_runners() {
} }
bool StableDiffusionGGML::reset_runners(const RunnerGroups& groups) { bool StableDiffusionGGML::reset_runners(const RunnerGroups& groups) {
conditioning_cache_->clear();
conditioning_loras_.clear();
end_runners(); end_runners();
clear_lora_adapters(); clear_lora_adapters();
runtime_lora_models.clear(); runtime_lora_models.clear();
@ -915,6 +920,11 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
return false; return false;
} }
} }
if (sd_ctx_params->conditioning_cache_size < 0) {
LOG_ERROR("conditioning_cache_size must be non-negative");
return false;
}
conditioning_cache_->set_capacity(static_cast<size_t>(sd_ctx_params->conditioning_cache_size));
auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params); auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
n_threads = sd_ctx_params->n_threads; n_threads = sd_ctx_params->n_threads;
tensor_executor = std::make_unique<sd::ParallelExecutor>(n_threads > 0 ? n_threads : sd_get_num_physical_cores()); tensor_executor = std::make_unique<sd::ParallelExecutor>(n_threads > 0 ? n_threads : sd_get_num_physical_cores());
@ -1765,8 +1775,19 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
int64_t t0 = ggml_time_ms(); int64_t t0 = ggml_time_ms();
end_runners(); end_runners();
clear_lora_adapters(); clear_lora_adapters();
if (!model_manager->prepare_lora_sources(all_loras)) if (!model_manager->prepare_lora_sources(all_loras)) {
conditioning_cache_->clear();
return false; return false;
}
if (!std::equal(all_loras.begin(), all_loras.end(),
conditioning_loras_.begin(), conditioning_loras_.end(),
[](const ModelManager::LoraSpec& a, const ModelManager::LoraSpec& b) {
return a.file_id == b.file_id && a.file_revision == b.file_revision &&
a.multiplier == b.multiplier && a.is_high_noise == b.is_high_noise &&
a.tensor_name_prefix_filter == b.tensor_name_prefix_filter;
})) {
conditioning_cache_->clear();
}
runtime_lora_models.erase(std::remove_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) { runtime_lora_models.erase(std::remove_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) {
return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) { return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) {
return entry.matches(spec); return entry.matches(spec);
@ -1776,6 +1797,7 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
const bool success = apply_lora_immediately ? apply_loras_immediately(all_loras) const bool success = apply_lora_immediately ? apply_loras_immediately(all_loras)
: apply_loras_at_runtime(all_loras); : apply_loras_at_runtime(all_loras);
if (!success) { if (!success) {
conditioning_cache_->clear();
clear_lora_adapters(); clear_lora_adapters();
runtime_lora_models.clear(); runtime_lora_models.clear();
return false; return false;
@ -1785,9 +1807,14 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
if (!all_loras.empty()) { if (!all_loras.empty()) {
LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
} }
conditioning_loras_ = std::move(all_loras);
return true; return true;
} }
SDCondition StableDiffusionGGML::get_learned_condition(const ConditionerParams& params) {
return conditioning_cache_->get(*cond_stage_model, n_threads, params);
}
void StableDiffusionGGML::reset_generation_extensions() { void StableDiffusionGGML::reset_generation_extensions() {
for (auto& extension : generation_extensions) { for (auto& extension : generation_extensions) {
extension->reset_runtime_condition(); extension->reset_runtime_condition();
@ -1972,6 +1999,8 @@ void StableDiffusionGGML::preview_image(int step,
int patch_sz = 1; int patch_sz = 1;
const float(*latent_rgb_proj)[3] = nullptr; const float(*latent_rgb_proj)[3] = nullptr;
float* latent_rgb_bias = nullptr; float* latent_rgb_bias = nullptr;
const float* latent_alpha_proj = nullptr;
float latent_alpha_bias = 1.f;
if (channels == 128) { if (channels == 128) {
if (sd_version_uses_flux2_vae(version)) { if (sd_version_uses_flux2_vae(version)) {
@ -1985,6 +2014,16 @@ void StableDiffusionGGML::preview_image(int step,
LOG_WARN("No latent to RGB projection known for this model"); LOG_WARN("No latent to RGB projection known for this model");
return; return;
} }
} else if (channels == 64) {
if (version == VERSION_QWEN_IMAGE_2_1) {
latent_rgb_proj = qwen21_latent_rgb_proj;
latent_rgb_bias = qwen21_latent_rgb_bias;
latent_alpha_proj = qwen21_latent_alpha_proj;
latent_alpha_bias = qwen21_latent_alpha_bias;
} else {
LOG_WARN("No latent to RGB projection known for this model");
return;
}
} else if (channels == 48) { } else if (channels == 48) {
if (sd_version_is_wan(version)) { if (sd_version_is_wan(version)) {
latent_rgb_proj = wan_22_latent_rgb_proj; latent_rgb_proj = wan_22_latent_rgb_proj;
@ -2035,13 +2074,14 @@ void StableDiffusionGGML::preview_image(int step,
uint32_t img_width = static_cast<uint32_t>(_latents.shape()[0]) * patch_sz; uint32_t img_width = static_cast<uint32_t>(_latents.shape()[0]) * patch_sz;
uint32_t img_height = static_cast<uint32_t>(_latents.shape()[1]) * patch_sz; uint32_t img_height = static_cast<uint32_t>(_latents.shape()[1]) * patch_sz;
uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * 3 * sizeof(uint8_t)); uint32_t img_channels = latent_alpha_proj != nullptr ? 4 : 3;
uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * img_channels * sizeof(uint8_t));
GGML_ASSERT(data != nullptr); GGML_ASSERT(data != nullptr);
preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz); preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz, latent_alpha_proj, latent_alpha_bias);
sd_image_t* images = (sd_image_t*)malloc(frames * sizeof(sd_image_t)); sd_image_t* images = (sd_image_t*)malloc(frames * sizeof(sd_image_t));
GGML_ASSERT(images != nullptr); GGML_ASSERT(images != nullptr);
for (uint32_t i = 0; i < frames; i++) { for (uint32_t i = 0; i < frames; i++) {
images[i] = {img_width, img_height, 3, data + i * img_width * img_height * 3}; images[i] = {img_width, img_height, img_channels, data + i * img_width * img_height * img_channels};
} }
step_callback(step, frames, images, is_noisy, step_callback_data); step_callback(step, frames, images, is_noisy, step_callback_data);
free(data); free(data);
@ -2216,6 +2256,15 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
}; };
RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()}; RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()};
// These inputs are immutable for this sampling run. Extensions may replace or
// modify them per step, so those paths need an explicit stability contract first.
const bool cache_qwen_prefix = version == VERSION_QWEN_IMAGE_2_1 &&
std::none_of(generation_extensions.begin(), generation_extensions.end(),
[](const auto& extension) { return extension->is_enabled(); });
using QwenPrefixInputs = std::tuple<const sd::Tensor<float>*, const sd::Tensor<int32_t>*,
const std::vector<sd::Tensor<float>>*>;
std::vector<QwenPrefixInputs> qwen_prefix_inputs;
RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
const bool apply_denoise_mask = !denoise_mask.empty() && const bool apply_denoise_mask = !denoise_mask.empty() &&
@ -2392,6 +2441,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
timesteps_tensor, timesteps_tensor,
cond, cond,
&controls); &controls);
bool uncond_controls_ready = false;
static const std::vector<sd::Tensor<float>> empty_ref_latents; static const std::vector<sd::Tensor<float>> empty_ref_latents;
bool uncond_without_ref_latents = !img_uncond.empty() && bool uncond_without_ref_latents = !img_uncond.empty() &&
@ -2481,10 +2531,33 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
return std::move(cached_output); return std::move(cached_output);
} }
// A re-enabled condition can miss the cache even when the positive pass was reused.
if (!uncond_controls_ready && !uncond.empty() &&
(&condition == &uncond || &condition == &img_uncond)) {
compute_sample_controls(control_image,
noised_input,
timesteps_tensor,
uncond,
&controls);
uncond_controls_ready = true;
}
for (const auto& extension : generation_extensions) { for (const auto& extension : generation_extensions) {
extension->before_diffusion(diffusion_params, step); extension->before_diffusion(diffusion_params, step);
} }
if (cache_qwen_prefix) {
auto* extra = std::get_if<QwenImage21DiffusionExtra>(&diffusion_params.extra);
if (extra != nullptr) {
auto key = std::make_tuple(diffusion_params.context, extra->image_slots,
diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr);
auto entry = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key);
extra->prefix_id = static_cast<uint64_t>(entry - qwen_prefix_inputs.begin()) + 1;
if (entry == qwen_prefix_inputs.end()) {
qwen_prefix_inputs.push_back(key);
}
}
}
auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params); auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params);
if (output_opt.empty()) { if (output_opt.empty()) {
LOG_ERROR("diffusion model compute failed"); LOG_ERROR("diffusion model compute failed");
@ -2508,41 +2581,69 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
} }
} }
float effective_guidance_scale = guidance_schedule.empty()
? cfg_scale
: guidance_schedule[guidance_schedule.size() - 1 - step];
float image_guidance_scale = img_cfg_scale;
constexpr float kEpsilon = 1e-5f;
bool skip_uncond = false;
if (!uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
if (!img_uncond.empty()) {
skip_uncond = std::abs(image_guidance_scale - effective_guidance_scale) < kEpsilon;
} else {
skip_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
}
}
bool skip_img_uncond = false;
if (!img_uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
if (!uncond.empty()) {
skip_img_uncond = std::abs(image_guidance_scale - 1.0f) < kEpsilon;
} else {
skip_img_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
}
}
cond_out = run_condition(*positive_condition, c_concat_override); cond_out = run_condition(*positive_condition, c_concat_override);
if (cond_out.empty()) { if (cond_out.empty()) {
return {}; return {};
} }
if (!uncond.empty()) { if (!uncond.empty()) {
if (!step_cache.is_step_skipped()) { if (!skip_uncond) {
compute_sample_controls(control_image, const std::vector<int>* uncond_skip_layers = nullptr;
noised_input, if (is_skiplayer_step && slg_uncond) {
timesteps_tensor, LOG_VERBOSE("Skipping layers at uncond step %d\n", step);
uncond, uncond_skip_layers = &skip_layer_guidance.layers();
&controls); }
} uncond_out = run_condition(uncond,
const std::vector<int>* uncond_skip_layers = nullptr; uncond.c_concat.empty() ? nullptr : &uncond.c_concat,
if (is_skiplayer_step && slg_uncond) { uncond_skip_layers,
LOG_VERBOSE("Skipping layers at uncond step %d\n", step); nullptr,
uncond_skip_layers = &skip_layer_guidance.layers(); true);
} if (uncond_out.empty()) {
uncond_out = run_condition(uncond, return {};
uncond.c_concat.empty() ? nullptr : &uncond.c_concat, }
uncond_skip_layers, } else {
nullptr, step_cache.invalidate_condition(&uncond);
true);
if (uncond_out.empty()) {
return {};
} }
} }
if (!img_uncond.empty()) { if (!img_uncond.empty()) {
img_uncond_out = run_condition(img_uncond, if (!skip_img_uncond) {
img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat, img_uncond_out = run_condition(img_uncond,
nullptr, img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat,
uncond_without_ref_latents ? &empty_ref_latents : nullptr, nullptr,
true); uncond_without_ref_latents ? &empty_ref_latents : nullptr,
if (img_uncond_out.empty()) { true);
return {}; if (img_uncond_out.empty()) {
return {};
}
} else {
step_cache.invalidate_condition(&img_uncond);
} }
} }
sd::guidance::GuidanceInput guidance_input; sd::guidance::GuidanceInput guidance_input;
@ -2552,7 +2653,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
guidance_input.pred_uncond = uncond_out.empty() ? nullptr : &uncond_out; guidance_input.pred_uncond = uncond_out.empty() ? nullptr : &uncond_out;
guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out; guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out;
sd::guidance::GuiderOutput guided = guidance_schedule.empty() ? primary_guidance.forward(guidance_input, {}) : primary_guidance.forward(guidance_input, {}, guidance_schedule[guidance_schedule.size() - 1 - step]); sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {}, effective_guidance_scale);
if (guided.pred.empty()) { if (guided.pred.empty()) {
return {}; return {};
} }

View File

@ -26,6 +26,7 @@ class RNG;
struct Denoiser; struct Denoiser;
struct LoraModel; struct LoraModel;
struct ConditionerParams; struct ConditionerParams;
class ConditioningCache;
struct SDCondition; struct SDCondition;
struct RefImageParams; struct RefImageParams;
namespace Wav2Vec2 { namespace Wav2Vec2 {
@ -178,6 +179,8 @@ public:
std::recursive_mutex execution_mutex; std::recursive_mutex execution_mutex;
std::unique_ptr<ModelConfig> config_; std::unique_ptr<ModelConfig> config_;
RunnerState runner_state_; RunnerState runner_state_;
std::unique_ptr<ConditioningCache> conditioning_cache_;
std::vector<ModelManager::LoraSpec> conditioning_loras_;
bool executing_ = false; bool executing_ = false;
std::shared_ptr<Denoiser> denoiser; std::shared_ptr<Denoiser> denoiser;
@ -361,6 +364,8 @@ public:
bool apply_loras(const sd_lora_t* loras, uint32_t lora_count); bool apply_loras(const sd_lora_t* loras, uint32_t lora_count);
SDCondition get_learned_condition(const ConditionerParams& params);
void reset_generation_extensions(); void reset_generation_extensions();
void prepare_generation_extensions(const sd_pm_params_t& pm_params, void prepare_generation_extensions(const sd_pm_params_t& pm_params,

View File

@ -10,6 +10,7 @@
#include "model/vae/vae.hpp" #include "model/vae/vae.hpp"
#include "request.h" #include "request.h"
#include "runtime/denoiser.hpp" #include "runtime/denoiser.hpp"
#include "runtime/image_preprocess.h"
#include "upscaler.h" #include "upscaler.h"
namespace sd::pipeline { namespace sd::pipeline {
@ -440,8 +441,7 @@ namespace sd::pipeline {
sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength); sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength);
int64_t prepare_start_ms = ggml_time_ms(); int64_t prepare_start_ms = ggml_time_ms();
condition_params.zero_out_masked = false; condition_params.zero_out_masked = false;
auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, auto cond = sd->get_learned_condition(condition_params);
condition_params);
if (cond.empty()) { if (cond.empty()) {
LOG_ERROR("failed to encode prompt"); LOG_ERROR("failed to encode prompt");
return std::nullopt; return std::nullopt;
@ -479,8 +479,7 @@ namespace sd::pipeline {
// LLaDA-Image CFG keeps the source latent but drops its SigVQ features. // LLaDA-Image CFG keeps the source latent but drops its SigVQ features.
condition_params.ref_images = nullptr; condition_params.ref_images = nullptr;
} }
uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, uncond = sd->get_learned_condition(condition_params);
condition_params);
if (uncond.empty()) { if (uncond.empty()) {
LOG_ERROR("failed to encode negative prompt"); LOG_ERROR("failed to encode negative prompt");
return std::nullopt; return std::nullopt;
@ -508,8 +507,7 @@ namespace sd::pipeline {
if (use_ref_latent_img_cfg) { if (use_ref_latent_img_cfg) {
condition_params.ref_images = &empty_ref_images; condition_params.ref_images = &empty_ref_images;
} }
img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, img_uncond = sd->get_learned_condition(condition_params);
condition_params);
if (img_uncond.empty()) { if (img_uncond.empty()) {
LOG_ERROR("failed to encode image guidance prompt"); LOG_ERROR("failed to encode image guidance prompt");
return std::nullopt; return std::nullopt;
@ -800,6 +798,12 @@ namespace sd::pipeline {
int64_t t0 = ggml_time_ms(); int64_t t0 = ggml_time_ms();
sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params;
GenerationRequest request(sd, sd_img_gen_params); GenerationRequest request(sd, sd_img_gen_params);
sd::ImagePreprocessor preprocessing(sd_img_gen_params->image_preprocess.rules);
sd_img_gen_params_t processed_params = *sd_img_gen_params;
if (!preprocessing.prepare_inputs(processed_params, request.width, request.height))
return false;
sd_img_gen_params = &processed_params;
request.pm_params = processed_params.pm_params;
LOG_INFO("generate_image %dx%d", request.width, request.height); LOG_INFO("generate_image %dx%d", request.width, request.height);
sd->rng->manual_seed(request.seed); sd->rng->manual_seed(request.seed);

View File

@ -291,7 +291,8 @@ namespace sd::model_builders {
result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map, tensor_storage_map,
"model.diffusion_model", "model.diffusion_model",
weight_manager); weight_manager,
sd_ctx_params->model_args);
} else { } else {
result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map, tensor_storage_map,

View File

@ -15,6 +15,7 @@
#include "model/vae/vae.hpp" #include "model/vae/vae.hpp"
#include "request.h" #include "request.h"
#include "runtime/denoiser.hpp" #include "runtime/denoiser.hpp"
#include "runtime/image_preprocess.h"
namespace sd::pipeline { namespace sd::pipeline {
@ -1161,8 +1162,7 @@ namespace sd::pipeline {
} }
int64_t prepare_start_ms = ggml_time_ms(); int64_t prepare_start_ms = ggml_time_ms();
embeds.cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, embeds.cond = sd->get_learned_condition(condition_params);
condition_params);
if (embeds.cond.empty()) { if (embeds.cond.empty()) {
LOG_ERROR("failed to encode video prompt"); LOG_ERROR("failed to encode video prompt");
return std::nullopt; return std::nullopt;
@ -1187,8 +1187,7 @@ namespace sd::pipeline {
} }
if (request.use_uncond) { if (request.use_uncond) {
condition_params.text = request.negative_prompt; condition_params.text = request.negative_prompt;
embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, embeds.uncond = sd->get_learned_condition(condition_params);
condition_params);
if (embeds.uncond.empty()) { if (embeds.uncond.empty()) {
LOG_ERROR("failed to encode negative video prompt"); LOG_ERROR("failed to encode negative video prompt");
return std::nullopt; return std::nullopt;
@ -1526,6 +1525,7 @@ namespace sd::pipeline {
img_gen_params.qwen_image_layers = 0; img_gen_params.qwen_image_layers = 0;
img_gen_params.circular_x = sd_vid_gen_params->circular_x; img_gen_params.circular_x = sd_vid_gen_params->circular_x;
img_gen_params.circular_y = sd_vid_gen_params->circular_y; img_gen_params.circular_y = sd_vid_gen_params->circular_y;
img_gen_params.image_preprocess = sd_vid_gen_params->image_preprocess;
sd->animatediff_num_frames = n_frames; sd->animatediff_num_frames = n_frames;
bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out);
@ -1556,6 +1556,11 @@ namespace sd::pipeline {
sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; 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); sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
GenerationRequest request(sd, sd_vid_gen_params); GenerationRequest request(sd, sd_vid_gen_params);
sd::ImagePreprocessor preprocessing(sd_vid_gen_params->image_preprocess.rules);
sd_vid_gen_params_t processed_params = *sd_vid_gen_params;
if (!preprocessing.prepare_inputs(processed_params, request.width, request.height))
return false;
sd_vid_gen_params = &processed_params;
if (fps_out != nullptr) { if (fps_out != nullptr) {
*fps_out = request.fps; *fps_out = request.fps;
} }

View File

@ -728,12 +728,20 @@ inline float flux_time_shift(float mu, float sigma, float t) {
// https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289 // https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289
struct FluxScheduler : SigmaScheduler { struct FluxScheduler : SigmaScheduler {
int image_seq_len = 0; int image_seq_len = 0;
float base_shift = 0.5f; int base_image_seq_len = 256;
float max_shift = 1.15f; int max_image_seq_len = 4096;
float base_shift = 0.5f;
float max_shift = 1.15f;
float shift_terminal = 0.0f;
explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr) FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr)
: image_seq_len(image_seq_len) { : image_seq_len(image_seq_len) {
if (version == VERSION_QWEN_IMAGE_2_1) {
max_image_seq_len = 8192;
max_shift = 0.9f;
shift_terminal = 0.02f;
}
parse_extra_sample_args(extra_sample_args); parse_extra_sample_args(extra_sample_args);
} }
@ -752,10 +760,8 @@ struct FluxScheduler : SigmaScheduler {
} }
float compute_mu() const { float compute_mu() const {
constexpr float base_shift_anchor = 256.0f; float m = (max_shift - base_shift) / static_cast<float>(max_image_seq_len - base_image_seq_len);
constexpr float max_shift_anchor = 4096.0f; float b = base_shift - m * static_cast<float>(base_image_seq_len);
float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
float b = base_shift - m * base_shift_anchor;
return static_cast<float>(image_seq_len) * m + b; return static_cast<float>(image_seq_len) * m + b;
} }
@ -764,7 +770,7 @@ struct FluxScheduler : SigmaScheduler {
sigmas.reserve(n + 1); sigmas.reserve(n + 1);
float mu = compute_mu(); float mu = compute_mu();
LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f, shift_terminal=%.3f", image_seq_len, n, mu, shift_terminal);
if (n == 0) { if (n == 0) {
sigmas.push_back(1.0f); sigmas.push_back(1.0f);
@ -780,6 +786,16 @@ struct FluxScheduler : SigmaScheduler {
} }
} }
if (shift_terminal > 0.0f && n > 1) {
// The terminal shift applies to the last model evaluation, not the final zero sigma.
float scale_factor = (1.0f - sigmas[n - 1]) / (1.0f - shift_terminal);
if (std::isfinite(scale_factor) && scale_factor > 0.0f) {
for (uint32_t i = 0; i < n; ++i) {
sigmas[i] = 1.0f - (1.0f - sigmas[i]) / scale_factor;
}
}
}
sigmas[n] = 0.0f; sigmas[n] = 0.0f;
return sigmas; return sigmas;
} }
@ -1178,7 +1194,7 @@ struct Denoiser {
} }
case FLUX_SCHEDULER: { case FLUX_SCHEDULER: {
LOG_INFO("get_sigmas with Flux scheduler"); LOG_INFO("get_sigmas with Flux scheduler");
scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args); scheduler = std::make_shared<FluxScheduler>(image_seq_len, version, extra_sample_args);
break; break;
} }
default: default:

View File

@ -0,0 +1,447 @@
#include "image_preprocess.h"
#include "core/util.h"
#include <climits>
#include <set>
namespace sd {
static constexpr std::pair<const char*, ImageTarget> image_targets[] = {
{"init", ImageTarget::Init},
{"end", ImageTarget::End},
{"mask", ImageTarget::Mask},
{"control", ImageTarget::Control},
{"ref", ImageTarget::Ref},
{"ip-adapter", ImageTarget::IPAdapter},
{"id", ImageTarget::ID},
{"control-frame", ImageTarget::ControlFrame},
};
static constexpr std::pair<const char*, ImageResizeMode> image_resize_modes[] = {
{"auto", ImageResizeMode::Auto},
{"none", ImageResizeMode::None},
{"stretch", ImageResizeMode::Stretch},
{"crop", ImageResizeMode::Crop},
{"crop-resize", ImageResizeMode::CropResize},
{"fit-pad", ImageResizeMode::FitPad},
};
template <typename T, size_t N>
static bool parse_enum(const std::string& text, const std::pair<const char*, T> (&names)[N], T& value) {
for (const auto& entry : names) {
if (text == entry.first) {
value = entry.second;
return true;
}
}
return false;
}
template <typename T, size_t N>
static const char* enum_name(T value, const std::pair<const char*, T> (&names)[N]) {
for (const auto& entry : names) {
if (value == entry.second)
return entry.first;
}
return "unknown";
}
template <typename T>
static bool one_of(T value, std::initializer_list<T> choices) {
return std::find(choices.begin(), choices.end(), value) != choices.end();
}
static bool one_of(const std::string& value, std::initializer_list<const char*> choices) {
for (const char* choice : choices) {
if (value == choice)
return true;
}
return false;
}
static ImageResizeMode resolve_mode(const std::map<std::string, std::string>& options, ImageResizeMode default_mode) {
auto it = options.find("mode");
ImageResizeMode mode = ImageResizeMode::Auto;
if (it != options.end())
parse_enum(it->second, image_resize_modes, mode);
if (mode != ImageResizeMode::Auto)
return mode;
return options.count("width") && default_mode == ImageResizeMode::None ? ImageResizeMode::Stretch : default_mode;
}
bool ImagePreprocessor::fail(const std::string& message) const {
LOG_ERROR("image preprocessing: %s", message.c_str());
valid_ = false;
return false;
}
ImagePreprocessor::ImagePreprocessor(const char* text) {
if (text == nullptr || trim(text).empty())
return;
for (const auto& part : split_string(text, ';')) {
ImagePreprocessRule rule;
std::set<std::string> keys;
if (trim(part).empty()) {
fail("empty rule");
return;
}
for (const auto& entry : split_string(part, ',')) {
size_t equal = entry.find('=');
if (equal == std::string::npos) {
fail("expected key=value: " + entry);
return;
}
std::string key = trim(entry.substr(0, equal));
std::string value = trim(entry.substr(equal + 1));
bool ok = !value.empty() && keys.insert(key).second;
int number = 0;
if (key == "target") {
ok &= parse_enum(value, image_targets, rule.target);
} else if (key == "index") {
ok &= parse_strict_int(value, rule.index) && rule.index >= 0;
} else {
if (key == "mode") {
ImageResizeMode mode;
ok &= parse_enum(value, image_resize_modes, mode);
} else if (key == "filter") {
ok &= one_of(value, {"auto", "nearest", "nearest-exact", "bilinear", "bicubic", "lanczos"});
} else if (key == "antialias") {
ok &= one_of(value, {"auto", "true", "false"});
} else if (key == "canny") {
ok &= one_of(value, {"true", "false"});
} else if (key == "anchor") {
ok &= one_of(value, {"center", "top", "bottom", "left", "right"});
} else if (key == "width" || key == "height") {
ok &= parse_strict_int(value, number) && number > 0;
} else if (key == "pad_color") {
ok &= value.size() == 7 || value.size() == 9;
ok &= !value.empty() && value[0] == '#';
for (size_t i = 1; i < value.size(); ++i)
ok &= std::isxdigit(static_cast<unsigned char>(value[i])) != 0;
} else {
ok = false;
}
rule.options[key] = value;
}
if (!ok) {
fail("invalid or duplicate option: " + entry);
return;
}
}
if (!keys.count("target") || rule.options.empty() ||
rule.options.count("width") != rule.options.count("height") ||
(rule.index >= 0 && !one_of(rule.target, {ImageTarget::Ref, ImageTarget::ID, ImageTarget::ControlFrame}))) {
fail("invalid target, index, or incomplete dimensions: " + part);
return;
}
rules_.push_back(std::move(rule));
}
for (const auto& rule : rules_) {
const auto options = resolve_options(rule.target, std::max(0, rule.index));
if (options.count("antialias") && options.at("antialias") == "true" && options.count("filter") &&
one_of(options.at("filter"), {"nearest", "nearest-exact"})) {
fail("antialias requires bilinear, bicubic, or lanczos");
return;
}
}
}
std::map<std::string, std::string> ImagePreprocessor::resolve_options(ImageTarget target, int index) const {
std::map<std::string, std::string> options;
for (int specificity = 0; specificity < 2; ++specificity) {
for (const auto& rule : rules_) {
if (rule.target == target &&
rule.index == (specificity == 0 ? -1 : index)) {
for (const auto& entry : rule.options)
options[entry.first] = entry.second;
}
}
}
return options;
}
bool ImagePreprocessor::validate_inputs(const sd_img_gen_params_t& params) const {
const std::map<ImageTarget, int> counts = {
{ImageTarget::Init, params.init_image.data != nullptr},
{ImageTarget::Mask, params.mask_image.data != nullptr},
{ImageTarget::Control, params.control_image.data != nullptr},
{ImageTarget::IPAdapter, params.ip_adapter_image.data != nullptr},
{ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0},
{ImageTarget::ID, params.pm_params.id_images != nullptr ? params.pm_params.id_images_count : 0},
};
for (const auto& rule : rules_) {
auto it = counts.find(rule.target);
int count = it == counts.end() ? 0 : it->second;
if (count <= 0 || rule.index >= count) {
return fail(std::string("rule targets an unavailable image: ") + enum_name(rule.target, image_targets));
}
}
return valid_;
}
bool ImagePreprocessor::validate_inputs(const sd_vid_gen_params_t& params) const {
const std::map<ImageTarget, int> counts = {
{ImageTarget::Init, params.init_image.data != nullptr},
{ImageTarget::End, params.end_image.data != nullptr},
{ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0},
{ImageTarget::ControlFrame, params.control_frames != nullptr ? params.control_frames_size : 0},
};
for (const auto& rule : rules_) {
auto it = counts.find(rule.target);
int count = it == counts.end() ? 0 : it->second;
if (count <= 0 || rule.index >= count)
return fail(std::string("rule targets an unavailable video input: ") + enum_name(rule.target, image_targets));
}
return valid_;
}
static int anchor_offset(int remaining, const std::string& anchor, bool horizontal) {
if (anchor == (horizontal ? "left" : "top"))
return 0;
if (anchor == (horizontal ? "right" : "bottom"))
return remaining;
return remaining / 2;
}
Tensor<float> ImagePreprocessor::apply_transform(const Tensor<float>& image, const std::map<std::string, std::string>& options, ImageTransform p, const std::string& label, ops::InterpolateMode default_filter) const {
auto value = [&](const char* key, const char* fallback) {
auto it = options.find(key);
return it == options.end() ? std::string(fallback) : it->second;
};
std::string filter = value("filter", "auto");
ops::InterpolateMode mode = default_filter;
if (filter == "nearest")
mode = ops::InterpolateMode::Nearest;
if (filter == "nearest-exact")
mode = ops::InterpolateMode::NearestExact;
if (filter == "bilinear")
mode = ops::InterpolateMode::Bilinear;
if (filter == "bicubic")
mode = ops::InterpolateMode::Bicubic;
if (filter == "lanczos")
mode = ops::InterpolateMode::Lanczos;
bool filtered = ops::is_2d_filter_interpolate_mode(mode);
bool antialias = value("antialias", "auto") == "true" ||
(value("antialias", "auto") == "auto" && filtered &&
(p.resize_width < p.crop_width || p.resize_height < p.crop_height));
if (antialias && !filtered) {
fail(label + ": antialias requires bilinear, bicubic, or lanczos");
return {};
}
auto cropped = ops::slice(ops::slice(image, 0, p.x, p.x + p.crop_width), 1, p.y, p.y + p.crop_height);
int channels = static_cast<int>(image.shape()[2]);
bool resize = p.resize_width != p.crop_width || p.resize_height != p.crop_height;
if (resize && channels == 4 && filtered) {
for (int64_t i = 0, pixels = cropped.shape()[0] * cropped.shape()[1]; i < pixels; ++i) {
for (int c = 0; c < 3; ++c)
cropped[i + c * pixels] *= cropped[i + 3 * pixels];
}
}
auto resized = ops::interpolate(cropped, {p.resize_width, p.resize_height, channels, 1}, mode, false, antialias);
if (resize && channels == 4 && filtered) {
for (int64_t i = 0, pixels = resized.shape()[0] * resized.shape()[1]; i < pixels; ++i) {
float alpha = std::clamp(resized[i + 3 * pixels], 0.f, 1.f);
for (int c = 0; c < 3; ++c)
resized[i + c * pixels] = alpha > 1e-6f ? resized[i + c * pixels] / alpha : 0.f;
}
}
resized = ops::clamp(resized, 0.f, 1.f);
Tensor<float> output({p.width, p.height, channels, 1});
std::string color = value("pad_color", "#000000ff");
if (color.size() == 7)
color += "ff";
uint8_t rgba[4];
for (int c = 0; c < 4; ++c)
rgba[c] = static_cast<uint8_t>(std::strtoul(color.substr(1 + c * 2, 2).c_str(), nullptr, 16));
for (int c = 0; c < channels; ++c) {
float fill = rgba[channels == 1 ? 0 : c] / 255.f;
for (int y = 0; y < p.height; ++y) {
for (int x = 0; x < p.width; ++x) {
output.index(x, y, c, 0) = x >= p.pad_x && x < p.pad_x + p.resize_width && y >= p.pad_y && y < p.pad_y + p.resize_height
? resized.index(x - p.pad_x, y - p.pad_y, c, 0)
: fill;
}
}
}
LOG_INFO("preprocess %s: %dx%d crop=(%d,%d,%d,%d) resize=%dx%d pad=(%d,%d) output=%dx%d filter=%s(%d) antialias=%s",
label.c_str(), p.source_width, p.source_height, p.x, p.y, p.crop_width, p.crop_height,
p.resize_width, p.resize_height, p.pad_x, p.pad_y, p.width, p.height, filter.c_str(), static_cast<int>(mode), BOOL_STR(antialias));
return output;
}
Tensor<float> ImagePreprocessor::apply_geometry(const Tensor<float>& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode, ops::InterpolateMode default_filter, ImageTransform* plan_out) const {
if (!valid_ || image.empty())
return {};
const std::string label = std::string(enum_name(target, image_targets)) + "[" + std::to_string(index) + "]";
auto options = resolve_options(target, index);
if (image.dim() != 4 || image.shape()[3] != 1 || image.shape()[2] < 1 || image.shape()[2] > 4) {
fail(label + ": expected one image with 1 to 4 channels");
return {};
}
ImageTransform p;
p.source_width = p.crop_width = static_cast<int>(image.shape()[0]);
p.source_height = p.crop_height = static_cast<int>(image.shape()[1]);
int target_width = width > 0 ? width : p.source_width;
int target_height = height > 0 ? height : p.source_height;
if (options.count("width")) {
parse_strict_int(options.at("width"), target_width);
parse_strict_int(options.at("height"), target_height);
}
ImageResizeMode mode = resolve_mode(options, default_mode);
std::string anchor = options.count("anchor") ? options.at("anchor") : "center";
p.width = p.resize_width = target_width;
p.height = p.resize_height = target_height;
if (mode == ImageResizeMode::None) {
if (options.count("width") && (target_width != p.source_width || target_height != p.source_height)) {
fail(label + ": mode=none conflicts with requested dimensions");
return {};
}
p.width = p.resize_width = p.source_width;
p.height = p.resize_height = p.source_height;
} else if (mode == ImageResizeMode::Crop || mode == ImageResizeMode::CropResize) {
if (mode == ImageResizeMode::Crop) {
p.crop_width = target_width;
p.crop_height = target_height;
} else if (int64_t(p.source_width) * target_height > int64_t(p.source_height) * target_width) {
p.crop_width = std::max(1, static_cast<int>(int64_t(p.source_height) * target_width / target_height));
} else {
p.crop_height = std::max(1, static_cast<int>(int64_t(p.source_width) * target_height / target_width));
}
if (p.crop_width > p.source_width || p.crop_height > p.source_height) {
fail(label + ": crop exceeds source dimensions");
return {};
}
p.x = anchor_offset(p.source_width - p.crop_width, anchor, true);
p.y = anchor_offset(p.source_height - p.crop_height, anchor, false);
} else if (mode == ImageResizeMode::FitPad) {
double scale = std::min(double(target_width) / p.source_width, double(target_height) / p.source_height);
p.resize_width = std::max(1, std::min(target_width, static_cast<int>(std::round(p.source_width * scale))));
p.resize_height = std::max(1, std::min(target_height, static_cast<int>(std::round(p.source_height * scale))));
p.pad_x = anchor_offset(target_width - p.resize_width, anchor, true);
p.pad_y = anchor_offset(target_height - p.resize_height, anchor, false);
}
if (p.width <= 0 || p.height <= 0) {
fail(label + ": invalid output dimensions");
return {};
}
uint64_t max_pixels = std::min<uint64_t>(INT64_MAX, SIZE_MAX / sizeof(float)) / static_cast<uint64_t>(image.shape()[2]);
if (uint64_t(p.width) * p.height > max_pixels || uint64_t(p.resize_width) * p.resize_height > max_pixels) {
fail(label + ": image allocation size overflows");
return {};
}
if (plan_out != nullptr)
*plan_out = p;
return apply_transform(image, options, p, label, default_filter);
}
Tensor<float> ImagePreprocessor::preprocess_input(sd_image_t image, ImageTarget target, int index, int width, int height) {
if (image.data == nullptr || image.width == 0 || image.height == 0 || image.width > INT_MAX || image.height > INT_MAX || image.channel < 1 || image.channel > 4) {
fail(std::string(enum_name(target, image_targets)) + ": invalid input image");
return {};
}
auto tensor = sd_image_to_tensor(image);
if (target == ImageTarget::Mask && has_init_transform_) {
auto options = resolve_options(target, index);
if (image.width != init_transform_.source_width || image.height != init_transform_.source_height) {
fail("mask and init source dimensions must match");
return {};
}
bool geometry_override = options.count("width") || options.count("anchor") ||
(options.count("mode") && options.at("mode") != "auto");
if (geometry_override) {
ImageTransform p;
auto init_options = resolve_options(ImageTarget::Init, 0);
ImageResizeMode default_mode = resolve_mode(init_options, ImageResizeMode::CropResize);
auto result = apply_geometry(tensor, target, index, init_transform_.width, init_transform_.height, default_mode, ops::InterpolateMode::NearestExact, &p);
if (result.empty())
return {};
const auto& q = init_transform_;
if (p.x != q.x || p.y != q.y || p.crop_width != q.crop_width || p.crop_height != q.crop_height ||
p.resize_width != q.resize_width || p.resize_height != q.resize_height || p.pad_x != q.pad_x || p.pad_y != q.pad_y || p.width != q.width || p.height != q.height) {
fail("mask geometry conflicts with init; configure geometry on init and filter on mask");
return {};
}
return result;
}
return apply_transform(tensor, options, init_transform_, "mask[0]", ops::InterpolateMode::NearestExact);
}
auto result = apply_geometry(tensor, target, index, width, height, width > 0 ? ImageResizeMode::CropResize : ImageResizeMode::None,
target == ImageTarget::Mask ? ops::InterpolateMode::NearestExact : ops::InterpolateMode::Nearest,
target == ImageTarget::Init ? &init_transform_ : nullptr);
if (target == ImageTarget::Init)
has_init_transform_ = !result.empty();
return result;
}
ImagePreprocessor::~ImagePreprocessor() {
for (const auto& image : owned_images_)
std::free(image.data);
}
bool ImagePreprocessor::prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height) {
if (image.data == nullptr)
return true;
auto options = resolve_options(target, index);
bool canny = options.count("canny") && options.at("canny") == "true";
auto tensor = preprocess_input(image, target, index, width, height);
if (tensor.empty())
return false;
auto output = tensor_to_sd_image(tensor);
if (output.data == nullptr)
return fail("could not allocate input preprocessing buffer");
owned_images_.push_back(output);
if (canny && !preprocess_canny(output, 0.08f, 0.08f, 0.8f, 1.f, false))
return fail("Canny preprocessing failed");
image = output;
return true;
}
bool ImagePreprocessor::prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector<sd_image_t>& storage, int width, int height) {
if (count < 0 || (count > 0 && images == nullptr))
return fail(std::string("invalid image array: ") + enum_name(target, image_targets));
if (count == 0)
return true;
storage.assign(images, images + count);
for (int i = 0; i < count; ++i) {
if (storage[i].data == nullptr)
return fail(std::string("empty image in array: ") + enum_name(target, image_targets));
if (!prepare_image(storage[i], target, i, width, height))
return false;
}
images = storage.data();
return true;
}
bool ImagePreprocessor::prepare_inputs(sd_img_gen_params_t& params, int width, int height) {
if (prepared_)
return fail("inputs have already been prepared");
prepared_ = true;
if (!valid_ || !validate_inputs(params))
return false;
if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) ||
!prepare_image(params.mask_image, ImageTarget::Mask, 0, width, height) ||
!prepare_image(params.control_image, ImageTarget::Control, 0, width, height) ||
!prepare_image(params.ip_adapter_image, ImageTarget::IPAdapter, 0, -1, -1) ||
!prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) ||
!prepare_array(params.pm_params.id_images, params.pm_params.id_images_count, ImageTarget::ID, id_images_))
return false;
params.image_preprocess = {};
return true;
}
bool ImagePreprocessor::prepare_inputs(sd_vid_gen_params_t& params, int width, int height) {
if (prepared_)
return fail("inputs have already been prepared");
prepared_ = true;
if (!valid_ || !validate_inputs(params))
return false;
if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) ||
!prepare_image(params.end_image, ImageTarget::End, 0, width, height) ||
!prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) ||
!prepare_array(params.control_frames, params.control_frames_size, ImageTarget::ControlFrame, control_frames_, width, height))
return false;
params.image_preprocess = {};
return true;
}
} // namespace sd

View File

@ -0,0 +1,88 @@
#ifndef __SD_RUNTIME_IMAGE_PREPROCESS_H__
#define __SD_RUNTIME_IMAGE_PREPROCESS_H__
#include <map>
#include <string>
#include <vector>
#include "core/tensor.hpp"
#include "stable-diffusion.h"
namespace sd {
enum class ImageTarget {
Init,
End,
Mask,
Control,
Ref,
IPAdapter,
ID,
ControlFrame,
};
enum class ImageResizeMode {
Auto,
None,
Stretch,
Crop,
CropResize,
FitPad,
};
struct ImageTransform {
int source_width = 0;
int source_height = 0;
int x = 0;
int y = 0;
int crop_width = 0;
int crop_height = 0;
int resize_width = 0;
int resize_height = 0;
int width = 0;
int height = 0;
int pad_x = 0;
int pad_y = 0;
};
struct ImagePreprocessRule {
ImageTarget target = ImageTarget::Init;
int index = -1;
std::map<std::string, std::string> options;
};
class ImagePreprocessor {
std::vector<ImagePreprocessRule> rules_;
mutable bool valid_ = true;
ImageTransform init_transform_;
bool has_init_transform_ = false;
bool prepared_ = false;
std::vector<sd_image_t> owned_images_;
std::vector<sd_image_t> ref_images_;
std::vector<sd_image_t> id_images_;
std::vector<sd_image_t> control_frames_;
bool fail(const std::string& message) const;
std::map<std::string, std::string> resolve_options(ImageTarget target, int index) const;
Tensor<float> apply_transform(const Tensor<float>& image, const std::map<std::string, std::string>& options, ImageTransform plan, const std::string& label, ops::InterpolateMode default_filter) const;
bool prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height);
bool prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector<sd_image_t>& storage, int width = -1, int height = -1);
public:
explicit ImagePreprocessor(const char* rules = nullptr);
~ImagePreprocessor();
ImagePreprocessor(const ImagePreprocessor&) = delete;
ImagePreprocessor& operator=(const ImagePreprocessor&) = delete;
bool prepare_inputs(sd_img_gen_params_t& params, int width, int height);
bool prepare_inputs(sd_vid_gen_params_t& params, int width, int height);
bool is_valid() const { return valid_; }
bool validate_inputs(const sd_img_gen_params_t& params) const;
bool validate_inputs(const sd_vid_gen_params_t& params) const;
Tensor<float> apply_geometry(const Tensor<float>& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode = ImageResizeMode::Stretch, ops::InterpolateMode default_filter = ops::InterpolateMode::Nearest, ImageTransform* plan_out = nullptr) const;
Tensor<float> preprocess_input(sd_image_t image, ImageTarget target, int index = 0, int width = -1, int height = -1);
};
} // namespace sd
#endif // __SD_RUNTIME_IMAGE_PREPROCESS_H__

View File

@ -4,6 +4,86 @@
#include "core/tensor.hpp" #include "core/tensor.hpp"
#include "ggml.h" #include "ggml.h"
// RGB is projected to [-1, 1]; alpha is projected directly to [0, 1].
const float qwen21_latent_rgb_proj[64][3] = {
{0.00860495522f, 0.01219501462f, -0.00321337196f},
{0.01889233090f, 0.01246581216f, 0.01074959482f},
{0.1255941446f, 0.1176879344f, -0.0332212352f},
{0.0418238528f, 0.1043427754f, 0.0121666316f},
{0.02025338f, 0.01453670296f, -0.000224336044f},
{-0.01896720702f, -0.0206099030f, -0.0322728584f},
{0.00438984796f, -0.01374969766f, 0.02849196f},
{-0.0374495856f, -0.0286777126f, -0.0693192810f},
{0.01511914734f, 0.0242979386f, 0.0553878870f},
{-0.1138629518f, -0.020391466f, 0.001550520522f},
{-0.0233650696f, -0.0417292018f, -0.0362361182f},
{-0.0351603342f, -0.0243595924f, -0.00216261038f},
{0.01093355288f, -0.0373466924f, 0.00241315350f},
{0.01778704744f, -0.00401984678f, -0.0343259192f},
{0.0486059334f, 0.0253144f, 0.0672564966f},
{0.0309463558f, 0.0277963166f, 0.0520869622f},
{0.0374485008f, 0.0551753676f, 0.0225853902f},
{-0.0090809962f, -0.004756176f, 0.00636443612f},
{-0.0270455652f, -0.0384966954f, -0.00905908082f},
{-0.00553493756f, 0.01484553684f, -0.0211502468f},
{0.01319502562f, 0.00948005666f, 0.0483789212f},
{-0.00931847104f, -0.00276452734f, -0.01011985302f},
{0.0180478258f, 0.01614954356f, -0.0209424690f},
{-0.0214530434f, -0.00272961176f, 0.0217887476f},
{-0.0636772304f, -0.0208893548f, 0.0479167742f},
{-0.0250321236f, -0.0286715676f, 0.0530110146f},
{-0.01853078078f, 0.01647272818f, -0.00207747588f},
{0.0023101082f, 0.01228800748f, 0.01303505006f},
{-0.01243671408f, -0.0258638728f, -0.0379116264f},
{0.00598934710f, 0.00642563550f, -0.01234514304f},
{-0.0296733996f, -0.0234698050f, 0.00060018212f},
{-0.0322019498f, -0.0529200462f, -0.00344987414f},
{-0.00205026458f, -0.00846599446f, 0.00455971038f},
{-0.01082227064f, 0.0315661948f, -0.0677753362f},
{0.0645553474f, 0.1109666998f, 0.0674744864f},
{0.01036801108f, -0.00484841210f, -0.001529168474f},
{0.01264353566f, 0.01548126338f, -0.00966374324f},
{-0.0223892408f, -0.00871751526f, -0.000306421670f},
{0.0271322742f, 0.03496524f, -0.0089692858f},
{0.0512178672f, 0.0173080034f, 0.00804227746f},
{0.01210987192f, 0.00758025926f, -0.00281712586f},
{0.1897278390f, 0.1210261828f, 0.062603892f},
{0.0208058822f, 0.00547548182f, 0.01262955638f},
{0.00813332858f, 0.01015930914f, 0.01301771290f},
{-0.000927236014f, -0.00152540594f, -0.00599213302f},
{0.01663314616f, -0.00582789626f, 0.0163958132f},
{-0.0252546342f, -0.0604193732f, -0.1606919922f},
{-0.091722686f, -0.0409201224f, -0.0959576198f},
{0.0282963112f, -0.01387223872f, -0.01648814464f},
{0.0552316818f, 0.0967547788f, 0.0413586632f},
{0.00922849292f, 0.00451467542f, -0.0529172378f},
{0.0558600768f, 0.0122988308f, -0.01445942422f},
{0.000210660902f, -0.01295958782f, -0.01804761764f},
{0.0358136250f, -0.0472505970f, -0.1156405142f},
{-0.0506390696f, -0.0471914842f, 0.0349791468f},
{-0.0480143168f, 0.00628389868f, -0.0545163826f},
{0.0315499582f, 0.0564846606f, -0.0430850488f},
{-0.0362330316f, -0.01267788554f, 0.0061024772f},
{0.0038627542f, 0.00911055916f, -0.00758526008f},
{-0.0447103298f, -0.00835411408f, 0.01545872328f},
{-0.015006738f, 0.00270612302f, -0.00784361356f},
{-0.0221755048f, -0.0513344748f, -0.0475317424f},
{-0.01036656294f, -0.00422146068f, -0.0213499052f},
{0.01788952706f, 0.01191944190f, 0.0397205238f},
};
float qwen21_latent_rgb_bias[3] = {-0.043293118f, -0.02695978f, -0.11986706f};
const float qwen21_latent_alpha_proj[64] = {
-0.0416241114f, -0.00678954612f, -0.0169095515f, -0.0230551401f, 0.0100882595f, 0.00655586802f, 0.0401166874f, -0.0055510216f,
0.0224234441f, -0.0389640963f, -0.0114492163f, -0.00721128977f, -0.0029064082f, 0.0150300547f, -0.00321615308f, -0.0498856338f,
-0.0215251401f, 0.0240220482f, 0.0117338008f, -0.0460420624f, 0.0387872889f, 0.0131517207f, 0.0147100836f, 0.0266985286f,
0.0153097324f, -0.0418119757f, 0.0421013917f, 0.0401724499f, 0.00972515915f, 0.011718495f, 0.0117622291f, 0.0136505134f,
-0.0350017363f, -0.0100692606f, -0.0131582529f, -0.00660639315f, 0.00253611396f, -0.0195736368f, -0.04240184f, 0.0321299262f,
0.0106089414f, -0.0179845306f, -0.00806212708f, 0.0135889057f, -0.0157393098f, -0.0267791344f, 0.0109068534f, 0.0283931966f,
-0.0435370078f, 0.00187883536f, -0.0108995378f, -0.0450757676f, -0.0699481501f, 0.0123562106f, -0.0222592249f, 0.0216155907f,
0.0563755424f, -0.0073379912f, 0.0160012921f, 0.0411637742f, 0.0189607258f, -0.024025029f, -0.0161487905f, -0.016913203f};
const float qwen21_latent_alpha_bias = 0.871268134f;
const float minimax_latent_rgb_proj[24][3] = { const float minimax_latent_rgb_proj[24][3] = {
{0.19819857f, 0.11584999f, 0.07929777f}, {0.19819857f, 0.11584999f, 0.07929777f},
{-0.16047224f, -0.10601170f, -0.15996324f}, {-0.16047224f, -0.10601170f, -0.15996324f},
@ -324,7 +404,7 @@ const float sd_latent_rgb_proj[4][3] = {
{-0.178022f, -0.200862f, -0.678514f}}; {-0.178022f, -0.200862f, -0.678514f}};
float sd_latent_rgb_bias[3] = {-0.017478f, -0.055834f, -0.105825f}; float sd_latent_rgb_bias[3] = {-0.017478f, -0.055834f, -0.105825f};
void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) { void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size, const float* latent_alpha_proj = nullptr, float latent_alpha_bias = 1.f) {
size_t buffer_head = 0; size_t buffer_head = 0;
uint32_t latent_width = static_cast<uint32_t>(latents->ne[0]); uint32_t latent_width = static_cast<uint32_t>(latents->ne[0]);
@ -338,7 +418,8 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
uint32_t rgb_width = latent_width * patch_size; uint32_t rgb_width = latent_width * patch_size;
uint32_t rgb_height = latent_height * patch_size; uint32_t rgb_height = latent_height * patch_size;
uint32_t unpatched_dim = dim / (patch_size * patch_size); uint32_t unpatched_dim = dim / (patch_size * patch_size);
const uint32_t output_channels = latent_alpha_proj != nullptr ? 4 : 3;
for (uint32_t k = 0; k < frames; k++) { for (uint32_t k = 0; k < frames; k++) {
for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) { for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) {
@ -356,13 +437,16 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
// should be incremented by 1 for each pixel // should be incremented by 1 for each pixel
size_t pixel_id = k * rgb_width * rgb_height + rgb_y * rgb_width + rgb_x; size_t pixel_id = k * rgb_width * rgb_height + rgb_y * rgb_width + rgb_x;
float r = 0, g = 0, b = 0; float r = 0, g = 0, b = 0, a = 0;
if (latent_rgb_proj != nullptr) { if (latent_rgb_proj != nullptr) {
for (uint32_t d = 0; d < unpatched_dim; d++) { for (uint32_t d = 0; d < unpatched_dim; d++) {
float value = *(float*)((char*)latents->data + latent_id + (d * patch_size * patch_size + channel_offset) * latents->nb[ggml_n_dims(latents) - 1]); float value = *(float*)((char*)latents->data + latent_id + (d * patch_size * patch_size + channel_offset) * latents->nb[ggml_n_dims(latents) - 1]);
r += value * latent_rgb_proj[d][0]; r += value * latent_rgb_proj[d][0];
g += value * latent_rgb_proj[d][1]; g += value * latent_rgb_proj[d][1];
b += value * latent_rgb_proj[d][2]; b += value * latent_rgb_proj[d][2];
if (latent_alpha_proj != nullptr) {
a += value * latent_alpha_proj[d];
}
} }
} else { } else {
// interpret first 3 channels as RGB // interpret first 3 channels as RGB
@ -386,9 +470,13 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
g = g >= 0 ? g <= 1 ? g : 1 : 0; g = g >= 0 ? g <= 1 ? g : 1 : 0;
b = b >= 0 ? b <= 1 ? b : 1 : 0; b = b >= 0 ? b <= 1 ? b : 1 : 0;
buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255); buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255);
buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255); buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 255);
buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255); buffer[pixel_id * output_channels + 2] = (uint8_t)(b * 255);
if (latent_alpha_proj != nullptr) {
a = std::min(1.0f, std::max(0.0f, a + latent_alpha_bias));
buffer[pixel_id * output_channels + 3] = (uint8_t)(a * 255);
}
} }
} }
} }
@ -398,16 +486,17 @@ static inline bool preview_latent_tensor_is_video(const sd::Tensor<float>& laten
return latents.dim() == 5; return latents.dim() == 5;
} }
void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) { void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size, const float* latent_alpha_proj = nullptr, float latent_alpha_bias = 1.f) {
uint32_t latent_width = static_cast<uint32_t>(latents.shape()[0]); uint32_t latent_width = static_cast<uint32_t>(latents.shape()[0]);
uint32_t latent_height = static_cast<uint32_t>(latents.shape()[1]); uint32_t latent_height = static_cast<uint32_t>(latents.shape()[1]);
bool is_video = preview_latent_tensor_is_video(latents); bool is_video = preview_latent_tensor_is_video(latents);
uint32_t frames = is_video ? static_cast<uint32_t>(latents.shape()[2]) : 1; uint32_t frames = is_video ? static_cast<uint32_t>(latents.shape()[2]) : 1;
uint32_t dim = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]); uint32_t dim = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
uint32_t rgb_width = latent_width * patch_size; uint32_t rgb_width = latent_width * patch_size;
uint32_t rgb_height = latent_height * patch_size; uint32_t rgb_height = latent_height * patch_size;
uint32_t unpatched_dim = dim / (patch_size * patch_size); uint32_t unpatched_dim = dim / (patch_size * patch_size);
const uint32_t output_channels = latent_alpha_proj != nullptr ? 4 : 3;
for (uint32_t k = 0; k < frames; k++) { for (uint32_t k = 0; k < frames; k++) {
for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) { for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) {
@ -427,7 +516,7 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, con
: latents.values()[latent_x + latent_width * (latent_y + latent_height * latent_channel)]; : latents.values()[latent_x + latent_width * (latent_y + latent_height * latent_channel)];
}; };
float r = 0.f, g = 0.f, b = 0.f; float r = 0.f, g = 0.f, b = 0.f, a = 0.f;
if (latent_rgb_proj != nullptr) { if (latent_rgb_proj != nullptr) {
for (uint32_t d = 0; d < unpatched_dim; d++) { for (uint32_t d = 0; d < unpatched_dim; d++) {
uint32_t latent_channel = d * patch_size * patch_size + channel_offset; uint32_t latent_channel = d * patch_size * patch_size + channel_offset;
@ -435,6 +524,9 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, con
r += value * latent_rgb_proj[d][0]; r += value * latent_rgb_proj[d][0];
g += value * latent_rgb_proj[d][1]; g += value * latent_rgb_proj[d][1];
b += value * latent_rgb_proj[d][2]; b += value * latent_rgb_proj[d][2];
if (latent_alpha_proj != nullptr) {
a += value * latent_alpha_proj[d];
}
} }
} else { } else {
r = latent_value(0); r = latent_value(0);
@ -450,9 +542,13 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, con
g = std::min(1.0f, std::max(0.0f, g * .5f + .5f)); g = std::min(1.0f, std::max(0.0f, g * .5f + .5f));
b = std::min(1.0f, std::max(0.0f, b * .5f + .5f)); b = std::min(1.0f, std::max(0.0f, b * .5f + .5f));
buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255); buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255);
buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255); buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 255);
buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255); buffer[pixel_id * output_channels + 2] = (uint8_t)(b * 255);
if (latent_alpha_proj != nullptr) {
a = std::min(1.0f, std::max(0.0f, a + latent_alpha_bias));
buffer[pixel_id * output_channels + 3] = (uint8_t)(a * 255);
}
} }
} }
} }

View File

@ -165,16 +165,18 @@ static inline sd::Tensor<float> convolve_tensor(const sd::Tensor<float>& input,
return output; return output;
} }
static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& rgb_img) { static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& image) {
GGML_ASSERT(rgb_img.dim() == 4); GGML_ASSERT(image.dim() == 4);
GGML_ASSERT(rgb_img.shape()[2] >= 3); GGML_ASSERT(image.shape()[2] >= 1);
sd::Tensor<float> grayscale({rgb_img.shape()[0], rgb_img.shape()[1], 1, rgb_img.shape()[3]}); sd::Tensor<float> grayscale({image.shape()[0], image.shape()[1], 1, image.shape()[3]});
for (int64_t iy = 0; iy < rgb_img.shape()[1]; ++iy) { for (int64_t iy = 0; iy < image.shape()[1]; ++iy) {
for (int64_t ix = 0; ix < rgb_img.shape()[0]; ++ix) { for (int64_t ix = 0; ix < image.shape()[0]; ++ix) {
float r = preprocessing_get_4d(rgb_img, ix, iy, 0, 0); float gray = preprocessing_get_4d(image, ix, iy, 0, 0);
float g = preprocessing_get_4d(rgb_img, ix, iy, 1, 0); if (image.shape()[2] >= 3) {
float b = preprocessing_get_4d(rgb_img, ix, iy, 2, 0); float g = preprocessing_get_4d(image, ix, iy, 1, 0);
float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b; float b = preprocessing_get_4d(image, ix, iy, 2, 0);
gray = 0.2989f * gray + 0.5870f * g + 0.1140f * b;
}
preprocessing_set_4d(grayscale, gray, ix, iy, 0, 0); preprocessing_set_4d(grayscale, gray, ix, iy, 0, 0);
} }
} }
@ -317,11 +319,12 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold,
image_gray = non_max_supression(G, theta); image_gray = non_max_supression(G, theta);
threshold_hystersis(&image_gray, high_threshold, low_threshold, weak, strong); threshold_hystersis(&image_gray, high_threshold, low_threshold, weak, strong);
const uint32_t color_channels = img.channel == 2 || img.channel == 4 ? img.channel - 1 : img.channel;
for (uint32_t iy = 0; iy < img.height; ++iy) { for (uint32_t iy = 0; iy < img.height; ++iy) {
for (uint32_t ix = 0; ix < img.width; ++ix) { for (uint32_t ix = 0; ix < img.width; ++ix) {
float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0); float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0);
gray = inverse ? 1.0f - gray : gray; gray = inverse ? 1.0f - gray : gray;
for (uint32_t c = 0; c < img.channel; ++c) { for (uint32_t c = 0; c < color_channels; ++c) {
preprocessing_set_4d(image, gray, ix, iy, c, 0); preprocessing_set_4d(image, gray, ix, iy, c, 0);
} }
} }

View File

@ -275,6 +275,26 @@ namespace sd_sample {
} }
} }
void SampleStepCacheDispatcher::invalidate_condition(const void* condition) {
if (condition == nullptr) {
return;
}
switch (runtime.mode) {
case SampleCacheMode::EASYCACHE:
runtime.easycache.cache_diffs.erase(condition);
break;
case SampleCacheMode::UCACHE:
runtime.ucache.cache_diffs.erase(condition);
break;
case SampleCacheMode::CACHEDIT:
runtime.cachedit.cache_diffs.erase(condition);
break;
case SampleCacheMode::NONE:
break;
}
}
bool SampleStepCacheDispatcher::is_step_skipped() const { bool SampleStepCacheDispatcher::is_step_skipped() const {
switch (runtime.mode) { switch (runtime.mode) {
case SampleCacheMode::EASYCACHE: case SampleCacheMode::EASYCACHE:

View File

@ -46,6 +46,7 @@ namespace sd_sample {
bool before_condition(const void* condition, const sd::Tensor<float>& input, sd::Tensor<float>* output); bool before_condition(const void* condition, const sd::Tensor<float>& input, sd::Tensor<float>* output);
void after_condition(const void* condition, const sd::Tensor<float>& input, const sd::Tensor<float>& output); void after_condition(const void* condition, const sd::Tensor<float>& input, const sd::Tensor<float>& output);
void invalidate_condition(const void* condition);
bool is_step_skipped() const; bool is_step_skipped() const;
}; };

View File

@ -326,6 +326,7 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) {
void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
*sd_ctx_params = {}; *sd_ctx_params = {};
sd_ctx_params->n_threads = sd_get_num_physical_cores(); sd_ctx_params->n_threads = sd_get_num_physical_cores();
sd_ctx_params->conditioning_cache_size = 4;
sd_ctx_params->wtype = SD_TYPE_COUNT; sd_ctx_params->wtype = SD_TYPE_COUNT;
sd_ctx_params->rng_type = CUDA_RNG; sd_ctx_params->rng_type = CUDA_RNG;
sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT;
@ -378,6 +379,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"pulid_weights_path: %s\n" "pulid_weights_path: %s\n"
"tensor_type_rules: %s\n" "tensor_type_rules: %s\n"
"n_threads: %d\n" "n_threads: %d\n"
"conditioning_cache_size: %d\n"
"wtype: %s\n" "wtype: %s\n"
"rng_type: %s\n" "rng_type: %s\n"
"sampler_rng_type: %s\n" "sampler_rng_type: %s\n"
@ -418,6 +420,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
SAFE_STR(sd_ctx_params->pulid_weights_path), SAFE_STR(sd_ctx_params->pulid_weights_path),
SAFE_STR(sd_ctx_params->tensor_type_rules), SAFE_STR(sd_ctx_params->tensor_type_rules),
sd_ctx_params->n_threads, sd_ctx_params->n_threads,
sd_ctx_params->conditioning_cache_size,
sd_type_name(sd_ctx_params->wtype), sd_type_name(sd_ctx_params->wtype),
sd_rng_type_name(sd_ctx_params->rng_type), sd_rng_type_name(sd_ctx_params->rng_type),
sd_rng_type_name(sd_ctx_params->sampler_rng_type), sd_rng_type_name(sd_ctx_params->sampler_rng_type),
@ -751,28 +754,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
int* num_frames_out, int* num_frames_out,
sd_audio_t** audio_out, sd_audio_t** audio_out,
int* fps_out) { int* fps_out) {
if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { if (frames_out != nullptr)
if (fps_out != nullptr) {
*fps_out = 0;
}
return false;
}
if (frames_out != nullptr) {
*frames_out = nullptr; *frames_out = nullptr;
} if (audio_out != nullptr)
if (audio_out != nullptr) {
*audio_out = nullptr; *audio_out = nullptr;
} if (num_frames_out != nullptr)
if (num_frames_out != nullptr) {
*num_frames_out = 0; *num_frames_out = 0;
if (fps_out != nullptr)
*fps_out = 0;
if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) {
return false;
} }
StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd);
if (!execution.ready) { if (!execution.ready) {
if (fps_out != nullptr) {
*fps_out = 0;
}
return false; return false;
} }

View File

@ -111,10 +111,22 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) { sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
sd::ParallelScope tensor_scope(&tensor_executor); sd::ParallelScope tensor_scope(&tensor_executor);
if (input_tensor.empty() || input_tensor.dim() != 4 ||
(input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) {
LOG_ERROR("esrgan expects a 4D RGB or RGBA image tensor");
return {};
}
const bool has_alpha = input_tensor.shape()[2] == 4;
sd::Tensor<float> rgb;
if (has_alpha) {
rgb = sd::ops::slice(input_tensor, 2, 0, 3);
}
const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
sd::Tensor<float> upscaled; sd::Tensor<float> upscaled;
const int scale = esrgan_upscaler->config.scale; const int scale = esrgan_upscaler->config.scale;
if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
upscaled = esrgan_upscaler->compute(n_threads, input_tensor); upscaled = esrgan_upscaler->compute(n_threads, model_input);
} else { } else {
auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> { auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
auto output_tile = esrgan_upscaler->compute(n_threads, input_tile); auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
@ -125,7 +137,7 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
return output_tile; return output_tile;
}; };
upscaled = process_tiles_2d(input_tensor, upscaled = process_tiles_2d(model_input,
static_cast<int>(input_tensor.shape()[0] * scale), static_cast<int>(input_tensor.shape()[0] * scale),
static_cast<int>(input_tensor.shape()[1] * scale), static_cast<int>(input_tensor.shape()[1] * scale),
scale, scale,
@ -141,6 +153,14 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
LOG_ERROR("esrgan compute failed"); LOG_ERROR("esrgan compute failed");
return {}; return {};
} }
if (has_alpha) {
auto alpha = sd::ops::slice(input_tensor, 2, 3, 4);
auto alpha_shape = alpha.shape();
alpha_shape[0] = upscaled.shape()[0];
alpha_shape[1] = upscaled.shape()[1];
alpha = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear);
upscaled = sd::ops::concat(upscaled, alpha, 2);
}
return upscaled; return upscaled;
} }
@ -242,6 +262,30 @@ int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {
return upscaler_ctx->upscaler->esrgan_upscaler->config.scale; return upscaler_ctx->upscaler->esrgan_upscaler->config.scale;
} }
int get_upscaler_model_scale(const char* model_path) {
if (model_path == nullptr || model_path[0] == '\0') {
return 0;
}
try {
ModelLoader loader;
if (!loader.init_from_file_and_convert_name(model_path, "", VERSION_ESRGAN)) {
return 0;
}
const auto& tensors = loader.get_tensor_storage_map();
auto first = tensors.find("conv_first.weight");
auto last = tensors.find("conv_last.weight");
if (first == tensors.end() || last == tensors.end() ||
tensors.count("body.0.rdb1.conv1.weight") == 0 ||
first->second.n_dims != 4 || last->second.n_dims != 4 ||
first->second.ne[2] != 3 || last->second.ne[3] != 3) {
return 0;
}
return ESRGANConfig::detect_from_weights(tensors).scale;
} catch (const std::exception&) {
return 0;
}
}
void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) { void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) {
if (upscaler_ctx->upscaler != nullptr) { if (upscaler_ctx->upscaler != nullptr) {
delete upscaler_ctx->upscaler; delete upscaler_ctx->upscaler;