diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 06485ca16c..5838cf3b03 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -564,6 +564,18 @@ jobs: if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }} uses: pr-mpt/actions-commit-hash@01d19a83c242e1851c9aa6cf9625092ecd095d09 # v2 + - name: Bundle HIP runtime libraries + if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }} + run: | + $ErrorActionPreference = "Stop" + $binPath = (rocm-sdk path --bin).Trim() + if (-not $binPath) { throw "rocm-sdk path --bin returned empty" } + foreach ($dll in @("amdhip64_7.dll", "rocm_kpack.dll", "amd_comgr.dll")) { + $f = Get-ChildItem -Path $binPath -Filter $dll -ErrorAction SilentlyContinue + if (-not $f) { throw "no match for $dll in $binPath" } + Copy-Item $f.FullName -Destination .\build\bin\ -Force + } + - name: Pack artifacts if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }} run: | diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d9975cf60d..776626d7e8 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -12,6 +12,10 @@ If you want to update a third-party dependency, please open an issue first inste ## 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. Do not include test code or test scripts in commits or PRs. Keep them local and report verification results in the PR description. diff --git a/README.md b/README.md index f184c51b0f..832ee339f8 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,7 @@ API and command-line option may change frequently.*** - [PiD](./docs/pid.md) - [LongCat Image](./docs/longcat_image.md) - [Z-Image](./docs/z_image.md) + - [Z-Image L2P](./docs/z_image_l2p.md) - [MiniT2I](./docs/minit2i.md) - [SenseNova U1.5](./docs/sensenova_u1.md) - [Ovis-Image](./docs/ovis_image.md) @@ -64,6 +65,9 @@ API and command-line option may change frequently.*** - [HiDream-O1-Image](./docs/hidream_o1_image.md) - [Ideogram4](./docs/ideogram4.md) - [LLaDA-Image](./docs/llada_image.md) + - [Ming-Image Design](./docs/ming_image.md) + - [PixArt](./docs/pixart.md) + - [Iris-3B](./docs/iris.md) - [Image Edit Models](./docs/edit.md) - [FLUX.1-Kontext-dev](./docs/kontext.md) - [Qwen Image Edit series](./docs/qwen_image_edit.md) diff --git a/assets/iris/example.png b/assets/iris/example.png new file mode 100644 index 0000000000..af3a4fd75b Binary files /dev/null and b/assets/iris/example.png differ diff --git a/assets/qwen/qwen-image-2.1-alpha-in1.png b/assets/qwen/qwen-image-2.1-alpha-in1.png new file mode 100644 index 0000000000..e65fcf0719 Binary files /dev/null and b/assets/qwen/qwen-image-2.1-alpha-in1.png differ diff --git a/assets/qwen/qwen-image-2.1-alpha-out1.png b/assets/qwen/qwen-image-2.1-alpha-out1.png new file mode 100644 index 0000000000..6841b7f787 Binary files /dev/null and b/assets/qwen/qwen-image-2.1-alpha-out1.png differ diff --git a/assets/qwen/qwen-image-2.1-alpha-out2.png b/assets/qwen/qwen-image-2.1-alpha-out2.png new file mode 100644 index 0000000000..8bef08a213 Binary files /dev/null and b/assets/qwen/qwen-image-2.1-alpha-out2.png differ diff --git a/assets/wan/Wan2.2_A14B_vace_r2v.mp4 b/assets/wan/Wan2.2_A14B_vace_r2v.mp4 new file mode 100644 index 0000000000..caf19c2151 Binary files /dev/null and b/assets/wan/Wan2.2_A14B_vace_r2v.mp4 differ diff --git a/docker/Dockerfile b/docker/Dockerfile index 0ef3d56acc..3a80ae9bf8 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -40,4 +40,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ chmod +x /sd-cli /sd-server +WORKDIR /sd.cpp + ENTRYPOINT [ "/sd-cli" ] diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 3e51bda60c..625dc60e70 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -57,4 +57,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ chmod +x /sd-cli /sd-server +WORKDIR /sd.cpp + ENTRYPOINT [ "/sd-cli" ] diff --git a/docker/Dockerfile.musa b/docker/Dockerfile.musa index 1c1d4d28c3..7dda431aaf 100644 --- a/docker/Dockerfile.musa +++ b/docker/Dockerfile.musa @@ -42,4 +42,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ chmod +x /sd-cli /sd-server +WORKDIR /sd.cpp + ENTRYPOINT [ "/sd-cli" ] diff --git a/docker/Dockerfile.sycl b/docker/Dockerfile.sycl index 9927ae7e58..9225e6eb2f 100644 --- a/docker/Dockerfile.sycl +++ b/docker/Dockerfile.sycl @@ -29,4 +29,6 @@ FROM intel/oneapi-basekit:${SYCL_VERSION}-devel-ubuntu24.04 AS runtime COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli COPY --from=build /sd.cpp/build/bin/sd-server /sd-server +WORKDIR /sd.cpp + ENTRYPOINT [ "/sd-cli" ] \ No newline at end of file diff --git a/docker/Dockerfile.vulkan b/docker/Dockerfile.vulkan index 2c40c21a77..c7810a1fa5 100644 --- a/docker/Dockerfile.vulkan +++ b/docker/Dockerfile.vulkan @@ -41,4 +41,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ chmod +x /sd-cli /sd-server +WORKDIR /sd.cpp + ENTRYPOINT [ "/sd-cli" ] diff --git a/docs/backend.md b/docs/backend.md index 846aafb858..70eef34f8f 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -156,8 +156,14 @@ the runner's graph-cut capacity checks. Runtime capacity checks also leave 512 MiB of currently free device memory for backend scratch buffers and pipelines, including with explicit backend assignments. -They cap stale free-memory reports by the device's total memory minus tracked -resident allocations and reject reports that exceed the device's total memory. +They cap free-memory reports by the device's total memory minus tracked +resident allocations. Vulkan reports exceeding total memory are rejected because +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 +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 used diffusion weights have priority. Each component's weights use the first diff --git a/docs/caching.md b/docs/caching.md index 01f0197443..f69cfbeca8 100644 --- a/docs/caching.md +++ b/docs/caching.md @@ -2,6 +2,16 @@ 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 | Mode | Target | Description | diff --git a/docs/edit.md b/docs/edit.md index 9f8a9c3f3b..0bd3b56200 100644 --- a/docs/edit.md +++ b/docs/edit.md @@ -26,6 +26,9 @@ Stable-diffusion.spp also supports basic Unet-based editing models like instruct ## 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). 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. diff --git a/docs/esrgan.md b/docs/esrgan.md index 39a9760502..708f49bad7 100644 --- a/docs/esrgan.md +++ b/docs/esrgan.md @@ -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. +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: ```bash diff --git a/docs/image_preprocessing.md b/docs/image_preprocessing.md new file mode 100644 index 0000000000..719781750e --- /dev/null +++ b/docs/image_preprocessing.md @@ -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 +`{...}` 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(¶ms); +/* 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, ¶ms, &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. diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index 3b65113d2d..dcbda83763 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -84,17 +84,18 @@ The floating-point output is reconstructed as Y[r, o] ~= A[r, o] * s_x[r] * s_w[o] + b[o] ``` -The packed runtime activation tensor contains the I8 activation rows and their floating-point row scales. Linear layers that share the same input and convrot group size reuse this packed tensor, avoiding repeated rotation and activation quantization within the graph. +The packed runtime activation tensor contains the I8 activation rows and their floating-point row scales. Linear layers with an input scale of `1` that share the same input and convrot group size reuse this packed tensor, avoiding repeated rotation and activation quantization within the graph. For other input scales, activations are scaled before explicit convrot quantization, and the linear output is unscaled before adding bias. ## Backend support - CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations. - NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. It uses cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition. -- Vulkan and other GPU backends do not currently have dedicated INT8 convrot kernels. They use the backend scheduler to fall back to CPU, which is expected to be substantially slower than the CUDA path. +- AMD HIP devices in the CDNA, RDNA3 (including RDNA3.5), and RDNA4 families use the same INT8/H256 kernels with hipBLAS for I8 x I8 to I32 GEMM. Other AMD architectures fall back to CPU for INT8 matrix multiplication. +- Vulkan provides native H256 activation quantization and INT8 matrix multiplication when the build and device support accelerated packed INT8 dot products. Unsupported configurations and GPU backends without these kernels use the backend scheduler to fall back to CPU. LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result. -The dedicated CUDA convrot activation path currently requires a group size of `256`; other supported group sizes use CPU execution. +The dedicated CUDA and HIP convrot activation paths currently require a group size of `256`; other supported group sizes use CPU execution. ## Example diff --git a/docs/iris.md b/docs/iris.md new file mode 100644 index 0000000000..940fd1eaae --- /dev/null +++ b/docs/iris.md @@ -0,0 +1,24 @@ +# How to Use + +Iris-3B generates images directly in pixel space and uses Qwen3-VL-4B-Instruct as the text encoder. No VAE is required. + +## Download weights + +- Download Iris-3B + - safetensors: https://huggingface.co/speridlabs/iris-3b/tree/main (`model.safetensors` in the root directory) +- Download Qwen3-VL-4B-Instruct + - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders + - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main + +## Examples + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\iris-3b.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf -p "a lovely cat" --cfg-scale 3 -H 1024 -W 1024 --diffusion-fa +``` + +iris-3b example + +## Notes + +- Width and height must be multiples of 16. Do not pass `--vae`. +- Captions are limited to 300 tokens including the assistant-turn suffix. Positive and negative prompts use the same template. diff --git a/docs/mage_flow.md b/docs/mage_flow.md index 1d8c9a4cd7..6cd5e7d64a 100644 --- a/docs/mage_flow.md +++ b/docs/mage_flow.md @@ -4,20 +4,8 @@ ## Download weights -- Download Mage-Flow - - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/transformer -- Download Mage-Flow-Base - - safetensors: https://huggingface.co/microsoft/Mage-Flow-Base/tree/main/transformer -- Download Mage-Flow-Turbo - - safetensors: https://huggingface.co/microsoft/Mage-Flow-Turbo/tree/main/transformer -- Download Mage-Flow-Edit - - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit/tree/main/transformer -- Download Mage-Flow-Edit-Turbo - - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo/tree/main/transformer -- Download Mage-Flow-Edit-Base - - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Base/tree/main/transformer -- Download Mage-Flow vae - - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/vae +- Download Mage-Flow diffusion from https://huggingface.co/Comfy-Org/Mage-Flow/tree/main/diffusion_models +- Download Mage-Flow vae from https://huggingface.co/Comfy-Org/Mage-Flow/tree/main/vae - Download Qwen3-VL 4B - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main diff --git a/docs/ming_image.md b/docs/ming_image.md new file mode 100644 index 0000000000..6077548c84 --- /dev/null +++ b/docs/ming_image.md @@ -0,0 +1,24 @@ +# Ming-Image + +[Ming-Image](https://github.com/inclusionAI/Ming-Image) 0.1 Design uses a 6B diffusion transformer (DiT), Ling-mini-2.0 for text conditioning, and the Ming-Image VAE. Text-to-image generation with RGBA output is supported. + +## Download weights + +- Download Ming-Image 0.1 Design DiT + - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/diffusion_models +- Download Ling-mini-2.0 BF16 + - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/text_encoders +- Download Ming-Image VAE + - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/vae +- Download Ling tokenizer + - tokenizer.json: https://huggingface.co/inclusionAI/Ming-Image-0.1-Design/blob/main/mllm/tokenizer.json + +The example below uses `ming_image_0.1_design_bf16.safetensors` for the DiT. You can also use `ming_image_0.1_design_int8_convrot.safetensors` with [INT8 convrot support](int8_convrot.md). Use the BF16 text encoder. + +## Text-to-image + +Pass the Ling `tokenizer.json` with `--tokenizer` and use the matching Ming-Image VAE. Image dimensions must be multiples of 16. Save the output as PNG to preserve the alpha channel. + +```bash +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\ming_image_0.1_design_bf16.safetensors --llm ..\models\text_encoders\ming_image_0.1_ling_mini_2.0_bf16.safetensors --vae ..\models\vae\ming_image_vae_bf16.safetensors --tokenizer ..\models\text_encoders\tokenizer.json -p "A cheerful orange cat sticker, transparent background" --width 1024 --height 1024 --steps 12 --cfg-scale 1 --sampling-method euler --diffusion-fa -v --offload-to-cpu -o ming_image.png +``` diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index d50b9f8b18..b8c5334d50 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -80,9 +80,10 @@ and audio differ slightly from the previous kernels. `GGML_CUDA_CUDNN_ATTN=0` re `GGML_CUDA_CUDNN_SAGE=0` / `=1` does the same for the sage op only, and `GGML_CUDA_CUDNN_ATTN_BF16=1` runs cuDNN in BF16 (less accurate, same speed). -The video VAE decodes one 16x16 latent tile per decoder graph by default. `SD_H3_VAE_TILE=N` -uses N x N latent tiles instead (20 decodes about 0.7 s faster at 960x544x124 on B200); the tile -seams move, so the frames differ from the default (36 dB PSNR at 20) and it is opt-in. +The video VAE decodes one 16x16 latent tile (256x256 pixels) per decoder graph by default. +`SD_H3_VAE_TILE=N` uses N x N latent tiles instead, overriding `--vae-tile-size` (20 decodes +about 0.7 s faster at 960x544x124 on B200); the tile seams move, so the frames differ from the +default (36 dB PSNR at 20) and it is opt-in. `SD_H3_VAE_TILE_BATCH=auto` puts several tiles into one graph, sized from free device memory (at most 4 unless `SD_H3_VAE_TILE_BATCH_MAX` raises it), and `SD_H3_VAE_TILE_BATCH=N` forces N; the batched projections can round differently from the per-tile decode on some GPUs, so it is diff --git a/docs/performance.md b/docs/performance.md index 2166720200..028a6409a5 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -21,6 +21,60 @@ CPU fallback. It excludes weights and cache buffers. Within a runner lifecycle, the summary is printed only on the first graph or when backend capacities or the segment count change. +## Use VAE tiling to reduce encode and decode memory usage. + +`--vae-tiling` enables spatial tiling for both VAE encoding and decoding. The +default tile size is 256x256 **image pixels**, independent of the VAE scale factor: + +```shell +--vae-tiling --vae-tile-size 256x256 --vae-tile-overlap 0.5 +``` + +`--vae-tile-size` accepts one size or `WIDTHxHEIGHT`. A zero dimension uses the +256-pixel default. Sizes are rounded down to a multiple of the VAE scale factor +and capped at the current input dimensions. Explicit sizes below four latent +pixels per axis (or the full axis when it is smaller) are rejected. Encoding and +decoding use the same spatial sizes, without an additional encoding multiplier. +Inputs that fit within a tile are processed as one tile. + +For a 512x512 image with the default 50% overlap, both encoding and decoding use +3x3 tiles. A 256-pixel tile corresponds to 32 latent pixels for an 8x VAE, 16 for +a 16x VAE, and 8 for a 32x VAE. Smaller tiles reduce each graph's memory demand, +but overlapping work can increase processing time and tiling can affect image +quality, especially during encoding. Use larger tiles when more context is needed. + +`--vae-relative-tile-size` overrides the absolute size on each axis with a positive +value. Values up to and including 1 specify a fraction of the current input size; +values greater than 1 specify a target number of tiles per axis, accounting for +overlap. For example, `0.5x0.5` uses half the width and height in both encode and +decode. The target overlap is clamped to 0 through 0.5 and the actual overlap is +adjusted to fit the image. On an axis that does not wrap, a tile size that would leave +neighboring tiles overlapping by less than half the target (or by fewer than 2 latent +pixels) is adjusted to the nearest size that does not: smaller first, larger only up to +twice the requested size. Without it, a tile a little over half of the axis gives two +tiles that barely overlap and the tile edge shows as a line. Size and overlap options +require `--vae-tiling`. + +**Migration:** `--vae-tile-size` and the C/JSON fields `tile_size_w` and +`tile_size_h` now use image pixels instead of latent units. The C/JSON fields +`tile_size_x/y` have been renamed to `tile_size_w/h`, and `rel_size_x/y` to +`rel_size_w/h`. The command-line option names are unchanged. For example, an old +decode tile size of 32 corresponds to 256 pixels for an 8x VAE or 512 pixels for a +16x VAE. Encoding no longer enlarges explicit or relative tile sizes. + +The main VAE decode and encode paths retry allocation failures with smaller tiles, +even without `--vae-tiling`. Supported video VAEs first try temporal tiling; spatial +retries use at most 256-pixel tiles initially and then halve the effective tile +dimensions down to the minimum size. Each spatial retry must reduce the effective +tile size. These runtime adjustments do not change the caller's parameters. +Execution failures are not retried. Encoding retries the same way, without the +temporal step. + +`--temporal-tiling` remains independent of spatial tiling. MiniMax H3 always uses +spatial tiling (256x256 pixels and 25% overlap by default) and its own temporal +windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap`; explicit +spatial sizes are honored. + ## Offload weights to the CPU to save VRAM without reducing generation speed. Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM without reducing generation speed. @@ -63,6 +117,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 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 ` 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-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments. diff --git a/docs/pixart.md b/docs/pixart.md new file mode 100644 index 0000000000..f9386abf16 --- /dev/null +++ b/docs/pixart.md @@ -0,0 +1,46 @@ +# How to Use + +You can run PixArt-α / PixArt-Σ with stable-diffusion.cpp. + +PixArt is a DiT-based text-to-image model family conditioned by a T5-XXL text +encoder and a 4-channel VAE: SDXL-style for PixArt-Σ and SD1.x-style for PixArt-α. + +## Download weights + +- Download the transformer (diffusion model) + - PixArt-Σ XL-2 1024-MS: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/transformer + - PixArt-α XL-2 1024-MS: https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS/tree/main/transformer +- Download the T5-XXL text encoder + - safetensors: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/text_encoder +- Download the VAE + - PixArt-Σ: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/vae + - PixArt-α: https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS/tree/main/vae + - Use the VAE matching the checkpoint's latent space. For TAE decoding or + preview, use TAESDXL for PixArt-Σ and TAESD for PixArt-α. +- Tokenizer: the T5 vocabulary is embedded; no extra tokenizer file is needed. + +## Examples + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\pixart_sigma_xl2_1024_ms.safetensors --t5xxl ..\models\text_encoders\t5xxl.safetensors --vae ..\models\vae\pixart_vae.safetensors -p "a lovely cat" --cfg-scale 4.5 -W 1024 -H 1024 --steps 20 -v +``` + +## Notes + +- The VAE scaling factor defaults to `0.13025` for PixArt-Σ. PixArt-α + checkpoints with resolution micro-condition weights use `0.18215`. + PixArt-α 512 has the same tensor layout as PixArt-Σ, so it requires an + explicit override: `--model-args "pixart_vae_scale_factor=0.18215"`. + This argument can also override the scale for other compatible checkpoints. +- PixArt-Σ checkpoints compute 2D sincos positional embeddings at runtime; + the trained grid is 64x64 patches with an interpolation scale of 2. + For checkpoints trained at a different resolution, the positional embedding + parameters can be adjusted via model args: + `--model-args "pixart_pos_embed_base_size=,pixart_interpolation_scale="` + (e.g. `pixart_pos_embed_base_size=32,pixart_interpolation_scale=1,pixart_vae_scale_factor=0.18215` for + PixArt-α XL-2 512). +- Checkpoints carrying resolution/aspect-ratio micro-condition weights are + detected but those conditions are not applied yet; a warning is logged and + generation proceeds with the timestep embedding only. +- The transformer predicts 8 channels (noise + learned variance); only the + noise half is used for sampling, matching the reference implementation. diff --git a/docs/qwen_image_2.1.md b/docs/qwen_image_2.1.md index 08a50875ec..56999bcc88 100644 --- a/docs/qwen_image_2.1.md +++ b/docs/qwen_image_2.1.md @@ -23,7 +23,7 @@ Run the following commands from the build directory. Use image dimensions divisi ### Text to image ```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 ``` Qwen Image 2.1 example @@ -35,7 +35,53 @@ 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`. ```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`. + +### 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 + +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 is an RGBA image with transparency. . The image has alpha channel and the background is transparent.` + +Since transparency is decided by the prompt rather than by the input or an explicit switch, the same format applies equally to editing, whether or not the reference image itself has an alpha channel. Note that alpha is kept only in `.png` and `.webp` outputs; saving as `.jpg` drops the transparency. + +Here are some examples ran with Q6_K quantization: +| Input | Prompt | Output | +| --- | --- | --- | +| ![Qwen Image 2.1 alpha input example 1](../assets/qwen/qwen-image-2.1-alpha-in1.png) | This is an RGBA image with transparency. Replace the text "BLOOM" with "Qwen Image 2.1", keeping the same font of the original text. The image has alpha channel and the background is transparent. | ![Qwen Image 2.1 alpha output example 1](../assets/qwen/qwen-image-2.1-alpha-out1.png) | +| ![Qwen Image 2.1 alpha input example 2](../assets/logo.png) | This is an RGBA image with transparency. Remove the background of the image, keeping only the text and cat. The image has alpha channel and the background is transparent. | ![Qwen Image 2.1 alpha output example 2](../assets/qwen/qwen-image-2.1-alpha-out2.png) | + +### Other features + +Other features of the model could be found on the [model card from QwenLM/Qwen-Image-2.1 repo](https://github.com/QwenLM/Qwen-Image-2.1), including 2 finetuned prompt rewriting Qwen3.5-9B model. diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 22007db2f2..86d4292581 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -1,5 +1,14 @@ # Troubleshooting +## Video model used in image generation mode + +If generation reports that a model cannot be run with `generate_image()`, add +`--mode vid_gen` to the CLI command. `--video-frames` alone does not select video +mode. Video models require this mode even when generating a single frame. +Library callers must use `generate_video()` for these models; use +`sd_ctx_supports_image_generation()` and `sd_ctx_supports_video_generation()` to +check the available generation modes. + ## Completely black or white images or videos / NaNs Some ggml backends can encounter numerical overflow during inference, producing diff --git a/docs/wan.md b/docs/wan.md index c7cf1dd9e8..1b38bcbc6a 100644 --- a/docs/wan.md +++ b/docs/wan.md @@ -1,5 +1,7 @@ # How to Use +Wan models require `-M vid_gen`, including single-frame generation. `--video-frames` alone does not select video mode. Library callers must use `generate_video()` instead of `generate_image()`. + ## Download weights - Download Wan @@ -38,6 +40,9 @@ - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models - gguf: https://huggingface.co/QuantStack/Wan2.2-S2V-14B-GGUF/tree/main - int8_convrot safetensors: https://huggingface.co/noctrex/Wan2.2-S2V-14B-int8_convrot + - Wan2.2 VACE-Fun A14B + - safetensors: https://huggingface.co/alibaba-pai/Wan2.2-VACE-Fun-A14B + - gguf: https://huggingface.co/QuantStack/Wan2.2-VACE-Fun-A14B-GGUF/tree/main - Download vae - wan_2.1_vae (for all the wan model except Wan2.2 TI2V 5B) - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors @@ -254,3 +259,25 @@ ffmpeg -i ..\..\ComfyUI\input\post+depth.mp4 -qscale:v 1 -vf fps=8 post+depth\fr ``` + +### Wan2.2 VACE-Fun A14B + +VACE-Fun runs as a MoE pair: `--diffusion-model` takes the low-noise expert and +`--high-noise-diffusion-model` the high-noise one. Reference-to-video uses `-i` +for the reference image, same as Wan2.1 VACE. + +#### R2V + +``` +.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\Wan2.2-VACE-Fun-A14B-low-noise-Q8_0.gguf --high-noise-diffusion-model ..\models\diffusion_models\Wan2.2-VACE-Fun-A14B-high-noise-Q8_0.gguf --vae ..\models\vae\wan_2.1_vae.safetensors --t5xxl ..\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf -p "a lovely cat" --cfg-scale 3.5 --sampling-method euler --steps 10 --high-noise-cfg-scale 3.5 --high-noise-sampling-method euler --high-noise-steps 8 -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa -i ..\assets\cat_with_sd_cpp_42.png --video-frames 33 --offload-to-cpu +``` + + + +#### T2V + +Same command without `-i` (VACE context is synthesized from an empty control +video, like Wan2.1 VACE t2v). + +> On GPUs with ~12 GB VRAM, VACE also needs `--vae-tiling` — the control-video +> encode can exceed the budget otherwise. diff --git a/docs/z_image_l2p.md b/docs/z_image_l2p.md new file mode 100644 index 0000000000..6c56b5e4e6 --- /dev/null +++ b/docs/z_image_l2p.md @@ -0,0 +1,21 @@ +# Z-Image L2P + +[L2P](https://nju-pcalab.github.io/projects/L2P/) (Latent-to-Pixel) transfers Z-Image-Turbo to pixel space. The VAE is replaced by 16x16 patch tokenization on RGB pixels and a small convolutional local decoder that turns the last DiT hidden states back into pixels, so no VAE is needed. The text encoder is the Qwen3-4B used by Z-Image-Turbo. + +## Download weights + +- Download L2P (1K) + - safetensors: https://huggingface.co/zhen-nan/L2P/tree/main +- Download Qwen3 4b + - safetensors: https://huggingface.co/Comfy-Org/z_image_turbo/tree/main/split_files/text_encoders + - gguf: https://huggingface.co/unsloth/Qwen3-4B-Instruct-2507-GGUF/tree/main + +## Text-to-image + +Do not pass `--vae`. Image dimensions must be multiples of 16. The reference pipeline uses 30 steps with a CFG scale of 2.0. + +The released checkpoint keeps part of the transformer in F32. Pass `--type bf16` to load it in BF16, as the reference pipeline runs it; this reduces the diffusion model from about 18.6 GB to 11.8 GB. + +```bash +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\model-1k-merge.safetensors --llm ..\models\text_encoders\qwen_3_4b.safetensors -p "an origami pig on fire in the middle of a dark room with a pentagram on the floor" --cfg-scale 2.0 --steps 30 --sampling-method euler -W 1024 -H 1024 --diffusion-fa --type bf16 -v -o z_image_l2p.png +``` diff --git a/examples/cli/README.md b/examples/cli/README.md index fe47122b63..63df3e91bd 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -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 [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: ```bash diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index eb96d9b5fe..21c1ce71da 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -41,7 +41,6 @@ struct SDCliParams { std::string metadata_format = "text"; sd_log_level_t log_level = SD_LOG_INFO; - bool canny_preprocess = false; bool convert_name = false; preview_t preview_method = PREVIEW_NONE; @@ -107,10 +106,6 @@ struct SDCliParams { }; options.bool_options = { - {"", - "--canny", - "apply canny preprocessor (edge detection)", - true, &canny_preprocess}, {"", "--convert-name", "convert tensor name (for convert mode)", @@ -268,7 +263,6 @@ struct SDCliParams { << " metadata_format: \"" << metadata_format << "\",\n" << " log_level: " << log_level_name(log_level) << ",\n" << " color: " << (color ? "true" : "false") << ",\n" - << " canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n" << " convert_name: " << (convert_name ? "true" : "false") << ",\n" << " preview_method: " << previews_str[preview_method] << ",\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, std::vector& images, - int expected_width = 0, - int expected_height = 0, - int max_image_num = 0) { + int max_image_num = 0) { if (!fs::exists(dir) || !fs::is_directory(dir)) { LOG_ERROR("'%s' is not a valid directory\n", dir.c_str()); return false; @@ -357,7 +349,8 @@ bool load_images_from_dir(const std::string dir, LOG_VERBOSE("load image %zu from '%s'", images.size(), path.c_str()); int width = 0; int height = 0; - uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height); + int loaded_channel = 0; + uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, 0, 0); if (image_buffer == nullptr) { LOG_ERROR("load image from '%s' failed", path.c_str()); return false; @@ -365,7 +358,7 @@ bool load_images_from_dir(const std::string dir, images.emplace_back(sd_image_t{(uint32_t)width, (uint32_t)height, - 3, + (uint32_t)loaded_channel, image_buffer}); if (max_image_num > 0 && static_cast(images.size()) >= max_image_num) { @@ -651,10 +644,11 @@ int main(int argc, const char* argv[]) { SDCliParams cli_params; SDContextParams ctx_params; + ctx_params.conditioning_cache_size = 0; SDGenerationParams gen_params; - parse_args(argc, argv, cli_params, ctx_params, gen_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) { MetadataReadOptions options; @@ -722,6 +716,15 @@ int main(int argc, const char* argv[]) { } if (cli_params.mode == CONVERT) { + std::vector loras; + loras.reserve(gen_params.lora_map.size() + gen_params.high_noise_lora_map.size()); + for (const auto& [path, multiplier] : gen_params.lora_map) { + loras.push_back({false, multiplier, path.c_str()}); + } + for (const auto& [path, multiplier] : gen_params.high_noise_lora_map) { + loras.push_back({true, multiplier, path.c_str()}); + } + bool success = convert_with_components(ctx_params.model_path.c_str(), ctx_params.clip_l_path.c_str(), ctx_params.clip_g_path.c_str(), @@ -732,7 +735,9 @@ int main(int argc, const char* argv[]) { ctx_params.wtype, ctx_params.tensor_type_rules.c_str(), cli_params.convert_name, - ctx_params.n_threads); + ctx_params.n_threads, + loras.empty() ? nullptr : loras.data(), + (int)loras.size()); if (!success) { LOG_ERROR("convert '%s'/'%s' to '%s' failed", ctx_params.model_path.c_str(), @@ -750,16 +755,8 @@ int main(int argc, const char* argv[]) { auto load_image_and_update_size = [&](const std::string& path, SDImageOwner& image, - bool resize_image = true, int expected_channel = 3) -> bool { - int expected_width = 0; - 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)) { + if (!load_sd_image_from_file(image.put(), path.c_str(), 0, 0, expected_channel)) { LOG_ERROR("load image from '%s' failed", path.c_str()); return false; } @@ -781,7 +778,8 @@ int main(int argc, const char* argv[]) { }; if (gen_params.init_image_path.size() > 0) { - if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) { + 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, native_init ? 0 : 3)) { return 1; } } @@ -795,8 +793,8 @@ int main(int argc, const char* argv[]) { if (gen_params.ref_image_paths.size() > 0) { gen_params.ref_images.clear(); for (auto& path : gen_params.ref_image_paths) { - SDImageOwner ref_image({0, 0, 3, nullptr}); - if (!load_image_and_update_size(path, ref_image, false)) { + SDImageOwner ref_image({0, 0, 0, nullptr}); + if (!load_image_and_update_size(path, ref_image, 0)) { return 1; } gen_params.ref_images.push_back(std::move(ref_image)); @@ -837,41 +835,22 @@ int main(int argc, const char* argv[]) { if (gen_params.mask_image_path.size() > 0) { if (!load_sd_image_from_file(gen_params.mask_image.put(), gen_params.mask_image_path.c_str(), - gen_params.get_resolved_width(), - gen_params.get_resolved_height(), + 0, + 0, 1)) { LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str()); 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 (!load_sd_image_from_file(gen_params.control_image.put(), gen_params.control_image_path.c_str(), - gen_params.get_resolved_width(), - gen_params.get_resolved_height())) { + 0, + 0)) { LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str()); 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) { @@ -888,8 +867,6 @@ int main(int argc, const char* argv[]) { gen_params.control_frames.clear(); if (!load_images_from_dir(gen_params.control_video_path, gen_params.control_frames, - gen_params.get_resolved_width(), - gen_params.get_resolved_height(), gen_params.video_frames)) { return 1; } @@ -898,10 +875,7 @@ int main(int argc, const char* argv[]) { if (!gen_params.pm_id_images_dir.empty()) { gen_params.pm_id_images.clear(); if (!load_images_from_dir(gen_params.pm_id_images_dir, - gen_params.pm_id_images, - 0, - 0, - 0)) { + gen_params.pm_id_images)) { return 1; } } diff --git a/examples/common/common.cpp b/examples/common/common.cpp index afea3afa09..89b226f3c1 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -518,7 +518,8 @@ ArgOptions SDContextParams::get_options() { {"", "--model-args", "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)',', &model_args}, {"", @@ -571,6 +572,10 @@ ArgOptions SDContextParams::get_options() { "number of threads to use during computation (default: -1). " "If threads <= 0, then threads will be set to the number of CPU physical cores", &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 = { @@ -822,6 +827,10 @@ bool SDContextParams::resolve(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) { const bool has_convert_input = model_path.length() != 0 || clip_l_path.length() != 0 || @@ -898,6 +907,7 @@ std::string SDContextParams::to_string() const { std::ostringstream oss; oss << "SDContextParams {\n" << " n_threads: " << n_threads << ",\n" + << " conditioning_cache_size: " << conditioning_cache_size << ",\n" << " model_path: \"" << model_path << "\",\n" << " clip_l_path: \"" << clip_l_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.tensor_type_rules = tensor_type_rules.c_str(); sd_ctx_params.n_threads = n_threads; + sd_ctx_params.conditioning_cache_size = conditioning_cache_size; sd_ctx_params.wtype = wtype; sd_ctx_params.rng_type = 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)", (int)',', &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 = { @@ -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).", true, &increase_ref_index}, - {"", - "--disable-auto-resize-ref-image", - "disable auto resize of ref images", - false, - &auto_resize_ref_image}, {"", "--circular", "enable circular padding on both axes for tileable output", @@ -1332,7 +1341,7 @@ ArgOptions SDGenerationParams::get_options() { &embed_image_metadata}, {"", "--vae-tiling", - "process vae in tiles to reduce memory usage", + "process vae encode and decode in spatial tiles to reduce memory usage (default: 256x256 image pixels)", true, &vae_tiling_params.enabled}, {"", @@ -1596,12 +1605,12 @@ ArgOptions SDGenerationParams::get_options() { size_t x_pos = tile_size_str.find('x'); try { if (x_pos != std::string::npos) { - std::string tile_x_str = tile_size_str.substr(0, x_pos); - std::string tile_y_str = tile_size_str.substr(x_pos + 1); - vae_tiling_params.tile_size_x = std::stoi(tile_x_str); - vae_tiling_params.tile_size_y = std::stoi(tile_y_str); + std::string tile_w_str = tile_size_str.substr(0, x_pos); + std::string tile_h_str = tile_size_str.substr(x_pos + 1); + vae_tiling_params.tile_size_w = std::stoi(tile_w_str); + vae_tiling_params.tile_size_h = std::stoi(tile_h_str); } else { - vae_tiling_params.tile_size_x = vae_tiling_params.tile_size_y = std::stoi(tile_size_str); + vae_tiling_params.tile_size_w = vae_tiling_params.tile_size_h = std::stoi(tile_size_str); } } catch (const std::invalid_argument&) { return -1; @@ -1619,12 +1628,12 @@ ArgOptions SDGenerationParams::get_options() { size_t x_pos = rel_size_str.find('x'); try { if (x_pos != std::string::npos) { - std::string rel_x_str = rel_size_str.substr(0, x_pos); - std::string rel_y_str = rel_size_str.substr(x_pos + 1); - vae_tiling_params.rel_size_x = std::stof(rel_x_str); - vae_tiling_params.rel_size_y = std::stof(rel_y_str); + std::string rel_w_str = rel_size_str.substr(0, x_pos); + std::string rel_h_str = rel_size_str.substr(x_pos + 1); + vae_tiling_params.rel_size_w = std::stof(rel_w_str); + vae_tiling_params.rel_size_h = std::stof(rel_h_str); } else { - vae_tiling_params.rel_size_x = vae_tiling_params.rel_size_y = std::stof(rel_size_str); + vae_tiling_params.rel_size_w = vae_tiling_params.rel_size_h = std::stof(rel_size_str); } } catch (const std::invalid_argument&) { return -1; @@ -1754,11 +1763,11 @@ ArgOptions SDGenerationParams::get_options() { on_scm_policy_arg}, {"", "--vae-tile-size", - "tile size for vae tiling, format [X]x[Y] (default: 32x32)", + "tile size for vae encode and decode in image pixels, format [W]x[H] or [S] (default: 256x256; requires --vae-tiling)", on_tile_size_arg}, {"", "--vae-relative-tile-size", - "relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)", + "relative tile size for vae encode and decode, format [W]x[H] or [S]: <=1 is a dimension fraction, >1 a target tile count (overrides --vae-tile-size; requires --vae-tiling)", on_relative_tile_size_arg}, {"", "--prompt-file", @@ -1848,28 +1857,28 @@ bool decode_base64_image(const std::string& encoded_input, return false; } - int decoded_width = 0; - int decoded_height = 0; - uint8_t* raw_data = load_image_from_memory(reinterpret_cast(image_bytes.data()), - static_cast(image_bytes.size()), - decoded_width, - decoded_height, - expected_width, - expected_height, - target_channels); + int decoded_width = 0; + int decoded_height = 0; + int resolved_channel = target_channels; + uint8_t* raw_data = load_image_from_memory(reinterpret_cast(image_bytes.data()), + static_cast(image_bytes.size()), + decoded_width, + decoded_height, + resolved_channel, + expected_width, + expected_height, + target_channels); if (raw_data == nullptr) { return false; } - out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)target_channels, raw_data}); + out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)resolved_channel, raw_data}); return true; } static bool parse_image_json_field(const json& parent, const char* key, int channels, - int expected_width, - int expected_height, SDImageOwner& out_image) { if (!parent.contains(key)) { return true; @@ -1881,14 +1890,12 @@ static bool parse_image_json_field(const json& parent, if (!parent.at(key).is_string()) { return false; } - return decode_base64_image(parent.at(key).get(), channels, expected_width, expected_height, out_image); + return decode_base64_image(parent.at(key).get(), channels, 0, 0, out_image); } static bool parse_image_array_json_field(const json& parent, const char* key, int channels, - int expected_width, - int expected_height, std::vector& out_images) { if (!parent.contains(key)) { return true; @@ -1907,7 +1914,7 @@ static bool parse_image_array_json_field(const json& parent, return false; } SDImageOwner image; - if (!decode_base64_image(item.get(), channels, expected_width, expected_height, image)) { + if (!decode_base64_image(item.get(), channels, 0, 0, image)) { return false; } out_images.push_back(std::move(image)); @@ -2006,6 +2013,29 @@ static bool resolve_model_file_from_dir(const std::string& model_name, 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(); + } 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(); + } + } 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( const std::string& json_str, const std::function& lora_path_resolver) { @@ -2017,6 +2047,9 @@ bool SDGenerationParams::from_json_str( 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) { if (j.contains(key)) { using T = std::decay_t; @@ -2054,6 +2087,7 @@ bool SDGenerationParams::from_json_str( load_if_exists("cache_mode", cache_mode); load_if_exists("cache_option", cache_option); 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("width", width); @@ -2071,7 +2105,6 @@ bool SDGenerationParams::from_json_str( load_if_exists("moe_boundary", moe_boundary); 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("embed_image_metadata", embed_image_metadata); @@ -2191,20 +2224,20 @@ bool SDGenerationParams::from_json_str( if (tiling_json.contains("temporal_tiling") && tiling_json["temporal_tiling"].is_boolean()) { vae_tiling_params.temporal_tiling = tiling_json["temporal_tiling"]; } - if (tiling_json.contains("tile_size_x") && tiling_json["tile_size_x"].is_number_integer()) { - vae_tiling_params.tile_size_x = tiling_json["tile_size_x"]; + if (tiling_json.contains("tile_size_w") && tiling_json["tile_size_w"].is_number_integer()) { + vae_tiling_params.tile_size_w = tiling_json["tile_size_w"]; } - if (tiling_json.contains("tile_size_y") && tiling_json["tile_size_y"].is_number_integer()) { - vae_tiling_params.tile_size_y = tiling_json["tile_size_y"]; + if (tiling_json.contains("tile_size_h") && tiling_json["tile_size_h"].is_number_integer()) { + vae_tiling_params.tile_size_h = tiling_json["tile_size_h"]; } if (tiling_json.contains("target_overlap") && tiling_json["target_overlap"].is_number()) { vae_tiling_params.target_overlap = tiling_json["target_overlap"]; } - if (tiling_json.contains("rel_size_x") && tiling_json["rel_size_x"].is_number()) { - vae_tiling_params.rel_size_x = tiling_json["rel_size_x"]; + if (tiling_json.contains("rel_size_w") && tiling_json["rel_size_w"].is_number()) { + vae_tiling_params.rel_size_w = tiling_json["rel_size_w"]; } - if (tiling_json.contains("rel_size_y") && tiling_json["rel_size_y"].is_number()) { - vae_tiling_params.rel_size_y = tiling_json["rel_size_y"]; + if (tiling_json.contains("rel_size_h") && tiling_json["rel_size_h"].is_number()) { + vae_tiling_params.rel_size_h = tiling_json["rel_size_h"]; } if (tiling_json.contains("extra_tiling_args") && tiling_json["extra_tiling_args"].is_string()) { extra_tiling_args = tiling_json["extra_tiling_args"].get(); @@ -2215,46 +2248,38 @@ bool SDGenerationParams::from_json_str( LOG_ERROR("invalid lora"); return false; } - if (!parse_image_json_field(j, "init_image", 3, width, height, init_image)) { - LOG_ERROR("invalid init_image"); - return false; - } - if (!parse_image_json_field(j, "end_image", 3, width, height, end_image)) { - LOG_ERROR("invalid end_image"); - return false; - } - if (!parse_image_array_json_field(j, "ref_images", 3, width, height, 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"); + auto load_image = [&](const char* key, int channels, SDImageOwner& 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; } - if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) { - LOG_ERROR("invalid ip_adapter_image"); + if (!parse_image_array_json_field(j, "ref_images", 0, ref_images) || + !parse_image_array_json_field(j, "control_frames", 3, control_frames)) { + LOG_ERROR("invalid input image array"); return false; } return true; } -void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_dir) { +bool SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_dir) { if (lora_model_dir.empty()) { - return; + return true; } static const std::regex re(R"(]+):([^>]+)>)"); static const std::vector valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"}; std::smatch m; + bool missing = false; std::string tmp = prompt; while (std::regex_search(tmp, m, re)) { @@ -2296,8 +2321,9 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d } if (!found) { LOG_WARN("can not found lora %s", final_path.lexically_normal().string().c_str()); - tmp = m.suffix().str(); - prompt = std::regex_replace(prompt, re, "", std::regex_constants::format_first_only); + tmp = m.suffix().str(); + prompt = std::regex_replace(prompt, re, "", std::regex_constants::format_first_only); + missing = true; continue; } } @@ -2313,6 +2339,8 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d tmp = m.suffix().str(); } + + return !missing; } bool SDGenerationParams::width_and_height_are_set() const { @@ -2434,7 +2462,7 @@ bool SDGenerationParams::initialize_cache_params() { return true; } -bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) { +bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict, bool validate_missing_loras) { vae_tiling_params.extra_tiling_args = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str(); if (high_noise_sample_params.sample_steps <= 0) { @@ -2478,12 +2506,18 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s prompt_with_lora = prompt; if (!lora_model_dir.empty()) { - extract_and_remove_lora(lora_model_dir); + if (!extract_and_remove_lora(lora_model_dir) && validate_missing_loras) { + return false; + } } return true; } 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) { LOG_ERROR("error: batch_count must be greater than 0"); return false; @@ -2602,7 +2636,7 @@ bool SDGenerationParams::resolve_and_validate(SDMode mode, const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) { - if (!resolve(lora_model_dir, hires_upscalers_dir, strict)) { + if (!resolve(lora_model_dir, hires_upscalers_dir, strict, mode == CONVERT)) { return false; } if (!validate(mode)) { @@ -2659,14 +2693,6 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { 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 (!ref_image_args.empty()) { ref_image_args += ","; @@ -2714,6 +2740,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { params.hires.custom_sigmas_count = static_cast(hires_custom_sigmas.size()); params.circular_x = circular || circular_x; params.circular_y = circular || circular_y; + params.image_preprocess = {image_preprocess.c_str()}; return params; } @@ -2816,6 +2843,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() { params.hires.custom_sigmas_count = static_cast(hires_custom_sigmas.size()); params.circular_x = circular || circular_x; params.circular_y = circular || circular_y; + params.image_preprocess = {image_preprocess.c_str()}; return params; } @@ -2872,7 +2900,8 @@ std::string SDGenerationParams::to_string() const { << " ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n" << " ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n" << " control_video_path: \"" << control_video_path << "\",\n" - << " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n" + << " image_preprocess: " << image_preprocess << ",\n" + << " ref_image_args: " << ref_image_args << ",\n" << " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n" << " pm_id_images_dir: \"" << pm_id_images_dir << "\",\n" << " pm_id_embed_path: \"" << pm_id_embed_path << "\",\n" @@ -2911,11 +2940,11 @@ std::string SDGenerationParams::to_string() const { << " vae_tiling_params: { " << vae_tiling_params.enabled << ", " << vae_tiling_params.temporal_tiling << ", " - << vae_tiling_params.tile_size_x << ", " - << vae_tiling_params.tile_size_y << ", " + << vae_tiling_params.tile_size_w << ", " + << vae_tiling_params.tile_size_h << ", " << vae_tiling_params.target_overlap << ", " - << vae_tiling_params.rel_size_x << ", " - << vae_tiling_params.rel_size_y << ", " + << vae_tiling_params.rel_size_w << ", " + << vae_tiling_params.rel_size_h << ", " << "\"" << extra_tiling_args << "\" },\n" << "}"; return oss.str(); @@ -3023,12 +3052,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); root["models"] = std::move(models); - root["clip_skip"] = gen_params.clip_skip; - root["strength"] = gen_params.strength; - root["control_strength"] = gen_params.control_strength; - root["ip_adapter_strength"] = gen_params.ip_adapter_strength; - root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image; - root["increase_ref_index"] = gen_params.increase_ref_index; + root["clip_skip"] = gen_params.clip_skip; + root["strength"] = gen_params.strength; + root["control_strength"] = gen_params.control_strength; + root["ip_adapter_strength"] = gen_params.ip_adapter_strength; + root["ref_image_args"] = gen_params.ref_image_args; + root["image_preprocess"] = gen_params.image_preprocess; + root["increase_ref_index"] = gen_params.increase_ref_index; if (mode == VID_GEN) { root["video"] = { {"frame_count", gen_params.video_frames}, @@ -3116,11 +3146,11 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, root["vae_tiling"] = { {"enabled", gen_params.vae_tiling_params.enabled}, {"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling}, - {"tile_size_x", gen_params.vae_tiling_params.tile_size_x}, - {"tile_size_y", gen_params.vae_tiling_params.tile_size_y}, + {"tile_size_w", gen_params.vae_tiling_params.tile_size_w}, + {"tile_size_h", gen_params.vae_tiling_params.tile_size_h}, {"target_overlap", gen_params.vae_tiling_params.target_overlap}, - {"rel_size_x", gen_params.vae_tiling_params.rel_size_x}, - {"rel_size_y", gen_params.vae_tiling_params.rel_size_y}, + {"rel_size_w", gen_params.vae_tiling_params.rel_size_w}, + {"rel_size_h", gen_params.vae_tiling_params.rel_size_h}, {"extra_tiling_args", gen_params.extra_tiling_args}, }; } diff --git a/examples/common/common.h b/examples/common/common.h index 0e5c50e3cb..8b60c76019 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -116,7 +116,8 @@ bool decode_base64_image(const std::string& encoded_input, SDImageOwner& out_image); struct SDContextParams { - int n_threads = -1; + int n_threads = -1; + int conditioning_cache_size = 4; std::string model_path; std::string clip_l_path; std::string clip_g_path; @@ -200,18 +201,17 @@ struct SDGenerationParams { std::string ad_prompt; std::string ad_negative_prompt; std::string extra_ad_args; - int clip_skip = -1; // <= 0 represents unspecified - int width = -1; - int height = -1; - int batch_count = 1; - int qwen_image_layers = 3; - int64_t seed = 42; - float strength = 0.75f; - float control_strength = 0.9f; - float ip_adapter_strength = 1.0f; - bool auto_resize_ref_image = true; - bool increase_ref_index = false; - bool embed_image_metadata = true; + int clip_skip = -1; // <= 0 represents unspecified + int width = -1; + int height = -1; + int batch_count = 1; + int qwen_image_layers = 3; + int64_t seed = 42; + float strength = 0.75f; + float control_strength = 0.9f; + float ip_adapter_strength = 1.0f; + bool increase_ref_index = false; + bool embed_image_metadata = true; std::string init_image_path; std::string end_image_path; @@ -247,6 +247,7 @@ struct SDGenerationParams { std::string extra_tiling_args; std::string ref_image_args; + std::string image_preprocess; std::string pm_id_images_dir; std::string pm_id_embed_path; @@ -310,13 +311,14 @@ struct SDGenerationParams { ArgOptions get_options(); bool from_json_str(const std::string& json_str, const std::function& lora_path_resolver = {}); + bool parse_image_preprocess_json(const std::string& json_str); bool initialize_cache_params(); - void extract_and_remove_lora(const std::string& lora_model_dir); + bool extract_and_remove_lora(const std::string& lora_model_dir); bool width_and_height_are_set() const; void set_width_and_height_if_unset(int w, int h); int get_resolved_width() const; int get_resolved_height() const; - bool resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict = false); + bool resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict = false, bool validate_missing_loras = false); bool validate(SDMode mode); bool resolve_and_validate(SDMode mode, const std::string& lora_model_dir, diff --git a/examples/common/log.cpp b/examples/common/log.cpp index c272c2376d..12d0683274 100644 --- a/examples/common/log.cpp +++ b/examples/common/log.cpp @@ -91,34 +91,34 @@ void log_print(enum sd_log_level_t level, const char* log, sd_log_level_t min_le switch (level) { case SD_LOG_DEBUG: tag_color = 37; - level_str = "DEBUG"; + level_str = "D"; break; case SD_LOG_VERBOSE: tag_color = 37; - level_str = "VERBOSE"; + level_str = "V"; break; case SD_LOG_INFO: tag_color = 34; - level_str = "INFO"; + level_str = "I"; break; case SD_LOG_WARN: tag_color = 35; - level_str = "WARN"; + level_str = "W"; break; case SD_LOG_ERROR: tag_color = 31; - level_str = "ERROR"; + level_str = "E"; break; default: tag_color = 33; - level_str = "?????"; + level_str = "?"; break; } if (color) { - fprintf(out_stream, "\033[%d;1m[%-7s]\033[0m ", tag_color, level_str); + fprintf(out_stream, "\033[%d;1m[%s]\033[0m ", tag_color, level_str); } else { - fprintf(out_stream, "[%-7s] ", level_str); + fprintf(out_stream, "[%s] ", level_str); } fflush(out_stream); print_utf8(out_stream, log); @@ -131,18 +131,15 @@ void example_log_printf(sd_log_level_t level, const char* file, int line, const va_list args; va_start(args, format); - static char log_buffer[LOG_BUFFER_SIZE + 1]; - int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", sd_basename(file).c_str(), line); + char log_buffer[LOG_BUFFER_SIZE + 1] = {}; + vsnprintf(log_buffer, sizeof(log_buffer), format, args); + va_end(args); - if (written >= 0 && written < static_cast(LOG_BUFFER_SIZE)) { - vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args); - } - size_t len = strlen(log_buffer); - if (len == 0 || log_buffer[len - 1] != '\n') { - strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len); + std::string message = log_buffer; + while (!message.empty() && (message.back() == '\n' || message.back() == '\r')) { + message.pop_back(); } + message += " --- " + sd_basename(file) + ":" + std::to_string(line) + '\n'; - log_print(level, log_buffer, log_level, log_color); - - va_end(args); + log_print(level, message.c_str(), log_level, log_color); } diff --git a/examples/common/media_io.cpp b/examples/common/media_io.cpp index 345b6f1523..ba0e8c5b29 100644 --- a/examples/common/media_io.cpp +++ b/examples/common/media_io.cpp @@ -261,6 +261,10 @@ uint8_t* decode_webp_image_to_buffer(const uint8_t* data, height = features.height; source_channel_count = features.has_alpha ? 4 : 3; + if (expected_channel == 0) { + expected_channel = source_channel_count; + } + const size_t pixel_count = static_cast(width) * static_cast(height); if (expected_channel == 1) { @@ -481,7 +485,8 @@ uint8_t* load_image_common(bool from_memory, int& height, int expected_width, int expected_height, - int expected_channel) { + int expected_channel, + int& out_channel) { const char* image_path; FreeUniquePtr image_buffer; int source_channel_count = 0; @@ -538,6 +543,32 @@ uint8_t* load_image_common(bool from_memory, LOG_ERROR("load image from '%s' failed", image_path); return nullptr; } + if (expected_channel == 0) { + expected_channel = source_channel_count == 2 ? 4 : (source_channel_count == 1 ? 3 : source_channel_count); + if (expected_channel != source_channel_count) { + FreeUniquePtr promoted((uint8_t*)malloc((size_t)width * height * expected_channel)); + if (promoted == nullptr) { + LOG_ERROR("error: allocate memory for channel promotion, image_path = %s", image_path); + return nullptr; + } + const size_t pixel_count = (size_t)width * (size_t)height; + for (size_t i = 0; i < pixel_count; ++i) { + if (source_channel_count == 1) { + promoted.get()[i * 3 + 0] = image_buffer.get()[i]; + promoted.get()[i * 3 + 1] = image_buffer.get()[i]; + promoted.get()[i * 3 + 2] = image_buffer.get()[i]; + } else { + promoted.get()[i * 4 + 0] = image_buffer.get()[i * 2]; + promoted.get()[i * 4 + 1] = image_buffer.get()[i * 2]; + promoted.get()[i * 4 + 2] = image_buffer.get()[i * 2]; + promoted.get()[i * 4 + 3] = image_buffer.get()[i * 2 + 1]; + } + } + image_buffer = std::move(promoted); + source_channel_count = expected_channel; + } + } + // stb reports the source channel count even when it converts the output. if (source_channel_count < expected_channel) { fprintf(stderr, "the number of channels for the input image must be >= %d," @@ -597,7 +628,7 @@ uint8_t* load_image_common(bool from_memory, } stbir_resize(image_buffer.get(), width, height, 0, resized_image_buffer.get(), expected_width, expected_height, 0, STBIR_TYPE_UINT8, - expected_channel, STBIR_ALPHA_CHANNEL_NONE, 0, + expected_channel, expected_channel == 4 ? 3 : STBIR_ALPHA_CHANNEL_NONE, 0, STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_FILTER_BOX, STBIR_FILTER_BOX, STBIR_COLORSPACE_SRGB, nullptr); @@ -605,6 +636,7 @@ uint8_t* load_image_common(bool from_memory, height = expected_height; image_buffer = std::move(resized_image_buffer); } + out_channel = expected_channel; return image_buffer.release(); } @@ -777,10 +809,11 @@ bool write_image_to_file(const std::string& path, uint8_t* load_image_from_file(const char* image_path, int& width, int& height, + int& out_channel, int expected_width, int expected_height, int expected_channel) { - return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel); + return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, out_channel); } bool load_sd_image_from_file(sd_image_t* image, @@ -790,13 +823,14 @@ bool load_sd_image_from_file(sd_image_t* image, int expected_channel) { int width; int height; - image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel); + int resolved_channel = expected_channel; + image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, resolved_channel); if (image->data == nullptr) { return false; } image->width = width; image->height = height; - image->channel = expected_channel; + image->channel = resolved_channel; return true; } @@ -804,10 +838,11 @@ uint8_t* load_image_from_memory(const char* image_bytes, int len, int& width, int& height, + int& out_channel, int expected_width, int expected_height, int expected_channel) { - return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel); + return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel, out_channel); } static void append_avi_metadata(std::vector& data, const std::string& parameters) { diff --git a/examples/common/media_io.h b/examples/common/media_io.h index 6adbe98758..b2b56eccde 100644 --- a/examples/common/media_io.h +++ b/examples/common/media_io.h @@ -32,9 +32,12 @@ bool write_image_to_file(const std::string& path, const std::string& parameters = "", int quality = 90); +// expected_channel == 0 preserves native channels (grayscale -> RGB, gray+alpha -> RGBA). +// out_channel receives the output channel count. uint8_t* load_image_from_file(const char* image_path, int& width, int& height, + int& out_channel, int expected_width = 0, int expected_height = 0, int expected_channel = 3); @@ -49,6 +52,7 @@ uint8_t* load_image_from_memory(const char* image_bytes, int len, int& width, int& height, + int& out_channel, int expected_width = 0, int expected_height = 0, int expected_channel = 3); diff --git a/examples/server/CMakeLists.txt b/examples/server/CMakeLists.txt index 5c462258bc..302a8d6099 100644 --- a/examples/server/CMakeLists.txt +++ b/examples/server/CMakeLists.txt @@ -21,21 +21,21 @@ if(SD_SERVER_BUILD_FRONTEND AND EXISTS "${FRONTEND_DIR}") set(HAVE_FRONTEND_BUILD ON) add_custom_target(${TARGET}_frontend_install - COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" install + COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" install WORKING_DIRECTORY "${FRONTEND_DIR}" COMMENT "Installing frontend dependencies" VERBATIM ) add_custom_target(${TARGET}_frontend_build - COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build + COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" run build WORKING_DIRECTORY "${FRONTEND_DIR}" COMMENT "Building frontend" VERBATIM ) add_custom_target(${TARGET}_frontend_header - COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build:header + COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" run build:header WORKING_DIRECTORY "${FRONTEND_DIR}" COMMENT "Generating gen_index_html.h" VERBATIM diff --git a/examples/server/api.md b/examples/server/api.md index 3ec9a1da8f..cff4e84559 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -37,6 +37,7 @@ Current generation-related endpoints include: - `POST /sdapi/v1/txt2img` - `POST /sdapi/v1/img2img` +- `GET /sdapi/v1/progress` - `GET /sdapi/v1/loras` - `GET /sdapi/v1/upscalers` - `GET /sdapi/v1/latent-upscale-modes` @@ -56,6 +57,7 @@ Current endpoints include: - `GET /sdcpp/v1/jobs/{id}` - `POST /sdcpp/v1/jobs/{id}/cancel` - `POST /sdcpp/v1/vid_gen` +- `POST /sdcpp/v1/upscale` ## `sd_cpp_extra_args` @@ -148,6 +150,19 @@ Native extension fields: - 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 {"image_preprocess":"target=ref,mode=none","ref_image_args":"resize_before_vae=false"} +``` + Response fields: | Field | Type | Notes | @@ -264,6 +279,26 @@ Response fields: | `parameters` | `object` | Echo of the parsed outer request body | | `info` | `string` | Currently empty string | +#### `GET /sdapi/v1/progress` + +Poll this endpoint while a synchronous SDAPI generation request is running. +An optional `id_task` query parameter selects the task with the same `id_task` +provided in the generation request (default: `"sdapi"`). Set +`skip_current_image=true` or `1` to omit the preview image. + +`current_image` contains the latest base64 JPEG latent projection, or `null` +when unavailable. `state.sampling_step` is the positive logical step and +`state.sampling_steps` is the actual step count for the current sampling pass, +including schedule and img2img strength adjustments. `state.job_no` is the +zero-based pass index; `state.job_count` includes one pass per batch image and +an additional pass per image when highres fix is enabled. `progress` weights +these passes equally. Sampling progress can reach `1` before final decoding +finishes; `eta_relative` is currently always `0`. + +When no SDAPI task is active, or `id_task` does not match, the endpoint returns +`progress=0`, `state.job=""`, zero job/step counts, and `current_image=null`. +Completed and failed requests do not retain an active preview. + #### Discovery / Compatibility Endpoints Currently exposed: @@ -396,6 +431,20 @@ Field types: | `queue_position` | `integer` | | `result` | `object \| null` | | `error` | `object \| null` | +| `preview` | `object \| null` | + +`preview` sub-fields: + +| Field | Type | Notes | +| --- | --- | --- | +| `preview.pass` | `integer` | Sampling pass number, starting at 1 within the job | +| `preview.step` | `integer` | Positive logical sample step within this pass | +| `preview.total_steps` | `integer` | Actual sample steps in this pass, after schedule and strength adjustments | +| `preview.b64_json` | `string` | Base64-encoded PNG of the current preview frame | + +The preview is updated in place as sampling progresses; poll `GET /sdcpp/v1/jobs/{id}` to retrieve the latest frame. Only the most recent preview is retained. + +`step / total_steps` describes the current sampling pass, not overall job completion. A new batch item, high/low-noise stage, or highres pass increments `pass` and restarts `step`. Samplers with multiple denoiser evaluations per logical step may update a preview more than once at the same step. The frame and its pass/step metadata are published together. ### Endpoints @@ -421,7 +470,9 @@ Top-level fields: | `samplers` | `array` | Available sampling methods | | `schedulers` | `array` | Available schedulers | | `loras` | `array` | Available LoRA entries | -| `upscalers` | `array` | Available model-backed highres upscalers | +| `upscalers` | `array` | Available highres upscalers, built-in and model-backed | +| `upscale` | `boolean` | Whether a compatible RGB ESRGAN model is available for `POST /sdcpp/v1/upscale` | +| `preview_modes` | `array` | Available preview modes, e.g. `["none", "proj", "tae", "vae"]` | | `limits` | `object` | Shared queue and size limits | `model` @@ -463,6 +514,8 @@ Shared nested fields: | Field | Type | Notes | | --- | --- | --- | | `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. @@ -476,6 +529,8 @@ Built-in entries include `None`, `Lanczos`, `Nearest`, `Latent`, `Latent (neares | `limits.max_height` | `integer` | | `limits.max_batch_count` | `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`: @@ -505,11 +560,11 @@ Shared default fields used by both `img_gen` and `vid_gen`: | `vae_tiling_params` | `object` | | `vae_tiling_params.enabled` | `boolean` | | `vae_tiling_params.temporal_tiling` | `boolean` | -| `vae_tiling_params.tile_size_x` | `integer` | -| `vae_tiling_params.tile_size_y` | `integer` | +| `vae_tiling_params.tile_size_w` | `integer` | +| `vae_tiling_params.tile_size_h` | `integer` | | `vae_tiling_params.target_overlap` | `number` | -| `vae_tiling_params.rel_size_x` | `number` | -| `vae_tiling_params.rel_size_y` | `number` | +| `vae_tiling_params.rel_size_w` | `number` | +| `vae_tiling_params.rel_size_h` | `number` | | `vae_tiling_params.extra_tiling_args` | `string` | | `cache_mode` | `string` | | `cache_option` | `string` | @@ -518,6 +573,8 @@ Shared default fields used by both `img_gen` and `vid_gen`: | `output_format` | `string` | | `output_compression` | `integer` | +`vae_tiling_params.tile_size_w` and `tile_size_h` are in **image pixels**, with `0` selecting the 256-pixel default. Both encode and decode use these sizes without an encoding multiplier. Positive `rel_size_w`/`rel_size_h` values override the corresponding absolute size: values up to 1 are dimension fractions, and values greater than 1 are target tile counts. Set `enabled` to use spatial tiling. Sizes are aligned down to the VAE scale factor and capped at the input dimensions; explicit sizes below the minimum supported tile size are rejected. These fields previously used latent units; see [VAE tiling](../../docs/performance.md#use-vae-tiling-to-reduce-encode-and-decode-memory-usage) for migration and OOM retry behavior. + `vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`). LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEHV use overlap blending. MiniMax H3 keeps its model-specific fixed temporal windows because its latent-to-frame mapping is non-linear. @@ -526,7 +583,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH | Field | Type | | --- | --- | | `batch_count` | `integer` | -| `auto_resize_ref_image` | `boolean` | +| `ref_image_args` | `string` | | `increase_ref_index` | `boolean` | | `control_strength` | `number` | | `ip_adapter_strength` | `number` | @@ -577,6 +634,7 @@ Fields returned in `features_by_mode.img_gen`: - `cache` - `cancel_queued` - `cancel_generating` +- `preview` Fields returned in `features_by_mode.vid_gen`: @@ -589,6 +647,7 @@ Fields returned in `features_by_mode.vid_gen`: - `cache` - `cancel_queued` - `cancel_generating` +- `preview` #### `POST /sdcpp/v1/img_gen` @@ -628,6 +687,52 @@ Typical status codes: - `404 Not Found` - `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` | 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` Attempts to cancel an accepted job. @@ -653,7 +758,7 @@ Example: "strength": 0.75, "seed": -1, "batch_count": 1, - "auto_resize_ref_image": true, + "ref_image_args": "", "increase_ref_index": false, "control_strength": 0.9, "ip_adapter_strength": 1.0, @@ -702,11 +807,11 @@ Example: "vae_tiling_params": { "enabled": false, "temporal_tiling": false, - "tile_size_x": 0, - "tile_size_y": 0, + "tile_size_w": 0, + "tile_size_h": 0, "target_overlap": 0.5, - "rel_size_x": 0.0, - "rel_size_y": 0.0, + "rel_size_w": 0.0, + "rel_size_h": 0.0, "extra_tiling_args": "" }, @@ -715,6 +820,9 @@ Example: "scm_mask": "", "scm_policy_dynamic": true, + "preview": "none", + "preview_interval": 1, + "output_format": "png", "output_compression": 100 } @@ -728,6 +836,17 @@ Example: ### 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: - a raw base64 string, or @@ -735,12 +854,15 @@ Any image field accepts: Channel expectations: -- `init_image`: 3 channels -- `ref_images[]`: 3 channels +- `init_image`: native channels (3 or 4); alpha is preserved and applied per model +- `ref_images[]`: native channels (3 or 4); alpha is preserved and applied per model - `control_image`: 3 channels - `ip_adapter_image`: 3 channels - `mask_image`: 1 channel +Models that support RGBA (e.g. Qwen-Image 2.1) use the alpha channel of `init_image` +and `ref_images[]`. RGB-only models drop it, so sending RGBA is safe for every model. + If omitted or null: - single-image fields map to an empty `sd_image_t` @@ -760,11 +882,14 @@ Top-level scalar fields: | `strength` | `number` | | `seed` | `integer` | | `batch_count` | `integer` | -| `auto_resize_ref_image` | `boolean` | +| `ref_image_args` | `string` | +| `image_preprocess` | `string \| array` | | `increase_ref_index` | `boolean` | | `control_strength` | `number` | | `ip_adapter_strength` | `number` | | `embed_image_metadata` | `boolean` | +| `preview` | `string` | +| `preview_interval` | `integer` | Image fields: @@ -820,11 +945,11 @@ Other native fields: | `vae_tiling_params` | `object` | | `vae_tiling_params.enabled` | `boolean` | | `vae_tiling_params.temporal_tiling` | `boolean` | -| `vae_tiling_params.tile_size_x` | `integer` | -| `vae_tiling_params.tile_size_y` | `integer` | +| `vae_tiling_params.tile_size_w` | `integer` | +| `vae_tiling_params.tile_size_h` | `integer` | | `vae_tiling_params.target_overlap` | `number` | -| `vae_tiling_params.rel_size_x` | `number` | -| `vae_tiling_params.rel_size_y` | `number` | +| `vae_tiling_params.rel_size_w` | `number` | +| `vae_tiling_params.rel_size_h` | `number` | | `vae_tiling_params.extra_tiling_args` | `string` | | `cache_mode` | `string` | | `cache_option` | `string` | @@ -852,6 +977,10 @@ When omitted, backend defaults apply to these fields: - `sample_params.flow_shift` - `sample_params.guidance.img_cfg` +### Preview Interval Semantics + +`preview_interval` controls the period (in sample steps) at which preview frames are generated. The default is `1` (every step). Any non-positive values will be clamped to `1`. Note that if `preview` is set to `"none"`, no previews are produced regardless of the interval. + ### Completion Result Example completed job: @@ -966,7 +1095,8 @@ Response fields: Compared with `img_gen`, the `vid_gen` request body: - `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema -- `ref_images`, `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema +- `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema +- `ref_images` is accepted for MiniMax-H3 Ref2VA conditioning; other video model families currently ignore it - `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength` Example: @@ -987,6 +1117,7 @@ Example: "init_image": null, "end_image": null, + "ref_images": [], "control_frames": [], "sample_params": { @@ -1035,11 +1166,11 @@ Example: "vae_tiling_params": { "enabled": false, "temporal_tiling": false, - "tile_size_x": 0, - "tile_size_y": 0, + "tile_size_w": 0, + "tile_size_h": 0, "target_overlap": 0.5, - "rel_size_x": 0.0, - "rel_size_y": 0.0, + "rel_size_w": 0.0, + "rel_size_h": 0.0, "extra_tiling_args": "" }, @@ -1048,11 +1179,25 @@ Example: "scm_mask": "", "scm_policy_dynamic": true, + "preview": "none", + "preview_interval": 1, + "output_format": "webm", "output_compression": 100 } ``` +### Reference Image Rules + +- `ref_images` contains reference images for MiniMax-H3 Ref2VA conditioning. +- Images retain request order and correspond to ``, ``, and so on in the prompt. +- For MiniMax-H3, non-empty `ref_images` cannot be combined with `init_image` or `end_image`. A conflicting API request fails during generation; the WebUI checks this before submission. +- MiniMax-H3 does not support `control_frames`; leave that array empty. +- Other video model families currently ignore `ref_images`. +- `features_by_mode.vid_gen` does not currently advertise `ref_images`; its absence is not an indication that MiniMax-H3 reference images are unsupported. + +See [MiniMax-H3 reference conditioning](../../docs/minimax_h3.md#reference-to-audio-video-conditioning) for model requirements and prompt examples. + ### LoRA Rules - The server only accepts explicit LoRA entries from the `lora` field. @@ -1070,6 +1215,7 @@ Channel expectations: - `init_image`: 3 channels - `end_image`: 3 channels +- `ref_images[]`: decoded with native channels, then converted to RGB by MiniMax-H3 - `control_frames[]`: 3 channels Frame ordering rules: @@ -1099,6 +1245,8 @@ Top-level scalar fields: | `fps` | `integer` | | `moe_boundary` | `number` | | `vace_strength` | `number` | +| `preview` | `string` | +| `preview_interval` | `integer` | Image and frame fields: @@ -1106,6 +1254,7 @@ Image and frame fields: | --- | --- | | `init_image` | `string \| null` | | `end_image` | `string \| null` | +| `ref_images` | `array` | | `control_frames` | `array` | LoRA fields: @@ -1160,11 +1309,11 @@ Other native fields: | `vae_tiling_params` | `object` | | `vae_tiling_params.enabled` | `boolean` | | `vae_tiling_params.temporal_tiling` | `boolean` | -| `vae_tiling_params.tile_size_x` | `integer` | -| `vae_tiling_params.tile_size_y` | `integer` | +| `vae_tiling_params.tile_size_w` | `integer` | +| `vae_tiling_params.tile_size_h` | `integer` | | `vae_tiling_params.target_overlap` | `number` | -| `vae_tiling_params.rel_size_x` | `number` | -| `vae_tiling_params.rel_size_y` | `number` | +| `vae_tiling_params.rel_size_w` | `number` | +| `vae_tiling_params.rel_size_h` | `number` | | `vae_tiling_params.extra_tiling_args` | `string` | | `cache_mode` | `string` | | `cache_option` | `string` | diff --git a/examples/server/async_jobs.cpp b/examples/server/async_jobs.cpp index 287955a848..88408088e8 100644 --- a/examples/server/async_jobs.cpp +++ b/examples/server/async_jobs.cpp @@ -10,6 +10,55 @@ #include "common/media_io.h" #include "common/resource_owners.hpp" +static void preview_callback(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data) { + (void)is_noisy; + if (frame_count <= 0 || frames == nullptr || frames[0].data == nullptr) { + return; + } + + AsyncGenerationJob* job = static_cast(data); + const sd_image_t& frame = frames[0]; + const sd_preview_info_t info = sd_get_preview_info(); + + auto image_bytes = encode_image_to_vector(EncodedImageFormat::PNG, + frame.data, + frame.width, + frame.height, + frame.channel, + "", + 80); + if (image_bytes.empty()) { + return; + } + + std::string b64 = base64_encode(image_bytes); + std::lock_guard lock(job->preview_mutex); + job->preview_b64 = std::move(b64); + job->preview_step = step < 0 ? -step : step; + job->preview_pass = info.sample_pass; + job->preview_total_steps = info.total_steps; +} + +static void clear_preview_callback() { + sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr); +} + +static void set_preview_callback(AsyncGenerationJob& job) { + preview_t mode = str_to_preview(job.kind == AsyncJobKind::ImgGen + ? job.img_gen.preview_mode.c_str() + : job.vid_gen.preview_mode.c_str()); + if (mode == PREVIEW_COUNT || mode == PREVIEW_NONE) { + return; + } + int interval = job.kind == AsyncJobKind::ImgGen + ? job.img_gen.preview_interval + : job.vid_gen.preview_interval; + if (interval <= 0) { + interval = 1; + } + sd_set_preview_callback(preview_callback, mode, interval, true, false, &job); +} + const char* async_job_kind_name(AsyncJobKind kind) { switch (kind) { case AsyncJobKind::ImgGen: @@ -161,6 +210,18 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo result["error"] = nullptr; } + if (job.status == AsyncJobStatus::Generating) { + std::lock_guard lock(job.preview_mutex); + if (!job.preview_b64.empty()) { + result["preview"] = { + {"pass", job.preview_pass}, + {"step", job.preview_step}, + {"total_steps", job.preview_total_steps}, + {"b64_json", job.preview_b64}, + }; + } + } + return result; } @@ -174,12 +235,14 @@ bool execute_img_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); + set_preview_callback(job); sd_image_t* raw_results = nullptr; int num_results = 0; if (!generate_image(runtime.sd_ctx, ¶ms, &raw_results, &num_results)) { raw_results = nullptr; num_results = 0; } + clear_preview_callback(); results.adopt(raw_results, num_results); } @@ -247,10 +310,12 @@ bool execute_vid_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); + set_preview_callback(job); sd_image_t* raw_results = nullptr; if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio, &output_fps)) { raw_results = nullptr; } + clear_preview_callback(); results.adopt(raw_results, num_results); } diff --git a/examples/server/async_jobs.h b/examples/server/async_jobs.h index 89997a3b4a..2fe79bb242 100644 --- a/examples/server/async_jobs.h +++ b/examples/server/async_jobs.h @@ -44,6 +44,11 @@ struct AsyncGenerationJob { int result_fps = 0; std::string error_code; std::string error_message; + mutable std::mutex preview_mutex; + std::string preview_b64; + int preview_pass = 0; + int preview_step = 0; + int preview_total_steps = 0; }; struct AsyncJobManager { diff --git a/examples/server/frontend b/examples/server/frontend index c4bce3d6b3..8e7c115059 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit c4bce3d6b3f236614cca21014f076083b7270ba8 +Subproject commit 8e7c11505925c0dbf9aaea56a3a5058b8690d66b diff --git a/examples/server/main.cpp b/examples/server/main.cpp index 1f93e0dc7f..1768f0de73 100644 --- a/examples/server/main.cpp +++ b/examples/server/main.cpp @@ -75,10 +75,12 @@ int main(int argc, const char** argv) { } SDSvrParams svr_params; SDContextParams ctx_params; + // Preserve the CLI default while disabling cwd scans for sd-server. + ctx_params.lora_model_dir.clear(); 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); + parse_args(argc, argv, svr_params, ctx_params, default_gen_params); LOG_VERBOSE("version: %s", version_string().c_str()); LOG_VERBOSE("%s", sd_get_system_info()); diff --git a/examples/server/routes_openai.cpp b/examples/server/routes_openai.cpp index 0f122922ce..af0cdfaa17 100644 --- a/examples/server/routes_openai.cpp +++ b/examples/server/routes_openai.cpp @@ -157,42 +157,46 @@ static bool build_openai_edit_request(const httplib::Request& req, request.gen_params.height = height; 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) { - int img_w = 0; - int img_h = 0; - uint8_t* raw_pixels = load_image_from_memory( - reinterpret_cast(bytes.data()), - static_cast(bytes.size()), - img_w, img_h, - width, height, 3); + int img_w = 0; + int img_h = 0; + int resolved_channel = 0; + uint8_t* raw_pixels = load_image_from_memory( + reinterpret_cast(bytes.data()), + static_cast(bytes.size()), + img_w, img_h, resolved_channel, + 0, 0, + 0); if (raw_pixels == nullptr) { continue; } - SDImageOwner image_owner({(uint32_t)img_w, (uint32_t)img_h, 3, raw_pixels}); + const bool is_first_ref_image = request.gen_params.ref_images.empty(); + SDImageOwner image_owner({(uint32_t)img_w, (uint32_t)img_h, (uint32_t)resolved_channel, raw_pixels}); request.gen_params.set_width_and_height_if_unset(image_owner.get().width, image_owner.get().height); - request.gen_params.ref_images.push_back(std::move(image_owner)); - } - if (!request.gen_params.ref_images.empty()) { - request.gen_params.init_image = request.gen_params.ref_images.front(); + if (is_first_ref_image) { + request.gen_params.init_image = image_owner; + if (request.gen_params.init_image.get().data == nullptr) { + error_message = "could not allocate init image"; + return false; + } + } + + request.gen_params.ref_images.push_back(std::move(image_owner)); } 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_h = 0; + int mask_w = 0; + int mask_h = 0; + int mask_channel = 0; uint8_t* mask_raw = load_image_from_memory( reinterpret_cast(mask_bytes.data()), static_cast(mask_bytes.size()), - mask_w, mask_h, - expected_width, expected_height, 1); + mask_w, mask_h, mask_channel, + 0, 0, 1); 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(); request.gen_params.set_width_and_height_if_unset(mask_image.width, mask_image.height); @@ -205,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)) { error_message = "invalid sd_cpp_extra_args"; return false; diff --git a/examples/server/routes_sdapi.cpp b/examples/server/routes_sdapi.cpp index 699ba022d8..22a710b7ea 100644 --- a/examples/server/routes_sdapi.cpp +++ b/examples/server/routes_sdapi.cpp @@ -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; } -static void assign_solid_mask(SDImageOwner& mask_owner, int width, int height) { - const size_t pixel_count = static_cast(width) * static_cast(height); - uint8_t* raw_mask = static_cast(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, ServerRuntime& runtime, bool img2img, @@ -193,15 +182,25 @@ static bool build_sdapi_img_gen_request(const json& j, } } - if (img2img) { - const int expected_width = request.gen_params.width_and_height_are_set() ? request.gen_params.width : 0; - const int expected_height = request.gen_params.width_and_height_are_set() ? request.gen_params.height : 0; + if (j.contains("ref_image_args")) { + if (!j["ref_image_args"].is_string()) { + error_message = "ref_image_args must be a string"; + return false; + } + request.gen_params.ref_image_args = j["ref_image_args"].get(); + } + + 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 (decode_base64_image(j["init_images"][0].get(), - 3, - expected_width, - expected_height, + 0, + 0, + 0, request.gen_params.init_image)) { const sd_image_t& image = request.gen_params.init_image.get(); 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 (decode_base64_image(j["mask"].get(), 1, - expected_width, - expected_height, + 0, + 0, request.gen_params.mask_image)) { const sd_image_t& image = request.gen_params.mask_image.get(); 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 { - const int resolved_width = request.gen_params.get_resolved_width(); - const int resolved_height = request.gen_params.get_resolved_height(); - assign_solid_mask(request.gen_params.mask_image, resolved_width, resolved_height); + request.gen_params.mask_image.reset({0, 0, 1, nullptr}); } float denoising_strength = j.value("denoising_strength", -1.f); @@ -243,9 +240,8 @@ static bool build_sdapi_img_gen_request(const json& j, } SDImageOwner image_owner; if (decode_base64_image(extra_image.get(), - 3, - request.gen_params.width_and_height_are_set() ? request.gen_params.width : 0, - request.gen_params.width_and_height_are_set() ? request.gen_params.height : 0, + 0, + 0, 0, image_owner)) { const sd_image_t& image = image_owner.get(); request.gen_params.set_width_and_height_if_unset(image.width, image.height); @@ -305,6 +301,76 @@ static nlohmann::json prepare_info_field(const SDContextParams& ctx_params, return jsoninfo; } +struct SDAPIPreviewState { + std::mutex mutex; + std::string b64; + std::string id_task; + std::string job_type; + int step = 0; + int total = 0; + bool active = false; + int64_t job_count = 0; + int job_no = 0; +}; + +static SDAPIPreviewState g_preview; + +static void preview_callback(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data) { + (void)is_noisy; + (void)data; + if (frame_count <= 0 || frames == nullptr || frames[0].data == nullptr) { + return; + } + + const sd_image_t& frame = frames[0]; + const sd_preview_info_t info = sd_get_preview_info(); + + auto image_bytes = encode_image_to_vector(EncodedImageFormat::JPEG, + frame.data, + frame.width, + frame.height, + frame.channel, + "", + 80); + if (image_bytes.empty()) { + return; + } + + std::string b64 = base64_encode(image_bytes); + std::lock_guard lock(g_preview.mutex); + g_preview.b64 = std::move(b64); + g_preview.step = static_cast(std::min(step < 0 ? -static_cast(step) : step, + std::max(0, info.total_steps))); + g_preview.total = std::max(0, info.total_steps); + g_preview.job_no = std::max(0, info.sample_pass - 1); +} + +struct SDAPIPreviewGuard { + SDAPIPreviewGuard(const std::string& id_task, const sd_img_gen_params_t& params, const std::string& job_type) { + std::lock_guard lock(g_preview.mutex); + g_preview.id_task = id_task; + g_preview.total = 0; + g_preview.job_type = job_type; + g_preview.b64.clear(); + g_preview.step = 0; + g_preview.active = true; + // Each batch item has a base pass and, when enabled, a hires pass. + g_preview.job_count = static_cast(std::max(1, params.batch_count)) * (params.hires.enabled ? 2 : 1); + g_preview.job_no = 0; + sd_set_preview_callback(preview_callback, PREVIEW_PROJ, 1, true, false, nullptr); + } + + ~SDAPIPreviewGuard() { + sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr); + std::lock_guard lock(g_preview.mutex); + g_preview.active = false; + g_preview.b64.clear(); + } + + SDAPIPreviewGuard(const SDAPIPreviewGuard&) = delete; + SDAPIPreviewGuard& operator=(const SDAPIPreviewGuard&) = delete; +}; + void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { ServerRuntime* runtime = &rt; @@ -321,7 +387,9 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { return; } - json j = json::parse(req.body); + std::string id_task = "sdapi"; + json j = json::parse(req.body); + id_task = j.value("id_task", id_task); ImgGenJobRequest request; std::string error_message; if (!build_sdapi_img_gen_request(j, *runtime, img2img, request, error_message)) { @@ -338,6 +406,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { { std::lock_guard lock(*runtime->sd_ctx_mutex); + SDAPIPreviewGuard preview_guard(id_task, img_gen_params, img2img ? "img2img" : "txt2img"); sd_image_t* raw_results = nullptr; if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) { raw_results = nullptr; @@ -532,4 +601,62 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { r["sd_model_checkpoint"] = model_path.stem(); res.set_content(r.dump(), "application/json"); }); + + svr.Get("/sdapi/v1/progress", [](const httplib::Request& req, httplib::Response& res) { + bool skip_current_image = false; + bool has_id_param = req.has_param("id_task"); + std::string req_id_task = has_id_param ? req.get_param_value("id_task") : ""; + + if (req.has_param("skip_current_image")) { + std::string skip_str = req.get_param_value("skip_current_image"); + skip_current_image = (skip_str == "true" || skip_str == "1"); + } + + json r; + r["eta_relative"] = 0.0f; + json state; + state["skipped"] = false; + state["interrupted"] = false; + state["job_timestamp"] = "0"; + r["textinfo"] = nullptr; + + { + std::lock_guard lock(g_preview.mutex); + + bool is_matching_task = g_preview.active && (!has_id_param || req_id_task == g_preview.id_task); + + int step = is_matching_task ? g_preview.step : 0; + int steps = is_matching_task ? g_preview.total : 0; + std::string b64 = is_matching_task ? g_preview.b64 : ""; + std::string job_type = is_matching_task ? g_preview.job_type : ""; + int64_t job_count = is_matching_task ? g_preview.job_count : 0; + int job_no = is_matching_task ? g_preview.job_no : 0; + + float batch_progress = (steps > 0) ? (static_cast(step) / static_cast(steps)) : 0.0f; + if (batch_progress > 1.0f) + batch_progress = 1.0f; + + float overall_progress = job_count > 0 + ? (static_cast(job_no) + batch_progress) / static_cast(job_count) + : 0.0f; + if (overall_progress > 1.0f) + overall_progress = 1.0f; + + r["progress"] = overall_progress; + state["job"] = job_type; + state["sampling_step"] = step; + state["sampling_steps"] = steps; + state["job_count"] = job_count; + state["job_no"] = job_no; + r["state"] = state; + + if (!skip_current_image && !b64.empty()) { + r["current_image"] = b64; + } else { + r["current_image"] = nullptr; + } + } + + res.set_content(r.dump(), "application/json"); + }); } diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index 7253dbf496..3139682f90 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/examples/server/routes_sdcpp.cpp @@ -3,12 +3,33 @@ #include #include #include +#include #include "async_jobs.h" #include "common/common.h" +#include "common/media_io.h" +#include "common/resource_owners.hpp" 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) { if (mode_str == "disabled") { mode_out = SD_CACHE_DISABLED; @@ -57,11 +78,11 @@ static json make_vae_tiling_json(const sd_tiling_params_t& params) { return { {"enabled", params.enabled}, {"temporal_tiling", params.temporal_tiling}, - {"tile_size_x", params.tile_size_x}, - {"tile_size_y", params.tile_size_y}, + {"tile_size_w", params.tile_size_w}, + {"tile_size_h", params.tile_size_h}, {"target_overlap", params.target_overlap}, - {"rel_size_x", params.rel_size_x}, - {"rel_size_y", params.rel_size_y}, + {"rel_size_w", params.rel_size_w}, + {"rel_size_h", params.rel_size_h}, {"extra_tiling_args", params.extra_tiling_args ? params.extra_tiling_args : ""}, }; } @@ -127,7 +148,8 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const {"seed", defaults.seed}, {"batch_count", defaults.batch_count}, {"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}, {"control_strength", defaults.control_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}, {"seed", defaults.seed}, {"video_frames", defaults.video_frames}, + {"image_preprocess", defaults.image_preprocess}, {"fps", defaults.fps}, {"moe_boundary", defaults.moe_boundary}, {"vace_strength", defaults.vace_strength}, @@ -182,6 +205,7 @@ static json make_img_gen_features_json() { {"cache", true}, {"cancel_queued", true}, {"cancel_generating", false}, + {"preview", true}, }; } @@ -196,9 +220,18 @@ static json make_vid_gen_features_json() { {"cache", true}, {"cancel_queued", true}, {"cancel_generating", false}, + {"preview", true}, }; } +static json make_preview_modes_json() { + json modes = json::array(); + for (int i = 0; i < PREVIEW_COUNT; ++i) { + modes.push_back(sd_preview_name((preview_t)i)); + } + return modes; +} + static json make_capabilities_json(ServerRuntime& runtime) { refresh_lora_cache(runtime); refresh_upscaler_cache(runtime); @@ -239,37 +272,59 @@ static json make_capabilities_json(ServerRuntime& runtime) { available_upscalers.push_back({ {"name", "None"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Lanczos"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Nearest"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Latent"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Latent (nearest)"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Latent (nearest-exact)"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Latent (antialiased)"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Latent (bicubic)"}, + {"model", false}, + {"image_upscale", false}, }); available_upscalers.push_back({ {"name", "Latent (bicubic antialiased)"}, + {"model", false}, + {"image_upscale", false}, }); + bool have_upscaler_models = false; { std::lock_guard lock(*runtime.upscaler_mutex); for (const auto& entry : *runtime.upscaler_cache) { available_upscalers.push_back({ {"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 +394,8 @@ static json make_capabilities_json(ServerRuntime& runtime) { {"max_height", 4096}, {"max_batch_count", 8}, {"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["schedulers"] = schedulers; @@ -348,6 +405,8 @@ static json make_capabilities_json(ServerRuntime& runtime) { result["features_by_mode"] = features_by_mode; result["loras"] = available_loras; result["upscalers"] = available_upscalers; + result["upscale"] = have_upscaler_models; + result["preview_modes"] = make_preview_modes_json(); return result; } @@ -370,6 +429,8 @@ static bool parse_img_gen_request(const json& body, if (!assign_output_options(request, output_format, output_compression, true, error_message)) { return false; } + request.preview_mode = body.value("preview", std::string("none")); + request.preview_interval = body.value("preview_interval", 1); // Intentionally disable prompt-embedded LoRA tag parsing for server APIs. if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) { error_message = "invalid generation parameters"; @@ -397,6 +458,8 @@ static bool parse_vid_gen_request(const json& body, if (!assign_output_options(request, output_format, output_compression, error_message)) { return false; } + request.preview_mode = body.value("preview", std::string("none")); + request.preview_interval = body.value("preview_interval", 1); // Intentionally disable prompt-embedded LoRA tag parsing for server APIs. if (!request.gen_params.resolve_and_validate(VID_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) { error_message = "invalid generation parameters"; @@ -413,6 +476,171 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { 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() <= static_cast(std::numeric_limits::max()) + : value.is_number_integer() && value.get() >= std::numeric_limits::min() && + value.get() <= std::numeric_limits::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 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 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) { try { if (req.body.empty()) { diff --git a/examples/server/runtime.cpp b/examples/server/runtime.cpp index cd0e2e722b..1a71cf76fe 100644 --- a/examples/server/runtime.cpp +++ b/examples/server/runtime.cpp @@ -254,28 +254,55 @@ std::string SDSvrParams::to_string() const { void refresh_lora_cache(ServerRuntime& rt) { std::vector new_cache; + if (rt.ctx_params->lora_model_dir.empty()) { + std::lock_guard lock(*rt.lora_mutex); + rt.lora_cache->clear(); + return; + } + fs::path lora_dir = rt.ctx_params->lora_model_dir; - if (fs::exists(lora_dir) && fs::is_directory(lora_dir)) { - for (auto& entry : fs::recursive_directory_iterator(lora_dir, fs::directory_options::skip_permission_denied)) { - if (!entry.is_regular_file()) { - continue; - } - const fs::path& p = entry.path(); - if (!is_supported_model_ext(p)) { - continue; + std::error_code ec; + if (fs::exists(lora_dir, ec) && !ec && fs::is_directory(lora_dir, ec) && !ec) { + try { + auto it = fs::recursive_directory_iterator( + lora_dir, + fs::directory_options::skip_permission_denied, + ec); + auto end = fs::recursive_directory_iterator(); + while (!ec && it != end) { + std::error_code entry_ec; + bool is_reg = it->is_regular_file(entry_ec); + if (entry_ec || !is_reg) { + it.increment(ec); + continue; + } + const fs::path& p = it->path(); + if (!is_supported_model_ext(p)) { + it.increment(ec); + continue; + } + + LoraEntry lora_entry; + lora_entry.name = p.stem().u8string(); + lora_entry.fullpath = p.u8string(); + std::string rel = p.lexically_relative(lora_dir).u8string(); + std::replace(rel.begin(), rel.end(), '\\', '/'); + lora_entry.path = rel; + + new_cache.push_back(std::move(lora_entry)); + it.increment(ec); } - - LoraEntry lora_entry; - lora_entry.name = p.stem().u8string(); - lora_entry.fullpath = p.u8string(); - std::string rel = p.lexically_relative(lora_dir).u8string(); - std::replace(rel.begin(), rel.end(), '\\', '/'); - lora_entry.path = rel; - - new_cache.push_back(std::move(lora_entry)); + } catch (const std::exception& e) { + LOG_WARN("error while scanning lora directory '%s': %s", lora_dir.string().c_str(), e.what()); + return; } } + if (ec) { + LOG_WARN("error while scanning lora directory '%s': %s", lora_dir.string().c_str(), ec.message().c_str()); + return; + } + std::sort(new_cache.begin(), new_cache.end(), [](const LoraEntry& a, const LoraEntry& b) { return a.path < b.path; }); @@ -295,28 +322,78 @@ std::string get_lora_full_path(ServerRuntime& rt, const std::string& path) { void refresh_upscaler_cache(ServerRuntime& rt) { std::vector new_cache; + std::vector previous_cache; + { + std::lock_guard lock(*rt.upscaler_mutex); + previous_cache = *rt.upscaler_cache; + } fs::path upscaler_dir = rt.ctx_params->hires_upscalers_dir; - if (fs::exists(upscaler_dir) && fs::is_directory(upscaler_dir)) { - for (auto& entry : fs::directory_iterator(upscaler_dir)) { - if (!entry.is_regular_file()) { - continue; - } - const fs::path& p = entry.path(); - if (!is_supported_model_ext(p)) { - continue; + std::error_code ec; + if (fs::exists(upscaler_dir, ec) && !ec && fs::is_directory(upscaler_dir, ec) && !ec) { + try { + auto it = fs::directory_iterator( + upscaler_dir, + fs::directory_options::skip_permission_denied, + ec); + auto end = fs::directory_iterator(); + while (!ec && it != end) { + std::error_code entry_ec; + bool is_reg = it->is_regular_file(entry_ec); + if (entry_ec || !is_reg) { + it.increment(ec); + continue; + } + const fs::path& p = it->path(); + if (!is_supported_model_ext(p)) { + it.increment(ec); + continue; + } + + UpscalerEntry upscaler_entry; + upscaler_entry.name = p.stem().u8string(); + std::error_code abs_ec; + fs::path abs_path = fs::absolute(p, abs_ec); + upscaler_entry.fullpath = (abs_ec ? p : abs_path.lexically_normal()).u8string(); + upscaler_entry.model_name = "ESRGAN_4x"; + upscaler_entry.path = p.filename().u8string(); + + std::error_code size_ec; + upscaler_entry.file_size = it->file_size(size_ec); + if (size_ec) { + it.increment(ec); + continue; + } + + std::error_code time_ec; + upscaler_entry.last_modified = it->last_write_time(time_ec); + 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 && + (!time_ec && 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)); + it.increment(ec); } - - UpscalerEntry upscaler_entry; - upscaler_entry.name = p.stem().u8string(); - upscaler_entry.fullpath = fs::absolute(p).lexically_normal().u8string(); - upscaler_entry.model_name = "ESRGAN_4x"; - upscaler_entry.path = p.filename().u8string(); - - new_cache.push_back(std::move(upscaler_entry)); + } catch (const std::exception& e) { + LOG_WARN("error while scanning upscalers directory '%s': %s", upscaler_dir.string().c_str(), e.what()); + return; } } + if (ec) { + LOG_WARN("error while scanning upscalers directory '%s': %s", upscaler_dir.string().c_str(), ec.message().c_str()); + return; + } + std::sort(new_cache.begin(), new_cache.end(), [](const UpscalerEntry& a, const UpscalerEntry& b) { return a.name < b.name; }); diff --git a/examples/server/runtime.h b/examples/server/runtime.h index 677bb0ba0c..5baa633a6b 100644 --- a/examples/server/runtime.h +++ b/examples/server/runtime.h @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -43,6 +44,9 @@ struct UpscalerEntry { std::string fullpath; std::string model_name; int scale = 4; + int image_upscale_factor = 0; + uintmax_t file_size = 0; + std::filesystem::file_time_type last_modified; }; struct ServerRuntime { @@ -62,6 +66,8 @@ struct ImgGenJobRequest { SDGenerationParams gen_params; std::string output_format = "png"; int output_compression = 100; + std::string preview_mode = "none"; + int preview_interval = 1; sd_img_gen_params_t to_sd_img_gen_params_t() { return gen_params.to_sd_img_gen_params_t(); @@ -72,6 +78,8 @@ struct VidGenJobRequest { SDGenerationParams gen_params; std::string output_format = "webm"; int output_compression = 100; + std::string preview_mode = "none"; + int preview_interval = 1; sd_vid_gen_params_t to_sd_vid_gen_params_t() { return gen_params.to_sd_vid_gen_params_t(); diff --git a/ggml b/ggml index f583f393cd..d25b121ce1 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit f583f393cd5dfdc129360bbf75cb3d49ccc837a8 +Subproject commit d25b121ce1d6ae0c7e62ad080460d9fc1e46dd93 diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 9173240be6..a2e1575560 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -173,11 +173,13 @@ enum lora_apply_mode_t { typedef struct { bool enabled; bool temporal_tiling; - int tile_size_x; - int tile_size_y; + // Spatial tile dimensions in image pixels for both encode and decode; 0 uses 256. + int tile_size_w; + int tile_size_h; float target_overlap; - float rel_size_x; - float rel_size_y; + // Positive values override tile_size: <= 1 is a dimension fraction, > 1 a target tile count. + float rel_size_w; + float rel_size_h; const char* extra_tiling_args; } sd_tiling_params_t; @@ -247,6 +249,7 @@ typedef struct { 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 bool sage_attn; + int conditioning_cache_size; // Maximum cached conditioning entries per context; 0 disables caching (default: 4) } sd_ctx_params_t; typedef struct { @@ -263,6 +266,11 @@ typedef struct { uint8_t* data; } sd_image_t; +typedef struct { + // Semicolon-separated target=...,key=value rules. NULL preserves defaults. + const char* rules; +} sd_image_preprocess_params_t; + typedef struct { sd_image_t* frames; int frame_count; @@ -410,6 +418,7 @@ typedef struct { int qwen_image_layers; bool circular_x; bool circular_y; + sd_image_preprocess_params_t image_preprocess; } sd_img_gen_params_t; typedef struct { @@ -443,6 +452,7 @@ typedef struct { sd_hires_params_t hires; bool circular_x; bool circular_y; + sd_image_preprocess_params_t image_preprocess; } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; @@ -451,6 +461,11 @@ struct ggml_tensor; typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data); typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data); typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data); + +typedef struct { + int sample_pass; + int total_steps; +} sd_preview_info_t; typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data); SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data); @@ -459,6 +474,10 @@ SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data); // negative interval previews only completed logical step -interval. Zero previews the final // completed step of the first sampling pass (base-resolution or high-noise). SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data); +// Query from a preview callback. Passes are numbered from 1 since the last +// sd_set_preview_callback call; total_steps is the actual count in that pass. +// Both fields are zero before sampling starts. +SD_API sd_preview_info_t sd_get_preview_info(); SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data); SD_API int32_t sd_get_num_physical_cores(); SD_API const char* sd_get_system_info(); @@ -550,6 +569,8 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx, int* num_images_out); 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; @@ -589,7 +610,9 @@ SD_API bool convert_with_components(const char* model_path, enum sd_type_t output_type, const char* tensor_type_rules, bool convert_name, - int n_threads); + int n_threads, + const sd_lora_t* loras, + int lora_count); SD_API bool preprocess_canny(sd_image_t image, float high_threshold, diff --git a/scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch b/scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch index f8b379f189..04cc944502 100644 --- a/scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch +++ b/scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch @@ -1,4 +1,4 @@ -From 22b560daa5e90bdb72c492f0b55667468ad3fd6b Mon Sep 17 00:00:00 2001 +From fa18cddc3f2a87bec12af3a00d4be0e28f9d514b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 30 Sep 2026 22:28:21 -0700 Subject: [PATCH 1/2] ggml-cuda: BF16 cuBLAS for large-batch quantized matmul, @@ -21,10 +21,10 @@ in one tiled H3 VAE decode. The flag is now read once at registration. 1 file changed, 47 insertions(+), 8 deletions(-) diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index e15fdc5e..517d0894 100644 +index 5e9c323b..a153de57 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -2427,6 +2427,25 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { +@@ -2431,6 +2431,25 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } @@ -50,7 +50,7 @@ index e15fdc5e..517d0894 100644 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_TENSOR_BINARY_OP_LOCALS -@@ -2492,6 +2511,23 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor +@@ -2496,6 +2515,23 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); return; } @@ -74,7 +74,7 @@ index e15fdc5e..517d0894 100644 if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); return; -@@ -5200,6 +5236,10 @@ struct ggml_backend_cuda_device_context { +@@ -5593,6 +5629,10 @@ struct ggml_backend_cuda_device_context { std::string description; std::string pci_bus_id; int op_offload_min_batch_size; @@ -85,9 +85,9 @@ index e15fdc5e..517d0894 100644 }; static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) { -@@ -5303,12 +5343,9 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * +@@ -5696,12 +5736,9 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * // ref: https://github.com/ggml-org/llama.cpp/pull/17368 - #if defined(__linux__) + #if defined(__linux__) && !defined(GGML_USE_HIP) // Check if this is a UMA (Unified Memory Architecture) system - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device))); @@ -99,7 +99,7 @@ index e15fdc5e..517d0894 100644 if (is_uma) { // For UMA systems (like DGX Spark), use system memory info -@@ -5332,10 +5369,7 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * +@@ -5725,10 +5762,7 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) { ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *) dev->context; @@ -111,7 +111,7 @@ index e15fdc5e..517d0894 100644 ? GGML_BACKEND_DEVICE_TYPE_IGPU : GGML_BACKEND_DEVICE_TYPE_GPU; } -@@ -6095,6 +6129,11 @@ ggml_backend_reg_t ggml_backend_cuda_reg() { +@@ -6503,6 +6537,11 @@ ggml_backend_reg_t ggml_backend_cuda_reg() { dev_ctx->device = i; dev_ctx->name = GGML_CUDA_NAME + std::to_string(i); dev_ctx->description = ggml_cuda_device_description(i); @@ -127,7 +127,7 @@ index e15fdc5e..517d0894 100644 2.43.0 -From 0bf02a4e324ef837685c63ebaeddf2b4a4c0f317 Mon Sep 17 00:00:00 2001 +From 09f9ed593f044105ff5e206b1003327275c08b9b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 30 Sep 2026 22:54:01 -0700 Subject: [PATCH 2/2] ggml-cuda: vectorized row copy for strided F32->F32/F16 @@ -145,7 +145,7 @@ cpy_scalar. 1 file changed, 83 insertions(+) diff --git a/src/ggml-cuda/cpy.cu b/src/ggml-cuda/cpy.cu -index 4c6fb4a2..c26a081c 100644 +index c32bb36b..4aa183af 100644 --- a/src/ggml-cuda/cpy.cu +++ b/src/ggml-cuda/cpy.cu @@ -1,4 +1,5 @@ diff --git a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch index 32b2c5378f..ca9141e547 100644 --- a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch +++ b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch @@ -1,4 +1,4 @@ -From ed3f840655078aef3c2bb4e6b9a1d26f8e868e64 Mon Sep 17 00:00:00 2001 +From dc60887c4e3eec372301d069f4dd90874745dc91 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 05:41:03 -0700 Subject: [PATCH 1/3] ggml-cuda: pipelined long-sequence FlashAttention for @@ -34,7 +34,7 @@ tiles, several Q tiles and GQA broadcast. create mode 100644 src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu diff --git a/src/ggml-cuda/fattn-mma-f16.cuh b/src/ggml-cuda/fattn-mma-f16.cuh -index 7f4cfd55..b4496734 100644 +index 449a77c5..308218f1 100644 --- a/src/ggml-cuda/fattn-mma-f16.cuh +++ b/src/ggml-cuda/fattn-mma-f16.cuh @@ -61,6 +61,10 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co @@ -48,7 +48,7 @@ index 7f4cfd55..b4496734 100644 GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 8, 64, 4, 64, 96, 64, 64, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 16, 64, 4, 32, 96, 64, 64, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2, 32, 96, 64, 64, 2, true); -@@ -358,6 +362,24 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, c +@@ -387,6 +391,24 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages( #endif // CP_ASYNC_AVAILABLE } @@ -72,47 +72,47 @@ index 7f4cfd55..b4496734 100644 + // ------------------------------------------------------------------------------------------------------------------ - template -@@ -529,7 +551,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( + template +@@ -581,7 +603,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( template + typename T_A_KQ, typename T_B_KQ, typename T_C_KQ, typename T_A_VKQ, typename T_B_VKQ, typename T_C_VKQ, int nstages_override = -1> static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, -@@ -566,7 +588,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( +@@ -619,7 +641,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols); constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); -- constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); -+ constexpr int nstages = nstages_override >= 0 ? nstages_override : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); +- constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); ++ constexpr int nstages = nstages_override >= 0 ? nstages_override : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, use_sparse); - constexpr int stride_tile_K = nbatch_K2 + 4; - -@@ -1113,7 +1135,8 @@ template struct mma_tile_sizes { + constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols1, ncols2); + constexpr int stride_tile_K = ggml_cuda_fattn_mma_get_stride_tile(nbatch_K2, swz); +@@ -1195,7 +1217,8 @@ template struct mma_tile_sizes { }; #endif // defined(TURING_MMA_AVAILABLE) --template -+template ++template static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, -@@ -1158,7 +1181,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( +@@ -1241,7 +1264,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); -- constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); +- constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); + constexpr int nstages = longseq ? ggml_cuda_fattn_mma_get_nstages_longseq(DKQ, DV, ncols) : -+ ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); -+ static_assert(!longseq || (ncols2 == 1 && !V_is_K_view), "bad longseq variant"); ++ ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); ++ static_assert(!longseq || (ncols2 == 1 && !V_is_K_view && !use_sparse), "bad longseq variant"); if (cols_per_warp > ncols) { NO_DEVICE_CODE; -@@ -1255,22 +1280,63 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( +@@ -1340,6 +1365,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( int kb0 = kb0_start; @@ -121,36 +121,37 @@ index 7f4cfd55..b4496734 100644 + // Preload mask and K data for first iteration when using cp_async with multiple stages: if constexpr (nstages > 1) { - static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); + static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); +@@ -1347,16 +1375,54 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool use_cp_async = true; constexpr bool oob_check = false; constexpr int k_VKQ_sup = nbatch_fa; - if (ncols2 > 1 || mask_h) { -- flash_attn_ext_f16_load_mask -- (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); +- flash_attn_ext_f16_load_mask +- (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr); + if (!longseq || kb0 < kb0_stop_pipe) { + if (ncols2 > 1 || mask_h) { -+ flash_attn_ext_f16_load_mask -+ (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); ++ flash_attn_ext_f16_load_mask ++ (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr); + } -+ flash_attn_ext_f16_load_tile -+ (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); ++ flash_attn_ext_f16_load_tile ++ (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr); } -- flash_attn_ext_f16_load_tile -- (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); +- flash_attn_ext_f16_load_tile +- (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr); } // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. -- if constexpr (ncols2 == 1) { +- if constexpr (ncols2 == 1 || use_sparse) { + if constexpr (longseq && nstages > 1) { + for (; kb0 < kb0_stop_pipe-1; ++kb0) { + constexpr bool last_iter = false; + constexpr bool oob_check = false; + constexpr int k_VKQ_sup = nbatch_fa; + flash_attn_ext_f16_iter -+ -+ (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ++ (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, + ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + } @@ -159,9 +160,9 @@ index 7f4cfd55..b4496734 100644 + constexpr bool oob_check = false; + constexpr int k_VKQ_sup = nbatch_fa; + flash_attn_ext_f16_iter -+ -+ (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ++ (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, + ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + ++kb0; @@ -172,52 +173,52 @@ index 7f4cfd55..b4496734 100644 + constexpr bool oob_check = true; + const int k_VKQ_sup = ne11 - kb0*nbatch_fa; + flash_attn_ext_f16_iter -+ -+ (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ++ (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, + ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + } -+ } else if constexpr (ncols2 == 1) { ++ } else if constexpr (ncols2 == 1 || use_sparse) { constexpr bool oob_check = true; for (; kb0 < kb0_stop-1; ++kb0) { constexpr bool last_iter = false; -@@ -1700,7 +1766,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( - #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) +@@ -1802,7 +1868,7 @@ static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_u + (DKQ == 256 && DV == 256 && ncols1 == 8 && ncols2 == 8); } --template -+template +-template ++template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, -@@ -1829,12 +1895,12 @@ static __global__ void flash_attn_ext_f16( +@@ -1945,12 +2011,12 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. -- flash_attn_ext_f16_process_tile -+ flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, +- flash_attn_ext_f16_process_tile ++ flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. -- flash_attn_ext_f16_process_tile -+ flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, +- flash_attn_ext_f16_process_tile ++ flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } -@@ -1875,7 +1941,7 @@ static __global__ void flash_attn_ext_f16( +@@ -1994,7 +2060,7 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; -- flash_attn_ext_f16_process_tile -+ flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, +- flash_attn_ext_f16_process_tile ++ flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else -@@ -1892,6 +1958,47 @@ static __global__ void flash_attn_ext_f16( - #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) - } +@@ -2013,6 +2079,47 @@ static __global__ void flash_attn_ext_f16( + + bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1, const int ncols2); +// Whether the long-sequence variant (see ggml_cuda_fattn_mma_get_nstages_longseq) can and should be used for dst. +// GGML_CUDA_FA_LONGSEQ=0 disables it, =1 enables it on every NVIDIA GPU with cp.async, unset = per-architecture default. @@ -263,7 +264,7 @@ index 7f4cfd55..b4496734 100644 template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; -@@ -1900,13 +2007,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml +@@ -2021,13 +2128,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml constexpr int ncols = ncols1 * ncols2; @@ -287,14 +288,14 @@ index 7f4cfd55..b4496734 100644 const int cols_per_warp = std::min(ncols, get_cols_per_warp(cc)); const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size; -@@ -1935,6 +2051,21 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml - using fattn_kernel_ptr_t = fattn_kernel_t; +@@ -2061,6 +2177,21 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml #endif // defined(GGML_USE_HIP) fattn_kernel_t fattn_kernel; + bool use_sparse = false; + if constexpr (longseq_instantiated) { + if (longseq) { + GGML_ASSERT(logit_softcap == 0.0f); -+ fattn_kernel = flash_attn_ext_f16; ++ fattn_kernel = flash_attn_ext_f16; + + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { @@ -302,14 +303,14 @@ index 7f4cfd55..b4496734 100644 + shared_memory_limit_raised[id] = true; + } + launch_fattn -+ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); ++ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host); + return; + } + } if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - fattn_kernel = flash_attn_ext_f16; -@@ -2003,6 +2134,12 @@ DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 64) + #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +@@ -2157,6 +2288,12 @@ DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 64) DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128, 64) DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256, 64) @@ -323,10 +324,10 @@ index 7f4cfd55..b4496734 100644 extern DECL_FATTN_MMA_F16_CASE(512, 512, 8, 2); extern DECL_FATTN_MMA_F16_CASE(512, 512, 16, 2); diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu -index ab7a3b29..00e9554a 100644 +index d1fcf58c..e3e34605 100644 --- a/src/ggml-cuda/fattn.cu +++ b/src/ggml-cuda/fattn.cu -@@ -30,6 +30,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con +@@ -187,6 +187,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con return; } @@ -357,12 +358,12 @@ index 00000000..b9bf54a1 +DECL_FATTN_MMA_F16_CASE(128, 128, 128, 1); +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index fdf505cd..dd290caf 100644 +index e198d092..6a6954fe 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp -@@ -9588,6 +9588,21 @@ static std::vector> make_test_cases_eval() { - } - } +@@ -10840,6 +10840,21 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_flash_attn_ext(72, 64, 4, {1, 1}, 256, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 72, 4, {1, 1}, 256, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // diffusion-shaped cases without a mask (CUDA long-sequence FA variant): KV lengths that are not a multiple + // of the KV tile (partial last tile), several Q tiles per head, and with/without GQA broadcast. @@ -386,7 +387,7 @@ index fdf505cd..dd290caf 100644 2.43.0 -From 7f7c132c57dbfd8fd44a13dd29031c6724589efc Mon Sep 17 00:00:00 2001 +From 277581a8840cfd5bbf2335f379ef9be1f70e4754 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:27:44 -0700 Subject: [PATCH 2/3] ggml-cuda: per-architecture tile width for the @@ -403,10 +404,10 @@ GGML_CUDA_FA_LONGSEQ_NCOLS=64|128 overrides. 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/ggml-cuda/fattn-mma-f16.cuh b/src/ggml-cuda/fattn-mma-f16.cuh -index b4496734..57f284c6 100644 +index 308218f1..c2d3196d 100644 --- a/src/ggml-cuda/fattn-mma-f16.cuh +++ b/src/ggml-cuda/fattn-mma-f16.cuh -@@ -1965,9 +1965,11 @@ static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { +@@ -2086,9 +2086,11 @@ static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { return true; } @@ -421,7 +422,7 @@ index b4496734..57f284c6 100644 static const int env = [] { const char * e = getenv("GGML_CUDA_FA_LONGSEQ_NCOLS"); return e != nullptr ? atoi(e) : -1; -@@ -1975,8 +1977,9 @@ static int ggml_cuda_fattn_mma_longseq_ncols(const int cc) { +@@ -2096,8 +2098,9 @@ static int ggml_cuda_fattn_mma_longseq_ncols(const int cc) { if (env == 64 || env == 128) { return env; } @@ -434,10 +435,10 @@ index b4496734..57f284c6 100644 static bool ggml_cuda_fattn_mma_use_longseq(const int cc, const ggml_tensor * dst) { diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu -index 00e9554a..b8be30a0 100644 +index e3e34605..f0fe5577 100644 --- a/src/ggml-cuda/fattn.cu +++ b/src/ggml-cuda/fattn.cu -@@ -32,7 +32,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con +@@ -189,7 +189,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) if constexpr (ncols2 == 1 && DKQ == DV && (DKQ == 64 || DKQ == 128)) { @@ -450,7 +451,7 @@ index 00e9554a..b8be30a0 100644 2.43.0 -From e6e78b0ba7477be2b0e6869e19bf86ac89b0381b Mon Sep 17 00:00:00 2001 +From ad669abc9a620559691dab67db22a2ec41518127 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 07:36:53 -0700 Subject: [PATCH 3/3] ggml-cuda: long-sequence FlashAttention takes the stock @@ -471,22 +472,22 @@ on any architecture; GGML_CUDA_FA_LONGSEQ_NCOLS=128 opts in. 2 files changed, 33 insertions(+), 14 deletions(-) diff --git a/src/ggml-cuda/fattn-common.cuh b/src/ggml-cuda/fattn-common.cuh -index e67cc7fd..8a8d4023 100644 +index 6d1ce52d..bc50e31f 100644 --- a/src/ggml-cuda/fattn-common.cuh +++ b/src/ggml-cuda/fattn-common.cuh -@@ -972,7 +972,10 @@ static __global__ void flash_attn_combine_results( - template +@@ -976,7 +976,10 @@ template void launch_fattn( ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared, -- const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE -+ const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE, + const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse, +- const int warp_size = WARP_SIZE ++ const int warp_size = WARP_SIZE, + // Optional: size the grid from the occupancy of this kernel instead of fattn_kernel. A drop-in kernel variant launched + // with the grid of the kernel it replaces splits the work (stream-k seams) identically and so gives identical results. + fattn_kernel_t occupancy_kernel = nullptr, const size_t occupancy_nbytes_shared = 0 ) { constexpr int ncols = ncols1 * ncols2; -@@ -1110,7 +1113,8 @@ void launch_fattn( +@@ -1128,7 +1131,8 @@ void launch_fattn( const dim3 block_dim(warp_size, nwarps, 1); int max_blocks_per_sm = 1; // Max. number of active blocks limited by occupancy. @@ -497,10 +498,10 @@ index e67cc7fd..8a8d4023 100644 int parallel_blocks = max_blocks_per_sm; diff --git a/src/ggml-cuda/fattn-mma-f16.cuh b/src/ggml-cuda/fattn-mma-f16.cuh -index 57f284c6..d5fd8eea 100644 +index c2d3196d..dfc3182d 100644 --- a/src/ggml-cuda/fattn-mma-f16.cuh +++ b/src/ggml-cuda/fattn-mma-f16.cuh -@@ -1965,21 +1965,18 @@ static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { +@@ -2086,21 +2086,18 @@ static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { return true; } @@ -529,12 +530,12 @@ index 57f284c6..d5fd8eea 100644 } static bool ggml_cuda_fattn_mma_use_longseq(const int cc, const ggml_tensor * dst) { -@@ -2064,8 +2061,26 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml +@@ -2190,8 +2187,26 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); shared_memory_limit_raised[id] = true; } - launch_fattn -- (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); +- (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host); + if (ncols1 <= 64) { + // Same grid as the stock kernel of this tile size: identical stream-k seams, so results are bit-identical. + const int nstages_stock = ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, cc); @@ -549,11 +550,11 @@ index 57f284c6..d5fd8eea 100644 + stock_limit_raised[id] = true; + } + launch_fattn -+ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host, ++ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host, + stock_kernel, nbytes_total_stock); + } else { + launch_fattn -+ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); ++ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host); + } return; } diff --git a/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch b/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch index 1e14c53584..e227077e25 100644 --- a/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch +++ b/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch @@ -1,4 +1,4 @@ -From 76b49e564c23b99758642d5bc5013d09774a1e24 Mon Sep 17 00:00:00 2001 +From 54d3121a3e39ddf79c025a87271c9e349fb76b5d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 05:46:18 -0700 Subject: [PATCH] ggml-cuda: fuse bias epilogues into F16 cuBLAS matmuls, @@ -24,10 +24,10 @@ the 256-thread launch. 2 files changed, 262 insertions(+), 1 deletion(-) diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index 517d089..0e61996 100644 +index a153de57..72d090e3 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -2017,8 +2017,73 @@ struct batched_mul_mat_traits { +@@ -2014,8 +2014,73 @@ struct batched_mul_mat_traits { static inline auto convert_nc(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } }; @@ -102,7 +102,7 @@ index 517d089..0e61996 100644 using traits = batched_mul_mat_traits; using cuda_t = typename traits::cuda_type; -@@ -2226,8 +2291,17 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const +@@ -2223,8 +2288,17 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const // Convert output back to F32 if needed if (cu_data_type != CUDA_R_32F) { @@ -120,7 +120,7 @@ index 517d089..0e61996 100644 } } -@@ -3838,6 +3912,123 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, +@@ -4111,6 +4185,123 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } // try and fuse nodes and return the number of nodes to skip @@ -244,8 +244,8 @@ index 517d089..0e61996 100644 static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); -@@ -4523,6 +4714,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph - return fused_node_count - 1; +@@ -4825,6 +5016,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph + return 2; } + if (node->op == GGML_OP_MUL_MAT) { @@ -259,7 +259,7 @@ index 517d089..0e61996 100644 ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); return 2; diff --git a/src/ggml-cuda/norm.cu b/src/ggml-cuda/norm.cu -index c3758cd..9c9385f 100644 +index c3758cd5..9c9385f8 100644 --- a/src/ggml-cuda/norm.cu +++ b/src/ggml-cuda/norm.cu @@ -301,10 +301,68 @@ static void group_norm_f32_cuda( diff --git a/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch index 689dfa5f67..29592cf662 100644 --- a/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch +++ b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch @@ -1,4 +1,4 @@ -From 04e9fe6c85b680c2f8e29c22e084dab6012babbe Mon Sep 17 00:00:00 2001 +From cc2ac8fad7ecb73b19cb7971c874883903393a96 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:48:01 -0700 Subject: [PATCH] ggml : fused value-preserving ops for adaLN DiT blocks @@ -43,22 +43,22 @@ partial and full rotation, F32/F16 outputs, chained segments and the scales. create mode 100644 src/ggml-cuda/rope-pe.cuh diff --git a/include/ggml.h b/include/ggml.h -index 991cc4fc..153eed2a 100644 +index 761d3f28..8c92967a 100644 --- a/include/ggml.h +++ b/include/ggml.h -@@ -594,6 +594,9 @@ extern "C" { - - GGML_OP_QUANTIZE_I8_CONVROT, +@@ -607,6 +607,9 @@ extern "C" { GGML_OP_SAGE_ATTN, + GGML_OP_SOL_ATTN, + GGML_OP_ROPE_APPLY, + GGML_OP_ROPE_PE_PERMUTE, + GGML_OP_MODULATE_ROWS, + GGML_OP_SWIGLU_SCALED, GGML_OP_COUNT, }; -@@ -2461,6 +2464,65 @@ extern "C" { - float scale, - enum ggml_sage_attn_mode mode); +@@ -2551,6 +2554,65 @@ extern "C" { + float scale, + float tau); + // Rotary embedding with a precomputed cos/sin table, fused with the head-major relayout + // attention kernels consume, an optional scale and an optional cast to F16. @@ -119,28 +119,28 @@ index 991cc4fc..153eed2a 100644 + float in_scale, + float out_scale); + - GGML_API void ggml_flash_attn_ext_set_prec( + GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec( struct ggml_tensor * a, - enum ggml_prec prec); + enum ggml_prec prec), diff --git a/src/ggml-backend-meta.cpp b/src/ggml-backend-meta.cpp -index 539b4f70..5a34e748 100644 +index 98ba5020..426eaa86 100644 --- a/src/ggml-backend-meta.cpp +++ b/src/ggml-backend-meta.cpp -@@ -151,7 +151,8 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type( +@@ -153,7 +153,8 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type( static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { GGML_ASSERT(ggml_backend_dev_is_meta(dev)); -- if (op->op == GGML_OP_SAGE_ATTN) { -+ if (op->op == GGML_OP_SAGE_ATTN || op->op == GGML_OP_ROPE_PE_PERMUTE || +- if (op->op == GGML_OP_SAGE_ATTN || op->op == GGML_OP_SOL_ATTN) { ++ if (op->op == GGML_OP_SAGE_ATTN || op->op == GGML_OP_SOL_ATTN || op->op == GGML_OP_ROPE_PE_PERMUTE || + op->op == GGML_OP_MODULATE_ROWS || op->op == GGML_OP_SWIGLU_SCALED) { return false; } const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; diff --git a/src/ggml-cpu/ggml-cpu.c b/src/ggml-cpu/ggml-cpu.c -index 350777c1..d83f15f9 100644 +index 92de0631..d4dcefdf 100644 --- a/src/ggml-cpu/ggml-cpu.c +++ b/src/ggml-cpu/ggml-cpu.c -@@ -2190,6 +2190,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm +@@ -2224,6 +2224,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_glu(params, tensor); } break; @@ -159,9 +159,9 @@ index 350777c1..d83f15f9 100644 case GGML_OP_GET_REL_POS: { ggml_compute_forward_get_rel_pos(params, tensor); -@@ -2522,6 +2534,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { - case GGML_OP_ROPE: +@@ -2558,6 +2570,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_ROPE_BACK: + case GGML_OP_ROPE_APPLY: case GGML_OP_ADD_REL_POS: + case GGML_OP_ROPE_PE_PERMUTE: + case GGML_OP_MODULATE_ROWS: @@ -170,7 +170,7 @@ index 350777c1..d83f15f9 100644 n_tasks = n_threads; } break; diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp -index 727a4058..40fe1949 100644 +index b68afce1..f4ca82e7 100644 --- a/src/ggml-cpu/ops.cpp +++ b/src/ggml-cpu/ops.cpp @@ -12,6 +12,7 @@ @@ -181,7 +181,7 @@ index 727a4058..40fe1949 100644 struct ggml_fp8_e4m3_cpu { uint8_t value; -@@ -12136,3 +12137,156 @@ void ggml_compute_forward_quantize_i8_convrot(const ggml_compute_params * params +@@ -12380,3 +12381,156 @@ void ggml_compute_forward_quantize_i8_convrot(const ggml_compute_params * params memset(dst_data + rows * dst->nb[1], 0, (size_t)(rows_padded - rows) * dst->nb[1]); } } @@ -339,7 +339,7 @@ index 727a4058..40fe1949 100644 + } +} diff --git a/src/ggml-cpu/ops.h b/src/ggml-cpu/ops.h -index 3b579167..e09cc11d 100644 +index 19d7075e..d731a958 100644 --- a/src/ggml-cpu/ops.h +++ b/src/ggml-cpu/ops.h @@ -98,6 +98,9 @@ void ggml_compute_forward_win_part(const struct ggml_compute_params * params, st @@ -353,10 +353,10 @@ index 3b579167..e09cc11d 100644 void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index 517d0894..2c0984a0 100644 +index 72d090e3..7e98d729 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -44,6 +44,8 @@ +@@ -47,6 +47,8 @@ #include "ggml-cuda/rope.cuh" #include "ggml-cuda/roll.cuh" #include "ggml-cuda/scale.cuh" @@ -365,9 +365,9 @@ index 517d0894..2c0984a0 100644 #include "ggml-cuda/snake.cuh" #include "ggml-cuda/softcap.cuh" #include "ggml-cuda/softmax.cuh" -@@ -3002,6 +3004,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg - case GGML_OP_SAGE_ATTN: - ggml_cuda_sage_attn(ctx, dst); +@@ -3123,6 +3125,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg + case GGML_OP_SOL_ATTN: + ggml_cuda_sol_attn(ctx, dst); break; + case GGML_OP_ROPE_PE_PERMUTE: + ggml_cuda_op_rope_pe_permute(ctx, dst); @@ -381,10 +381,10 @@ index 517d0894..2c0984a0 100644 case GGML_OP_CROSS_ENTROPY_LOSS: ggml_cuda_cross_entropy_loss(ctx, dst); break; -@@ -5886,6 +5897,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g - return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); - case GGML_OP_SAGE_ATTN: +@@ -6492,6 +6503,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return ggml_cuda_sage_attn_supported(dev_ctx->device, op); + case GGML_OP_SOL_ATTN: + return ggml_cuda_sol_attn_supported(dev_ctx->device, op); + case GGML_OP_ROPE_PE_PERMUTE: + return ggml_cuda_rope_pe_permute_supported(op); + case GGML_OP_MODULATE_ROWS: @@ -701,38 +701,38 @@ index 00000000..b68b6af4 + +bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op); diff --git a/src/ggml.c b/src/ggml.c -index dc4a29f9..fec933bf 100644 +index 562b6eb9..165514d9 100644 --- a/src/ggml.c +++ b/src/ggml.c -@@ -1115,9 +1115,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { - - "QUANTIZE_I8_CONVROT", +@@ -1118,9 +1118,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "SAGE_ATTN", + "SOL_ATTN", + "ROPE_APPLY", + "ROPE_PE_PERMUTE", + "MODULATE_ROWS", + "SWIGLU_SCALED", }; --static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); -+static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); +-static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); ++static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT != 108"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", -@@ -1233,9 +1236,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { - - "quantize_i8_convrot(x)", +@@ -1238,9 +1241,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "sage_attn(x,y,z)", + "sol_attn(x,y,z)", + "rope_apply(x,pe)", + "rope_pe_permute(x,pe)", + "modulate_rows(x,a,b)", + "swiglu_scaled(x,y)", }; --static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); -+static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); +-static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); ++static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT != 108"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); -@@ -5557,6 +5563,125 @@ struct ggml_tensor * ggml_sage_attn( +@@ -5699,6 +5705,125 @@ struct ggml_tensor* ggml_sol_attn( return result; } @@ -859,10 +859,10 @@ index dc4a29f9..fec933bf 100644 struct ggml_tensor * a, enum ggml_prec prec) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index dd290caf..1e46361c 100644 +index 6a6954fe..6b49765d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp -@@ -2125,6 +2125,118 @@ struct test_glu : public test_case { +@@ -2188,6 +2188,118 @@ struct test_glu : public test_case { } }; @@ -981,7 +981,7 @@ index dd290caf..1e46361c 100644 struct test_glu_split : public test_case { const ggml_glu_op op; const ggml_type type; -@@ -8101,6 +8213,33 @@ static std::vector> make_test_cases_eval() { +@@ -9106,6 +9218,33 @@ static std::vector> make_test_cases_eval() { } } diff --git a/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch b/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch index 65b336de54..ce6492e9ae 100644 --- a/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch +++ b/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch @@ -1,4 +1,4 @@ -From 092f5d51d0844f7cf282ca115a81fccdde9fafad Mon Sep 17 00:00:00 2001 +From b8666f9e06e9e07255b8bfc92c42988716b7b865 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 23:16:25 -0700 Subject: [PATCH] ggml : rope_pe_permute with the per-head RMS norm folded in @@ -19,10 +19,10 @@ reported unsupported and callers keep the separate norm. test-backend-ops cases 5 files changed, 174 insertions(+), 5 deletions(-) diff --git a/include/ggml.h b/include/ggml.h -index 153eed2a..8134626f 100644 +index 8c92967a..41ac28bb 100644 --- a/include/ggml.h +++ b/include/ggml.h -@@ -2483,6 +2483,18 @@ extern "C" { +@@ -2573,6 +2573,18 @@ extern "C" { float scale, enum ggml_type type); @@ -42,10 +42,10 @@ index 153eed2a..8134626f 100644 // different modulation vectors. Writes rows [row0, row1) along dim 1 (all of dims 2 and 3) of the // result and leaves every other row as it is in `acc`: chain one call per segment, passing the diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp -index 40fe1949..1d67b482 100644 +index f4ca82e7..bdd36aea 100644 --- a/src/ggml-cpu/ops.cpp +++ b/src/ggml-cpu/ops.cpp -@@ -12151,6 +12151,8 @@ void ggml_compute_forward_rope_pe_permute( +@@ -12395,6 +12395,8 @@ void ggml_compute_forward_rope_pe_permute( const int n_rot = ggml_get_op_params_i32(dst, 0); const float scale = ggml_get_op_params_f32(dst, 1); @@ -54,7 +54,7 @@ index 40fe1949..1d67b482 100644 const int half = n_rot / 2; const int64_t D = src0->ne[0]; -@@ -12165,7 +12167,7 @@ void ggml_compute_forward_rope_pe_permute( +@@ -12409,7 +12411,7 @@ void ggml_compute_forward_rope_pe_permute( // Same rounding points as the unfused graph: each product, the sum, the scale and the // conversion are separate passes, so no contraction into fused multiply-adds can happen. @@ -63,7 +63,7 @@ index 40fe1949..1d67b482 100644 for (int64_t ir = ir0; ir < ir1; ++ir) { const int64_t h = ir % H; -@@ -12174,6 +12176,19 @@ void ggml_compute_forward_rope_pe_permute( +@@ -12418,6 +12420,19 @@ void ggml_compute_forward_rope_pe_permute( const float * x = (const float *) ((const char *) src0->data + h * src0->nb[1] + l * src0->nb[2] + n * src0->nb[3]); @@ -228,10 +228,10 @@ index 3afe4423..b6dd7a69 100644 return pe == nullptr || (pe->type == GGML_TYPE_F32 && ggml_is_contiguous(pe)); } diff --git a/src/ggml.c b/src/ggml.c -index fec933bf..62a157c6 100644 +index 165514d9..1fe59a0e 100644 --- a/src/ggml.c +++ b/src/ggml.c -@@ -5594,6 +5594,21 @@ struct ggml_tensor * ggml_rope_pe_permute( +@@ -5736,6 +5736,21 @@ struct ggml_tensor * ggml_rope_pe_permute( return result; } @@ -254,10 +254,10 @@ index fec933bf..62a157c6 100644 struct ggml_context * ctx, struct ggml_tensor * acc, diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index 1e46361c..b58c815c 100644 +index 6b49765d..af772ae9 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp -@@ -2132,15 +2132,16 @@ struct test_rope_pe_permute : public test_case { +@@ -2195,15 +2195,16 @@ struct test_rope_pe_permute : public test_case { const int n_rot; const float scale; const bool view; // a is a strided view (one third of a fused qkv projection) @@ -277,7 +277,7 @@ index 1e46361c..b58c815c 100644 ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * a; -@@ -2159,7 +2160,8 @@ struct test_rope_pe_permute : public test_case { +@@ -2222,7 +2223,8 @@ struct test_rope_pe_permute : public test_case { pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, n_rot / 2, ne_a[2]); ggml_set_name(pe, "pe"); } @@ -287,7 +287,7 @@ index 1e46361c..b58c815c 100644 ggml_set_name(out, "out"); return out; } -@@ -8221,6 +8223,11 @@ static std::vector> make_test_cases_eval() { +@@ -9226,6 +9228,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 0, scale, view)); test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 32, scale, view)); test_cases.emplace_back(new test_rope_pe_permute(type, {200, 2, 3, 1}, 100, scale, view)); diff --git a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch index 2b23434e88..68aecc56f6 100644 --- a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch +++ b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch @@ -1,4 +1,4 @@ -From 784b441f8b098863343e69e0345ef62d2ff580af Mon Sep 17 00:00:00 2001 +From 1602a3095b12f20b584066ac88c37a9b2ad07f1d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 3 Oct 2026 18:27:29 -0700 Subject: [PATCH] ggml-cuda: optional cuDNN attention for unmasked @@ -50,11 +50,11 @@ lengths, GQA with batch, F32 K/V, permuted Q) and unmasked perf cases. create mode 100644 src/ggml-cuda/cudnn-attn.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt -index b6843578..1c32cf3b 100644 +index 752dabbb..91e94ba3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -208,6 +208,9 @@ option(GGML_CUDA_FA "ggml: compile ggml FlashAttention C - option(GGML_CUDA_FA_ALL_QUANTS "ggml: compile all quants for FlashAttention" OFF) +@@ -208,6 +208,9 @@ set (GGML_CUDA_FA_QUANTS "q4_0-q4_0;q8_0-q8_0;f16-f16;bf16-bf16" CACHE STRING + "ggml: FlashAttention K-V type combinations to compile, \"all\" or a list such as \"q8_0-q8_0;q8_0-q4_0\"") option(GGML_CUDA_GRAPHS "ggml: use CUDA graphs (llama.cpp only)" ${GGML_CUDA_GRAPHS_DEFAULT}) option(GGML_CUDA_NCCL "ggml: use NVIDIA Collective Comm. Library" ON) +option(GGML_CUDA_CUDNN "ggml: cuDNN attention (cudnn-frontend, libcudnn opened at run time)" OFF) @@ -64,12 +64,12 @@ index b6843578..1c32cf3b 100644 "ggml: cuda link binary compression mode; requires cuda 12.8+") set_property(CACHE GGML_CUDA_COMPRESSION_MODE PROPERTY STRINGS "none;speed;balance;size") diff --git a/src/ggml-cuda/CMakeLists.txt b/src/ggml-cuda/CMakeLists.txt -index 05b81487..0ddeb81e 100644 +index d272f6f7..ba2a2080 100644 --- a/src/ggml-cuda/CMakeLists.txt +++ b/src/ggml-cuda/CMakeLists.txt -@@ -124,11 +124,60 @@ if (CUDAToolkit_FOUND) - template-instances/fattn-vec-instance-bf16-bf16.cu) - endif() +@@ -115,11 +115,60 @@ if (CUDAToolkit_FOUND) + ggml_cuda_fattn_vec_instances(${CMAKE_CURRENT_SOURCE_DIR} SRCS) + list(APPEND GGML_SOURCES_CUDA ${SRCS}) + if (GGML_CUDA_CUDNN AND WIN32) + message(WARNING "GGML_CUDA_CUDNN: the run-time cuDNN loader is Linux only, building without it") @@ -852,7 +852,7 @@ index 00000000..067bea63 +// so the caller runs its own kernel. +bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu -index b8be30a0..ddc93e30 100644 +index f0fe5577..2eda0ccd 100644 --- a/src/ggml-cuda/fattn.cu +++ b/src/ggml-cuda/fattn.cu @@ -4,6 +4,7 @@ @@ -861,9 +861,9 @@ index b8be30a0..ddc93e30 100644 #include "fattn.cuh" +#include "cudnn-attn.cuh" - template - static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { -@@ -573,11 +574,15 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d + #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + // one list per group of ncols1 queries: a column is selected if any query of the group can see it +@@ -758,11 +759,15 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d const ggml_cuda_flash_attn_ext_f16_extra_data f16_extra = ggml_cuda_flash_attn_ext_get_f16_extra_data(dst, need_f16_K, need_f16_V); @@ -881,18 +881,18 @@ index b8be30a0..ddc93e30 100644 case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index 86a6cc5b..749aead9 100644 +index 7e98d729..831d0de8 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -919,7 +919,7 @@ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_t +@@ -922,7 +922,7 @@ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_t ? ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor) : ggml_nbytes(tensor); if (tensor->op == GGML_OP_SAGE_ATTN) { - size = ggml_cuda_sage_attn_get_alloc_size(tensor); + size = ggml_cuda_sage_attn_get_alloc_size(buft_ctx->device, tensor); } - int64_t ne0 = tensor->ne[0]; - + if (tensor->op == GGML_OP_SOL_ATTN) { + size = ggml_cuda_sol_attn_get_alloc_size(tensor); diff --git a/src/ggml-cuda/sage-attn.cu b/src/ggml-cuda/sage-attn.cu index 2c103eb0..ceb29ed2 100644 --- a/src/ggml-cuda/sage-attn.cu @@ -947,12 +947,12 @@ index 126c7a74..e2648857 100644 #endif diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index b58c815c..55db14a0 100644 +index af772ae9..778bed20 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp -@@ -9773,6 +9773,21 @@ static std::vector> make_test_cases_eval() { - GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); - } +@@ -11099,6 +11099,21 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // unmasked F16 K/V with the MiniMax-H3 DiT head layout (42 heads x 128), the shape class the CUDA + // cuDNN attention path takes: odd and tile-aligned lengths, GQA with batch, F32 K/V, head size 64. @@ -972,9 +972,9 @@ index b58c815c..55db14a0 100644 test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); -@@ -10124,6 +10139,16 @@ static std::vector> make_test_cases_perf() { - // Qwen3-VL-8B https://github.com/ggml-org/llama.cpp/issues/17012 - test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); +@@ -11527,6 +11542,16 @@ static std::vector> make_test_cases_perf() { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); + } + // unmasked diffusion attention: MiniMax-H3 DiT (19315 tokens x 56 heads at 960x544x124; 22440 x 42), + // its token refiner (31 tokens) and video VAE decoder (1797 x 32 x 64), FLUX-class 1024px (4096 x 24), diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index db5ad31980..0062c4e3e9 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -16,6 +16,7 @@ #include "model/te/clip.hpp" #include "model/te/llada_image_te.hpp" #include "model/te/llm.hpp" +#include "model/te/ming_image_te.hpp" #include "model/te/t5.hpp" #include "model_loader.h" #include "tokenizers/sensenova_u1_tokenizer.h" @@ -398,7 +399,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -816,7 +817,7 @@ struct SD3CLIPEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -908,6 +909,31 @@ struct SD3CLIPEmbedder : public Conditioner { size_t chunk_count = std::max(std::max(clip_l_tokens.size(), clip_g_tokens.size()), t5_tokens.size()) / chunk_len; + // Each active encoder must cover every chunk of the longest token sequence. + Tokenizer* tokenizers[] = { + clip_l ? clip_l_tokenizer.get() : nullptr, + clip_g ? clip_g_tokenizer.get() : nullptr, + t5 ? &t5_tokenizer : nullptr, + }; + const size_t target_length = chunk_count * chunk_len; + + for (size_t encoder_idx = 0; encoder_idx < 3; encoder_idx++) { + auto& tokens = token_and_weights[encoder_idx].first; + auto& weights = token_and_weights[encoder_idx].second; + if (!tokenizers[encoder_idx] || tokens.size() >= target_length) { + continue; + } + + std::vector empty_tokens; + std::vector empty_weights; + tokenizers[encoder_idx]->pad_tokens(empty_tokens, &empty_weights, nullptr, chunk_len, chunk_len, false); + + while (tokens.size() < target_length) { + tokens.insert(tokens.end(), empty_tokens.begin(), empty_tokens.end()); + weights.insert(weights.end(), empty_weights.begin(), empty_weights.end()); + } + } + for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) { // clip_l sd::Tensor chunk_hidden_states_l; @@ -1206,7 +1232,7 @@ struct FluxCLIPEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -1483,7 +1509,7 @@ struct T5CLIPEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -1872,7 +1898,7 @@ struct AnimaConditioner : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str()); } std::vector qwen_tokens; @@ -1986,9 +2012,10 @@ struct LLMEmbedder : public Conditioner { sd_version_is_sefi_image(version) || sd_version_is_krea2(version) || sd_version_is_minimax_h3(version) || - sd_version_is_mage_flow(version)) { + sd_version_is_mage_flow(version) || + version == VERSION_IRIS) { arch = LLM::LLMArch::QWEN3_VL; - } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) { + } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) { arch = LLM::LLMArch::QWEN3; } llm = std::make_shared(arch, @@ -2112,9 +2139,10 @@ struct LLMEmbedder : public Conditioner { std::tuple, std::vector, std::vector> tokenize(std::string text, const std::pair& attn_range, - size_t min_length = 0, - size_t max_length = 100000000, - bool spell_quotes = false) { + size_t min_length = 0, + size_t max_length = 100000000, + bool spell_quotes = false, + const std::string& suffix = "") { std::vector> parsed_attention; if (attn_range.first >= 0 && attn_range.second > 0) { if (attn_range.first > 0) { @@ -2143,7 +2171,7 @@ struct LLMEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str()); } std::vector tokens; @@ -2159,6 +2187,20 @@ struct LLMEmbedder : public Conditioner { weights.insert(weights.end(), curr_tokens.size(), curr_weight); } + if (!suffix.empty()) { + std::vector suffix_tokens; + if (!tokenizer->encode(suffix, suffix_tokens) || suffix_tokens.size() >= max_length) { + return {}; + } + // Reserve the assistant-turn marker before truncating the caption. + if (tokens.size() > max_length - suffix_tokens.size()) { + tokens.resize(max_length - suffix_tokens.size()); + weights.resize(tokens.size()); + } + tokens.insert(tokens.end(), suffix_tokens.begin(), suffix_tokens.end()); + weights.insert(weights.end(), suffix_tokens.size(), 1.f); + } + std::vector mask; tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length); @@ -2181,8 +2223,10 @@ struct LLMEmbedder : public Conditioner { bool spell_quotes = false, int max_length = 100000000, const LLM::DeepStackImageEmbeds& deepstack_image_embeds = {}, - const std::vector& image_grids = {}) { - auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes); + const std::vector& image_grids = {}, + const std::string& prompt_suffix = "", + sd::Tensor* output_mask = nullptr) { + auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix); auto& tokens = std::get<0>(tokens_weights_mask); auto& weights = std::get<1>(tokens_weights_mask); auto& mask = std::get<2>(tokens_weights_mask); @@ -2219,7 +2263,10 @@ struct LLMEmbedder : public Conditioner { false, deepstack_image_embeds, image_grids); - GGML_ASSERT(!hidden_states.empty()); + if (hidden_states.empty()) { + LOG_ERROR("LLM prompt encoding failed"); + return {}; + } hidden_states = apply_token_weights(std::move(hidden_states), weights); GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx); @@ -2242,6 +2289,12 @@ struct LLMEmbedder : public Conditioner { 1); } + if (output_mask != nullptr) { + *output_mask = sd::Tensor::zeros({new_hidden_states.shape()[1]}); + for (size_t i = prompt_template_encode_start_idx; i < mask.size(); ++i) { + output_mask->data()[i - prompt_template_encode_start_idx] = mask[i]; + } + } return new_hidden_states; } @@ -2308,6 +2361,8 @@ struct LLMEmbedder : public Conditioner { int hidden_states_min_length = 0; // zero pad hidden_states bool spell_quotes = false; std::set out_layers; + std::string prompt_suffix; + sd::Tensor text_mask; int64_t t0 = ggml_time_ms(); RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode; @@ -2949,7 +3004,7 @@ struct LLMEmbedder : public Conditioner { prompt_attn_range.second = static_cast(prompt.size()); prompt += "<|end|><|start|>assistant<|channel|>analysis<|message|>Need to generate one image according to the description.<|end|><|start|>assistant<|channel|>final<|message|>"; - } else if (sd_version_is_z_image(version)) { + } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) { prompt_template_encode_start_idx = 0; out_layers = {35}; // -2 @@ -3057,6 +3112,23 @@ struct LLMEmbedder : public Conditioner { SDCondition result; result.c_crossattn = std::move(hidden_states); return result; + } else if (version == VERSION_IRIS) { + prompt = + "<|im_start|>system\n" + "Describe the image by detailing the color, shape, size, texture, quantity, text, spatial " + "relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"; + auto prefix_tokens = std::get<0>(tokenize(prompt, {0, 0})); + if (prefix_tokens.empty()) { + return {}; + } + prompt_template_encode_start_idx = static_cast(prefix_tokens.size()); + hidden_states_min_length = 300; + max_length = prompt_template_encode_start_idx + hidden_states_min_length; + out_layers = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35}; + prompt_attn_range.first = static_cast(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(prompt.size()); + prompt_suffix = "<|im_end|>\n<|im_start|>assistant\n"; } else { GGML_ABORT("unknown version %d", version); } @@ -3072,7 +3144,9 @@ struct LLMEmbedder : public Conditioner { spell_quotes, max_length, deepstack_image_embeds, - image_grids); + image_grids, + prompt_suffix, + version == VERSION_IRIS ? &text_mask : nullptr); if (hidden_states.empty()) { return {}; } @@ -3137,6 +3211,7 @@ struct LLMEmbedder : public Conditioner { SDCondition result; result.c_crossattn = std::move(hidden_states); result.extra_c_crossattns = std::move(extra_hidden_states_vec); + result.c_vector = std::move(text_mask); if (version == VERSION_QWEN_IMAGE_2_1) { auto slots = sd::Tensor::zeros({result.c_crossattn.shape()[1]}); for (size_t i = 0; i < image_embeds.size(); ++i) { @@ -3163,6 +3238,66 @@ struct LLMEmbedder : public Conditioner { int64_t tag_count = static_cast(tags.size()); result.c_token_types = sd::Tensor({tag_count}, std::move(tags)); } + + return result; + } +}; + +struct MingImageEmbedder : public Conditioner { + std::shared_ptr tokenizer; + std::shared_ptr text_model; + const std::string prefix = "text_encoders.llm"; + + MingImageEmbedder(ggml_backend_t backend, const String2TensorStorage& tensors, std::shared_ptr weight_manager, const TokenizerConfig& tokenizers) { + if (!tokenizers.has(TokenizerConfig::MAIN)) { + throw std::runtime_error("Ming-Image requires the Ling tokenizer.json; pass --tokenizer FILE"); + } + text_model = std::make_shared(backend, tensors, prefix, weight_manager); + tokenizer = tokenizers.create(TokenizerConfig::MAIN, text_model->config.backbone.vocab_size, 156895); + } + + void get_param_tensors(std::map& tensors) override { + text_model->get_param_tensors(tensors, prefix); + } + void get_param_tensor_ops(std::map& ops) override { + text_model->get_param_tensor_ops(ops); + } + void get_layer_split_param_tensors(std::map& tensors) override { + text_model->get_param_tensors(tensors, prefix); + } + void set_flash_attention_enabled(bool enabled) override { text_model->set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t bytes) override { text_model->set_max_graph_vram_bytes(bytes); } + void set_runtime_backends(const std::vector& backends) override { text_model->set_runtime_backends(backends); } + void set_graph_cut_layer_split_enabled(bool enabled) override { text_model->set_graph_cut_layer_split_enabled(enabled); } + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { text_model->set_graph_cut_layer_split_backend_vram_limits(limits); } + void set_scale_overrides(float linear, float attention) override { text_model->set_scale_overrides(linear, attention); } + void set_weight_adapter(const std::shared_ptr& adapter) override { text_model->set_weight_adapter(adapter); } + void runner_end() override { text_model->runner_end(); } + + SDCondition get_learned_condition(int n_threads, const ConditionerParams& input) override { + if (input.ref_images != nullptr && !input.ref_images->empty()) { + LOG_ERROR("Ming-Image currently supports text-to-image only"); + return {}; + } + std::string prompt = + "SYSTEM你是一个友好的AI助手。\n\ndetailed thinking off<|role_end|>" + "HUMAN" + + input.text + "<|role_end|>ASSISTANT"; + std::vector tokens; + if (!tokenizer->encode(prompt, tokens, nullptr)) { + return {}; + } + if (tokens.size() + 258 > 32768) { + LOG_ERROR("Ming-Image prompt exceeds the text encoder context length"); + return {}; + } + auto output = text_model->compute(n_threads, tokens); + if (output.empty()) { + return {}; + } + SDCondition result; + result.c_crossattn = sd::ops::slice(sd::ops::slice(output, 1, 0, 256), 0, 0, text_model->config.caption_dim); + result.extra_c_crossattns.push_back(sd::ops::slice(output, 1, 256, output.shape()[1])); return result; } }; diff --git a/src/conditioning/conditioning_cache.h b/src/conditioning/conditioning_cache.h new file mode 100644 index 0000000000..6e6cbcaf77 --- /dev/null +++ b/src/conditioning/conditioning_cache.h @@ -0,0 +1,107 @@ +#ifndef __SD_CONDITIONING_CONDITIONING_CACHE_H__ +#define __SD_CONDITIONING_CONDITIONING_CACHE_H__ + +#include +#include +#include + +#include "conditioning/conditioner.hpp" + +class ConditioningCache { + struct Entry { + ConditionerParams params; + std::vector> ref_images; + std::vector 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 entries_; + + static bool same_images(const std::vector>& a, + const std::vector>& b) { + return std::equal(a.begin(), a.end(), b.begin(), b.end(), + [](const sd::Tensor& x, const sd::Tensor& 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__ diff --git a/src/convert.cpp b/src/convert.cpp index c82d2557ac..7fb0825f5c 100644 --- a/src/convert.cpp +++ b/src/convert.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -11,11 +12,15 @@ #include #include +#include "core/ggml_extend_backend.h" #include "core/util.h" #include "model_io/gguf_io.h" #include "model_io/safetensors_io.h" #include "model_io/streaming_writer.h" #include "model_loader.h" +#include "model_manager.h" +#include "name_conversion.h" +#include "runtime/imatrix.h" struct TensorExportInfo { TensorStorage storage; @@ -75,9 +80,9 @@ static size_t export_tensor_nbytes(const TensorExportInfo& info) { return static_cast(output_storage.nbytes()); } -static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo& info) { +static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo& info, const std::string& name) { TensorWritePlan plan; - plan.name = info.storage.name; + plan.name = name; plan.type = info.type; plan.n_dims = info.storage.n_dims; for (int i = 0; i < SD_MAX_DIMS; i++) { @@ -86,11 +91,14 @@ static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo return plan; } -static std::vector tensor_write_plans_from_export_infos(const std::vector& tensors) { +static std::vector tensor_write_plans_from_export_infos(const std::vector& tensors, + bool convert_name, + SDVersion version) { std::vector plans; plans.reserve(tensors.size()); for (const TensorExportInfo& info : tensors) { - plans.push_back(tensor_write_plan_from_export_info(info)); + const std::string name = convert_name ? convert_tensor_name(info.storage.name, version) : info.storage.name; + plans.push_back(tensor_write_plan_from_export_info(info, name)); } return plans; } @@ -158,16 +166,42 @@ static bool load_tensor_for_export(ModelLoader& model_loader, TensorExportJob& j return true; } +static bool export_tensor_from_memory(const TensorExportInfo& info, + const std::map& mem, + SDVersion version, + std::vector& out) { + auto found = mem.find(info.storage.name); + if (found == mem.end() || found->second == nullptr || found->second->data == nullptr) { + return false; + } + const TensorStorage& storage = info.storage; + out.resize(export_tensor_nbytes(info)); + if (out.size() > 0) { + std::vector imatrix = get_imatrix_collector().get_values(convert_tensor_name(info.storage.name, version)); + convert_tensor(found->second->data, + found->second->type, + out.data(), + info.type, + (int)(storage.nelements() / storage.ne[0]), + (int)storage.ne[0], + std::move(imatrix)); + } + return true; +} + static bool stream_tensor_data(ModelLoader& model_loader, const std::string& output_path, const std::vector& tensors, const StreamingModelWriter& writer, int n_threads, + const std::map* mem, std::string* error) { n_threads = n_threads > 0 ? n_threads : sd_get_num_physical_cores(); n_threads = std::max(1, n_threads); LOG_INFO("streaming convert with %d threads", n_threads); + const SDVersion version = model_loader.get_sd_version(); + int64_t start_time = ggml_time_ms(); uint64_t bytes_written = 0; size_t tensors_written = 0; @@ -241,7 +275,10 @@ static bool stream_tensor_data(ModelLoader& model_loader, TensorExportJob job; job.info = tensors[tensor_index]; try { - job.success = load_tensor_for_export(model_loader, job); + job.success = (mem != nullptr) && export_tensor_from_memory(job.info, *mem, version, job.data); + if (!job.success) { + job.success = load_tensor_for_export(model_loader, job); + } } catch (const std::exception& e) { job.error = e.what(); job.success = false; @@ -298,8 +335,10 @@ static bool write_model_file_streaming(ModelLoader& model_loader, const std::vector& tensors, StreamingModelWriter& writer, int n_threads, + bool convert_name, + const std::map* mem, std::string* error) { - std::vector plans = tensor_write_plans_from_export_infos(tensors); + std::vector plans = tensor_write_plans_from_export_infos(tensors, convert_name, model_loader.get_sd_version()); if (!writer.write_metadata(output_path, plans, error)) { return false; } @@ -307,7 +346,7 @@ static bool write_model_file_streaming(ModelLoader& model_loader, return false; } model_loader.process_model_files(false, false); - return stream_tensor_data(model_loader, output_path, tensors, writer, n_threads, error); + return stream_tensor_data(model_loader, output_path, tensors, writer, n_threads, mem, error); } static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) { @@ -326,7 +365,9 @@ static bool export_loaded_model(ModelLoader& model_loader, const char* output_path, sd_type_t output_type, const char* tensor_type_rules, - int n_threads) { + int n_threads, + bool convert_name, + const std::map* mem = nullptr) { ggml_type type = sd_type_to_ggml_type(output_type); bool output_is_safetensors = ends_with(output_path, ".safetensors"); TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules); @@ -341,7 +382,7 @@ static bool export_loaded_model(ModelLoader& model_loader, } else { writer = std::make_unique(); } - success = write_model_file_streaming(model_loader, output_path, tensors, *writer, n_threads, &error); + success = write_model_file_streaming(model_loader, output_path, tensors, *writer, n_threads, convert_name, mem, &error); } if (!success && !error.empty()) { @@ -351,6 +392,174 @@ static bool export_loaded_model(ModelLoader& model_loader, return success; } +static bool has_active_loras(const sd_lora_t* loras, int lora_count) { + if (loras == nullptr) { + return false; + } + for (int i = 0; i < lora_count; i++) { + if (loras[i].multiplier != 0.0f) { + return true; + } + } + return false; +} + +static std::vector build_lora_specs(const sd_lora_t* loras, int lora_count) { + std::vector specs; + specs.reserve(lora_count); + for (int i = 0; i < lora_count; i++) { + ModelManager::LoraSpec spec; + spec.path = loras[i].path != nullptr ? loras[i].path : ""; + spec.multiplier = loras[i].multiplier; + spec.is_high_noise = loras[i].is_high_noise; + spec.required = true; + specs.push_back(std::move(spec)); + } + return specs; +} + +static ModelComponent component_for_tensor_name(const std::string& name) { + if (starts_with(name, "text_encoders.") || starts_with(name, "cond_stage_model.")) { + return ModelComponent::Conditioner; + } + if (starts_with(name, "model.high_noise_diffusion_model.")) { + return ModelComponent::HighNoiseDiffusion; + } + if (starts_with(name, "model.diffusion_model.")) { + return ModelComponent::Diffusion; + } + if (starts_with(name, "vae.") || starts_with(name, "first_stage_model.")) { + return ModelComponent::VAE; + } + return ModelComponent::Diffusion; +} + +static bool load_model_into_memory(ModelLoader& model_loader, + const sd_lora_t* loras, + int lora_count, + int n_threads, + ModelManager& manager, + ggml_backend_t cpu, + ggml_context* ctx, + std::map& mem, + std::vector& all_tensors, + ggml_type output_type, + const TensorTypeRules& type_rules) { + const String2TensorStorage& storage_map = model_loader.get_tensor_storage_map(); + if (storage_map.empty()) { + LOG_ERROR("no tensors to load into memory for convert"); + return false; + } + + const SDVersion version = model_loader.get_sd_version(); + ModelLoader merge_loader = model_loader; + merge_loader.convert_tensors_name(); + + std::map> groups; + size_t total_bytes = 0; + for (const auto& [name, storage] : storage_map) { + ggml_type load_type = storage.type; + const ggml_type export_type = get_export_tensor_type(model_loader, storage, output_type, type_rules); + const bool quant_src = ggml_is_quantized(storage.type); + const bool quant_dst = ggml_is_quantized(export_type); + if (quant_src && quant_dst) { + load_type = GGML_TYPE_F16; + } else if (quant_src || quant_dst) { + load_type = quant_src ? export_type : storage.type; + } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32) { + load_type = GGML_TYPE_F32; + } + ggml_tensor* tensor = ggml_new_tensor(ctx, load_type, storage.n_dims, storage.ne); + if (tensor == nullptr) { + LOG_ERROR("failed to create tensor '%s' for in-memory convert", name.c_str()); + return false; + } + const std::string merge_name = convert_tensor_name(name, version); + auto& group = groups[component_for_tensor_name(merge_name)]; + if (!group.emplace(merge_name, tensor).second) { + LOG_ERROR("duplicate tensor name '%s' after name conversion", merge_name.c_str()); + return false; + } + mem[name] = tensor; + all_tensors.push_back(tensor); + total_bytes += ggml_nbytes(tensor); + } + LOG_INFO("loading %zu tensors (%.2f GB) into memory for convert", + all_tensors.size(), total_bytes / (1024.0 * 1024.0 * 1024.0)); + + manager.set_n_threads(n_threads); + manager.set_enable_mmap(false); + if (!manager.set_loader(std::move(merge_loader))) { + LOG_ERROR("failed to set model loader for in-memory convert"); + return false; + } + + for (const auto& [component, tensors] : groups) { + if (!manager.register_param_tensors(component, tensors, ModelManager::ResidencyMode::ParamBackend, cpu, cpu)) { + LOG_ERROR("failed to register %s tensors for in-memory convert", model_component_name(component)); + return false; + } + } + + std::vector lora_specs = build_lora_specs(loras, lora_count); + if (!manager.prepare_lora_sources(lora_specs)) { + LOG_ERROR("failed to prepare LoRA sources for convert"); + return false; + } + if (!manager.set_loras(lora_specs, version)) { + LOG_ERROR("failed to set LoRAs for convert"); + return false; + } + if (!manager.prepare_params(all_tensors)) { + LOG_ERROR("failed to load model into memory for convert"); + return false; + } + + return true; +} + +static bool convert_model_in_memory(ModelLoader& model_loader, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name, + int n_threads, + const sd_lora_t* loras, + int lora_count) { + auto cpu = std::unique_ptr(sd_backend_cpu_init(), ggml_backend_free); + if (cpu == nullptr) { + LOG_ERROR("failed to init CPU backend for in-memory convert"); + return false; + } + + ggml_init_params ctx_params; + ctx_params.mem_size = model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead(); + ctx_params.mem_buffer = nullptr; + ctx_params.no_alloc = true; + ggml_context* ctx = ggml_init(ctx_params); + if (ctx == nullptr) { + LOG_ERROR("ggml_init failed for in-memory convert"); + return false; + } + + ggml_type type = sd_type_to_ggml_type(output_type); + TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules); + + std::vector all_tensors; + std::map mem; + bool success = false; + { + // The manager owns the buffers backing `mem` and must outlive the export; it is + // destroyed (freeing those buffers) before the tensor context is freed below. + ModelManager manager; + if (load_model_into_memory(model_loader, loras, lora_count, n_threads, manager, cpu.get(), ctx, mem, all_tensors, type, type_rules)) { + success = export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads, convert_name, &mem); + } + } + ggml_free(ctx); + return success; +} + bool convert_with_components(const char* model_path, const char* clip_l_path, const char* clip_g_path, @@ -361,10 +570,22 @@ bool convert_with_components(const char* model_path, sd_type_t output_type, const char* tensor_type_rules, bool convert_name, - int n_threads) { + int n_threads, + const sd_lora_t* loras, + int lora_count) { if (!validate_tensor_types(output_type, tensor_type_rules)) { return false; } + + if (loras != nullptr) { + for (int i = 0; i < lora_count; i++) { + LOG_INFO("lora %d: '%s' (multiplier %.2f%s)", + i + 1, loras[i].path != nullptr ? loras[i].path : "", + loras[i].multiplier, + loras[i].is_high_noise ? ", high noise" : ""); + } + } + ModelLoader model_loader; bool loaded_any = false; @@ -382,11 +603,17 @@ bool convert_with_components(const char* model_path, return false; } + if (has_active_loras(loras, lora_count)) { + LOG_INFO("loading the full model into memory to apply LoRAs"); + return convert_model_in_memory(model_loader, output_path, output_type, tensor_type_rules, convert_name, n_threads, + loras, lora_count); + } + if (convert_name) { model_loader.convert_tensors_name(); } - return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads); + return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads, false, nullptr); } bool convert(const char* input_path, @@ -405,5 +632,7 @@ bool convert(const char* input_path, output_type, tensor_type_rules, convert_name, + 0, + nullptr, 0); } diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 0eddb568ad..6698fedf74 100644 --- a/src/core/backend_fit.cpp +++ b/src/core/backend_fit.cpp @@ -478,57 +478,61 @@ namespace sd::backend_fit { return true; } - bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { - return prepare_vae_retry_tiling(tiling_params, prefer_temporal_tiling, false); + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, + bool prefer_temporal_tiling, + ggml_status status, + int latent_tile_size_w, + int latent_tile_size_h, + int scale_factor) { + return prepare_vae_retry_tiling(tiling_params, prefer_temporal_tiling, status, + latent_tile_size_w, latent_tile_size_h, scale_factor, false); } - bool prepare_vae_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, bool encode) { - // Spatial retry tiles are a fraction of the latent: half of it for decode. Image VAE encode tiles are - // scaled up 2x by get_tile_sizes, so encode starts a step lower to actually split the image. - // get_tile_sizes keeps the tiles overlapping (sd_tiling_seam_safe_tile_size), so a fractional tile never - // collapses into two tiles that meet edge to edge. If a retry still fails, the tiles are halved again - // down to an eighth of the latent before giving up. - const float first_rel_size = encode ? 0.25f : 0.5f; - const float min_rel_size = 0.125f; - const char* stage = encode ? "encode" : "decode"; - const char* retry_mode = nullptr; + bool prepare_vae_retry_tiling(sd_tiling_params_t& tiling_params, + bool prefer_temporal_tiling, + ggml_status status, + int latent_tile_size_w, + int latent_tile_size_h, + int scale_factor, + bool encode) { + // Execution failures can leave the device unusable; tiling only helps with allocation failures. + if (status != GGML_STATUS_ALLOC_FAILED) { + return false; + } + const char* stage = encode ? "encode" : "decode"; + const char* retry_mode = nullptr; if (prefer_temporal_tiling && !tiling_params.temporal_tiling) { tiling_params.temporal_tiling = true; retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal"; - } else if (!tiling_params.enabled) { - tiling_params.enabled = true; - tiling_params.rel_size_x = first_rel_size; - tiling_params.rel_size_y = first_rel_size; - if (tiling_params.tile_size_x <= 0) { - tiling_params.tile_size_x = 256; - } - if (tiling_params.tile_size_y <= 0) { - tiling_params.tile_size_y = 256; - } - if (tiling_params.target_overlap <= 0.f) { - tiling_params.target_overlap = 0.5f; + } else { + if (latent_tile_size_w <= 0 || latent_tile_size_h <= 0 || scale_factor <= 0) { + return false; } - retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial"; - } else if (tiling_params.rel_size_x > 0.f && tiling_params.rel_size_x <= 1.f && - tiling_params.rel_size_y > 0.f && tiling_params.rel_size_y <= 1.f && - std::max(tiling_params.rel_size_x, tiling_params.rel_size_y) > min_rel_size) { - // An axis already below the floor keeps its size: clamping it up would grow the tile that just failed. - auto shrink = [&](float rel) { return rel > min_rel_size ? std::max(rel * 0.5f, min_rel_size) : rel; }; - tiling_params.rel_size_x = shrink(tiling_params.rel_size_x); - tiling_params.rel_size_y = shrink(tiling_params.rel_size_y); - if (tiling_params.target_overlap <= 0.f) { - tiling_params.target_overlap = 0.5f; + auto smaller_tile = [&](int size) { + int next_size = size / 2; + if (!tiling_params.enabled) { + next_size = std::min(next_size, 256 / scale_factor); + } + return std::min(size, std::max(4, next_size)); + }; + const int tile_size_w = smaller_tile(latent_tile_size_w); + const int tile_size_h = smaller_tile(latent_tile_size_h); + if (tile_size_w == latent_tile_size_w && tile_size_h == latent_tile_size_h) { + return false; } - LOG_WARN("VAE %s failed (likely out of memory); retrying with smaller tiles (%.3g x %.3g of the latent)", + tiling_params.enabled = true; + tiling_params.rel_size_w = 0.0f; + tiling_params.rel_size_h = 0.0f; + tiling_params.tile_size_w = tile_size_w * scale_factor; + tiling_params.tile_size_h = tile_size_h * scale_factor; + retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial"; + LOG_WARN("Reducing VAE %s tiles from %dx%d to %dx%d image pixels", stage, - tiling_params.rel_size_x, - tiling_params.rel_size_y); - return true; - } else { - return false; + latent_tile_size_w * scale_factor, latent_tile_size_h * scale_factor, + tiling_params.tile_size_w, tiling_params.tile_size_h); } - LOG_WARN("VAE %s failed (likely out of memory); retrying with %s tiling", + LOG_WARN("VAE %s ran out of memory; retrying with %s tiling", stage, retry_mode); return true; diff --git a/src/core/backend_fit.h b/src/core/backend_fit.h index 6db5315847..40b5bbdf52 100644 --- a/src/core/backend_fit.h +++ b/src/core/backend_fit.h @@ -16,12 +16,19 @@ namespace sd::backend_fit { std::string& params_spec); bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, - bool prefer_temporal_tiling); + bool prefer_temporal_tiling, + ggml_status status, + int latent_tile_size_w, + int latent_tile_size_h, + int scale_factor); - // Turns on (or tightens) VAE tiling after an encode/decode failed, most likely out of memory. - // Returns false once there is nothing left to try. + // Same as prepare_vae_decode_retry_tiling, for an encode when encode is true (only the log differs). bool prepare_vae_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, + ggml_status status, + int latent_tile_size_w, + int latent_tile_size_h, + int scale_factor, bool encode); } // namespace sd::backend_fit diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 58130582d4..585e542c4e 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -218,13 +218,13 @@ ggml_tensor* ggml_ext_linear_matmul(ggml_context* ctx, x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]); x = ggml_mul_mat(ctx, w, x); if (force_prec_f32) { - ggml_mul_mat_set_prec(x, GGML_PREC_F32); + ggml_prec_set_acc(x, GGML_PREC_F32); } x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3); } else { x = ggml_mul_mat(ctx, w, x); if (force_prec_f32) { - ggml_mul_mat_set_prec(x, GGML_PREC_F32); + ggml_prec_set_acc(x, GGML_PREC_F32); } } return x; @@ -263,14 +263,20 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, } ggml_tensor* fused_bias = scale == 1.f ? b : nullptr; + auto mul_mat = [&](ggml_tensor* input) { + if (input->type == GGML_TYPE_F32 && convrot_group_size > 0) { + input = ggml_quantize_i8_convrot(ctx, input, convrot_group_size); + } + return ggml_mul_mat_i8_tensorwise(ctx, w, input, weight_scale, fused_bias, convrot_group_size); + }; if (x->ne[2] * x->ne[3] > 1024) { int64_t ne2 = x->ne[2]; int64_t ne3 = x->ne[3]; x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]); - x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size); + x = mul_mat(x); x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3); } else { - x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size); + x = mul_mat(x); } if (scale != 1.f) { @@ -473,7 +479,11 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d1, int d2, bool force_prec_f32, - bool direct) { + bool direct, + float scale) { + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } if (direct) { int64_t OC = w->ne[3] / IC; int64_t N = x->ne[3] / IC; @@ -486,7 +496,7 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, x = ggml_mul_mat(ctx, ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2] * IC, OC)); - ggml_mul_mat_set_prec(x, GGML_PREC_F32); + ggml_prec_set_acc(x, GGML_PREC_F32); int64_t OD = im2col->ne[3] / N; x = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC); @@ -510,6 +520,9 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, } } + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, 1.f / scale); + } if (b != nullptr) { b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1] x = ggml_add_inplace(ctx, x, b); @@ -631,7 +644,11 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, bool skip_reshape, bool flash_attn, 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_k; int64_t C; @@ -674,8 +691,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, v_in->type == GGML_TYPE_F32 && sd_backend_supports_cuda_mma(backend); if (pad_head) { // CUDA FA MMA starts at 64 channels; keep the original head's attention scale. - q_in = ggml_pad(ctx, q_in, 64 - d_head, 0, 0, 0); - k_in = ggml_pad(ctx, k_in, 64 - d_head, 0, 0, 0); + q_in = ggml_pad(ctx, q_in, static_cast(64 - d_head), 0, 0, 0); + k_in = ggml_pad(ctx, k_in, static_cast(64 - d_head), 0, 0, 0); } if (kv_scale != 1.0f) { k_in = ggml_ext_scale(ctx, k_in, kv_scale); @@ -685,7 +702,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3)); v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N); if (pad_head) { - v_in = ggml_pad(ctx, v_in, 64 - d_head, 0, 0, 0); + v_in = ggml_pad(ctx, v_in, static_cast(64 - d_head), 0, 0, 0); } if (kv_scale != 1.0f) { v_in = ggml_ext_scale(ctx, v_in, kv_scale); @@ -712,7 +729,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, if (!ggml_backend_supports_op(backend, out)) { return nullptr; } - ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_prec_set_acc(out, GGML_PREC_F32); if (kv_scale != 1.0f) { out = ggml_ext_scale(ctx, out, 1.0f / kv_scale); } @@ -733,9 +750,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, } if (padded_head != d_head) { // Keep the original head's softmax scale when padding for the CUDA kernel. - q_in = ggml_pad(ctx, q_in, padded_head - d_head, 0, 0, 0); - k_in = ggml_pad(ctx, k_in, padded_head - d_head, 0, 0, 0); - v_in = ggml_pad(ctx, v_in, padded_head - d_head, 0, 0, 0); + q_in = ggml_pad(ctx, q_in, static_cast(padded_head - d_head), 0, 0, 0); + k_in = ggml_pad(ctx, k_in, static_cast(padded_head - d_head), 0, 0, 0); + v_in = ggml_pad(ctx, v_in, static_cast(padded_head - d_head), 0, 0, 0); } if (kv_scale != 1.0f) { k_in = ggml_ext_scale(ctx, k_in, kv_scale); @@ -763,6 +780,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, if (can_use_flash_attn) { kqv = build_kqv(q, k, v, mask); if (kqv != nullptr) { + if (used_flash_attn != nullptr) { + *used_flash_attn = true; + } kqv = ggml_view_4d(ctx, kqv, d_head, @@ -785,7 +805,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N); // [N * n_kv_head, d_head, L_k] auto kq = ggml_mul_mat(ctx, k, q); // [N * n_head, L_q, L_k] - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_prec_set_acc(kq, GGML_PREC_F32); kq = ggml_scale_inplace(ctx, kq, scale); if (mask) { kq = ggml_add_inplace(ctx, kq, mask); diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index a5076262ff..e94dd1451e 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -160,7 +160,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d1 = 1, int d2 = 1, bool force_prec_f32 = false, - bool direct = false); + bool direct = false, + float scale = 1.f); // w: [OC，IC, KD, 1 * 1] // x: [N, IC, ID, IH*IW] @@ -223,11 +224,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, ggml_tensor* k, ggml_tensor* v, int64_t n_head, - ggml_tensor* mask = nullptr, - bool skip_reshape = false, - bool flash_attn = false, - float kv_scale = 1.0f, - bool sage_attn = false); + ggml_tensor* mask = nullptr, + bool skip_reshape = false, + bool flash_attn = false, + float kv_scale = 1.0f, + bool sage_attn = false, + bool* used_flash_attn = nullptr); // Flash attention on inputs that are already in the head-major layout the kernel reads: // q: F32 [d_head, L_q, n_head * N], k/v: F16 [d_head, L_k, n_kv_head * N] with the kv_scale diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index f3049dc23f..b1bce98475 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -13,6 +13,7 @@ #endif #include "core/util.h" +#include "ggml-backend-impl.h" #include "ggml-impl.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; } +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) { #ifdef SD_USE_CUDA if (!sd_backend_is(backend, "CUDA")) { diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index b2ebd7eb28..df1d692400 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -88,6 +88,10 @@ class SDBackendManager { bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is_cpu(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(); 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, diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 3ac8e0b2dc..c0f7085b28 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, ggml_tensor* mask, bool skip_reshape, bool flash_attn, - float kv_scale) { + float kv_scale, + bool* used_flash_attn) { if (ctx->attn_scale > 0.f) { kv_scale = ctx->attn_scale; } - return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, 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() { @@ -515,9 +516,10 @@ GGMLRunner::~GGMLRunner() { free_params_ctx(); } -GGMLRunnerContext GGMLRunner::get_context() { +GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) { GGMLRunnerContext runner_ctx; runner_ctx.ggml_ctx = compute_ctx; + runner_ctx.graph = graph; runner_ctx.backend = runtime_backend; runner_ctx.flash_attn_enabled = flash_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) { return this->get_cache_tensor_by_name(name); }; - runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) { - this->cache(name, tensor); + runner_ctx.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) { + this->cache(name, tensor, graph); }; runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* 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) { tensor = ggml_cont(compute_ctx, tensor); } @@ -583,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) { ggml_set_output(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> GGMLRunner::compute(get_graph_cb_t get_graph, @@ -590,6 +596,7 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, bool auto_runner_end, bool no_return, const std::function& read_outputs) { + last_compute_status_ = GGML_STATUS_FAILED; if (graph_active_) { LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str()); return std::nullopt; @@ -613,7 +620,9 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, GGMLRunner& runner; const bool& success; ~GraphEndGuard() { - runner.workspace_.segment_end(); + if (!runner.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } runner.cache_.graph_end(false); runner.cut_cache_.clear(); runner.free_compute_ctx(); @@ -641,7 +650,12 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, std::optional> output; try { 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) { + last_compute_status_ = GGML_STATUS_FAILED; LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), ggml_backend_name(runtime_backend), error.what()); return std::nullopt; @@ -649,6 +663,7 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, success = output.has_value(); if (success) { cache_.graph_end(true); + last_compute_status_ = GGML_STATUS_SUCCESS; } return output; } @@ -766,6 +781,7 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) { } workspace_.synchronize(); if (status != GGML_STATUS_SUCCESS) { + last_compute_status_ = status; LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); return false; } @@ -818,13 +834,23 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n const auto& cached_plan = resolve_graph_cut_plan(graph); const auto full_measurement = measure(graph, cached_plan.compute_buffer_size); if (full_measurement.buffers.empty()) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; 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(); const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && manager != nullptr && manager->segmented_compute_enabled() && 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; if (!segmented) { monolithic_plan.segments.emplace_back(); @@ -888,7 +914,9 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n SegmentGraphBindings& bindings; ggml_context* context; ~SegmentCleanup() { - runner.workspace_.segment_end(); + if (!runner.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } bindings.restore(); weights.segment_end(); ggml_free(context); @@ -898,6 +926,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement; if (!workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace preparation"); } const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); @@ -912,7 +941,11 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n sync_runtime_residency(); requests = memory_requests(measurement.buffers, new_cache_bytes); } - return weights.ensure_segment_capacity(index, requests); + const bool ready = weights.ensure_segment_capacity(index, requests); + if (!ready && manager != nullptr) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + } + return ready; }; if (!weights.segment_start(index, ensure_capacity)) { return fail_segment("weight preparation"); @@ -921,12 +954,17 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n if (!workspace_.measurement_matches(segment_graph, measurement)) { measurement = measure(segment_graph, segment.compute_buffer_size); } - if (!workspace_.prepare(measurement) || !ensure_capacity()) { + if (!workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace preparation"); + } + if (!ensure_capacity()) { return fail_segment("workspace capacity check"); } if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { pin_multi_device_nodes(scheduler, current); })) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace allocation"); } for (const auto& size : measurement.buffers) { @@ -945,10 +983,16 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), index + 1, plan.segments.size(), segment.group_name.c_str()); - if (!execute_segment(segment_graph, n_threads) || - !cache_.capture(segment_graph) || - !cut_cache_.capture(graph, segment, get_desc().c_str())) { - return fail_segment("execution or output caching"); + if (!execute_segment(segment_graph, n_threads)) { + return fail_segment("execution"); + } + 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(); if (last) { @@ -964,6 +1008,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } } if (!workspace_.segment_end()) { + last_compute_status_ = GGML_STATUS_FAILED; return fail_segment("workspace synchronization"); } // Final outputs and their callbacks may still be views of consumed cuts. diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index c5e1a65164..5b963982a3 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -67,6 +67,7 @@ struct WeightAdapter { struct GGMLRunnerContext { ggml_backend_t backend = nullptr; ggml_context* ggml_ctx = nullptr; + ggml_cgraph* graph = nullptr; bool flash_attn_enabled = false; bool sage_attn_enabled = false; float linear_scale = 0.f; @@ -102,6 +103,12 @@ struct GGMLRunnerContext { 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 { if (!cache_tensor || tensor == nullptr) { return; @@ -122,15 +129,17 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, ggml_tensor* k, ggml_tensor* v, int64_t n_head, - ggml_tensor* mask = nullptr, - bool skip_reshape = false, - bool flash_attn = false, - float kv_scale = 1.f); + ggml_tensor* mask = nullptr, + bool skip_reshape = false, + bool flash_attn = false, + float kv_scale = 1.f, + bool* used_flash_attn = nullptr); struct GGMLRunner { private: std::map logged_compute_bytes_; - size_t logged_segment_count_ = 0; + size_t logged_segment_count_ = 0; + ggml_status last_compute_status_ = GGML_STATUS_SUCCESS; sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes); std::vector memory_requests(const std::vector& sizes, @@ -289,7 +298,8 @@ struct 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(); @@ -324,7 +334,7 @@ struct GGMLRunner { 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) { return cache_.get(name); @@ -336,6 +346,8 @@ struct GGMLRunner { bool no_return = false, const std::function& read_outputs = {}); + ggml_status last_compute_status() const { return last_compute_status_; } + void set_flash_attention_enabled(bool enabled) { flash_attn_enabled = enabled; } diff --git a/src/core/runner_cache.cpp b/src/core/runner_cache.cpp index 095621068b..b640bb45a6 100644 --- a/src/core/runner_cache.cpp +++ b/src/core/runner_cache.cpp @@ -26,10 +26,13 @@ namespace sd { std::unique_ptr CachedTensor::copy(ggml_backend_t backend, 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) { return nullptr; } + status = GGML_STATUS_ALLOC_FAILED; auto entry = std::make_unique(); entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); if (entry->context == nullptr) { @@ -50,6 +53,7 @@ namespace sd { } else { ggml_backend_tensor_copy(source, entry->tensor); } + status = GGML_STATUS_SUCCESS; return entry; } @@ -106,9 +110,9 @@ namespace sd { 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()) { - return true; + return GGML_STATUS_SUCCESS; } const auto tensors = cache_graph_tensors(graph); for (const auto& output : outputs_) { @@ -116,14 +120,15 @@ namespace sd { continue; } 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) { - return false; + return status; } pending_[output.first] = std::move(entry); } ggml_backend_synchronize(backend_); - return true; + return GGML_STATUS_SUCCESS; } void RunnerCache::graph_end(bool success) { @@ -180,9 +185,9 @@ namespace sd { } } - bool GraphCutTensorCache::capture(ggml_cgraph* graph, - const ggml_graph_cut::Segment& segment, - const char* log_desc) { + ggml_status GraphCutTensorCache::capture(ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment, + const char* log_desc) { size_t copied_bytes = 0; size_t copied_count = 0; for (int index : segment.output_node_indices) { @@ -191,10 +196,11 @@ namespace sd { !segment.future_cut_names.count(output->name)) { 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) { 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); 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_desc, copied_bytes / (1024.f * 1024.f), copied_count); } - return true; + return GGML_STATUS_SUCCESS; } } diff --git a/src/core/runner_cache.h b/src/core/runner_cache.h index bd331a4bbf..f1b5cf91c1 100644 --- a/src/core/runner_cache.h +++ b/src/core/runner_cache.h @@ -20,7 +20,8 @@ namespace sd { ~CachedTensor(); static std::unique_ptr copy(ggml_backend_t backend, const std::string& name, - ggml_tensor* source); + ggml_tensor* source, + ggml_status& status); }; using CachedTensors = std::map>; @@ -41,7 +42,8 @@ namespace sd { const std::map& outputs() const { return outputs_; } size_t pending_bytes(ggml_cgraph* graph) 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 clear(); }; @@ -57,7 +59,7 @@ namespace sd { size_t resident_bytes(ggml_backend_dev_t device) const; size_t estimate_output_bytes(ggml_cgraph* graph, 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& keep_names); void clear() { tensors_.clear(); } }; diff --git a/src/core/util.cpp b/src/core/util.cpp index b90f70ed3f..5292c38ba5 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -62,6 +62,21 @@ void replace_all_chars(std::string& str, char target, char replacement) { } } +std::string escape_newlines(const std::string& text) { + std::string escaped; + escaped.reserve(text.size()); + for (char ch : text) { + if (ch == '\n') { + escaped += "\\n"; + } else if (ch == '\r') { + escaped += "\\r"; + } else { + escaped += ch; + } + } + return escaped; +} + static std::string sd_vformat(const char* fmt, va_list ap) { char small[128]; va_list ap2; @@ -360,12 +375,13 @@ int32_t sd_get_num_physical_cores() { static sd_progress_cb_t sd_progress_cb = nullptr; void* sd_progress_cb_data = nullptr; -static sd_preview_cb_t sd_preview_cb = nullptr; -static void* sd_preview_cb_data = nullptr; -preview_t sd_preview_mode = PREVIEW_NONE; -int sd_preview_interval = 1; -bool sd_preview_denoised = true; -bool sd_preview_noisy = false; +static sd_preview_info_t sd_preview_info = {}; +static sd_preview_cb_t sd_preview_cb = nullptr; +static void* sd_preview_cb_data = nullptr; +preview_t sd_preview_mode = PREVIEW_NONE; +int sd_preview_interval = 1; +bool sd_preview_denoised = true; +bool sd_preview_noisy = false; static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr; static void* sd_backend_eval_cb_data = nullptr; @@ -644,9 +660,14 @@ void* sd_log_cb_data = nullptr; static void sd_log_dispatch(sd_log_level_t level, const std::string& origin, const std::string& text) { if (sd_log_cb == nullptr) return; - std::string message = origin + " - " + text; - if (message.back() != '\n') { - message += '\n'; + std::string message = text; + while (!message.empty() && (message.back() == '\n' || message.back() == '\r')) { + message.pop_back(); + } + if (origin == "ggml") { + message = origin + ": " + message + '\n'; + } else { + message += " --- " + origin + '\n'; } sd_log_cb(level, message.c_str(), sd_log_cb_data); } @@ -656,7 +677,7 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo va_start(args, format); std::string message = sd_vformat(format, args); va_end(args); - std::string origin = sd_format("%s:%-4d", sd_basename(file).c_str(), line); + std::string origin = sd_format("%s:%d", sd_basename(file).c_str(), line); sd_log_dispatch(level, origin, message); } @@ -691,6 +712,7 @@ void sd_set_progress_callback(sd_progress_cb_t cb, void* data) { sd_progress_cb_data = data; } void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, bool denoised, bool noisy, void* data) { + sd_preview_info = {}; sd_preview_cb = cb; sd_preview_cb_data = data; sd_preview_mode = mode; @@ -704,6 +726,17 @@ void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) { sd_backend_eval_cb_data = data; } +void sd_begin_preview_pass(int total_steps) { + if (sd_preview_cb != nullptr) { + ++sd_preview_info.sample_pass; + sd_preview_info.total_steps = total_steps; + } +} + +sd_preview_info_t sd_get_preview_info() { + return sd_preview_info; +} + sd_preview_cb_t sd_get_preview_callback() { return sd_preview_cb; } diff --git a/src/core/util.h b/src/core/util.h index 334224da26..0929f60e01 100644 --- a/src/core/util.h +++ b/src/core/util.h @@ -29,6 +29,7 @@ bool contains(const std::string& str, const std::string& substr); std::string sd_format(const char* fmt, ...); void replace_all_chars(std::string& str, char target, char replacement); +std::string escape_newlines(const std::string& text); int round_up_to(int value, int base); @@ -101,6 +102,7 @@ std::vector> split_quotation_attention( sd_progress_cb_t sd_get_progress_callback(); void* sd_get_progress_callback_data(); +void sd_begin_preview_pass(int total_steps); sd_preview_cb_t sd_get_preview_callback(); void* sd_get_preview_callback_data(); preview_t sd_get_preview_mode(); diff --git a/src/detailer.cpp b/src/detailer.cpp index bbf9c6880e..72d01dc6f0 100644 --- a/src/detailer.cpp +++ b/src/detailer.cpp @@ -970,6 +970,7 @@ bool adetail_image(adetailer_ctx_t* context, generation.pm_params = {}; generation.pulid_params = {}; generation.hires.enabled = false; + generation.image_preprocess = {}; if (params.steps > 0) { generation.sample_params.sample_steps = params.steps; generation.sample_params.custom_sigmas = nullptr; diff --git a/src/model.h b/src/model.h index 3fdf1a63ed..0d291ce439 100644 --- a/src/model.h +++ b/src/model.h @@ -62,6 +62,10 @@ enum SDVersion { VERSION_SENSENOVA_U1_5, VERSION_LLADA_IMAGE, VERSION_ESRGAN, + VERSION_PIXART, + VERSION_MING_IMAGE, + VERSION_Z_IMAGE_L2P, + VERSION_IRIS, VERSION_COUNT, }; @@ -174,6 +178,10 @@ static inline bool sd_version_is_z_image(SDVersion version) { return false; } +static inline bool sd_version_is_z_image_l2p(SDVersion version) { + return version == VERSION_Z_IMAGE_L2P; +} + static inline bool sd_version_is_llada_image(SDVersion version) { if (version == VERSION_LLADA_IMAGE) { return true; @@ -252,6 +260,18 @@ static inline bool sd_version_is_sensenova_u1(SDVersion version) { return version == VERSION_SENSENOVA_U1_5; } +static inline bool sd_version_is_pixart(SDVersion version) { + return version == VERSION_PIXART; +} + +static inline bool sd_version_supports_video_generation(SDVersion version) { + return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); +} + +static inline bool sd_version_supports_image_generation(SDVersion version) { + return !sd_version_supports_video_generation(version); +} + static inline bool sd_version_uses_flux_vae(SDVersion version) { if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) { return true; @@ -267,7 +287,7 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) { } static inline bool sd_version_uses_wan_vae(SDVersion version) { - if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) { + if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version) || version == VERSION_MING_IMAGE) { return true; } return false; @@ -301,6 +321,8 @@ static inline bool sd_version_is_dit(SDVersion version) { version == VERSION_HIDREAM_O1 || sd_version_is_anima(version) || sd_version_is_z_image(version) || + sd_version_is_z_image_l2p(version) || + version == VERSION_MING_IMAGE || sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_ernie_image(version) || @@ -312,7 +334,9 @@ static inline bool sd_version_is_dit(SDVersion version) { sd_version_is_sefi_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || - sd_version_is_sensenova_u1(version)) { + sd_version_is_sensenova_u1(version) || + sd_version_is_pixart(version) || + version == VERSION_IRIS) { return true; } return false; diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index cb368feade..e9d20a4ece 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -163,6 +163,27 @@ struct LoraModel : public GGMLRunner { lora_tensors = std::move(new_lora_tensors); } + + std::unordered_map new_lora_tensors; + for (const auto& [old_name, tensor] : lora_tensors) { + std::string new_name = old_name; + if (starts_with(old_name, "lora.model.diffusion_model.transformer_blocks.")) { + // Qwen Image 2.1 stores the gate before the projection in fused MLP weights. + for (const auto& suffix : {std::string(".img_mlp.gate_layer.weight."), std::string(".img_mlp.proj.weight.")}) { + size_t pos = old_name.find(suffix); + if (pos == std::string::npos) { + continue; + } + std::string fused_name = old_name.substr(5, pos - 5) + ".img_mlp.gate_up.weight"; + if (model_tensor_names.find(fused_name) != model_tensor_names.end()) { + new_name = "lora." + fused_name + (suffix == ".img_mlp.proj.weight." ? ".1." : ".") + old_name.substr(pos + suffix.size()); + } + break; + } + } + new_lora_tensors[new_name] = tensor; + } + lora_tensors = std::move(new_lora_tensors); } ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) { diff --git a/src/model/adapter/lora_ops.cpp b/src/model/adapter/lora_ops.cpp index 7dacfd164d..5a6f905de9 100644 --- a/src/model/adapter/lora_ops.cpp +++ b/src/model/adapter/lora_ops.cpp @@ -11,7 +11,9 @@ ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx, ggml_tensor* lora_mid) { ggml_tensor* updown; // flat lora tensors to multiply it - int64_t lora_up_rows = lora_up->ne[ggml_n_dims(lora_up) - 1]; + // ggml_n_dims drops the output dimension when a linear layer has one output. + int lora_up_n_dims = ggml_n_dims(lora_up); + int64_t lora_up_rows = lora_up_n_dims == 1 ? 1 : lora_up->ne[lora_up_n_dims - 1]; lora_up = ggml_reshape_2d(ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows); auto lora_down_n_dims = ggml_n_dims(lora_down); // assume n_dims should always be a multiple of 2 (otherwise rank 1 doesn't work) diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index b4c94b2d94..2142895777 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__ #define __SD_MODEL_COMMON_GGML_BLOCK_HPP__ +#include #include #include #include @@ -228,6 +229,11 @@ class Linear : public UnaryBlock { return force_prec_f32; } + bool has_quantized_weight() const { + auto it = params.find("weight"); + return it != params.end() && it->second != nullptr && ggml_is_quantized(it->second->type); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; const float scale = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale; @@ -238,7 +244,7 @@ class Linear : public UnaryBlock { if (ctx->backend != nullptr) { ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x); if (force_prec_f32) { - ggml_mul_mat_set_prec(fp8_matmul, GGML_PREC_F32); + ggml_prec_set_acc(fp8_matmul, GGML_PREC_F32); } supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul); } @@ -359,15 +365,75 @@ class Embedding : public UnaryBlock { protected: int64_t embedding_dim; int64_t num_embeddings; + bool is_int8_tensorwise = false; + int int8_convrot_group_size = 0; + + static void get_rows_i8(ggml_tensor* dst, int ith, int nth, void* userdata) { + const auto* embedding = static_cast(userdata); + const int group_size = embedding->int8_convrot_group_size; + const auto* weight = dst->src[0]; + const auto* input_ids = static_cast(dst->src[1]->data); + const auto* scales = static_cast(dst->src[2]->data); + const bool scalar_scale = ggml_nelements(dst->src[2]) == 1; + const float normalization = group_size > 0 ? 1.f / std::sqrt(static_cast(group_size)) : 1.f; + for (int64_t row = ith; row < dst->ne[1]; row += nth) { + const int32_t token = input_ids[row]; + GGML_ASSERT(token >= 0 && token < weight->ne[1]); + const auto* src = reinterpret_cast(static_cast(weight->data) + token * weight->nb[1]); + auto* out = reinterpret_cast(static_cast(dst->data) + row * dst->nb[1]); + const float scale = scales[scalar_scale ? 0 : token] * normalization; + for (int64_t i = 0; i < weight->ne[0]; ++i) { + out[i] = static_cast(src[i]) * scale; + } + // The regular Hadamard rotation is symmetric and its own inverse. + for (int stride = 1; stride < group_size; stride *= 4) { + for (int64_t base = 0; base < weight->ne[0]; base += 4 * stride) { + for (int j = 0; j < stride; ++j) { + float* values = out + base + j; + const float a = values[0]; + const float b = values[stride]; + const float c = values[2 * stride]; + const float d = values[3 * stride]; + values[0] = a + b + c - d; + values[stride] = a + b - c + d; + values[2 * stride] = a - b + c + d; + values[3 * stride] = -a + b + c + d; + } + } + } + } + } + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - if (!support_get_rows(wtype)) { + auto weight_storage = tensor_storage_map.find(prefix + "weight"); + is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise; + int8_convrot_group_size = 0; + enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); + if (is_int8_tensorwise) { + GGML_ASSERT(wtype == GGML_TYPE_I8); + auto scale_storage = tensor_storage_map.find(prefix + "weight_scale"); + GGML_ASSERT(scale_storage != tensor_storage_map.end()); + const int64_t scale_nelements = scale_storage->second.nelements(); + GGML_ASSERT(scale_nelements == 1 || scale_nelements == num_embeddings); + params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); + if (weight_storage->second.int8_convrot) { + int8_convrot_group_size = weight_storage->second.int8_convrot_group_size; + int remainder = int8_convrot_group_size; + while (remainder > 1 && remainder % 4 == 0) { + remainder /= 4; + } + GGML_ASSERT(remainder == 1 && embedding_dim % int8_convrot_group_size == 0); + } + } else if (!support_get_rows(wtype)) { wtype = GGML_TYPE_F32; } params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings); } enum ggml_op param_usage_op(const std::string& name) const override { + if (is_int8_tensorwise) { + return GGML_OP_CUSTOM; + } return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE; } @@ -387,9 +453,17 @@ class Embedding : public UnaryBlock { int64_t n = input_ids->ne[1]; input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]); - input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); - auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids); - embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); + ggml_tensor* embedding; + if (is_int8_tensorwise) { + GGML_ASSERT(input_ids->type == GGML_TYPE_I32); + ggml_tensor* args[] = {weight, input_ids, params["weight_scale"]}; + embedding = ggml_custom_4d(ctx->ggml_ctx, GGML_TYPE_F32, embedding_dim, input_ids->ne[0], 1, 1, + args, 3, get_rows_i8, GGML_N_TASKS_MAX, this); + } else { + input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); + embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids); + } + embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); // [N, n_token, embedding_dim] return embedding; diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index e447849366..93e457d4b0 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -16,6 +17,45 @@ namespace Rope { 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 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(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 values; + std::vector> ids; + PositionLayout positions; + std::vector frequencies; + EmbedNDLayout layout = EmbedNDLayout::Matrix; + int batch_size = 1; + }; + enum class RefIndexMode { FIXED, INCREASE, @@ -56,40 +96,25 @@ namespace Rope { return flat_vec; } - __STATIC_INLINE__ std::vector> rope(const std::vector& pos, - int dim, - float theta, - const std::vector& axis_wrap_dims = {}) { + __STATIC_INLINE__ std::vector rope_frequencies(int dim, float theta) { assert(dim % 2 == 0); - int half_dim = dim / 2; - + int half_dim = dim / 2; std::vector scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim); - std::vector omega(half_dim); for (int i = 0; i < half_dim; ++i) { omega[i] = 1.0f / ::powf(1.f * theta, scale[i]); } + return omega; + } + __STATIC_INLINE__ std::vector> rope(const std::vector& pos, + const std::vector& omega) { + int half_dim = static_cast(omega.size()); size_t pos_size = pos.size(); std::vector> out(pos_size, std::vector(half_dim)); for (size_t i = 0; i < pos_size; ++i) { for (size_t j = 0; j < half_dim; ++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; } @@ -108,6 +133,12 @@ namespace Rope { return result; } + __STATIC_INLINE__ std::vector> rope(const std::vector& pos, + int dim, + float theta) { + return rope(pos, rope_frequencies(dim, theta)); + } + // Generate IDs for image patches and text __STATIC_INLINE__ std::vector> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set arange_dims) { auto txt_ids = std::vector>(bs * context_len, std::vector(axes_dim_num, 0.0f)); @@ -136,12 +167,16 @@ namespace Rope { int patch_size, int bs, int axes_dim_num, - int index = 0, - int h_offset = 0, - int w_offset = 0, - bool scale_rope = false) { + int index = 0, + int h_offset = 0, + int w_offset = 0, + bool scale_rope = false, + PositionLayout* layout = nullptr) { int h_len = (h + (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> img_ids(h_len * w_len, std::vector(axes_dim_num, 0.0)); int h_start = h_offset; @@ -192,8 +227,8 @@ namespace Rope { int bs, const std::vector& axis_thetas, const std::vector& axes_dim, - const std::vector>& wrap_dims = {}, - EmbedNDLayout layout = EmbedNDLayout::Matrix) { + EmbedNDLayout layout = EmbedNDLayout::Matrix, + std::vector* frequencies = nullptr) { std::vector> trans_ids = transpose(ids); size_t pos_len = ids.size() / bs; size_t num_axes = axes_dim.size(); @@ -205,19 +240,25 @@ namespace Rope { for (int d : axes_dim) emb_dim += d / 2; + if (frequencies) { + frequencies->clear(); + frequencies->reserve(emb_dim); + } std::vector> emb(bs * pos_len, std::vector(emb_dim * 2 * 2, 0.0)); size_t offset = 0; for (size_t i = 0; i < num_axes; ++i) { - std::vector axis_wrap_dims; - if (!wrap_dims.empty() && i < (int)wrap_dims.size()) { - axis_wrap_dims = wrap_dims[i]; - } float axis_theta = 10000.0f; if (!axis_thetas.empty()) { 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> 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 j = 0; j < pos_len; ++j) { for (int k = 0; k < rope_emb[0].size(); ++k) { @@ -253,10 +294,10 @@ namespace Rope { int bs, float theta, const std::vector& axes_dim, - const std::vector>& wrap_dims = {}, - EmbedNDLayout layout = EmbedNDLayout::Matrix) { + EmbedNDLayout layout = EmbedNDLayout::Matrix, + std::vector* frequencies = nullptr) { std::vector 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 embed_interleaved_mrope(const std::vector>& ids, @@ -264,7 +305,7 @@ namespace Rope { float theta, int head_dim, const std::vector& mrope_section, - const std::vector>& axis_wrap_dims = {}) { + std::vector* frequencies = nullptr) { GGML_ASSERT(bs > 0); GGML_ASSERT(head_dim % 2 == 0); GGML_ASSERT(mrope_section.size() >= 3); @@ -273,20 +314,26 @@ namespace Rope { size_t pos_len = ids.size() / bs; 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>> axis_embs; axis_embs.reserve(3); for (int axis = 0; axis < 3; ++axis) { - std::vector axis_wrap; - if (axis < static_cast(axis_wrap_dims.size())) { - axis_wrap = axis_wrap_dims[axis]; - } - axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap)); + axis_embs.push_back(rope(trans_ids[axis], omega)); } std::vector> emb = axis_embs[0]; for (int axis = 1; axis < 3; ++axis) { int length = std::min(mrope_section[axis] * 3, half_dim); 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 (int k = 0; k < 4; ++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); } - __STATIC_INLINE__ std::vector embed_2d_interleaved(int height, - int width, - int dim, - float theta = 10000.f, - float scale = 16.f, - int ref_grid_h = 0, - int ref_grid_w = 0) { + __STATIC_INLINE__ Embedding embed_2d_interleaved(int height, + int width, + int dim, + float theta = 10000.f, + float scale = 16.f, + int ref_grid_h = 0, + int ref_grid_w = 0) { assert(dim % 4 == 0); int half_dim = dim / 2; int dim_axis = dim / 2; @@ -318,6 +365,10 @@ namespace Rope { w_ntk = std::pow(static_cast(width) / static_cast(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 x_pos; std::vector y_pos; x_pos.reserve(static_cast(height) * width); @@ -326,13 +377,20 @@ namespace Rope { float y = height == 1 ? 0.f : scale * static_cast(iy) / static_cast(height - 1); for (int ix = 0; ix < width; ++ix) { float x = width == 1 ? 0.f : scale * static_cast(ix) / static_cast(width - 1); + result.ids.push_back({0.f, y, x}); x_pos.push_back(x); y_pos.push_back(y); } } - auto x_emb = rope(x_pos, dim_axis, theta * w_ntk); - auto y_emb = rope(y_pos, dim_axis, theta * h_ntk); + auto x_freq = rope_frequencies(dim_axis, theta * w_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 out(static_cast(height) * width * half_dim * 4); 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> gen_refs_ids(int patch_size, @@ -359,7 +418,8 @@ namespace Rope { RefIndexMode ref_index_mode, float ref_index_scale, bool scale_rope, - int base_offset = 0) { + int base_offset = 0, + PositionLayout* layout = nullptr) { std::vector> ids; int curr_h_offset = 0; int curr_w_offset = 0; @@ -386,7 +446,8 @@ namespace Rope { static_cast(index * ref_index_scale), h_offset + base_offset, w_offset + base_offset, - scale_rope); + scale_rope, + layout); ids = concat_ids(ids, ref_ids, bs); if (ref_index_mode == RefIndexMode::INCREASE) { @@ -409,88 +470,53 @@ namespace Rope { const std::vector& ref_latents, RefIndexMode ref_index_mode, 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; 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; - 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); 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); } return ids; } // Generate flux positional embeddings - __STATIC_INLINE__ std::vector gen_flux_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - std::set txt_arange_dims, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - float ref_index_scale, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim, - bool is_longcat) { - std::vector> ids = gen_flux_ids(h, - w, - patch_size, - bs, - static_cast(axes_dim.size()), - context_len, - txt_arange_dims, - ref_latents, - ref_index_mode, - ref_index_scale, - is_longcat); - std::vector> wrap_dims; - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - 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(pos_len, 0)); - size_t cursor = context_len; // text first - const size_t img_tokens = static_cast(h_len) * static_cast(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(ref->ne[1]); - int ref_w = static_cast(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(ref_h_l) * static_cast(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(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_flux_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + std::set txt_arange_dims, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + float ref_index_scale, + int theta, + const std::vector& axes_dim, + bool is_longcat) { + Embedding result; + result.batch_size = bs; + result.ids = gen_flux_ids(h, + w, + patch_size, + bs, + static_cast(axes_dim.size()), + context_len, + txt_arange_dims, + ref_latents, + ref_index_mode, + ref_index_scale, + is_longcat, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> gen_vid_ids(int t, @@ -500,14 +526,18 @@ namespace Rope { int ph, int pw, int bs, - int t_offset = 0, - int h_offset = 0, - int w_offset = 0, - bool scale_rope = false) { + int t_offset = 0, + int h_offset = 0, + int w_offset = 0, + bool scale_rope = false, + PositionLayout* layout = nullptr) { int t_len = (t + (pt / 2)) / pt; int h_len = (h + (ph / 2)) / ph; int w_len = (w + (pw / 2)) / pw; + if (layout) { + layout->append_image(h_len, w_len, t_len); + } std::vector> vid_ids(t_len * h_len * w_len, std::vector(3, 0.0)); if (scale_rope) { @@ -573,7 +603,11 @@ namespace Rope { int bs, int context_len, const std::vector& 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 w_len = (w + (patch_size / 2)) / patch_size; int txt_id_start = std::max(h_len, w_len) / 2; @@ -585,90 +619,49 @@ namespace Rope { } } 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); if (ref_latents.size() > 0) { 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); } return ids; } // Generate qwen_image positional embeddings - __STATIC_INLINE__ std::vector gen_qwen_image_pe(int t, - int h, - int w, - int patch_size, - int bs, - int context_len, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode); - std::vector> wrap_dims; - // This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles - 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(total_tokens / bs, 0)); - size_t cursor = context_len; // ignore text tokens - const size_t img_tokens = static_cast(t) * static_cast(h_len) * static_cast(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(ref->ne[1]); - int ref_w = static_cast(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(ref_h_len) * static_cast(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(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_qwen_image_pe(int t, + int h, + int w, + int patch_size, + int bs, + int context_len, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } - __STATIC_INLINE__ std::vector gen_mage_flow_pe(int h, - int w, - int bs, - int context_len, - const std::vector& ref_latents, - int theta, - const std::vector& axes_dim) { + __STATIC_INLINE__ Embedding gen_mage_flow_pe(int h, + int w, + int bs, + int context_len, + const std::vector& ref_latents, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); const int axes_dim_num = static_cast(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> image_ids(static_cast(bs) * image_h * image_w, std::vector(axes_dim_num, 0.f)); + result.positions.append_image(image_h, image_w); int h_start = -(image_h - image_h / 2); int w_start = -(image_w - image_w / 2); for (int b = 0; b < bs; ++b) { @@ -692,15 +685,18 @@ namespace Rope { static_cast(i + 1)); ids = concat_ids(ids, ref_ids, bs); } - return embed_nd(ids, bs, static_cast(theta), axes_dim); + result.ids = std::move(ids); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> gen_lens_ids(int h, int w, int bs, int context_len, - bool scale_rope = true) { - auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope); + bool scale_rope = true, + 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; auto txt_ids = linspace(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); } - __STATIC_INLINE__ std::vector gen_lens_pe(int h, - int w, - int bs, - int context_len, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_lens_ids(h, w, bs, context_len, true); - std::vector> wrap_dims; - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - size_t pos_len = ids.size() / bs; - wrap_dims.assign(axes_dim.size(), std::vector(pos_len, 0)); - const size_t img_tokens = static_cast(h) * static_cast(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(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_lens_pe(int h, + int w, + int bs, + int context_len, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_lens_ids(h, w, bs, context_len, true, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> gen_ernie_image_ids(int h, int w, int patch_size, int bs, - int context_len) { + int context_len, + PositionLayout* layout = nullptr) { int h_len = h / patch_size; int w_len = w / patch_size; + if (layout) { + layout->append_image(h_len, w_len); + } std::vector> img_ids(h_len * w_len, std::vector(3, 0.0f)); std::vector h_ids = linspace(0.f, static_cast(h_len - 1), h_len); std::vector w_ids = linspace(0.f, static_cast(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); } - __STATIC_INLINE__ std::vector gen_ernie_image_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len); - std::vector> wrap_dims; - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - int h_len = h / patch_size; - 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(pos_len, 0)); - const size_t img_tokens = static_cast(h_len) * static_cast(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(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage); + __STATIC_INLINE__ Embedding gen_ernie_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.layout = EmbedNDLayout::ErnieImage; + result.ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } // Generate wan positional embeddings @@ -905,7 +880,8 @@ namespace Rope { int context_len, int seq_multi_of, const std::vector& ref_latents, - RefIndexMode ref_index_mode) { + RefIndexMode ref_index_mode, + PositionLayout* layout = nullptr) { SD_UNUSED(ref_index_mode); int padded_context_len = context_len + bound_mod(context_len, seq_multi_of); auto txt_ids = std::vector>(bs * padded_context_len, std::vector(3, 0.0f)); @@ -913,11 +889,17 @@ namespace Rope { txt_ids[i][0] = (i % padded_context_len) + 1.f; } + if (layout) { + layout->append_tokens(padded_context_len); + } int axes_dim_num = 3; 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(img_ids.size() / bs), seq_multi_of); + if (layout) { + layout->append_tokens(img_pad_len); + } if (img_pad_len > 0) { std::vector> img_pad_ids(bs * img_pad_len, std::vector(3, 0.f)); img_ids = concat_ids(img_ids, img_pad_ids, bs); @@ -936,7 +918,8 @@ namespace Rope { int patch_size, int bs, 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 padded_context_len = context_len + context_pad_len; auto txt_ids = std::vector>(bs * padded_context_len, std::vector(3, 0.0f)); @@ -947,11 +930,17 @@ namespace Rope { } } + if (layout) { + layout->append_tokens(padded_context_len); + } int axes_dim_num = 3; 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(img_ids.size() / bs), seq_multi_of); + if (layout) { + layout->append_tokens(img_pad_len); + } if (img_pad_len > 0) { std::vector> img_pad_ids(bs * img_pad_len, std::vector(3, 0.f)); img_ids = concat_ids(img_ids, img_pad_ids, bs); @@ -968,7 +957,8 @@ namespace Rope { int patch_size, int context_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 padded_context = context_len + context_pad; const int h_len = (h + (patch_size / 2)) / patch_size; @@ -994,11 +984,17 @@ namespace Rope { cursor += 2; } + if (layout) { + layout->append_tokens(cap_ids.size()); + } std::vector> img_ids; 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(), image_pad, std::vector(3, 0.f)); + if (layout) { + layout->append_tokens(image_pad); + } } const int sigvq_start = static_cast(cap_ids.size() + img_ids.size()) + 1; @@ -1016,101 +1012,77 @@ namespace Rope { ids.insert(ids.end(), cap_ids.begin(), cap_ids.end()); ids.insert(ids.end(), img_ids.begin(), img_ids.end()); ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end()); + if (layout) { + layout->append_tokens(sigvq_ids.size()); + } SD_UNUSED(padded_image); return ids; } - __STATIC_INLINE__ std::vector gen_llada_image_edit_pe(int h, - int w, - int patch_size, - int context_len, - int sigvq_len, - int seq_multi_of, - int theta, - const std::vector& axes_dim) { - auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of); - return embed_nd(ids, 1, static_cast(theta), axes_dim, {}); - } - - __STATIC_INLINE__ std::vector gen_llada_image_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - int seq_multi_of, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of); - std::vector> wrap_dims; - 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) { - size_t pos_len = ids.size() / bs; - wrap_dims.assign(axes_dim.size(), std::vector(pos_len, 0)); - size_t cursor = context_len + bound_mod(context_len, seq_multi_of); - size_t img_tokens = static_cast(h_len) * static_cast(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(theta), axes_dim, wrap_dims); - } - - // Generate z_image positional embeddings - __STATIC_INLINE__ std::vector gen_z_image_pe(int h, + __STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h, int w, int patch_size, - int bs, int context_len, + int sigvq_len, int seq_multi_of, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, int theta, - bool circular_h, - bool circular_w, const std::vector& axes_dim) { - std::vector> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode); - std::vector> wrap_dims; - 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) { - size_t pos_len = ids.size() / bs; - wrap_dims.assign(axes_dim.size(), std::vector(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(h_len) * static_cast(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; - } - } - } - } + Embedding result; + 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(theta), axes_dim, result.layout, &result.frequencies); + return result; + } + + __STATIC_INLINE__ Embedding gen_llada_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; + } - return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); + // Generate z_image positional embeddings + __STATIC_INLINE__ Embedding gen_z_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } - __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* apply_rope(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe, bool rope_interleaved = true) { + auto ctx = runner_ctx->ggml_ctx; +#ifndef SD_USE_UPSTREAM_GGML + auto backend = runner_ctx->backend; + if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 && + x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 && + pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) { + auto out = ggml_rope_apply(ctx, x, pe, rope_interleaved); + if (ggml_backend_supports_op(backend, out)) { + return out; + } + } +#endif // x: [N, L, n_head, d_head] // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]] int64_t d_head = x->ne[0]; @@ -1122,8 +1094,16 @@ namespace Rope { x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N); // [N * n_head, L, d_head/2, 2] x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2)); // [2, N * n_head, L, d_head/2] } else { - x = ggml_reshape_4d(ctx, x, d_head / 2, 2, L, n_head * N); // [N * n_head, L, 2, d_head/2] - x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 3, 1)); // [2, N * n_head, L, d_head/2] + auto x_0 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], 0); + auto x_1 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], (d_head / 2) * x->nb[0]); + + auto pe_2 = ggml_cont(ctx, ggml_permute(ctx, pe, 2, 3, 0, 1)); + auto cos = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], 0); + auto sin = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], pe_2->nb[3]); + + auto part_0 = ggml_sub(ctx, ggml_mul(ctx, x_0, cos), ggml_mul(ctx, x_1, sin)); + auto part_1 = ggml_add(ctx, ggml_mul(ctx, x_0, sin), ggml_mul(ctx, x_1, cos)); + return ggml_concat(ctx, part_0, part_1, 0); } int64_t offset = x->nb[2] * x->ne[2]; @@ -1141,10 +1121,7 @@ namespace Rope { auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1); // [L, d_head/2, 2] auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1)); // [N * n_head, L, d_head/2, 2] - if (!rope_interleaved) { - x_out = ggml_cont(ctx, ggml_permute(ctx, x_out, 1, 0, 2, 3)); // [N * n_head, L, x, d_head/2] - } - x_out = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N); // [N*n_head, L, d_head] + x_out = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N); // [N*n_head, L, d_head] return x_out; } @@ -1161,8 +1138,8 @@ namespace Rope { // return: [N, L, n_head*d_head] int64_t n_head = q->ne[1]; - q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head] - k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head] + q = apply_rope(ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head] + k = apply_rope(ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head] auto x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head] return x; diff --git a/src/model/common/rope_circular.hpp b/src/model/common/rope_circular.hpp new file mode 100644 index 0000000000..19dccb9593 --- /dev/null +++ b/src/model/common/rope_circular.hpp @@ -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(region.height_axis)) { + period = region.height_period; + } else if (circular_x && frequency.axis == static_cast(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__ diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index a8977c966f..a414d96e63 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -235,8 +235,8 @@ namespace Anima { if (pe_k == nullptr) { pe_k = pe_q; } - auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false); - auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false); + auto q_rope = Rope::apply_rope(ctx, q4, pe_q, false); + auto k_rope = Rope::apply_rope(ctx, k4, pe_k, false); attn_out = ggml_ext_attention_ext(ctx, q_rope, k_rope, @@ -603,34 +603,37 @@ namespace Anima { return std::pow(extrapolation_ratio, static_cast(axis_dim) / static_cast(axis_dim - 2)); } - static std::vector gen_anima_image_pe_vec(int bs, - int h, - int w, - int patch_size, - int theta, - const std::vector& axes_dim, - float h_extrapolation_ratio, - float w_extrapolation_ratio, - float t_extrapolation_ratio, - const std::vector& ref_latents) { - auto ids = Rope::gen_flux_ids(h, - w, - patch_size, - bs, - static_cast(axes_dim.size()), - 0, - {}, - ref_latents, - Rope::RefIndexMode::FIXED, - 1.0f, - false); + static Rope::Embedding gen_anima_image_pe_vec(int bs, + int h, + int w, + int patch_size, + int theta, + const std::vector& axes_dim, + float h_extrapolation_ratio, + float w_extrapolation_ratio, + float t_extrapolation_ratio, + const std::vector& ref_latents) { + Rope::Embedding result; + result.batch_size = bs; + result.ids = Rope::gen_flux_ids(h, + w, + patch_size, + bs, + static_cast(axes_dim.size()), + 0, + {}, + ref_latents, + Rope::RefIndexMode::FIXED, + 1.0f, + false, &result.positions); std::vector axis_thetas = { static_cast(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]), static_cast(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]), static_cast(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& x_tensor, @@ -657,16 +660,16 @@ namespace Anima { int64_t h_pad = x->ne[1] + pad_h; int64_t w_pad = x->ne[0] + pad_w; - image_pe_vec = gen_anima_image_pe_vec(1, - static_cast(h_pad), - static_cast(w_pad), - static_cast(config.patch_size), - config.theta, - config.axes_dim, - 4.0f, - 4.0f, - 1.0f, - ref_latents); + image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1, + static_cast(h_pad), + static_cast(w_pad), + static_cast(config.patch_size), + config.theta, + config.axes_dim, + 4.0f, + 4.0f, + 1.0f, + ref_latents)); int64_t image_pos_len = static_cast(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); set_backend_tensor_data(image_pe, image_pe_vec.data()); diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index 8c91bb55d2..1ca1370984 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -720,15 +720,18 @@ namespace Boogu { } } - __STATIC_INLINE__ std::vector gen_boogu_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - const std::vector& ref_latents, - int theta, - const std::vector& axes_dim) { - std::vector> ids; + __STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + const std::vector& ref_latents, + int theta, + const std::vector& axes_dim) { + Rope::Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); + auto& ids = result.ids; ids.reserve(static_cast(bs) * context_len); for (int b = 0; b < bs; b++) { for (int i = 0; i < context_len; i++) { @@ -741,15 +744,18 @@ namespace Boogu { for (ggml_tensor* ref : ref_latents) { int ref_h_tokens = patched_token_count(ref->ne[1], 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); pe_shift += std::max(ref_h_tokens, ref_w_tokens); } int h_tokens = patched_token_count(h, 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); - return Rope::embed_nd(ids, bs, static_cast(theta), axes_dim); + result.values = Rope::embed_nd(ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } struct BooguImageRunner : public DiffusionModelRunner { @@ -793,14 +799,14 @@ namespace Boogu { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = gen_boogu_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(gen_boogu_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); int pos_len = static_cast(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); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp index 34387d1b71..3b08ed1500 100644 --- a/src/model/diffusion/control.hpp +++ b/src/model/diffusion/control.hpp @@ -376,7 +376,7 @@ struct ControlNet : public GGMLRunner { hint = make_input(hint_tensor); } - auto runner_ctx = get_context(); + auto runner_ctx = get_context(gf); auto outs = control_net.forward(&runner_ctx, x, @@ -389,8 +389,7 @@ struct ControlNet : public GGMLRunner { if (guided_hint_input == nullptr && !outs.empty()) { guided_hint_output_ggml = outs[0]; ggml_set_output(guided_hint_output_ggml); - cache(guided_hint_cache_name(), guided_hint_output_ggml); - ggml_build_forward_expand(gf, guided_hint_output_ggml); + runner_ctx.persist_cache_tensor(guided_hint_cache_name(), guided_hint_output_ggml); } control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0); diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index 81bf142a07..b284cb83f7 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -415,15 +415,13 @@ namespace ErnieImage { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = Rope::gen_ernie_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); int pos_len = static_cast(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); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 67e79d3d96..825dabe13d 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -1548,20 +1548,18 @@ namespace Flux { } else if (version == VERSION_OVIS_IMAGE) { txt_arange_dims = {1, 2}; } - pe_vec = Rope::gen_flux_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - txt_arange_dims, - ref_latents, - ref_index_mode, - config.ref_index_scale, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim, - sd_version_is_longcat(version)); + pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + txt_arange_dims, + ref_latents, + ref_index_mode, + config.ref_index_scale, + config.theta, + config.axes_dim, + sd_version_is_longcat(version))); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index 9e2aa36122..3254571db2 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -149,18 +149,21 @@ namespace Ideogram4 { return std::make_shared(in_features, out_features, bias); } - __STATIC_INLINE__ std::vector gen_ideogram4_pe(int grid_h, - int grid_w, - int bs, - int context_len, - int head_dim, - int rope_theta, - const std::vector& mrope_section, - bool circular_x = false, - bool circular_y = false) { + __STATIC_INLINE__ Rope::Embedding gen_ideogram4_pe(int grid_h, + int grid_w, + int bs, + int context_len, + int head_dim, + int rope_theta, + const std::vector& mrope_section) { GGML_ASSERT(bs == 1); - std::vector> ids(static_cast(bs) * (context_len + grid_h * grid_w), - std::vector(3, 0.f)); + Rope::Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); + result.positions.append_image(grid_h, grid_w); + result.ids.assign(static_cast(bs) * (context_len + grid_h * grid_w), + std::vector(3, 0.f)); + auto& ids = result.ids; for (int i = 0; i < context_len; ++i) { ids[i] = {static_cast(i), static_cast(i), static_cast(i)}; @@ -175,29 +178,13 @@ namespace Ideogram4 { } } - std::vector> axis_wrap_dims(3); - if (circular_y || circular_x) { - size_t total_len = static_cast(bs) * (context_len + grid_h * grid_w); - axis_wrap_dims[1].assign(total_len, 0); - axis_wrap_dims[2].assign(total_len, 0); - if (circular_y) { - for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { - axis_wrap_dims[1][idx] = grid_h; - } - } - if (circular_x) { - for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { - axis_wrap_dims[2][idx] = grid_w; - } - } - } - - return Rope::embed_interleaved_mrope(ids, - bs, - static_cast(rope_theta), - head_dim, - mrope_section, - axis_wrap_dims); + result.values = Rope::embed_interleaved_mrope(ids, + bs, + static_cast(rope_theta), + head_dim, + mrope_section, + &result.frequencies); + return result; } class Ideogram4Attention : public GGMLBlock { @@ -509,15 +496,13 @@ namespace Ideogram4 { int64_t head_dim = config.emb_dim / config.num_heads; auto runner_ctx = get_context(); - pe_vec = gen_ideogram4_pe(static_cast(grid_h), - static_cast(grid_w), - static_cast(x->ne[3]), - static_cast(context_len), - static_cast(head_dim), - static_cast(config.rope_theta), - config.mrope_section, - runner_ctx.circular_x_enabled, - runner_ctx.circular_y_enabled); + pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast(grid_h), + static_cast(grid_w), + static_cast(x->ne[3]), + static_cast(context_len), + static_cast(head_dim), + static_cast(config.rope_theta), + config.mrope_section)); 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()); diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp new file mode 100644 index 0000000000..b90fbe9d44 --- /dev/null +++ b/src/model/diffusion/iris.hpp @@ -0,0 +1,450 @@ +#ifndef __SD_MODEL_DIFFUSION_IRIS_HPP__ +#define __SD_MODEL_DIFFUSION_IRIS_HPP__ + +#include +#include +#include "model/diffusion/pid.hpp" + +namespace Iris { + struct IrisConfig { + int64_t hidden_size = 2560; + int64_t depth = 24; + int64_t dual_depth = 8; + int64_t head_dim = 128; + int64_t num_heads = 20; + int64_t kv_heads = 5; + int64_t mlp_dim = 6826; + int64_t patch_size = 16; + int64_t text_dim = 2560; + int64_t text_len = 300; + int64_t text_layers = 12; + int64_t layer_heads = 32; + int64_t layer_mlp_dim = 3328; + int64_t pixel_dim = 16; + int64_t pixel_attn_dim = 1280; + int64_t pixel_heads = 10; + int64_t pixel_depth = 4; + + static IrisConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) { + IrisConfig config; + auto find = [&](const std::string& name) -> const TensorStorage& { + auto it = weights.find(prefix + "." + name); + if (it == weights.end()) { + throw std::runtime_error("Iris-3B: missing weight " + prefix + "." + name); + } + return it->second; + }; + config.hidden_size = find("s_embedder.proj.weight").ne[1]; + config.patch_size = static_cast(std::sqrt(find("s_embedder.proj.weight").ne[0] / 3)); + config.head_dim = find("blocks.0.attn.q_norm_x.weight").ne[0]; + config.num_heads = config.hidden_size / config.head_dim; + config.kv_heads = find("blocks.0.attn.k_proj_x.weight").ne[1] / config.head_dim; + config.mlp_dim = find("blocks.0.mlp_x.w1.weight").ne[1]; + config.text_dim = find("y_embedder.refiner.proj.weight").ne[0]; + config.text_len = find("y_pos_embedding").ne[1]; + config.text_layers = find("y_embedder.layer_pool.weight").ne[0]; + config.layer_mlp_dim = find("y_embedder.layer_blocks.0.mlp.0.weight").ne[1]; + config.pixel_dim = find("pixel_embedder.proj.weight").ne[1]; + config.pixel_attn_dim = find("pixel_blocks.0.compress.weight").ne[1]; + config.pixel_heads = config.pixel_attn_dim / find("pixel_blocks.0.attn.q_norm.weight").ne[0]; + config.depth = config.dual_depth = config.pixel_depth = 0; + for (const auto& [name, tensor] : weights) { + if (!starts_with(name, prefix + ".")) { + continue; + } + auto parts = split_string(name.substr(prefix.size() + 1), '.'); + if (parts.size() > 2 && parts[0] == "blocks") { + int64_t index = std::stoi(parts[1]); + config.depth = std::max(config.depth, index + 1); + if (parts[2] == "adaln_img") { + config.dual_depth = std::max(config.dual_depth, index + 1); + } + } else if (parts.size() > 2 && parts[0] == "pixel_blocks") { + config.pixel_depth = std::max(config.pixel_depth, int64_t(std::stoi(parts[1]) + 1)); + } + } + if (config.text_layers != 12 || config.text_dim != 2560 || config.text_len != 300 || + find("modulation_cores.adaln_img.weight").ne[1] != 6 * config.hidden_size || + find("pixel_blocks.0.adaln.weight").ne[1] != 4 * config.pixel_dim * config.patch_size * config.patch_size) { + throw std::runtime_error("Iris-3B: unsupported checkpoint configuration"); + } + LOG_VERBOSE("iris: hidden_size = %" PRId64 ", depth = %" PRId64 ", dual_depth = %" PRId64 ", heads = %" PRId64 "/%" PRId64 ", pixel_depth = %" PRId64, + config.hidden_size, config.depth, config.dual_depth, config.num_heads, config.kv_heads, config.pixel_depth); + return config; + } + }; + + struct SelfAttention : public GGMLBlock { + int64_t dim; + int64_t heads; + bool qk_norm; + + SelfAttention(int64_t dim, int64_t heads, bool qk_norm = true) + : dim(dim), heads(heads), qk_norm(qk_norm) { + blocks["qkv"] = std::make_shared(dim, 3 * dim, false); + blocks["proj"] = std::make_shared(dim, dim, true); + if (qk_norm) { + blocks["q_norm"] = std::make_shared(dim / heads, 1e-6f); + blocks["k_norm"] = std::make_shared(dim / heads, 1e-6f); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr, ggml_tensor* mask = nullptr) { + auto c = ctx->ggml_ctx; + auto qkv = split_qkv(c, std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x)); + auto q = ggml_reshape_4d(c, qkv[0], dim / heads, heads, x->ne[1], x->ne[2]); + auto k = ggml_reshape_4d(c, qkv[1], dim / heads, heads, x->ne[1], x->ne[2]); + auto v = ggml_reshape_4d(c, qkv[2], dim / heads, heads, x->ne[1], x->ne[2]); + if (qk_norm) { + q = std::dynamic_pointer_cast(blocks["q_norm"])->forward(ctx, q); + k = std::dynamic_pointer_cast(blocks["k_norm"])->forward(ctx, k); + } + if (pe) { + x = Rope::attention(ctx, q, k, v, pe, mask); + } else { + q = ggml_reshape_3d(c, q, dim, x->ne[1], x->ne[2]); + k = ggml_reshape_3d(c, k, dim, x->ne[1], x->ne[2]); + v = ggml_reshape_3d(c, v, dim, x->ne[1], x->ne[2]); + x = ggml_ext_attention_ext(ctx, q, k, v, heads, mask, false, ctx->flash_attn_enabled); + } + return std::dynamic_pointer_cast(blocks["proj"])->forward(ctx, x); + } + }; + + struct TextBlock : public GGMLBlock { + bool layerwise; + + TextBlock(int64_t dim, int64_t heads, int64_t mlp_dim, bool layerwise) + : layerwise(layerwise) { + blocks["norm1"] = std::make_shared(dim, 1e-6f); + blocks["norm2"] = std::make_shared(dim, 1e-6f); + blocks["attn"] = std::make_shared(dim, heads, !layerwise); + if (layerwise) { + blocks["mlp.0"] = std::make_shared(dim, mlp_dim, true); + blocks["mlp.2"] = std::make_shared(mlp_dim, dim, true); + } else { + blocks["mlp"] = std::make_shared(dim, mlp_dim); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask = nullptr) { + auto h = std::dynamic_pointer_cast(blocks["norm1"])->forward(ctx, x); + h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, nullptr, mask); + x = ggml_add(ctx->ggml_ctx, x, h); + h = std::dynamic_pointer_cast(blocks["norm2"])->forward(ctx, x); + if (layerwise) { + h = std::dynamic_pointer_cast(blocks["mlp.0"])->forward(ctx, h); + h = ggml_silu(ctx->ggml_ctx, h); + h = std::dynamic_pointer_cast(blocks["mlp.2"])->forward(ctx, h); + } else { + h = std::dynamic_pointer_cast(blocks["mlp"])->forward(ctx, h); + } + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct TextEmbedder : public GGMLBlock { + IrisConfig config; + + TextEmbedder(const IrisConfig& config) + : config(config) { + for (int i = 0; i < 2; ++i) { + blocks["layer_blocks." + std::to_string(i)] = std::make_shared(config.text_dim, config.layer_heads, config.layer_mlp_dim, true); + blocks["refiner.blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, config.num_heads, config.mlp_dim, false); + } + blocks["layer_pool"] = std::make_shared(config.text_layers, 1, true); + blocks["refiner.proj"] = std::make_shared(config.text_dim, config.hidden_size, true); + blocks["refiner.norm"] = std::make_shared(config.hidden_size, 1e-6f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask) { + auto c = ctx->ggml_ctx; + int64_t tokens = x->ne[1]; + int64_t batch = x->ne[2]; + // The encoder concatenates selected layers within each token's feature dimension. + x = ggml_reshape_3d(c, x, config.text_dim, config.text_layers, tokens * batch); + for (int i = 0; i < 2; ++i) { + x = std::dynamic_pointer_cast(blocks["layer_blocks." + std::to_string(i)])->forward(ctx, x); + } + x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3)); + x = std::dynamic_pointer_cast(blocks["layer_pool"])->forward(ctx, x); + x = ggml_reshape_3d(c, x, config.text_dim, tokens, batch); + x = std::dynamic_pointer_cast(blocks["refiner.proj"])->forward(ctx, x); + for (int i = 0; i < 2; ++i) { + x = std::dynamic_pointer_cast(blocks["refiner.blocks." + std::to_string(i)])->forward(ctx, x, mask); + } + return std::dynamic_pointer_cast(blocks["refiner.norm"])->forward(ctx, x); + } + }; + + struct TrunkBlock : public GGMLBlock { + IrisConfig config; + bool dual; + + TrunkBlock(const IrisConfig& config, bool dual) + : config(config), dual(dual) { + int64_t d = config.hidden_size; + for (const std::string stream : dual ? std::vector{"x", "y"} : std::vector{""}) { + std::string suffix = stream.empty() ? "" : "_" + stream; + std::string attn = dual ? "attn." : ""; + blocks[attn + "q_proj" + suffix] = std::make_shared(d, d, false); + blocks[attn + "k_proj" + suffix] = std::make_shared(d, config.kv_heads * config.head_dim, false); + blocks[attn + "v_proj" + suffix] = std::make_shared(d, config.kv_heads * config.head_dim, false); + blocks[attn + "q_norm" + suffix] = std::make_shared(config.head_dim, 1e-6f); + blocks[attn + "k_norm" + suffix] = std::make_shared(config.head_dim, 1e-6f); + blocks[dual ? "attn.proj" + suffix : "attn_proj"] = std::make_shared(d, d, true); + blocks["attn_gate" + suffix] = std::make_shared(d, d, false); + blocks["norm" + suffix + "1"] = std::make_shared(d, 1e-6f); + blocks["norm" + suffix + "2"] = std::make_shared(d, 1e-6f); + blocks["attn_post_norm" + suffix] = std::make_shared(d, 1e-6f); + blocks["mlp_post_norm" + suffix] = std::make_shared(d, 1e-6f); + blocks["mlp" + suffix] = std::make_shared(d, config.mlp_dim); + } + } + + void init_params(ggml_context* ctx, const String2TensorStorage& = {}, const std::string = "") override { + if (dual) { + params["adaln_img.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 6 * config.hidden_size); + params["adaln_txt.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 6 * config.hidden_size); + } else { + params["adaln.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 6 * config.hidden_size); + } + } + + ggml_tensor* norm(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); + } + + ggml_tensor* linear(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); + } + + std::array project(GGMLRunnerContext* ctx, ggml_tensor* h, const std::string& suffix) { + auto c = ctx->ggml_ctx; + std::string a = dual ? "attn." : ""; + auto q = linear(ctx, a + "q_proj" + suffix, h); + auto k = linear(ctx, a + "k_proj" + suffix, h); + auto v = linear(ctx, a + "v_proj" + suffix, h); + q = ggml_reshape_4d(c, q, config.head_dim, config.num_heads, h->ne[1], h->ne[2]); + k = ggml_reshape_4d(c, k, config.head_dim, config.kv_heads, h->ne[1], h->ne[2]); + v = ggml_reshape_4d(c, v, config.head_dim, config.kv_heads, h->ne[1], h->ne[2]); + return {norm(ctx, a + "q_norm" + suffix, q), norm(ctx, a + "k_norm" + suffix, k), v}; + } + + ggml_tensor* finish(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* h, ggml_tensor* out, const std::vector& mods, const std::string& suffix) { + auto c = ctx->ggml_ctx; + out = ggml_mul(c, out, ggml_sigmoid(c, linear(ctx, "attn_gate" + suffix, h))); + out = linear(ctx, dual ? "attn.proj" + suffix : "attn_proj", out); + out = norm(ctx, "attn_post_norm" + suffix, out); + x = ggml_add(c, x, ggml_mul(c, out, mods[2])); + h = Pid::apply_adaln(c, norm(ctx, "norm" + suffix + "2", x), mods[3], mods[4]); + h = std::dynamic_pointer_cast(blocks["mlp" + suffix])->forward(ctx, h); + h = norm(ctx, "mlp_post_norm" + suffix, h); + return ggml_add(c, x, ggml_mul(c, h, mods[5])); + } + + std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* y, ggml_tensor* img_mod, ggml_tensor* txt_mod, ggml_tensor* pe) { + auto c = ctx->ggml_ctx; + int64_t nt = y->ne[1]; + auto mx = ggml_ext_chunk(c, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0); + if (!dual) { + auto tokens = ggml_concat(c, y, x, 1); + auto h = Pid::apply_adaln(c, norm(ctx, "norm1", tokens), mx[0], mx[1]); + auto qkv = project(ctx, h, ""); + auto out = Rope::attention(ctx, qkv[0], qkv[1], qkv[2], pe, nullptr); + tokens = finish(ctx, tokens, h, out, mx, ""); + return {ggml_ext_slice(c, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(c, tokens, 1, 0, nt)}; + } + auto my = ggml_ext_chunk(c, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0); + auto hx = Pid::apply_adaln(c, norm(ctx, "norm_x1", x), mx[0], mx[1]); + auto hy = Pid::apply_adaln(c, norm(ctx, "norm_y1", y), my[0], my[1]); + auto qx = project(ctx, hx, "_x"); + auto qy = project(ctx, hy, "_y"); + auto out = Rope::attention(ctx, ggml_concat(c, qy[0], qx[0], 2), ggml_concat(c, qy[1], qx[1], 2), + ggml_concat(c, qy[2], qx[2], 2), pe, nullptr); + x = finish(ctx, x, hx, ggml_ext_slice(c, out, 1, nt, out->ne[1]), mx, "_x"); + y = finish(ctx, y, hy, ggml_ext_slice(c, out, 1, 0, nt), my, "_y"); + return {x, y}; + } + }; + + struct PixelBlock : public GGMLBlock { + IrisConfig config; + + PixelBlock(const IrisConfig& config) + : config(config) { + auto d = config.pixel_dim; + auto flat = d * config.patch_size * config.patch_size; + blocks["norm1"] = std::make_shared(d, 1e-6f); + blocks["norm2"] = std::make_shared(d, 1e-6f); + blocks["compress"] = std::make_shared(flat, config.pixel_attn_dim, true); + blocks["expand"] = std::make_shared(config.pixel_attn_dim, flat, true); + blocks["adaln"] = std::make_shared(config.hidden_size, 4 * flat, true); + blocks["attn"] = std::make_shared(config.pixel_attn_dim, config.pixel_heads); + blocks["mlp.fc1"] = std::make_shared(d, d * 4, true); + blocks["mlp.fc2"] = std::make_shared(d * 4, d, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cond, ggml_tensor* pe) { + auto c = ctx->ggml_ctx; + auto d = config.pixel_dim; + auto p2 = config.patch_size * config.patch_size; + auto patches = cond->ne[1]; + auto batch = cond->ne[2]; + auto mods = std::dynamic_pointer_cast(blocks["adaln"])->forward(ctx, cond); + mods = ggml_reshape_3d(c, mods, 4 * d, p2, patches * batch); + auto m = ggml_ext_chunk(c, mods, 4, 0); + auto h = std::dynamic_pointer_cast(blocks["norm1"])->forward(ctx, x); + h = ggml_reshape_3d(c, h, d * p2, patches, batch); + h = std::dynamic_pointer_cast(blocks["compress"])->forward(ctx, h); + h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, pe); + h = std::dynamic_pointer_cast(blocks["expand"])->forward(ctx, h); + h = ggml_reshape_3d(c, h, d, p2, patches * batch); + x = ggml_add(c, x, Pid::apply_adaln(c, h, m[1], m[0])); + h = std::dynamic_pointer_cast(blocks["norm2"])->forward(ctx, x); + h = std::dynamic_pointer_cast(blocks["mlp.fc1"])->forward(ctx, h); + h = ggml_gelu_erf(c, h); + h = std::dynamic_pointer_cast(blocks["mlp.fc2"])->forward(ctx, h); + return ggml_add(c, x, Pid::apply_adaln(c, h, m[3], m[2])); + } + }; + + struct IrisModel : public GGMLBlock { + IrisConfig config; + + IrisModel(const IrisConfig& config) + : config(config) { + blocks["s_embedder"] = std::make_shared(3 * config.patch_size * config.patch_size, config.hidden_size); + blocks["t_embedder"] = std::make_shared(config.hidden_size); + blocks["y_embedder"] = std::make_shared(config); + blocks["modulation_cores.adaln_img"] = std::make_shared(config.hidden_size, 6 * config.hidden_size, true); + blocks["modulation_cores.adaln_txt"] = std::make_shared(config.hidden_size, 6 * config.hidden_size, true); + for (int64_t i = 0; i < config.depth; ++i) { + blocks["blocks." + std::to_string(i)] = std::make_shared(config, i < config.dual_depth); + } + blocks["pixel_embedder"] = std::make_shared(3, config.pixel_dim); + for (int64_t i = 0; i < config.pixel_depth; ++i) { + blocks["pixel_blocks." + std::to_string(i)] = std::make_shared(config); + } + blocks["final_layer"] = std::make_shared(config.pixel_dim, 3); + } + + void init_params(ggml_context* ctx, const String2TensorStorage& = {}, const std::string = "") override { + params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config.hidden_size, config.text_len, 1); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image, ggml_tensor* t, ggml_tensor* text, ggml_tensor* mask, ggml_tensor* joint_pe, ggml_tensor* pixel_pe, ggml_tensor* pixel_pos) { + auto c = ctx->ggml_ctx; + int p = static_cast(config.patch_size); + auto x = DiT::patchify(c, image, p, p, true); + x = std::dynamic_pointer_cast(blocks["s_embedder"])->forward(ctx, x); + auto te = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); + te = ggml_reshape_3d(c, te, config.hidden_size, 1, image->ne[3]); + auto cond = ggml_silu(c, te); + auto im = std::dynamic_pointer_cast(blocks["modulation_cores.adaln_img"])->forward(ctx, cond); + auto tm = std::dynamic_pointer_cast(blocks["modulation_cores.adaln_txt"])->forward(ctx, cond); + auto y = std::dynamic_pointer_cast(blocks["y_embedder"])->forward(ctx, text, mask); + y = ggml_add(c, y, params["y_pos_embedding"]); + for (int64_t i = 0; i < config.depth; ++i) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + std::tie(x, y) = block->forward(ctx, x, y, im, tm, joint_pe); + sd::ggml_graph_cut::mark_graph_cut(x, "iris.blocks." + std::to_string(i), "x"); + sd::ggml_graph_cut::mark_graph_cut(y, "iris.blocks." + std::to_string(i), "y"); + } + cond = ggml_silu(c, ggml_add(c, x, te)); + x = std::dynamic_pointer_cast(blocks["pixel_embedder"])->forward(ctx, image, p, pixel_pos); + for (int64_t i = 0; i < config.pixel_depth; ++i) { + x = std::dynamic_pointer_cast(blocks["pixel_blocks." + std::to_string(i)])->forward(ctx, x, cond, pixel_pe); + sd::ggml_graph_cut::mark_graph_cut(x, "iris.pixel_blocks." + std::to_string(i), "x"); + } + x = std::dynamic_pointer_cast(blocks["final_layer"])->forward(ctx, x); + x = ggml_reshape_3d(c, x, 3 * p * p, cond->ne[1], image->ne[3]); + return DiT::unpatchify(c, x, image->ne[1] / p, image->ne[0] / p, p, p, false); + } + }; + + inline std::vector image_rope(int64_t height, int64_t width, int64_t dim) { + std::vector pe; + pe.reserve(height * width * dim * 2); + float step = 16.f / static_cast(std::max(std::max(height, width) - 1, 1)); + for (int64_t y = 0; y < height; ++y) { + for (int64_t x = 0; x < width; ++x) { + for (int64_t i = 0; i < dim / 4; ++i) { + float frequency = std::pow(10000.f, -4.f * static_cast(i) / static_cast(dim)); + for (int64_t pos : {x, y}) { + float angle = static_cast(pos) * step * frequency; + pe.insert(pe.end(), {std::cos(angle), -std::sin(angle), std::sin(angle), std::cos(angle)}); + } + } + } + } + return pe; + } + + struct IrisRunner : public DiffusionModelRunner { + IrisConfig config; + IrisModel model; + std::vector joint_pe_data; + std::vector pixel_pe_data; + std::vector pixel_pos_data; + std::vector text_mask_data; + + IrisRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(IrisConfig::detect_from_weights(weights, prefix)), + model(config) { + model.init(params_ctx, weights, prefix); + } + + std::string get_desc() override { return "Iris-3B"; } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const DiffusionParams& inputs) override { + if (!inputs.x || !inputs.timesteps || !inputs.context || !inputs.y || inputs.y->empty()) { + LOG_ERROR("Iris-3B requires text features and their padding mask"); + return {}; + } + auto get_graph = [&]() { + auto gf = new_graph_custom(196608); + auto x = make_input(*inputs.x); + auto t = make_input(*inputs.timesteps); + auto text = make_input(*inputs.context); + int h = static_cast(x->ne[1]); + int w = static_cast(x->ne[0]); + int nt = static_cast(config.text_len); + GGML_ASSERT(h % config.patch_size == 0 && w % config.patch_size == 0); + GGML_ASSERT(text->ne[0] == config.text_dim * config.text_layers && text->ne[1] == nt); + GGML_ASSERT(inputs.y->numel() == nt && x->ne[3] == 1); + joint_pe_data = Pid::make_rope_1d(nt, static_cast(config.head_dim), 10000.f); + auto img = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim); + joint_pe_data.insert(joint_pe_data.end(), img.begin(), img.end()); + pixel_pe_data = image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads); + pixel_pos_data = Pid::make_pixel_abs_pos(h, w, static_cast(config.pixel_dim)); + text_mask_data.resize(nt * nt); + for (int q = 0; q < nt; ++q) { + for (int k = 0; k < nt; ++k) { + text_mask_data[q * nt + k] = inputs.y->values()[k] != 0.f || q == k ? 0.f : -INFINITY; + } + } + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, nt + img.size() / (2 * config.head_dim)); + auto pp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.pixel_attn_dim / config.pixel_heads / 2, (h / config.patch_size) * (w / config.patch_size)); + auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.pixel_dim, h * w, 1); + auto mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nt, nt); + set_backend_tensor_data(pe, joint_pe_data.data()); + set_backend_tensor_data(pp, pixel_pe_data.data()); + set_backend_tensor_data(pos, pixel_pos_data.data()); + set_backend_tensor_data(mask, text_mask_data.data()); + auto ctx = get_context(); + auto out = model.forward(&ctx, x, t, text, mask, pe, pp, pos); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), inputs.x->dim()); + } + }; +} // namespace Iris + +#endif // __SD_MODEL_DIFFUSION_IRIS_HPP__ diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index 84ecfe0b1c..e58df41081 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -689,23 +689,28 @@ namespace Krea2 { } }; - __STATIC_INLINE__ std::vector gen_krea2_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - float theta, - const std::vector& axes_dim, - const std::vector& ref_latents, - Rope::RefIndexMode ref_index_mode) { + __STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + float theta, + const std::vector& axes_dim, + const std::vector& ref_latents, + 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 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); 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); } - 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 { @@ -749,15 +754,15 @@ namespace Krea2 { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = gen_krea2_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim, - ref_latents, - ref_image_params.ref_index_mode); + pe_vec = finish_rope_pe(gen_krea2_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim, + ref_latents, + ref_image_params.ref_index_mode)); int pos_len = static_cast(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); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 35b4b517d4..8fb1512b41 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -384,14 +384,12 @@ namespace Lens { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = Rope::gen_lens_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); int pos_len = static_cast(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); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp index d03383a559..e69c3ecd39 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -412,16 +412,14 @@ namespace LLaDAImage { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = Rope::gen_llada_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); int pos_len = static_cast(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); set_backend_tensor_data(pe, pe_vec.data()); @@ -461,14 +459,14 @@ namespace LLaDAImage { ggml_tensor* source = make_input(source_tensor); GGML_ASSERT(x->ne[3] == 1); - pe_vec = Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(context->ne[1]), - semantic != nullptr ? static_cast(semantic->ne[1]) : 0, - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(context->ne[1]), + semantic != nullptr ? static_cast(semantic->ne[1]) : 0, + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); int pos_len = static_cast(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); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index 31a65afa65..e10d5e829a 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -548,21 +548,22 @@ namespace LTXV { return build_rope_matrix_from_frequencies(freqs, dim); } - __STATIC_INLINE__ ggml_tensor* apply_hidden_rope(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* apply_hidden_rope(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe, int64_t heads, int64_t dim_head, bool rope_interleaved) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(x->ne[0] == heads * dim_head); auto x4 = ggml_reshape_4d(ctx, x, dim_head, heads, x->ne[1], x->ne[2]); if (pe != nullptr && pe->ne[3] == x->ne[1] * heads) { auto x_flat = ggml_reshape_4d(ctx, x4, dim_head, 1, x->ne[1] * heads, x->ne[2]); - auto out_flat = Rope::apply_rope(ctx, x_flat, pe, rope_interleaved); + auto out_flat = Rope::apply_rope(runner_ctx, x_flat, pe, rope_interleaved); auto out4 = ggml_reshape_4d(ctx, out_flat, dim_head, heads, x->ne[1], x->ne[2]); return ggml_reshape_3d(ctx, out4, heads * dim_head, x->ne[1], x->ne[2]); } - return Rope::apply_rope(ctx, x4, pe, rope_interleaved); + return Rope::apply_rope(runner_ctx, x4, pe, rope_interleaved); } struct TimestepEmbedder : public GGMLBlock { @@ -705,8 +706,8 @@ namespace LTXV { if (k_pe == nullptr) { k_pe = pe; } - q = apply_hidden_rope(ctx->ggml_ctx, q, pe, heads, dim_head, rope_interleaved); - k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved); + q = apply_hidden_rope(ctx, q, pe, heads, dim_head, rope_interleaved); + k = apply_hidden_rope(ctx, k, k_pe, heads, dim_head, rope_interleaved); } auto out = ggml_ext_attention_ext(ctx, diff --git a/src/model/diffusion/mage_flow.hpp b/src/model/diffusion/mage_flow.hpp index 89df0f4827..1b9b9f9be9 100644 --- a/src/model/diffusion/mage_flow.hpp +++ b/src/model/diffusion/mage_flow.hpp @@ -110,13 +110,13 @@ namespace MageFlow { } int batch_size = static_cast(x->ne[3]); - pe_vec = Rope::gen_mage_flow_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - batch_size, - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + batch_size, + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); int pos_len = static_cast(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); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/src/model/diffusion/ming_image.hpp b/src/model/diffusion/ming_image.hpp new file mode 100644 index 0000000000..f19be99033 --- /dev/null +++ b/src/model/diffusion/ming_image.hpp @@ -0,0 +1,135 @@ +#ifndef __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ + +#include "z_image.hpp" + +namespace MingImage { + struct MingImageConfig : ZImage::ZImageConfig { + bool split_qkv = true; + + static MingImageConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { + MingImageConfig config; + static_cast(config) = ZImage::ZImageConfig::detect_from_weights(tensors, prefix); + config.split_qkv = tensors.count(prefix + ".layers.0.attention.qkv.weight") == 0; + return config; + } + }; + + class MingImageModel : public GGMLBlock { + MingImageConfig config; + + public: + explicit MingImageModel(const MingImageConfig& config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); + blocks["t_embedder"] = std::make_shared(1024, 256, std::min(config.hidden_size, 256)); + blocks["cap_embedder.0"] = std::make_shared(config.cap_feat_dim, config.norm_eps); + blocks["cap_embedder.1"] = std::make_shared(config.cap_feat_dim, config.hidden_size); + auto add_blocks = [&](const std::string& prefix, int64_t count, bool modulation) { + for (int64_t i = 0; i < count; ++i) { + blocks[prefix + std::to_string(i)] = std::make_shared( + static_cast(i), config.hidden_size, config.head_dim, config.num_heads, + config.num_kv_heads, config.multiple_of, config.ffn_dim_multiplier, + config.norm_eps, config.qk_norm, modulation, true, config.split_qkv, 1e-5f); + } + }; + add_blocks("noise_refiner.", config.num_refiner_layers, true); + add_blocks("context_refiner.", config.num_refiner_layers, false); + add_blocks("layers.", config.num_layers, true); + blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* direct, ggml_tensor* pe) { + auto gctx = ctx->ggml_ctx; + const int64_t width = x->ne[0], height = x->ne[1]; + auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, false); + img = std::dynamic_pointer_cast(blocks["x_embedder"])->forward(ctx, img); + auto txt = std::dynamic_pointer_cast(blocks["cap_embedder.0"])->forward(ctx, context); + txt = std::dynamic_pointer_cast(blocks["cap_embedder.1"])->forward(ctx, txt); + txt = ggml_concat(gctx, txt, direct, 1); + auto t = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, timestep); + const int64_t n_txt = txt->ne[1], n_img = img->ne[1]; + auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, n_txt); + auto img_pe = ggml_ext_slice(gctx, pe, 3, n_txt, n_txt + n_img); + for (int64_t i = 0; i < config.num_refiner_layers; ++i) { + txt = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe); + sd::ggml_graph_cut::mark_graph_cut(txt, "ming_image.context_refiner." + std::to_string(i), "txt"); + } + for (int64_t i = 0; i < config.num_refiner_layers; ++i) { + img = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, img, img_pe, nullptr, t); + sd::ggml_graph_cut::mark_graph_cut(img, "ming_image.noise_refiner." + std::to_string(i), "img"); + } + auto combined = ggml_concat(gctx, txt, img, 1); + for (int64_t i = 0; i < config.num_layers; ++i) { + combined = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, combined, pe, nullptr, t); + sd::ggml_graph_cut::mark_graph_cut(combined, "ming_image.layers." + std::to_string(i), "combined"); + } + img = ggml_ext_slice(gctx, combined, 1, n_txt, n_txt + n_img); + img = std::dynamic_pointer_cast(blocks["final_layer"])->forward(ctx, img, t); + img = DiT::unpatchify_and_crop(gctx, img, height, width, config.patch_size, config.patch_size, false); + return ggml_scale(gctx, img, -1.f); + } + }; + + struct MingImageRunner : DiffusionModelRunner { + MingImageConfig config; + MingImageModel model; + std::vector pe_values; + + MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(MingImageConfig::detect_from_weights(tensors, prefix)), + model(config) { + model.init(params_ctx, tensors, prefix); + } + + std::string get_desc() override { return "ming_image"; } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const DiffusionParams& inputs) override { + const auto* extra = diffusion_extra_as(inputs); + if (inputs.ref_latents != nullptr && !inputs.ref_latents->empty()) { + LOG_ERROR("Ming-Image reference-image conditioning is not supported"); + return {}; + } + if (inputs.context == nullptr || extra->direct_context == nullptr) { + LOG_ERROR("Ming-Image requires both query and direct text conditions"); + return {}; + } + auto graph = [&]() { + auto gf = new_graph_custom(ZImage::Z_IMAGE_GRAPH_SIZE); + auto x = make_input(*inputs.x); + auto t = make_input(*inputs.timesteps); + auto context = make_input(*inputs.context); + auto direct = make_input(*extra->direct_context); + GGML_ASSERT(x->ne[3] == 1); + const int64_t n_txt = context->ne[1] + direct->ne[1]; + const int64_t n_img = ((x->ne[0] + config.patch_size - 1) / config.patch_size) * + ((x->ne[1] + config.patch_size - 1) / config.patch_size); + auto padded = finish_rope_pe(Rope::gen_z_image_pe( + static_cast(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, + static_cast(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED, + config.theta, config.axes_dim)); + // Zero-masked alignment tokens cannot affect valid queries. Omit them while + // retaining the padded caption length used to position image tokens. + const size_t stride = config.axes_dim_sum * 2; + const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast(n_txt), ZImage::SEQ_MULTI_OF); + pe_values.assign(padded.begin(), padded.begin() + n_txt * stride); + pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride, + padded.begin() + (padded_txt + n_img) * stride); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, n_txt + n_img); + set_backend_tensor_data(pe, pe_values.data()); + auto ctx = get_context(); + auto out = model.forward(&ctx, x, t, context, direct, pe); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), inputs.x->dim()); + } + }; +} + +#endif // __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index b065f02bd1..7caaeef7c2 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -256,12 +256,13 @@ namespace MiniMaxH3 { return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]); } - static ggml_tensor* apply_partial_rope(ggml_context* ctx, + static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) { + auto ctx = runner_ctx->ggml_ctx; int64_t rot_dim = pe->ne[2] * 2; GGML_ASSERT(rot_dim <= x->ne[0]); - auto rotated = Rope::apply_rope(ctx, + auto rotated = Rope::apply_rope(runner_ctx, ggml_ext_slice(ctx, x, 0, 0, rot_dim), pe, false); @@ -386,8 +387,8 @@ namespace MiniMaxH3 { q = q_norm->forward(ctx, q); k = k_norm->forward(ctx, k); if (pe != nullptr) { - q = apply_partial_rope(ctx->ggml_ctx, q, pe); - k = apply_partial_rope(ctx->ggml_ctx, k, pe); + q = apply_partial_rope(ctx, q, pe); + k = apply_partial_rope(ctx, k, pe); } else { q = attention_layout(ctx->ggml_ctx, q); k = attention_layout(ctx->ggml_ctx, k); @@ -441,6 +442,9 @@ namespace MiniMaxH3 { for (int64_t i = 0; i < num_layers; ++i) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(x, + "minimax_h3.token_refiner.blocks." + std::to_string(i), + "hidden_states"); } return std::dynamic_pointer_cast(blocks["final_norm"])->forward(ctx, x); } @@ -788,7 +792,11 @@ namespace MiniMaxH3 { GGML_ASSERT(context->ne[0] == config.text_dim); auto condition_proj = std::dynamic_pointer_cast(blocks["condition_proj"]); auto token_refiner = std::dynamic_pointer_cast(blocks["token_refiner"]); - return token_refiner->forward(ctx, condition_proj->forward(ctx, context)); + 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, diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index 06c4db5a5c..0ddacf02eb 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -154,18 +154,26 @@ namespace MiniT2I { return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), head_dim, 10000.f)); } - inline std::vector 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); int dim = head_dim / 2; int quarter = dim / 2; int length = side * side; + Rope::Embedding result; + result.positions.append_image(side, side); std::vector out(static_cast(length) * (head_dim / 2) * 4); std::vector freqs(quarter); for (int i = 0; i < quarter; ++i) { freqs[i] = 1.0f / std::pow(10000.0f, static_cast(2 * i) / static_cast(dim)); } + for (int axis : {1, 2}) { + for (float frequency : freqs) { + result.frequencies.push_back({static_cast(axis), frequency}); + } + } for (int y = 0; y < side; ++y) { for (int x = 0; x < side; ++x) { + result.ids.push_back({0.f, static_cast(y), static_cast(x)}); int pos = y * side + x; size_t base = static_cast(pos) * (head_dim / 2) * 4; 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 { @@ -475,6 +484,8 @@ namespace MiniT2I { int64_t cached_txt_len = -1; int64_t cached_hidden_size = -1; int64_t cached_head_dim = -1; + bool cached_circular_x = false; + bool cached_circular_y = false; MiniT2IRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -521,6 +532,8 @@ namespace MiniT2I { cached_txt_len == txt_len && cached_hidden_size == config.hidden_size && cached_head_dim == config.head_dim && + cached_circular_x == circular_x_enabled && + cached_circular_y == circular_y_enabled && cached_pos_embed != nullptr && cached_txt_pe != nullptr && cached_joint_pe != nullptr) { @@ -531,7 +544,7 @@ namespace MiniT2I { auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast(img_side), static_cast(config.hidden_size)); auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); - auto img_pe_vec = make_vision_rope(static_cast(img_side), static_cast(config.head_dim)); + auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(config.head_dim))); auto joint_pe_vec = txt_pe_vec; 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_hidden_size = config.hidden_size; 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& x_tensor, diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index fb4e362c86..957d575c1b 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -7,7 +7,7 @@ #include "core/ggml_runner.h" #include "core/tensor_ggml.hpp" -#include "model/common/rope.hpp" +#include "model/common/rope_circular.hpp" #include "model_manager.h" enum class RefImageResizeMode { @@ -71,6 +71,8 @@ struct AnimaDiffusionExtra { struct QwenImage21DiffusionExtra { const sd::Tensor* image_slots = nullptr; + // Nonzero IDs identify immutable prefix inputs within one sampling run. + uint64_t prefix_id = 0; }; struct WanDiffusionExtra { @@ -139,6 +141,10 @@ struct LLaDAImageDiffusionExtra { const sd::Tensor* semantic = nullptr; }; +struct MingImageDiffusionExtra { + const sd::Tensor* direct_context = nullptr; +}; + using DiffusionExtraParams = std::variant; + LLaDAImageDiffusionExtra, + MingImageDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; @@ -182,6 +189,11 @@ struct DiffusionModelRunner : public GGMLRunner { protected: std::string prefix; + std::vector finish_rope_pe(Rope::Embedding embedding) { + Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled); + return std::move(embedding.values); + } + public: DiffusionModelRunner(ggml_backend_t backend, const std::string& prefix, diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 59eac29503..6edc83422d 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -135,13 +135,13 @@ namespace Pid { return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), dim, theta)); } - inline std::vector make_rope_2d(int height, - int width, - int dim, - float theta = 10000.f, - float scale = 16.f, - int ref_grid_h = 0, - int ref_grid_w = 0) { + inline Rope::Embedding make_rope_2d(int height, + int width, + int dim, + float theta = 10000.f, + float scale = 16.f, + int ref_grid_h = 0, + int ref_grid_w = 0) { GGML_ASSERT(dim % 4 == 0); 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 Ws = Wp / config.patch_size; - pos_img_vec = make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.hidden_size / config.num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w)); + pos_img_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.hidden_size / config.num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); auto pos_img = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, @@ -904,13 +904,13 @@ namespace Pid { 1); set_backend_tensor_data(pixel_pos, pixel_pos_vec.data()); - pixel_pos_comp_vec = make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w)); + pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp new file mode 100644 index 0000000000..cc06291928 --- /dev/null +++ b/src/model/diffusion/pixart.hpp @@ -0,0 +1,399 @@ +#ifndef __SD_MODEL_DIFFUSION_PIXART_HPP__ +#define __SD_MODEL_DIFFUSION_PIXART_HPP__ + +#include +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/mmdit.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" + +// Ref: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/pixart_transformer_2d.py +// Ref: https://github.com/PixArt-alpha/PixArt-sigma + +namespace PixArt { + constexpr int PIXART_GRAPH_SIZE = 20480; + constexpr int ADALN_EMBED_DIM = 256; + + struct PixArtConfig { + int64_t in_channels = 4; + int64_t out_channels = 8; // learn_sigma: noise prediction + learned variance + int64_t hidden_size = 1152; + int64_t cross_attention_dim = 1152; + int64_t caption_channels = 4096; + int64_t num_heads = 16; + int64_t patch_size = 2; + int64_t ffn_dim = 4608; + int64_t pos_embed_base_size = 64; + float interpolation_scale = 2.f; + int num_layers = 28; + + static PixArtConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) { + PixArtConfig config; + auto find = [&](const std::string& suffix) -> const TensorStorage* { + auto it = weights.find(prefix + "." + suffix); + return it == weights.end() ? nullptr : &it->second; + }; + if (auto w = find("x_embedder.proj.weight")) { + config.hidden_size = w->ne[3]; + config.in_channels = w->ne[2]; + config.patch_size = w->ne[0]; + } + if (auto w = find("final_layer.linear.weight")) { + config.out_channels = w->ne[1] / (config.patch_size * config.patch_size); + } + if (auto w = find("y_embedder.y_proj.fc1.weight")) { + config.caption_channels = w->ne[0]; + } + if (auto w = find("blocks.0.cross_attn.kv_linear.weight")) { + config.cross_attention_dim = w->ne[0]; + } + if (auto w = find("blocks.0.mlp.fc1.weight")) { + config.ffn_dim = w->ne[1]; + } + if (find("csize_embedder.mlp.0.weight") != nullptr) { + LOG_WARN("pixart: resolution/aspect-ratio micro conditions are not supported; output may differ from the reference"); + } + int layers = 0; + const std::string block_prefix = prefix + ".blocks."; + for (const auto& [name, _] : weights) { + if (starts_with(name, block_prefix)) { + layers = std::max(layers, atoi(name.substr(block_prefix.size()).c_str()) + 1); + } + } + if (layers > 0) { + config.num_layers = layers; + LOG_VERBOSE("pixart: layers = %d, hidden_size = %" PRId64, + layers, config.hidden_size); + } + return config; + } + }; + + // Mirrors diffusers get_2d_sincos_pos_embed for a (gh, gw) patch grid. + static std::vector gen_2d_sincos_pos_embed(int64_t dim, + int64_t gh, + int64_t gw, + int64_t base_size, + float interpolation_scale) { + // diffusers: meshgrid(grid_w, grid_h, indexing="xy") -> grid[0]=w, grid[1]=h, + // embedding = concat(sincos(w), sincos(h)) + std::vector out(static_cast(gh) * gw * dim); + int64_t quarter = dim / 4; + for (int64_t h = 0; h < gh; ++h) { + float pos_h = static_cast(h) / (static_cast(gh) / base_size) / interpolation_scale; + for (int64_t w = 0; w < gw; ++w) { + float pos_w = static_cast(w) / (static_cast(gw) / base_size) / interpolation_scale; + float* dst_w = out.data() + (h * gw + w) * dim; + float* dst_h = dst_w + dim / 2; + for (int64_t i = 0; i < quarter; ++i) { + float omega = 1.f / powf(10000.f, static_cast(i) / quarter); + dst_w[i] = sinf(pos_w * omega); + dst_w[i + quarter] = cosf(pos_w * omega); + dst_h[i] = sinf(pos_h * omega); + dst_h[i + quarter] = cosf(pos_h * omega); + } + } + } + return out; + } + + class PixArtTimestepEmbedding : public GGMLBlock { + public: + PixArtTimestepEmbedding(int64_t in_channels, int64_t out_dim) { + blocks["mlp.0"] = std::make_shared(in_channels, out_dim); + blocks["mlp.2"] = std::make_shared(out_dim, out_dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + x = std::dynamic_pointer_cast(blocks["mlp.0"])->forward(ctx, x); + x = ggml_silu(ctx->ggml_ctx, x); + return std::dynamic_pointer_cast(blocks["mlp.2"])->forward(ctx, x); + } + }; + + class PixArtAttention : public GGMLBlock { + int64_t num_heads; + bool self_attention; + + public: + PixArtAttention(int64_t dim, int64_t num_heads, int64_t context_dim, bool self_attention) + : num_heads(num_heads), self_attention(self_attention) { + if (self_attention) { + blocks["qkv"] = std::make_shared(dim, 3 * dim); + } else { + blocks["q_linear"] = std::make_shared(dim, dim); + blocks["kv_linear"] = std::make_shared(context_dim, 2 * dim); + } + blocks["proj"] = std::make_shared(dim, dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* context, ggml_tensor* mask = nullptr) { + std::vector qkv; + if (self_attention) { + auto projected = std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x); + qkv = ggml_ext_chunk(ctx->ggml_ctx, projected, 3, 0); + } else { + auto q = std::dynamic_pointer_cast(blocks["q_linear"])->forward(ctx, x); + auto kv = std::dynamic_pointer_cast(blocks["kv_linear"])->forward(ctx, context); + auto parts = ggml_ext_chunk(ctx->ggml_ctx, kv, 2, 0); + qkv = {q, parts[0], parts[1]}; + } + auto out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, mask, false, ctx->flash_attn_enabled); + return std::dynamic_pointer_cast(blocks["proj"])->forward(ctx, out); + } + }; + + class PixArtBlock : public GGMLBlock { + int64_t dim; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ggml_type wtype = get_type(prefix + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32); + params["scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, dim, 6); + } + + public: + PixArtBlock(int64_t dim, int64_t num_heads, int64_t context_dim, int64_t ffn_dim) + : dim(dim) { + blocks["attn"] = std::make_shared(dim, num_heads, dim, true); + blocks["cross_attn"] = std::make_shared(dim, num_heads, context_dim, false); + blocks["mlp.fc1"] = std::make_shared(dim, ffn_dim); + blocks["mlp.fc2"] = std::make_shared(ffn_dim, dim); + } + + static ggml_tensor* norm(ggml_context* ctx, ggml_tensor* x) { + return ggml_norm(ctx, x, 1e-6f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mod, ggml_tensor* context, ggml_tensor* context_mask) { + // x: [N, n_token, dim] + // mod: [N, 6 * dim], shared adaLN-single output + int64_t N = x->ne[2]; + + auto table = params["scale_shift_table"]; + if (table->type != GGML_TYPE_F32) { + table = ggml_cast(ctx->ggml_ctx, table, GGML_TYPE_F32); + } + table = ggml_reshape_3d(ctx->ggml_ctx, table, dim, 6, 1); + auto m = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, mod, dim, 6, N), table); + auto mv = ggml_ext_chunk(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, m), dim * 6, N), 6, 0); + + auto attn1 = std::dynamic_pointer_cast(blocks["attn"]); + auto attn2 = std::dynamic_pointer_cast(blocks["cross_attn"]); + auto proj = std::dynamic_pointer_cast(blocks["mlp.fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["mlp.fc2"]); + + auto gate = [&](ggml_tensor* y, ggml_tensor* g) { + g = ggml_reshape_3d(ctx->ggml_ctx, g, dim, 1, N); + return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, y, g)); + }; + + auto h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]); + x = gate(attn1->forward(ctx, h, h), mv[2]); + // ada_norm_single: no norm before cross-attention (PixArtMS.py) + x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask)); + h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]); + h = proj->forward(ctx, h); + h = ggml_ext_gelu(ctx->ggml_ctx, h, true); + h = fc2->forward(ctx, h); + return gate(h, mv[5]); + } + }; + + class PixArtModel : public GGMLBlock { + PixArtConfig config; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ggml_type wtype = get_type(prefix + "final_layer.scale_shift_table", tensor_storage_map, GGML_TYPE_F32); + params["final_layer.scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, config.hidden_size, 2); + } + + public: + PixArtModel() = default; + PixArtModel(const PixArtConfig& config) + : config(config) { + blocks["x_embedder.proj"] = std::make_shared(config.in_channels, + config.hidden_size, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}); + blocks["t_embedder"] = std::make_shared(ADALN_EMBED_DIM, config.hidden_size); + blocks["t_block.1"] = std::make_shared(config.hidden_size, 6 * config.hidden_size); + blocks["y_embedder.y_proj.fc1"] = std::make_shared(config.caption_channels, config.hidden_size); + blocks["y_embedder.y_proj.fc2"] = std::make_shared(config.hidden_size, config.cross_attention_dim); + for (int i = 0; i < config.num_layers; ++i) { + blocks["blocks." + std::to_string(i)] = + std::make_shared(config.hidden_size, config.num_heads, config.cross_attention_dim, config.ffn_dim); + } + blocks["final_layer.norm_final"] = std::make_shared(config.hidden_size, 1e-6f, false); + blocks["final_layer.linear"] = std::make_shared(config.hidden_size, + config.patch_size * config.patch_size * config.out_channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timesteps, + ggml_tensor* context, + ggml_tensor* pos_embed, + ggml_tensor* context_mask) { + // x: [N, C, H, W] latent, context: [N, n_ctx, caption_channels] + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t N = x->ne[3]; + int64_t p = config.patch_size; + int64_t wp = W / p; + int64_t hp = H / p; + + auto h = std::dynamic_pointer_cast(blocks["x_embedder.proj"])->forward(ctx, x); // [N, hidden, hp, wp] + h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3)); // [N, hp, wp, hidden] -> [N, hp*wp, hidden] + h = ggml_reshape_3d(ctx->ggml_ctx, h, config.hidden_size, wp * hp, N); // [N, hp*wp, hidden] + h = ggml_add(ctx->ggml_ctx, h, pos_embed); + + auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, ADALN_EMBED_DIM, 10000); + auto emb = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); + + auto mod = std::dynamic_pointer_cast(blocks["t_block.1"]) + ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb)); // [N, 6 * hidden] + + auto ctx_emb = std::dynamic_pointer_cast(blocks["y_embedder.y_proj.fc1"])->forward(ctx, context); + ctx_emb = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true); + ctx_emb = std::dynamic_pointer_cast(blocks["y_embedder.y_proj.fc2"])->forward(ctx, ctx_emb); + + for (int i = 0; i < config.num_layers; ++i) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + h = block->forward(ctx, h, mod, ctx_emb, context_mask); + sd::ggml_graph_cut::mark_graph_cut(h, "pixart.blocks." + std::to_string(i), "h"); + } + + // scale_shift_table + emb -> (shift, scale) for the affine-free final norm + auto tail_table = params["final_layer.scale_shift_table"]; + if (tail_table->type != GGML_TYPE_F32) { + tail_table = ggml_cast(ctx->ggml_ctx, tail_table, GGML_TYPE_F32); + } + auto ss = ggml_add(ctx->ggml_ctx, + ggml_reshape_3d(ctx->ggml_ctx, tail_table, config.hidden_size, 2, 1), + ggml_reshape_3d(ctx->ggml_ctx, emb, config.hidden_size, 1, N)); // [2, hidden, N] + auto parts = ggml_ext_chunk(ctx->ggml_ctx, + ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N), + 2, 0); + h = std::dynamic_pointer_cast(blocks["final_layer.norm_final"])->forward(ctx, h); + h = modulate(ctx->ggml_ctx, h, parts[0], parts[1]); + h = std::dynamic_pointer_cast(blocks["final_layer.linear"])->forward(ctx, h); // [N, hp*wp, p*p*out_ch] + h = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast(p), static_cast(p), false); + return h; // [N, out_channels, H, W] + } + }; + + struct PixArtRunner : public DiffusionModelRunner { + PixArtConfig config; + PixArtModel model; + std::vector pos_vec; + + PixArtRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr, + const char* model_args = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(PixArtConfig::detect_from_weights(tensor_storage_map, prefix)) { + for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) { + if (key == "pixart_pos_embed_base_size") { + int parsed = 0; + if (parse_strict_int(value, parsed)) { + config.pos_embed_base_size = parsed; + } else { + LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "pixart_interpolation_scale") { + float parsed = 0.f; + if (parse_strict_float(value, parsed)) { + config.interpolation_scale = parsed; + } else { + LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + model = PixArtModel(config); + model.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "pixart"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& mask_tensor) { + ggml_cgraph* gf = new_graph_custom(PIXART_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + GGML_ASSERT(x->ne[3] == 1); + GGML_ASSERT(!context_tensor.empty()); + ggml_tensor* context = make_input(context_tensor); + + ggml_tensor* context_mask = nullptr; + if (!mask_tensor.empty()) { + // additive attention bias over context tokens: 0 keep / -inf discard + context_mask = ggml_reshape_4d(compute_ctx, make_input(mask_tensor), mask_tensor.shape()[0], 1, 1, 1); + } + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t wp = W / config.patch_size; + int64_t hp = H / config.patch_size; + + pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp, + config.pos_embed_base_size, config.interpolation_scale); + auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.hidden_size, wp * hp, 1); + set_backend_tensor_data(pos, pos_vec.data()); + + auto runner_ctx = get_context(); + ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask); + // learn_sigma: keep the noise prediction half of the output channels + out = ggml_ext_slice(compute_ctx, out, 2, 0, config.in_channels); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const sd::Tensor& context_mask) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, context_mask); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + auto context = tensor_or_empty(diffusion_params.context); + auto context_msk = tensor_or_empty(diffusion_params.y); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + context, + context_msk); + } + }; +} // namespace PixArt + +#endif // __SD_MODEL_DIFFUSION_PIXART_HPP__ diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index 07364958f6..82d1730059 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -635,18 +635,16 @@ namespace Qwen { ref_index_mode = Rope::RefIndexMode::DECREASE; } - pe_vec = Rope::gen_qwen_image_pe(time_len, - static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - batch_size, - static_cast(context->ne[1]), - ref_latents, - ref_index_mode, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len, + static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + batch_size, + static_cast(context->ne[1]), + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index cb044053fe..1d37d6afeb 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -68,6 +68,7 @@ namespace Qwen { std::vector segments; std::vector> positions; int64_t prefix_length = 0; + Rope::PositionLayout rope_layout; static QwenImage21Layout build(int64_t text_length, const sd::Tensor& image_slots, @@ -82,6 +83,7 @@ namespace Qwen { auto [height, width] = image_shapes[index]; int64_t start = static_cast(layout.positions.size()); layout.segments.push_back({start, start + height * width, context_start, index}); + layout.rope_layout.append_image(static_cast(height), static_cast(width)); for (int64_t h = 0; h < height; ++h) { for (int64_t w = 0; w < width; ++w) { layout.positions.push_back({static_cast(position), @@ -106,6 +108,7 @@ namespace Qwen { } else { int64_t start = static_cast(layout.positions.size()); 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) { float p = static_cast(position); 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 { public: using RMSNorm::RMSNorm; @@ -160,27 +177,70 @@ namespace Qwen { } } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector& segments, const std::vector& masks) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector& segments, const std::vector& masks, const QwenImage21PrefixCache& cache) { int64_t heads = x->ne[0] / dim_head; auto project = [&](const char* name) { auto h = std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]); }; - auto q = project("to_q"); - auto k = project("to_k"); - auto v = project("to_v"); - q = std::dynamic_pointer_cast(blocks["norm_q"])->forward(ctx, q); - k = std::dynamic_pointer_cast(blocks["norm_k"])->forward(ctx, k); - q = Rope::apply_rope(ctx->ggml_ctx, q, pe); - k = Rope::apply_rope(ctx->ggml_ctx, k, pe); + auto q = project("to_q"); + auto k = project("to_k"); + auto v = project("to_v"); + q = std::dynamic_pointer_cast(blocks["norm_q"])->forward(ctx, q); + k = std::dynamic_pointer_cast(blocks["norm_k"])->forward(ctx, k); + q = Rope::apply_rope(ctx, q, pe); + k = Rope::apply_rope(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; - 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 = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled); - result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); + if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) { + auto prefix_k = ctx->load_cache_tensor(cache.name + ".k"); + auto prefix_v = ctx->load_cache_tensor(cache.name + ".v"); + GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr); + if (prefix_k->type != k->type) { + prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type); + } + 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(blocks["to_out.0"]); 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); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks) { - auto h = std::dynamic_pointer_cast(blocks["img_norm1"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length); - h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks); - x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true)); - h = std::dynamic_pointer_cast(blocks["img_norm2"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length); + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks, const QwenImage21PrefixCache& cache) { + const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length; + auto h = std::dynamic_pointer_cast(blocks["img_norm1"])->forward(ctx, x); + h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length); + h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache); + x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true)); + h = std::dynamic_pointer_cast(blocks["img_norm2"])->forward(ctx, x); + h = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length); ggml_tensor* gate; auto fused = blocks.find("img_mlp.gate_up"); if (fused != blocks.end()) { @@ -239,7 +300,7 @@ namespace Qwen { } h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate)); h = std::dynamic_pointer_cast(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& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks, const QwenImage21PrefixCache& cache) { 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. 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); auto modulation = std::dynamic_pointer_cast(blocks["modulation.1"])->forward(ctx, time); auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0); - auto text = std::dynamic_pointer_cast(blocks["txt_in"])->forward(ctx, context); auto img_in = std::dynamic_pointer_cast(blocks["img_in"]); ggml_tensor* joint = nullptr; - for (const auto& segment : layout.segments) { - ggml_tensor* h; - if (segment.image_index < 0) { - 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(refs.size()) ? x : refs[segment.image_index]; - h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1)); + if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) { + joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1)); + } else { + auto text = std::dynamic_pointer_cast(blocks["txt_in"])->forward(ctx, context); + for (const auto& segment : layout.segments) { + ggml_tensor* h; + if (segment.image_index < 0) { + 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(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"); for (int i = 0; i < config.num_layers; ++i) { - auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); - joint = block->forward(ctx, joint, mod, pe, layout, masks); - sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint"); + const std::string layer = "transformer_blocks." + std::to_string(i); + auto layer_cache = cache; + layer_cache.name = cache.name + "." + std::to_string(i); + layer_cache.cut_group = "qwen_image_2_1." + layer; + auto block = std::dynamic_pointer_cast(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(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]); joint = std::dynamic_pointer_cast(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)); @@ -303,11 +374,61 @@ namespace Qwen { QwenImage21Model model; std::vector pe_data; std::vector> 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 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_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 weight_manager = nullptr, const char* model_args = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(QwenImage21Config::detect_from_weights(weights, prefix)), 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); } @@ -317,6 +438,20 @@ namespace Qwen { 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 compute(int n_threads, const DiffusionParams& inputs) override { const auto& x = tensor_or_empty(inputs.x); const auto& context = tensor_or_empty(inputs.context); @@ -345,38 +480,113 @@ namespace Qwen { LOG_ERROR("%s", error.what()); return {}; } - pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim); - mask_data.clear(); - for (const auto& segment : layout.segments) { - sd::Tensor mask; - if (segment.image_index < 0) { - mask = sd::Tensor::zeros({segment.end, segment.end - segment.start}); - 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; - } - } + if (!runner_started()) { + prefix_cache_disabled = false; + prefix_cache_auto_f32 = false; + } + QwenImage21PrefixCache cache; + if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) { + cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) + + ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled); + 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; } - mask_data.push_back(std::move(mask)); - } - auto build = [&]() { - 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, layout.positions.size()); - set_backend_tensor_data(pe, pe_data.data()); - std::vector masks, ref_inputs; - for (const auto& mask : mask_data) { - masks.push_back(mask.empty() ? nullptr : make_input(mask)); + 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; } - for (const auto& ref : refs) { - ref_inputs.push_back(make_input(ref)); + 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(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); + } } - auto ctx = get_context(); - auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context), - ref_inputs, pe, layout, masks); - ggml_build_forward_expand(graph, out); - return graph; + 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)); + mask_data.clear(); + if (!cached) { + for (const auto& segment : layout.segments) { + sd::Tensor mask; + if (segment.image_index < 0) { + mask = sd::Tensor::zeros({segment.end, segment.end - segment.start}); + 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)); + } + } + auto build = [&]() { + 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, + layout.positions.size() - (cached ? layout.prefix_length : 0)); + set_backend_tensor_data(pe, pe_data.data()); + std::vector masks, ref_inputs; + for (const auto& mask : mask_data) { + 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(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; } }; } diff --git a/src/model/diffusion/sensenova_u1.h b/src/model/diffusion/sensenova_u1.h index 22b789fac5..1702ed8546 100644 --- a/src/model/diffusion/sensenova_u1.h +++ b/src/model/diffusion/sensenova_u1.h @@ -442,16 +442,9 @@ namespace SenseNovaU1 { k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2); v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2); } else { - // Keep dedicated graph outputs alive until the runner copies them - // into its persistent cache buffer after graph execution. - auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k); - cache_k = ggml_cpy(ctx->ggml_ctx, k, cache_k); - ggml_set_output(cache_k); - auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v); - cache_v = ggml_cpy(ctx->ggml_ctx, v, cache_v); - ggml_set_output(cache_v); - ctx->persist_cache_tensor(layer_cache + ".k", cache_k); - ctx->persist_cache_tensor(layer_cache + ".v", cache_v); + ctx->expand_graph(q); + ctx->persist_cache_tensor(layer_cache + ".k", k); + ctx->persist_cache_tensor(layer_cache + ".v", v); } q = ggml_cont(ctx->ggml_ctx, @@ -687,7 +680,7 @@ namespace SenseNovaU1 { ggml_set_name(attention_mask, "snu15.prefix.attention_mask"); set_backend_tensor_data(attention_mask, attention_mask_vec.data()); - auto runner_ctx = get_context(); + auto runner_ctx = get_context(graph); auto text_model = model.text_model(); auto hidden = text_model->embed(&runner_ctx, ids); hidden = text_model->forward(&runner_ctx, diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index a1012cad84..ed2644047e 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -140,7 +140,8 @@ namespace ZImage { int64_t num_kv_heads, bool qk_norm, bool norm_elementwise_affine = true, - bool split_qkv = false) + bool split_qkv = false, + float qk_norm_eps = 1e-6f) : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm), split_qkv(split_qkv) { float scale = 1.f; if (split_qkv) { @@ -153,8 +154,8 @@ namespace ZImage { blocks["out"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); } if (qk_norm) { - blocks["q_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); - blocks["k_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); + blocks["q_norm"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + blocks["k_norm"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); } } @@ -167,13 +168,16 @@ namespace ZImage { int64_t N = x->ne[2]; auto out_proj = std::dynamic_pointer_cast(blocks[split_qkv ? "to_out.0" : "out"]); + if (out_proj->has_quantized_weight()) { + out_proj->set_scale(1.f / 16.f); + } + if (split_qkv) { auto q_proj = std::dynamic_pointer_cast(blocks["to_q"]); auto k_proj = std::dynamic_pointer_cast(blocks["to_k"]); auto v_proj = std::dynamic_pointer_cast(blocks["to_v"]); if (sd_backend_is(ctx->backend, "ROCm")) { - out_proj->set_scale(1.f / 16.f); out_proj->set_force_prec_f32(true); q_proj->set_force_prec_f32(true); k_proj->set_force_prec_f32(true); @@ -196,7 +200,6 @@ namespace ZImage { auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); if (sd_backend_is(ctx->backend, "ROCm")) { - out_proj->set_scale(1.f / 16.f); out_proj->set_force_prec_f32(true); qkv_proj->set_force_prec_f32(true); } @@ -318,9 +321,10 @@ namespace ZImage { bool qk_norm, bool modulation = true, bool norm_elementwise_affine = true, - bool split_qkv = false) + bool split_qkv = false, + float qk_norm_eps = 1e-6f) : modulation(modulation) { - blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv); + blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv, qk_norm_eps); blocks["feed_forward"] = std::make_shared(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier); blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); @@ -642,18 +646,16 @@ namespace ZImage { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = Rope::gen_z_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - SEQ_MULTI_OF, - ref_latents, - ref_index_mode, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + SEQ_MULTI_OF, + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); diff --git a/src/model/diffusion/z_image_l2p.hpp b/src/model/diffusion/z_image_l2p.hpp new file mode 100644 index 0000000000..0eb0b14d3d --- /dev/null +++ b/src/model/diffusion/z_image_l2p.hpp @@ -0,0 +1,247 @@ +#ifndef __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__ +#define __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__ + +#include + +#include "z_image.hpp" + +// Ref: https://github.com/TencentYoutuResearch/T2I-L2P/blob/main/diffsynth/models/z_image_dit_L2P.py + +namespace ZImageL2P { + struct ZImageL2PConfig : ZImage::ZImageConfig { + static ZImageL2PConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { + ZImageL2PConfig config; + static_cast(config) = ZImage::ZImageConfig::detect_from_weights(tensors, prefix); + config.in_channels = 3; + config.out_channels = 3; + config.patch_size = 16; + auto x_embedder = tensors.find(prefix + ".x_embedder.weight"); + if (x_embedder != tensors.end()) { + config.patch_size = static_cast(std::lround(std::sqrt(static_cast(x_embedder->second.ne[0] / config.in_channels)))); + } + // L2P ships split to_q/to_k/to_v, which name conversion maps to qkv.weight plus + // .1/.2 parts; the base detection only sees the q part and undercounts kv heads. + auto k_part = tensors.find(prefix + ".layers.0.attention.qkv.weight.1"); + if (k_part != tensors.end()) { + config.num_kv_heads = k_part->second.ne[1] / config.head_dim; + } + LOG_VERBOSE("z_image_l2p: patch_size = %d, in_channels = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64, + config.patch_size, + config.in_channels, + config.num_heads, + config.num_kv_heads); + return config; + } + }; + + // U-Net over the full-resolution noisy image. Its bottleneck sits at the DiT token + // grid, so the number of pooling levels is tied to patch_size == 16. + class LocalDecoder : public GGMLBlock { + static constexpr int LEVELS = 4; + static constexpr int64_t CHANNELS[LEVELS] = {64, 128, 256, 512}; + static constexpr int64_t BOTTLENECK_CHANNELS = 512; + + public: + LocalDecoder(int64_t in_channels, int64_t cond_channels) { + int64_t prev = in_channels; + for (int i = 0; i < LEVELS; ++i) { + blocks["enc" + std::to_string(i + 1) + ".0"] = std::make_shared(prev, CHANNELS[i], std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + prev = CHANNELS[i]; + } + blocks["bottleneck.0"] = std::make_shared(prev + cond_channels, BOTTLENECK_CHANNELS, std::pair{1, 1}); + prev = BOTTLENECK_CHANNELS; + for (int i = LEVELS - 1; i >= 0; --i) { + const std::string level = std::to_string(i + 1); + const int64_t out = i == 0 ? CHANNELS[0] : CHANNELS[i - 1]; + blocks["up" + level + ".1"] = std::make_shared(prev, prev, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + blocks["dec" + level + ".0"] = std::make_shared(prev + CHANNELS[i], out, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + prev = out; + } + blocks["out_conv"] = std::make_shared(prev, in_channels, std::pair{1, 1}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cond) { + // x: [N, C, H, W] + // cond: [N, cond_channels, H / 16, W / 16] + // return: [N, C, H, W] + auto gctx = ctx->ggml_ctx; + auto conv = [&](const std::string& name, ggml_tensor* h) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); + }; + + std::vector skips; + auto h = x; + for (int i = 0; i < LEVELS; ++i) { + h = ggml_silu(gctx, conv("enc" + std::to_string(i + 1) + ".0", h)); + skips.push_back(h); + h = ggml_pool_2d(gctx, h, GGML_OP_POOL_MAX, 2, 2, 2, 2, 0, 0); + } + + h = ggml_concat(gctx, h, cond, 2); + h = ggml_silu(gctx, conv("bottleneck.0", h)); + + for (int i = LEVELS - 1; i >= 0; --i) { + const std::string level = std::to_string(i + 1); + h = ggml_upscale(gctx, h, 2, GGML_SCALE_MODE_NEAREST); + h = conv("up" + level + ".1", h); + h = ggml_concat(gctx, h, skips[i], 2); + h = ggml_silu(gctx, conv("dec" + level + ".0", h)); + } + return conv("out_conv", h); + } + }; + + class ZImageL2PModel : public GGMLBlock { + ZImageL2PConfig config; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + } + + public: + explicit ZImageL2PModel(const ZImageL2PConfig& config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); + blocks["t_embedder"] = std::make_shared(MIN(config.hidden_size, 1024), 256, 256); + blocks["cap_embedder.0"] = std::make_shared(config.cap_feat_dim, config.norm_eps); + blocks["cap_embedder.1"] = std::make_shared(config.cap_feat_dim, config.hidden_size); + auto add_blocks = [&](const std::string& prefix, int64_t count, bool modulation) { + for (int64_t i = 0; i < count; ++i) { + blocks[prefix + std::to_string(i)] = std::make_shared( + static_cast(i), config.hidden_size, config.head_dim, config.num_heads, + config.num_kv_heads, config.multiple_of, config.ffn_dim_multiplier, + config.norm_eps, config.qk_norm, modulation, true, false, 1e-5f); + } + }; + add_blocks("noise_refiner.", config.num_refiner_layers, true); + add_blocks("context_refiner.", config.num_refiner_layers, false); + add_blocks("layers.", config.num_layers, true); + blocks["local_decoder"] = std::make_shared(config.in_channels, config.hidden_size); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* pe) { + // x: [N, C, H, W], H and W are multiples of patch_size + // context: [N, L, cap_feat_dim] + // return: [N, C, H, W] + auto gctx = ctx->ggml_ctx; + const int64_t W = x->ne[0]; + const int64_t H = x->ne[1]; + const int64_t N = x->ne[3]; + const int64_t w = W / config.patch_size; + const int64_t h = H / config.patch_size; + const int64_t n_img = w * h; + const int64_t n_txt = context->ne[1]; + + auto img = DiT::patchify(gctx, x, config.patch_size, config.patch_size, false); + img = std::dynamic_pointer_cast(blocks["x_embedder"])->forward(ctx, img); + auto txt = std::dynamic_pointer_cast(blocks["cap_embedder.0"])->forward(ctx, context); + txt = std::dynamic_pointer_cast(blocks["cap_embedder.1"])->forward(ctx, txt); + auto t = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, timestep); + sd::ggml_graph_cut::mark_graph_cut(txt, "z_image_l2p.prelude", "txt"); + sd::ggml_graph_cut::mark_graph_cut(img, "z_image_l2p.prelude", "img"); + sd::ggml_graph_cut::mark_graph_cut(t, "z_image_l2p.prelude", "t_emb"); + + const int64_t n_txt_pad = Rope::bound_mod(static_cast(n_txt), ZImage::SEQ_MULTI_OF); + if (n_txt_pad > 0) { + auto pad = params["cap_pad_token"]; + txt = ggml_concat(gctx, txt, ggml_repeat_4d(gctx, pad, pad->ne[0], n_txt_pad, N, 1), 1); + } + const int64_t n_img_pad = Rope::bound_mod(static_cast(n_img), ZImage::SEQ_MULTI_OF); + if (n_img_pad > 0) { + auto pad = params["x_pad_token"]; + img = ggml_concat(gctx, img, ggml_repeat_4d(gctx, pad, pad->ne[0], n_img_pad, N, 1), 1); + } + GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]); + + auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, txt->ne[1]); + auto img_pe = ggml_ext_slice(gctx, pe, 3, txt->ne[1], pe->ne[3]); + for (int64_t i = 0; i < config.num_refiner_layers; ++i) { + txt = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe); + sd::ggml_graph_cut::mark_graph_cut(txt, "z_image_l2p.context_refiner." + std::to_string(i), "txt"); + } + for (int64_t i = 0; i < config.num_refiner_layers; ++i) { + img = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, img, img_pe, nullptr, t); + sd::ggml_graph_cut::mark_graph_cut(img, "z_image_l2p.noise_refiner." + std::to_string(i), "img"); + } + + auto txt_img = ggml_concat(gctx, txt, img, 1); + for (int64_t i = 0; i < config.num_layers; ++i) { + txt_img = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, txt_img, pe, nullptr, t); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image_l2p.layers." + std::to_string(i), "txt_img"); + } + + // The local decoder consumes the raw last-layer hidden states: there is no final + // norm or adaLN modulation in front of it, unlike Z-Image's final_layer. + const int64_t img_start = n_txt + n_txt_pad; + auto feat = ggml_ext_slice(gctx, txt_img, 1, img_start, img_start + n_img); // [N, h*w, hidden_size] + feat = ggml_reshape_4d(gctx, feat, config.hidden_size, w, h, N); // [N, h, w, hidden_size] + feat = ggml_cont(gctx, ggml_permute(gctx, feat, 2, 0, 1, 3)); // [N, hidden_size, h, w] + auto out = std::dynamic_pointer_cast(blocks["local_decoder"])->forward(ctx, x, feat); // [N, C, H, W] + return ggml_ext_scale(gctx, out, -1.f); + } + }; + + struct ZImageL2PRunner : public DiffusionModelRunner { + ZImageL2PConfig config; + ZImageL2PModel model; + std::vector pe_vec; + + ZImageL2PRunner(ggml_backend_t backend, + const String2TensorStorage& tensors, + const std::string& prefix, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(ZImageL2PConfig::detect_from_weights(tensors, prefix)), + model(config) { + model.init(params_ctx, tensors, prefix); + } + + std::string get_desc() override { return "z_image_l2p"; } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const DiffusionParams& inputs) override { + GGML_ASSERT(inputs.x != nullptr); + GGML_ASSERT(inputs.timesteps != nullptr); + if (inputs.ref_latents != nullptr && !inputs.ref_latents->empty()) { + LOG_ERROR("Z-Image L2P reference-image conditioning is not supported"); + return {}; + } + if (inputs.context == nullptr || inputs.context->empty()) { + LOG_ERROR("Z-Image L2P requires a text condition"); + return {}; + } + auto graph = [&]() { + auto gf = new_graph_custom(ZImage::Z_IMAGE_GRAPH_SIZE); + auto x = make_input(*inputs.x); + auto t = make_input(*inputs.timesteps); + auto context = make_input(*inputs.context); + GGML_ASSERT(x->ne[3] == 1); + GGML_ASSERT(x->ne[0] % config.patch_size == 0 && x->ne[1] % config.patch_size == 0); + pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + {}, + Rope::RefIndexMode::FIXED, + config.theta, + config.axes_dim)); + int pos_len = static_cast(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); + set_backend_tensor_data(pe, pe_vec.data()); + auto ctx = get_context(); + auto out = model.forward(&ctx, x, t, context, pe); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), inputs.x->dim()); + } + }; +} // namespace ZImageL2P + +#endif // __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__ diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index da59cfab61..0d00eedbec 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -50,6 +50,8 @@ namespace LLM { GEMMA4_12B, GPT_OSS_20B, LLADA2_MOE, + BAILING_MOE, + QWEN2, ARCH_COUNT, }; @@ -64,6 +66,8 @@ namespace LLM { "gemma4_12b", "gpt_oss_20b", "llada2_moe", + "bailing_moe", + "qwen2", }; enum class MLPActivation { @@ -225,7 +229,7 @@ namespace LLM { config.intermediate_size = 9216; config.num_layers = 26; config.vocab_size = 256000; - } else if (arch == LLMArch::LLADA2_MOE) { + } else if (arch == LLMArch::LLADA2_MOE || arch == LLMArch::BAILING_MOE) { config.head_dim = 128; config.num_heads = 16; config.num_kv_heads = 4; @@ -240,7 +244,7 @@ namespace LLM { config.max_position_embeddings = 16384; config.rope_thetas = {600000.f}; config.qkv_fused = true; - config.bidirectional = true; + config.bidirectional = arch == LLMArch::LLADA2_MOE; config.partial_rotary = 0.5f; config.num_experts = 256; config.num_experts_per_tok = 8; @@ -250,6 +254,18 @@ namespace LLM { config.n_group = 8; config.topk_group = 4; config.routed_scaling_factor = 2.5f; + if (arch == LLMArch::BAILING_MOE) { + config.vocab_size = 157184; + config.max_position_embeddings = 32768; + } + } else if (arch == LLMArch::QWEN2) { + config.hidden_size = 1536; + config.intermediate_size = 8960; + config.num_heads = 12; + config.num_kv_heads = 2; + config.vocab_size = 151936; + config.max_position_embeddings = 32768; + config.bidirectional = true; } else if (arch == LLMArch::GPT_OSS_20B) { config.head_dim = 64; config.num_heads = 64; @@ -283,22 +299,22 @@ namespace LLM { if (contains(name, "attn.q_proj")) { config.llama_cpp_style = true; } - if (contains(name, "visual.patch_embed.proj.1.weight")) { + if (ends_with(name, "visual.patch_embed.proj.1.weight")) { config.vision.split_patch_embed = true; } - if (contains(name, "visual.patch_embed.proj.0.weight")) { + if (ends_with(name, "visual.patch_embed.proj.0.weight")) { config.vision.patch_size = static_cast(tensor_storage.ne[0]); config.vision.in_channels = tensor_storage.ne[2]; config.vision.hidden_size = tensor_storage.ne[3]; } // HF-format checkpoints keep the patch embed unsplit under a single name. - if (contains(name, "visual.patch_embed.proj.weight")) { + if (ends_with(name, "visual.patch_embed.proj.weight")) { config.vision.patch_size = static_cast(tensor_storage.ne[0]); } if (contains(name, "visual.patch_embed.bias") || contains(name, "visual.patch_embed.proj.bias")) { config.vision.hidden_size = tensor_storage.ne[0]; } - if (contains(name, "visual.pos_embed.weight")) { + if (ends_with(name, "visual.pos_embed.weight") && tensor_storage.n_dims == 2) { config.vision.hidden_size = tensor_storage.ne[0]; config.vision.num_position_embeddings = static_cast(tensor_storage.ne[1]); } @@ -332,7 +348,7 @@ namespace LLM { } } } - if (contains(name, "embed_tokens.weight")) { + if (ends_with(name, "embed_tokens.weight") && tensor_storage.n_dims == 2) { config.hidden_size = tensor_storage.ne[0]; config.vocab_size = tensor_storage.ne[1]; } @@ -471,6 +487,8 @@ namespace LLM { int64_t n_group; int64_t topk_group; float routed_scaling_factor; + bool image_router; + bool fused_experts = false; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, @@ -488,6 +506,10 @@ namespace LLM { // scores and the group sums match. params["gate.weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts); params["gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts); + if (image_router) { + params["image_gate.weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts); + params["image_gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts); + } ggml_type gate_type = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); ggml_type up_type = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); @@ -506,8 +528,14 @@ namespace LLM { } }; - declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size); - declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size); + fused_experts = tensor_storage_map.count(prefix + "experts.gate_up_proj.weight") != 0; + if (fused_experts) { + auto type = supported_type(get_type(prefix + "experts.gate_up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); + declare_experts("experts.gate_up_proj.weight", type, hidden_size, 2 * moe_intermediate_size); + } else { + declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size); + declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size); + } declare_experts("experts.down_proj.weight", down_type, moe_intermediate_size, hidden_size); } @@ -519,7 +547,8 @@ namespace LLM { num_experts_per_tok(config.num_experts_per_tok), n_group(config.n_group), topk_group(config.topk_group), - routed_scaling_factor(config.routed_scaling_factor) { + routed_scaling_factor(config.routed_scaling_factor), + image_router(config.arch == LLMArch::BAILING_MOE) { if (config.num_shared_experts > 0) { blocks["shared_experts"] = std::make_shared(config.hidden_size, config.moe_intermediate_size * config.num_shared_experts, @@ -582,7 +611,7 @@ namespace LLM { return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* image_mask = nullptr) { // x: [N, n_token, hidden_size] GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0); GGML_ASSERT(n_group > 0 && topk_group > 0 && num_experts % n_group == 0); @@ -596,10 +625,21 @@ namespace LLM { auto logits = ggml_mul_mat(gctx, params["gate.weight"], x); logits = ggml_reshape_2d(gctx, logits, num_experts, n_token_total); - auto scores = ggml_sigmoid(gctx, logits); // [num_experts, tokens] + auto bias = params["gate.expert_bias"]; + if (image_router) { + GGML_ASSERT(image_mask != nullptr); + auto mask = ggml_reshape_2d(gctx, image_mask, 1, n_token_total); + auto inverse_mask = ggml_scale_bias(gctx, mask, -1.f, 1.f); + auto image_logits = ggml_reshape_2d(gctx, ggml_mul_mat(gctx, params["image_gate.weight"], x), num_experts, n_token_total); + logits = ggml_add(gctx, ggml_mul(gctx, logits, inverse_mask), ggml_mul(gctx, image_logits, mask)); + bias = ggml_add(gctx, + ggml_mul(gctx, ggml_repeat_4d(gctx, bias, num_experts, n_token_total, 1, 1), inverse_mask), + ggml_mul(gctx, ggml_repeat_4d(gctx, params["image_gate.expert_bias"], num_experts, n_token_total, 1, 1), mask)); + } + auto scores = ggml_sigmoid(gctx, logits); // The bias steers selection only; the combine weights come from the unbiased scores. - auto routing = ggml_add(gctx, scores, params["gate.expert_bias"]); + auto routing = ggml_add(gctx, scores, bias); routing = ggml_add(gctx, routing, group_limited_mask(ctx, routing, n_token_total)); auto selected_experts = ggml_argsort_top_k(gctx, routing, (int)num_experts_per_tok); // [top_k, tokens] @@ -614,12 +654,17 @@ namespace LLM { weights = ggml_scale(gctx, weights, routed_scaling_factor); weights = ggml_reshape_3d(gctx, weights, 1, num_experts_per_tok, n_token_total); - auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total); - auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts); - auto up = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts); - auto activated = ggml_swiglu_split(gctx, gate, up); - auto experts = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts); - experts = ggml_mul(gctx, experts, weights); + auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total); + ggml_tensor* activated; + if (fused_experts) { + activated = ggml_swiglu(gctx, expert_linear(ctx, "experts.gate_up_proj.weight", xf, selected_experts)); + } else { + auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts); + auto up = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts); + activated = ggml_swiglu_split(gctx, gate, up); + } + auto experts = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts); + experts = ggml_mul(gctx, experts, weights); ggml_tensor* out = nullptr; for (int64_t i = 0; i < num_experts_per_tok; ++i) { @@ -1428,7 +1473,9 @@ namespace LLM { ggml_tensor* x, ggml_tensor* input_pos, ggml_tensor* attention_mask = nullptr, - int rope_index = 0) { + int rope_index = 0, + ggml_tensor* rope_cos = nullptr, + ggml_tensor* rope_sin = nullptr) { // x: [N, n_token, hidden_size] int64_t n_token = x->ne[1]; int64_t N = x->ne[2]; @@ -1471,13 +1518,28 @@ namespace LLM { v = ggml_rms_norm(ctx->ggml_ctx, v, rms_norm_eps); } - if (arch == LLMArch::MISTRAL_SMALL_3_2) { + if (rope_cos != nullptr) { + GGML_ASSERT(rope_sin != nullptr); + // Bailing video RoPE interleaves spatial frequencies within a partial NEOX head. + auto rotate = [&](ggml_tensor* input) { + auto gctx = ctx->ggml_ctx; + int64_t half = rope_cos->ne[0]; + auto first = ggml_ext_slice(gctx, input, 0, 0, half); + auto second = ggml_ext_slice(gctx, input, 0, half, 2 * half); + auto left = ggml_sub(gctx, ggml_mul(gctx, first, rope_cos), ggml_mul(gctx, second, rope_sin)); + auto right = ggml_add(gctx, ggml_mul(gctx, second, rope_cos), ggml_mul(gctx, first, rope_sin)); + auto rotated = ggml_concat(gctx, left, right, 0); + return ggml_concat(gctx, rotated, ggml_ext_slice(gctx, input, 0, 2 * half, input->ne[0]), 0); + }; + q = rotate(q); + k = rotate(k); + } else if (arch == LLMArch::MISTRAL_SMALL_3_2) { q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); } else if (arch == LLMArch::MINISTRAL_3_3B) { q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); - } else if (arch == LLMArch::QWEN3) { + } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN2) { q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); } else if (arch == LLMArch::GPT_OSS_20B) { @@ -1655,7 +1717,7 @@ namespace LLM { v_attn = ggml_reshape_3d(ctx->ggml_ctx, v_attn, n_token, head_dim, num_kv_heads * N); auto kq = ggml_mul_mat(ctx->ggml_ctx, k, q); - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_prec_set_acc(kq, GGML_PREC_F32); kq = ggml_scale_inplace(ctx->ggml_ctx, kq, 1.0f / std::sqrt(static_cast(head_dim))); if (attention_mask != nullptr) { kq = ggml_add_inplace(ctx->ggml_ctx, kq, attention_mask); @@ -1724,7 +1786,7 @@ namespace LLM { blocks["self_attn"] = std::make_shared(config, sliding_attention == 0); if (config.arch == LLMArch::GPT_OSS_20B) { blocks["mlp"] = std::make_shared(config); - } else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) { + } else if ((config.arch == LLMArch::LLADA2_MOE || config.arch == LLMArch::BAILING_MOE) && layer_index >= config.first_k_dense_replace) { blocks["mlp"] = std::make_shared(config); } else { blocks["mlp"] = std::make_shared(config.hidden_size, @@ -1746,7 +1808,10 @@ namespace LLM { ggml_tensor* x, ggml_tensor* input_pos, ggml_tensor* attention_mask = nullptr, - ggml_tensor* sliding_attention_mask = nullptr) { + ggml_tensor* sliding_attention_mask = nullptr, + ggml_tensor* image_mask = nullptr, + ggml_tensor* rope_cos = nullptr, + ggml_tensor* rope_sin = nullptr) { // x: [N, n_token, hidden_size] auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); auto input_layernorm = std::dynamic_pointer_cast(blocks["input_layernorm"]); @@ -1768,7 +1833,7 @@ namespace LLM { auto residual = x; x = input_layernorm->forward(ctx, x); - x = self_attn->forward(ctx, x, input_pos, block_attention_mask, rope_index); + x = self_attn->forward(ctx, x, input_pos, block_attention_mask, rope_index, rope_cos, rope_sin); if (post_attention_norm != nullptr) { x = post_attention_norm->forward(ctx, x); } @@ -1782,7 +1847,7 @@ namespace LLM { } else if (auto moe_mlp = std::dynamic_pointer_cast(blocks["mlp"])) { // LLaDA2 is dense for the first first_k_dense_replace layers and MoE afterwards, // so the block type varies per layer rather than per arch. - x = moe_mlp->forward(ctx, x); + x = moe_mlp->forward(ctx, x, image_mask); } else { auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); x = mlp->forward(ctx, x); @@ -1804,10 +1869,11 @@ namespace LLM { protected: int64_t num_layers; LLMConfig config; + std::string graph_cut_prefix; public: - TextModel(const LLMConfig& config) - : num_layers(config.num_layers), config(config) { + TextModel(const LLMConfig& config, const std::string& graph_cut_prefix = "llm.text") + : num_layers(config.num_layers), config(config), graph_cut_prefix(graph_cut_prefix) { blocks["embed_tokens"] = std::shared_ptr(new Embedding(config.vocab_size, config.hidden_size)); for (int i = 0; i < num_layers; i++) { blocks["layers." + std::to_string(i)] = std::shared_ptr(new TransformerBlock(config, i)); @@ -1831,7 +1897,10 @@ namespace LLM { std::set out_layers, const std::vector>>& deepstack_image_embeds = {}, ggml_tensor* sliding_attention_mask = nullptr, - bool return_all_hidden_states = false) { + bool return_all_hidden_states = false, + ggml_tensor* image_mask = nullptr, + ggml_tensor* rope_cos = nullptr, + ggml_tensor* rope_sin = nullptr) { auto norm = config.final_norm ? std::dynamic_pointer_cast(blocks["norm"]) : nullptr; std::vector intermediate_outputs; @@ -1843,18 +1912,18 @@ namespace LLM { intermediate_outputs.push_back(x); } - sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.prelude", "x"); + sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".prelude", "x"); for (int i = 0; i < num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); - x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask); + x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask, image_mask, rope_cos, rope_sin); if (i < static_cast(deepstack_image_embeds.size())) { x = add_deepstack_image_embeds(ctx, x, deepstack_image_embeds[static_cast(i)]); } if (return_all_hidden_states || out_layers.size() > 1) { x = ggml_cont(ctx->ggml_ctx, x); } - sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.layers." + std::to_string(i), "x"); + sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x"); if (return_all_hidden_states) { if (i + 1 < num_layers) { intermediate_outputs.push_back(x); diff --git a/src/model/te/ming_image_te.hpp b/src/model/te/ming_image_te.hpp new file mode 100644 index 0000000000..ab3b8e6dbc --- /dev/null +++ b/src/model/te/ming_image_te.hpp @@ -0,0 +1,159 @@ +#ifndef __SD_MODEL_TE_MING_IMAGE_TE_HPP__ +#define __SD_MODEL_TE_MING_IMAGE_TE_HPP__ + +#include "llm.hpp" + +namespace MingImageTE { + struct MingImageTEConfig { + LLM::LLMConfig backbone; + LLM::LLMConfig connector; + int64_t num_queries = 256; + int64_t caption_dim = 2560; + int64_t diffusion_dim = 3840; + + static MingImageTEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { + MingImageTEConfig config; + for (const auto& entry : tensors) { + if (starts_with(entry.first, prefix + ".backbone.") && + contains(entry.first, ".mlp.experts.") && entry.second.type == GGML_TYPE_I8) { + throw std::runtime_error("Ming-Image INT8/W4A8 text encoder experts are not supported; use the BF16 text encoder"); + } + } + bool vision = false; + config.backbone = LLM::LLMConfig::detect_from_weights(tensors, prefix + ".backbone.", LLM::LLMArch::BAILING_MOE, vision); + config.connector = LLM::LLMConfig::detect_from_weights(tensors, prefix + ".connector.", LLM::LLMArch::QWEN2, vision); + const auto query = tensors.find(prefix + ".query_tokens_dict.16x16"); + const auto projection = tensors.find(prefix + ".proj_out.weight"); + const auto direct = tensors.find(prefix + ".proj_directvlm.1.weight"); + if (query == tensors.end() || projection == tensors.end() || direct == tensors.end()) { + throw std::runtime_error("Ming-Image requires the learned queries, connector and both condition projections"); + } + config.num_queries = query->second.ne[1]; + config.caption_dim = projection->second.ne[1]; + config.diffusion_dim = direct->second.ne[1]; + if (config.num_queries != 256 || config.caption_dim != 2560 || config.diffusion_dim != 3840 || + config.backbone.num_layers != 20 || config.backbone.hidden_size != 2048 || + config.connector.num_layers != 28 || config.connector.hidden_size != 1536) { + throw std::runtime_error("unsupported Ming-Image text encoder configuration"); + } + LOG_VERBOSE("ming_image_te: queries = %" PRId64 ", caption_dim = %" PRId64 ", diffusion_dim = %" PRId64, + config.num_queries, config.caption_dim, config.diffusion_dim); + return config; + } + }; + + struct ConnectorModel : LLM::TextModel { + explicit ConnectorModel(const LLM::LLMConfig& config) + : LLM::TextModel(config, "ming_image.connector") { + blocks.erase("embed_tokens"); + } + }; + + class MingImageTextModel : public GGMLBlock { + MingImageTEConfig config; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensors = {}, const std::string prefix = "") override { + params["query_tokens_dict.16x16"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.backbone.hidden_size, config.num_queries); + } + + public: + explicit MingImageTextModel(const MingImageTEConfig& config) + : config(config) { + blocks["backbone"] = std::make_shared(config.backbone, "ming_image.backbone"); + blocks["connector"] = std::make_shared(config.connector); + blocks["proj_in"] = std::make_shared(config.backbone.hidden_size, config.connector.hidden_size); + blocks["proj_out"] = std::make_shared(config.connector.hidden_size, config.caption_dim); + blocks["proj_directvlm.0"] = std::make_shared(config.backbone.hidden_size * 3, 1e-5f); + blocks["proj_directvlm.1"] = std::make_shared(config.backbone.hidden_size * 3, config.diffusion_dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* ids, int64_t prompt_length, ggml_tensor* mask, ggml_tensor* image_mask, ggml_tensor* cos, ggml_tensor* sin, ggml_tensor* connector_positions) { + auto gctx = ctx->ggml_ctx; + auto backbone = std::dynamic_pointer_cast(blocks["backbone"]); + auto connector = std::dynamic_pointer_cast(blocks["connector"]); + auto x = backbone->embed(ctx, ids); + const int64_t query_start = prompt_length + 1; + auto before = ggml_ext_slice(gctx, x, 1, 0, query_start); + auto after = ggml_ext_slice(gctx, x, 1, query_start + config.num_queries, x->ne[1]); + x = ggml_concat(gctx, ggml_concat(gctx, before, params["query_tokens_dict.16x16"], 1), after, 1); + // HF hidden_states[20] includes final RMSNorm; sd.cpp selects it as num_layers + 1. + x = backbone->forward_embeds(ctx, x, nullptr, mask, {5, 12, 21}, {}, nullptr, false, image_mask, cos, sin); + auto direct = ggml_ext_slice(gctx, x, 1, 0, prompt_length); + direct = std::dynamic_pointer_cast(blocks["proj_directvlm.0"])->forward(ctx, direct); + direct = std::dynamic_pointer_cast(blocks["proj_directvlm.1"])->forward(ctx, direct); + auto queries = ggml_ext_slice(gctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3); + queries = ggml_ext_slice(gctx, queries, 1, query_start, query_start + config.num_queries); + queries = std::dynamic_pointer_cast(blocks["proj_in"])->forward(ctx, queries); + queries = connector->forward_embeds(ctx, queries, connector_positions, nullptr, {}); + queries = std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, queries); + queries = ggml_pad(gctx, queries, static_cast(config.diffusion_dim - config.caption_dim), 0, 0, 0); + return ggml_concat(gctx, queries, direct, 1); + } + }; + + struct MingImageTextRunner : GGMLRunner { + MingImageTEConfig config; + MingImageTextModel model; + + MingImageTextRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), config(MingImageTEConfig::detect_from_weights(tensors, prefix)), model(config) { + model.init(params_ctx, tensors, prefix); + } + + std::string get_desc() override { return "ming_image_text"; } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + model.get_param_tensors(tensors, prefix); + } + + void get_param_tensor_ops(std::map& ops) { + model.get_param_tensor_ops(ops); + } + + sd::Tensor compute(int n_threads, const std::vector& tokens) { + const int64_t prompt_length = tokens.size(); + const int64_t query_start = prompt_length + 1; + const int64_t total = prompt_length + config.num_queries + 2; + std::vector ids(tokens.begin(), tokens.end()); + ids.push_back(157158); + ids.insert(ids.end(), config.num_queries, 157157); + ids.push_back(157159); + auto input = sd::Tensor({total}, std::move(ids)); + sd::Tensor attention_mask({total, total}); + sd::Tensor image_mask({1, total}); + sd::Tensor cos({32, 1, total}), sin({32, 1, total}); + std::vector positions(config.num_queries); + std::iota(positions.begin(), positions.end(), 0); + auto connector_positions = sd::Tensor({config.num_queries}, std::move(positions)); + for (int64_t token = 0; token < total; ++token) { + const bool query = token >= query_start && token < query_start + config.num_queries; + image_mask[token] = query ? 1.f : 0.f; + for (int64_t key = 0; key < total; ++key) { + attention_mask[key + token * total] = key > token ? -INFINITY : 0.f; + } + // A 16x16 query bank is represented upstream as a [1, 2, 512] image grid, + // then spatially merged and centered to [1, 1, 256]. + int64_t temporal = query ? query_start : (token == total - 1 ? query_start + 1 : token); + int64_t width = query ? token - 127 : temporal; + for (int j = 0; j < 32; ++j) { + int64_t position = query && j < 24 && j % 2 ? width : temporal; + float frequency = 1.f / std::pow(600000.f, static_cast(2 * j) / 64.f); + float angle = static_cast(position) * frequency; + cos[token * 32 + j] = std::cos(angle); + sin[token * 32 + j] = std::sin(angle); + } + } + auto graph = [&]() { + auto gf = new_graph_custom(LLM::LLM_GRAPH_SIZE); + auto ctx = get_context(); + auto out = model.forward(&ctx, make_input(input), prompt_length, make_input(attention_mask), + make_input(image_mask), make_input(cos), make_input(sin), make_input(connector_positions)); + ggml_build_forward_expand(gf, out); + return gf; + }; + return take_or_empty(GGMLRunner::compute(graph, n_threads)); + } + }; +} + +#endif // __SD_MODEL_TE_MING_IMAGE_TE_HPP__ diff --git a/src/model/vae/auto_encoder_kl.hpp b/src/model/vae/auto_encoder_kl.hpp index 115ff34589..3a4733879a 100644 --- a/src/model/vae/auto_encoder_kl.hpp +++ b/src/model/vae/auto_encoder_kl.hpp @@ -544,7 +544,7 @@ class AutoEncoderKLModel : public GGMLBlock { const String2TensorStorage& tensor_storage_map = {}, const std::string& prefix = "") : version(version), decode_only(decode_only), use_video_decoder(use_video_decoder) { - if (sd_version_is_dit(version)) { + if (sd_version_is_dit(version) && version != VERSION_PIXART) { if (sd_version_uses_flux2_vae(version)) { dd_config.z_channels = 32; embed_dim = 32; @@ -678,7 +678,7 @@ struct AutoEncoderKL : public VAE { if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) { scale_factor = 0.18215f; shift_factor = 0.f; - } else if (sd_version_is_sdxl(version)) { + } else if (sd_version_is_sdxl(version) || sd_version_is_pixart(version)) { scale_factor = 0.13025f; shift_factor = 0.f; } else if (sd_version_is_sd3(version)) { diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp index ac22c09b45..f680f20081 100644 --- a/src/model/vae/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -1300,7 +1300,7 @@ struct LTXVideoVAE : public VAE { 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; ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx, z, @@ -1313,8 +1313,7 @@ struct LTXVideoVAE : public VAE { for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast(feat_map.size()); ++feat_idx) { ggml_tensor* feat_cache = feat_map[static_cast(feat_idx)]; if (feat_cache != nullptr) { - cache(temporal_feat_cache_name(static_cast(feat_idx)), feat_cache); - ggml_build_forward_expand(gf, feat_cache); + runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast(feat_idx)), feat_cache); } } diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 68cb9a9a27..2b1085d6be 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -233,11 +233,12 @@ namespace MiniMaxH3VAE { return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]); } - static ggml_tensor* apply_partial_rope(ggml_context* ctx, + static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) { + auto ctx = runner_ctx->ggml_ctx; int64_t rot_dim = pe->ne[2] * 2; - auto rotated = Rope::apply_rope(ctx, + auto rotated = Rope::apply_rope(runner_ctx, ggml_ext_slice(ctx, x, 0, 0, rot_dim), pe, false); @@ -404,8 +405,8 @@ namespace MiniMaxH3VAE { q = apply_partial_rope_tables(ctx->ggml_ctx, q, rope_a, rope_b, pe->ne[2] * 2); k = apply_partial_rope_tables(ctx->ggml_ctx, k, rope_a, rope_b, pe->ne[2] * 2); } else { - q = apply_partial_rope(ctx->ggml_ctx, q, pe); - k = apply_partial_rope(ctx->ggml_ctx, k, pe); + q = apply_partial_rope(ctx, q, pe); + k = apply_partial_rope(ctx, k, pe); } const int64_t tiles = v->ne[3]; if (tiles > 1) { @@ -739,17 +740,25 @@ namespace MiniMaxH3VAE { tensor.shape()[3]}); } - static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) { + sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const override { + if (!params.enabled) { + params.target_overlap = 0.25f; + } + if (params.tile_size_w == 0 && params.rel_size_w == 0.f) { + params.tile_size_w = 256; + } + if (params.tile_size_h == 0 && params.rel_size_h == 0.f) { + params.tile_size_h = 256; + } params.enabled = true; params.temporal_tiling = false; - params.tile_size_x = 16; - params.tile_size_y = 16; - params.target_overlap = 0.25f; - // SD_H3_VAE_TILE=N: opt-in N x N latent tiles; the seams move, so frames differ from the default + // SD_H3_VAE_TILE=N: opt-in N x N latent tiles (N * 16 image pixels); the seams move, so frames differ + // from the default if (const char* tile = getenv("SD_H3_VAE_TILE")) { const int n = atoi(tile); if (n >= 8) { - params.tile_size_x = params.tile_size_y = n; + params.tile_size_w = params.tile_size_h = n * 16; + params.rel_size_w = params.rel_size_h = 0.f; } } return params; @@ -799,7 +808,7 @@ namespace MiniMaxH3VAE { bool circular_x = false, bool circular_y = false) override { auto input = ensure_video_shape(x); - auto tiling = h3_tiling(tiling_params); + auto tiling = resolve_tiling_params(tiling_params); if (input.shape()[2] == 1) { auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y); if (!encoded.empty() && encoded.shape()[2] > 1) { @@ -821,8 +830,11 @@ namespace MiniMaxH3VAE { auto plan = make_vae_temporal_tile_plan(input.shape()[2], {17, 0}); auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& chunk, const VAETemporalTile& tile) { SD_UNUSED(tile); - return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y); + // keep the runner alive across chunks; ending it here would + // evict and reload the encoder weights every chunk + return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y, false); }); + runner_end(); if (result.empty()) { return {}; } @@ -840,7 +852,7 @@ namespace MiniMaxH3VAE { bool circular_y = false, bool silent = false) override { auto input = ensure_video_shape(x); - auto tiling = h3_tiling(tiling_params); + auto tiling = resolve_tiling_params(tiling_params); if (input.shape()[2] == 1) { auto decoded = VAE::decode(n_threads, input, @@ -1134,18 +1146,36 @@ namespace MiniMaxH3VAE { bool silent) { const int scale_factor = get_scale_factor(); float tile_overlap; - int tile_size_x, tile_size_y; - get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling, chunk.shape()[0], chunk.shape()[1]); + int tile_size_w, tile_size_h; + if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling, chunk.shape()[0], chunk.shape()[1], circular_x, circular_y)) { + return {}; + } + // Same circular handling as VAE::tiled_compute: full-width axes wrap in convolutions, split axes wrap + // between tiles. + const bool original_circular_x = circular_x_enabled; + const bool original_circular_y = circular_y_enabled; + circular_x = circular_x || original_circular_x; + circular_y = circular_y || original_circular_y; + set_circular_axes(circular_x && tile_size_w >= chunk.shape()[0], + circular_y && tile_size_h >= chunk.shape()[1]); + struct CircularAxesGuard { + VAE& vae; + bool x; + bool y; + ~CircularAxesGuard() { + vae.set_circular_axes(x, y); + } + } circular_axes_guard{*this, original_circular_x, original_circular_y}; return process_tiles_2d_batched( chunk, static_cast(chunk.shape()[0] * scale_factor), static_cast(chunk.shape()[1] * scale_factor), scale_factor, - tile_size_x, - tile_size_y, + tile_size_w, + tile_size_h, tile_overlap, - circular_x, - circular_y, + circular_x && tile_size_w < chunk.shape()[0], + circular_y && tile_size_h < chunk.shape()[1], [&](int) { return resolve_tile_batch(); }, [&](const std::vector>& tiles) { auto output = compute_tile_batch(n_threads, tiles); diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index e9257b7c82..173b40c5cb 100644 --- a/src/model/vae/tae.hpp +++ b/src/model/vae/tae.hpp @@ -701,7 +701,7 @@ class TAESD : public GGMLBlock { bool use_midblock_gn = false; taef2 = sd_version_uses_flux2_vae(version); - if (sd_version_is_dit(version)) { + if (sd_version_is_dit(version) && !sd_version_is_pixart(version)) { z_channels = 16; } if (taef2) { diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index 8383f7cf56..4e5248884e 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -1,6 +1,9 @@ #ifndef __SD_MODEL_VAE_VAE_HPP__ #define __SD_MODEL_VAE_VAE_HPP__ +#include +#include + #include "core/tensor_ggml.hpp" #include "model/common/block.hpp" #include "model/vae/vae_tiling.hpp" @@ -117,8 +120,8 @@ struct VAE : public GGMLRunner { int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y, @@ -138,17 +141,28 @@ struct VAE : public GGMLRunner { } return output_tile; }; - return ::process_tiles_2d(input, - output_width, - output_height, - scale, - p_tile_size_x, - p_tile_size_y, - tile_overlap_factor, - circular_x, - circular_y, - on_processing, - silent); + const bool original_circular_x = circular_x_enabled; + const bool original_circular_y = circular_y_enabled; + const int64_t latent_width = decode_graph ? input.shape()[0] : output_width; + const int64_t latent_height = decode_graph ? input.shape()[1] : output_height; + circular_x = circular_x || original_circular_x; + circular_y = circular_y || original_circular_y; + // Full-width axes wrap in convolutions; split axes wrap between tiles. + set_circular_axes(circular_x && p_tile_size_w >= latent_width, + circular_y && p_tile_size_h >= latent_height); + auto output = ::process_tiles_2d(input, + output_width, + output_height, + scale, + p_tile_size_w, + p_tile_size_h, + tile_overlap_factor, + circular_x && p_tile_size_w < latent_width, + circular_y && p_tile_size_h < latent_height, + on_processing, + silent); + set_circular_axes(original_circular_x, original_circular_y); + return output; } public: @@ -162,11 +176,20 @@ struct VAE : public GGMLRunner { int scale_factor = 8; if (version == VERSION_LTXAV) { scale_factor = 32; - } else if (version == VERSION_WAN2_2_TI2V || version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) { + } else if (version == VERSION_WAN2_2_TI2V || + version == VERSION_QWEN_IMAGE_2_1 || + sd_version_is_hunyuan_video(version) || + sd_version_is_mage_flow(version) || + sd_version_is_minimax_h3(version)) { scale_factor = 16; } else if (sd_version_uses_flux2_vae(version)) { scale_factor = 16; - } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) { + } else if (version == VERSION_CHROMA_RADIANCE || + version == VERSION_HIDREAM_O1 || + sd_version_is_minit2i(version) || + sd_version_is_sensenova_u1(version) || + sd_version_is_z_image_l2p(version) || + version == VERSION_IRIS) { scale_factor = 1; } return scale_factor; @@ -178,42 +201,60 @@ struct VAE : public GGMLRunner { return supports_temporal_tiling(VAETemporalDirection::DECODE); } - void get_tile_sizes(int& tile_size_x, - int& tile_size_y, + virtual sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const { + return params; + } + + bool get_tile_sizes(int& tile_size_w, + int& tile_size_h, float& tile_overlap, const sd_tiling_params_t& params, - int64_t latent_x, - int64_t latent_y, - float encoding_factor = 1.0f, - bool circular_x = false, - bool circular_y = false) { - tile_overlap = std::max(std::min(params.target_overlap, 0.5f), 0.0f); - auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) { - const int default_tile_size = 32; - const int min_tile_dimension = 4; - int tile_size = default_tile_size; - // factor <= 1 means simple fraction of the latent dimension - // factor > 1 means number of tiles across that dimension - if (factor > 0.f) { - if (factor > 1.0) - factor = 1 / (factor - factor * tile_overlap + tile_overlap); - tile_size = static_cast(std::round(latent_size * factor)); - } else if (requested_size >= min_tile_dimension) { - tile_size = requested_size; + int64_t latent_w, + int64_t latent_h, + bool circular_x = false, + bool circular_y = false) { + const auto tiling = resolve_tiling_params(params); + if (latent_w <= 0 || latent_h <= 0 || + latent_w > std::numeric_limits::max() || latent_h > std::numeric_limits::max() || + !std::isfinite(tiling.target_overlap)) { + LOG_ERROR("invalid VAE tiling dimensions or overlap"); + return false; + } + const int scale_factor = get_scale_factor(); + tile_overlap = std::max(std::min(tiling.target_overlap, 0.5f), 0.0f); + auto get_tile_size = [&](int requested_size, double factor, int64_t latent_size, int& tile_size) { + if (requested_size < 0 || !std::isfinite(factor) || factor < 0.0) { + LOG_ERROR("VAE tile sizes and relative sizes must be finite and non-negative"); + return false; } - tile_size = static_cast(tile_size * encoding_factor); - return std::max(std::min(tile_size, static_cast(latent_size)), min_tile_dimension); + const int min_tile_dimension = std::min(4, static_cast(latent_size)); + double size = (requested_size > 0 ? requested_size : 256) / scale_factor; + if (factor > 0.0) { + if (factor > 1.0) { + factor = 1.0 / (factor * (1.0 - tile_overlap) + tile_overlap); + } + size = std::floor(static_cast(latent_size) * factor); + } + if (size < min_tile_dimension && (requested_size > 0 || factor > 0.0)) { + LOG_ERROR("VAE tile size must be at least %d image pixels on this axis", min_tile_dimension * scale_factor); + return false; + } + tile_size = static_cast(std::min(static_cast(latent_size), std::max(min_tile_dimension, size))); + return true; }; - tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x); - tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y); + if (!get_tile_size(tiling.tile_size_w, tiling.rel_size_w, latent_w, tile_size_w) || + !get_tile_size(tiling.tile_size_h, tiling.rel_size_h, latent_h, tile_size_h)) { + return false; + } // Circular axes wrap their tiles instead of pinning the last one to the edge, so they keep the requested size. - if (!circular_x) { - tile_size_x = sd_tiling_seam_safe_tile_size(static_cast(latent_x), tile_size_x, tile_overlap); + if (!circular_x && !circular_x_enabled) { + tile_size_w = sd_tiling_seam_safe_tile_size(static_cast(latent_w), tile_size_w, tile_overlap); } - if (!circular_y) { - tile_size_y = sd_tiling_seam_safe_tile_size(static_cast(latent_y), tile_size_y, tile_overlap); + if (!circular_y && !circular_y_enabled) { + tile_size_h = sd_tiling_seam_safe_tile_size(static_cast(latent_h), tile_size_h, tile_overlap); } + return true; } virtual sd::Tensor encode(int n_threads, @@ -221,7 +262,17 @@ struct VAE : public GGMLRunner { sd_tiling_params_t tiling_params, bool circular_x = false, bool circular_y = false) { + return encode(n_threads, x, tiling_params, circular_x, circular_y, true); + } + + sd::Tensor encode(int n_threads, + const sd::Tensor& x, + sd_tiling_params_t tiling_params, + bool circular_x, + bool circular_y, + bool auto_end) { int64_t t0 = ggml_time_ms(); + tiling_params = resolve_tiling_params(tiling_params); sd::Tensor input = x; sd::Tensor output; if (scale_input) { @@ -233,21 +284,19 @@ struct VAE : public GGMLRunner { int64_t W = input.shape()[0] / scale_factor; int64_t H = input.shape()[1] / scale_factor; float tile_overlap; - int tile_size_x, tile_size_y; - // Image VAE encode is more sensitive to tile boundary context than decode. - // Keep the smaller legacy factor for video VAEs, but default image encode - // tiles to 64 latent pixels so a 512px SD image is encoded as one tile. - const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f - : 2.0f; - get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor, circular_x, circular_y); - LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + int tile_size_w, tile_size_h; + if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, W, H, circular_x, circular_y)) { + return {}; + } + LOG_VERBOSE("VAE encode tile size: %dx%d pixels (%dx%d latent)", + tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h); output = tiled_compute(input, n_threads, static_cast(W), static_cast(H), scale_factor, - tile_size_x, - tile_size_y, + tile_size_w, + tile_size_h, tile_overlap, circular_x, circular_y, @@ -261,7 +310,9 @@ struct VAE : public GGMLRunner { tiling_params); } - runner_end(); + if (auto_end) { + runner_end(); + } if (output.empty()) { LOG_ERROR("vae encode compute failed"); @@ -279,7 +330,19 @@ struct VAE : public GGMLRunner { bool circular_x = false, bool circular_y = false, bool silent = false) { + return decode(n_threads, x, tiling_params, decode_video, circular_x, circular_y, silent, true); + } + + sd::Tensor decode(int n_threads, + const sd::Tensor& x, + sd_tiling_params_t tiling_params, + bool decode_video, + bool circular_x, + bool circular_y, + bool silent, + bool auto_end) { int64_t t0 = ggml_time_ms(); + tiling_params = resolve_tiling_params(tiling_params); sd::Tensor input = x; sd::Tensor output; @@ -288,10 +351,13 @@ struct VAE : public GGMLRunner { int64_t W = input.shape()[0] * scale_factor; int64_t H = input.shape()[1] * scale_factor; float tile_overlap; - int tile_size_x, tile_size_y; - get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1], 1.0f, circular_x, circular_y); + int tile_size_w, tile_size_h; + if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, input.shape()[0], input.shape()[1], circular_x, circular_y)) { + return {}; + } if (!silent) { - LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + LOG_VERBOSE("VAE decode tile size: %dx%d pixels (%dx%d latent)", + tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h); } output = tiled_compute( input, @@ -299,8 +365,8 @@ struct VAE : public GGMLRunner { static_cast(W), static_cast(H), scale_factor, - tile_size_x, - tile_size_y, + tile_size_w, + tile_size_h, tile_overlap, circular_x, circular_y, @@ -315,7 +381,9 @@ struct VAE : public GGMLRunner { tiling_params); } - runner_end(); + if (auto_end) { + runner_end(); + } if (output.empty()) { LOG_ERROR("vae decode compute failed"); diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 035deae324..9dd30fdff6 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -24,6 +24,7 @@ namespace WAN { std::tuple padding; std::tuple dilation; bool bias; + float scale = 1.f; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { auto weight = tensor_storage_map.find(prefix + "weight"); @@ -60,6 +61,10 @@ namespace WAN { dilation(std::move(dilation)), bias(bias) {} + void set_scale(float scale_value) { + scale = scale_value; + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) { // x: [N*IC, ID, IH, IW] // result: x: [N*OC, ID, IH, IW] @@ -82,11 +87,25 @@ 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); + 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, false, false, scale); + 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, std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), 0, 0, 0, std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), - false, ctx->conv3d_direct_enabled); + false, ctx->conv3d_direct_enabled, scale); } }; @@ -1065,7 +1084,7 @@ namespace WAN { _conv_num = 34; _enc_conv_num = 26; - } else if (version == VERSION_QWEN_IMAGE_LAYERED) { + } else if (version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_MING_IMAGE) { input_channels = 4; } @@ -1103,6 +1122,19 @@ namespace WAN { } else { blocks["conv2"] = std::shared_ptr(new CausalConv3d(z_dim, z_dim, {1, 1, 1})); } + if (version == VERSION_QWEN_IMAGE_2_1) { + // Keep large VAE activations within the FP16 convolution range. + const float conv_scale = 1.f / 128.f; + std::vector all_blocks; + get_all_blocks(all_blocks); + for (auto block : all_blocks) { + if (auto conv = dynamic_cast(block)) { + conv->set_scale(conv_scale); + } else if (auto conv = dynamic_cast(block)) { + conv->set_scale(conv_scale); + } + } + } } static ggml_tensor* patchify(ggml_context* ctx, @@ -1391,11 +1423,17 @@ namespace WAN { } sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) override { + if (version == VERSION_MING_IMAGE) { + return latents / 8.0064f; + } auto [mean_tensor, std_tensor] = get_latents_mean_std(latents); return (latents * std_tensor) / scale_factor + mean_tensor; } sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) override { + if (version == VERSION_MING_IMAGE) { + return latents * 8.0064f; + } auto [mean_tensor, std_tensor] = get_latents_mean_std(latents); return ((latents - mean_tensor) * scale_factor) / std_tensor; } @@ -1429,15 +1467,14 @@ namespace WAN { 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); for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) { ggml_tensor* feat_cache = ae._feat_map[feat_idx]; if (feat_cache != nullptr) { - cache("feat_idx:" + std::to_string(feat_idx), feat_cache); - ggml_build_forward_expand(gf, feat_cache); + runner_ctx.persist_cache_tensor("feat_idx:" + std::to_string(feat_idx), feat_cache); } } diff --git a/src/model_io/gguf_io.cpp b/src/model_io/gguf_io.cpp index 081fcf978f..3e57448786 100644 --- a/src/model_io/gguf_io.cpp +++ b/src/model_io/gguf_io.cpp @@ -47,10 +47,14 @@ bool read_gguf_file(const std::string& file_path, gguf_context* ctx_gguf_ = nullptr; ggml_context* ctx_meta_ = nullptr; - ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_}); + GGUFReader gguf_reader; + bool probe_ok = gguf_reader.load(file_path); + + if (!probe_ok || !gguf_reader.has_tensors_beyond_ggml_limits()) { + ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_}); + } if (!ctx_gguf_) { - GGUFReader gguf_reader; - if (!gguf_reader.load(file_path)) { + if (!probe_ok && !gguf_reader.load(file_path)) { set_error(error, "failed to open '" + file_path + "' with GGUFReader"); return false; } diff --git a/src/model_io/gguf_reader_ext.h b/src/model_io/gguf_reader_ext.h index 7b372e4f59..b55d49f11b 100644 --- a/src/model_io/gguf_reader_ext.h +++ b/src/model_io/gguf_reader_ext.h @@ -35,23 +35,72 @@ enum class GGUFMetadataType : uint32_t { class GGUFReader { private: std::vector tensors_; + bool has_wide_tensors_ = false; + uint64_t remaining_bytes_ = 0; size_t data_offset_; size_t alignment_ = 32; // default alignment is 32 template bool safe_read(std::ifstream& fin, T& value) { - fin.read(reinterpret_cast(&value), sizeof(T)); - return fin.good(); + return safe_read(fin, reinterpret_cast(&value), sizeof(T)); } bool safe_read(std::ifstream& fin, char* buffer, size_t size) { + if (size > remaining_bytes_) + return false; fin.read(buffer, size); - return fin.good(); + if (!fin.good()) + return false; + remaining_bytes_ -= size; + return true; + } + + bool safe_skip(std::ifstream& fin, uint64_t count, uint64_t element_size = 1) { + if (count > remaining_bytes_ / element_size) + return false; + uint64_t size = count * element_size; + fin.seekg(static_cast(size), std::ios::cur); + if (!fin.good()) + return false; + remaining_bytes_ -= size; + return true; } - bool safe_seek(std::ifstream& fin, std::streamoff offset, std::ios::seekdir dir) { - fin.seekg(offset, dir); - return fin.good(); + bool skip_metadata_values(std::ifstream& fin, GGUFMetadataType type, uint64_t count) { + switch (type) { + case GGUFMetadataType::UINT8: + case GGUFMetadataType::INT8: + case GGUFMetadataType::BOOL: + return safe_skip(fin, count); + + case GGUFMetadataType::UINT16: + case GGUFMetadataType::INT16: + return safe_skip(fin, count, 2); + + case GGUFMetadataType::UINT32: + case GGUFMetadataType::INT32: + case GGUFMetadataType::FLOAT32: + return safe_skip(fin, count, 4); + + case GGUFMetadataType::UINT64: + case GGUFMetadataType::INT64: + case GGUFMetadataType::FLOAT64: + return safe_skip(fin, count, 8); + + case GGUFMetadataType::STRING: + if (count > remaining_bytes_ / sizeof(uint64_t)) + return false; + for (uint64_t i = 0; i < count; i++) { + uint64_t len = 0; + if (!safe_read(fin, len) || !safe_skip(fin, len)) + return false; + } + return true; + + default: + LOG_ERROR("Unknown metadata type=%u", static_cast(type)); + return false; + } } bool read_metadata(std::ifstream& fin) { @@ -84,52 +133,12 @@ class GGUFReader { return true; } - switch (static_cast(type)) { - case GGUFMetadataType::UINT8: - case GGUFMetadataType::INT8: - case GGUFMetadataType::BOOL: - return safe_seek(fin, 1, std::ios::cur); - - case GGUFMetadataType::UINT16: - case GGUFMetadataType::INT16: - return safe_seek(fin, 2, std::ios::cur); - - case GGUFMetadataType::UINT32: - case GGUFMetadataType::INT32: - case GGUFMetadataType::FLOAT32: - return safe_seek(fin, 4, std::ios::cur); - - case GGUFMetadataType::UINT64: - case GGUFMetadataType::INT64: - case GGUFMetadataType::FLOAT64: - return safe_seek(fin, 8, std::ios::cur); - - case GGUFMetadataType::STRING: { - uint64_t len = 0; - if (!safe_read(fin, len)) - return false; - return safe_seek(fin, len, std::ios::cur); - } - - case GGUFMetadataType::ARRAY: { - uint32_t elem_type = 0; - uint64_t len = 0; - if (!safe_read(fin, elem_type)) - return false; - if (!safe_read(fin, len)) - return false; - - for (uint64_t i = 0; i < len; i++) { - if (!read_metadata(fin)) - return false; - } - return true; - } - - default: - LOG_ERROR("Unknown metadata type=%u", type); + uint64_t count = 1; + if (type == static_cast(GGUFMetadataType::ARRAY)) { + if (!safe_read(fin, type) || !safe_read(fin, count)) return false; } + return skip_metadata_values(fin, static_cast(type), count); } GGUFTensorInfo read_tensor_info(std::ifstream& fin) { @@ -154,6 +163,7 @@ class GGUFReader { } if (n_dims > GGML_MAX_DIMS) { + has_wide_tensors_ = true; for (uint32_t i = GGML_MAX_DIMS; i < n_dims; i++) { info.shape[GGML_MAX_DIMS - 1] *= info.shape[i]; // stack to last dim; } @@ -174,12 +184,20 @@ class GGUFReader { public: bool load(const std::string& file_path) { - std::ifstream fin(file_path, std::ios::binary); + std::ifstream fin(file_path, std::ios::binary | std::ios::ate); if (!fin) { LOG_ERROR("failed to open '%s'", file_path.c_str()); return false; } + std::streamoff file_size = fin.tellg(); + if (file_size < 0) + return false; + remaining_bytes_ = static_cast(file_size); + fin.seekg(0, std::ios::beg); + if (!fin.good()) + return false; + // --- Header --- char magic[4]; if (!safe_read(fin, magic, 4) || strncmp(magic, "GGUF", 4) != 0) { @@ -228,6 +246,8 @@ class GGUFReader { } const std::vector& tensors() const { return tensors_; } + + bool has_tensors_beyond_ggml_limits() const { return has_wide_tensors_; } size_t data_offset() const { return data_offset_; } }; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 57cc11ff28..34490be4a0 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -115,7 +115,7 @@ void convert_tensor(void* src, ggml_type dst_type, int nrows, int n_per_row, - std::vector imatrix = {}) { + std::vector imatrix) { int n = nrows * n_per_row; if (src_type == dst_type) { size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type); @@ -322,7 +322,13 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, } for (const std::string& shard_path : shard_paths) { - if (!parse_file(shard_path, prefix)) { + FileStamp stamp; + if (!read_file_stamp(shard_path, stamp)) { + return false; + } + parsed_dependencies_.push_back(stamp); + LOG_INFO("load %s using safetensors format", shard_path.c_str()); + if (!init_from_safetensors_file(shard_path, prefix)) { return false; } } @@ -518,6 +524,12 @@ SDVersion ModelLoader::get_sd_version() const { return VERSION_LLADA_IMAGE; } if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) { + if (tensor_storage_map.find("text_encoders.llm.connector.layers.0.self_attn.q_proj.weight") != tensor_storage_map.end()) { + return VERSION_MING_IMAGE; + } + if (tensor_storage_map.find("model.diffusion_model.local_decoder.out_conv.weight") != tensor_storage_map.end()) { + return VERSION_Z_IMAGE_L2P; + } return VERSION_Z_IMAGE; } if (tensor_storage.name.find("double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") != std::string::npos) { @@ -526,13 +538,28 @@ SDVersion ModelLoader::get_sd_version() const { if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) { return VERSION_ERNIE_IMAGE; } + if (tensor_storage.name.find("model.diffusion_model.t_block.1.weight") != std::string::npos && + tensor_storage_map.find("model.diffusion_model.x_embedder.proj.weight") != tensor_storage_map.end() && + tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) { + return VERSION_PIXART; + } if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) { + // PixArt shares this timestep embedding with LTX-AV. + if (tensor_storage_map.find("model.diffusion_model.pos_embed.proj.weight") != tensor_storage_map.end() && + tensor_storage_map.find("model.diffusion_model.adaln_single.linear.weight") != tensor_storage_map.end() && + tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) { + return VERSION_PIXART; + } return VERSION_LTXAV; } if (tensor_storage.name.find("model.diffusion_model.video_patch_proj.weight") != std::string::npos && tensor_storage_map.find("model.diffusion_model.audio_patch_proj.weight") != tensor_storage_map.end()) { return VERSION_MINIMAX_H3; } + if (name == "model.diffusion_model.y_embedder.layer_pool.weight" && + tensor_storage_map.count("model.diffusion_model.modulation_cores.adaln_img.weight")) { + return VERSION_IRIS; + } if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) { is_wan = true; } @@ -874,7 +901,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) { std::vector ModelLoader::mmap_tensors(std::map& tensors, std::set ignore_tensors, - bool writable_mmap) { + bool writable_mmap, + ggml_backend_dev_t device) { std::set names; for (const auto& entry : tensors) { names.insert(entry.first); @@ -896,6 +924,39 @@ std::vector ModelLoader::mmap_tensors(std::map 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(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(buf, ggml_backend_buffer_free); + fdata.device_mmbuffers[device] = file_buffer; + return file_buffer.get(); + }; + const std::vector& file_tensors = fdata.tensors; size_t file_mapped_bytes = 0; @@ -944,10 +1005,13 @@ std::vector ModelLoader::mmap_tensors(std::map(ggml_backend_buffer_get_base(buf_mmap)); - dst_tensor->buffer = buf_mmap; - dst_tensor->data = mmap_data + tensor_offset; + ggml_backend_buffer_t buf_mmap = buffer_for_file(); + if (buf_mmap == nullptr) { + break; + } + uint8_t* mmap_data = static_cast(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_tensors++; @@ -956,7 +1020,7 @@ std::vector ModelLoader::mmap_tensors(std::map 0) { mapped_tensors += file_mapped_tensors; mapped_bytes += file_mapped_bytes; - result.push_back({fdata.mmapped, fdata.mmbuffer}); + result.push_back({fdata.mmapped, file_buffer}); } } @@ -972,6 +1036,16 @@ std::vector ModelLoader::mmap_tensors(std::map ModelLoader::get_device_mmap_buffers() const { + std::vector 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 enable_mmap, const std::set* target_tensor_names, @@ -1037,10 +1111,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, if (tensors_to_process.empty()) { continue; } - LOG_VERBOSE("loading %zu/%zu tensors from %s", - tensors_to_process.size(), - file_tensors.size(), - file_path.c_str()); + if (log_progress) { + LOG_VERBOSE("loading %zu/%zu tensors from %s", + tensors_to_process.size(), + file_tensors.size(), + file_path.c_str()); + } bool is_zip = fdata.is_zip; @@ -1115,6 +1191,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()) { 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(); @@ -1551,6 +1632,9 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage // Pass, do not convert. For Unet } else if (contains(name, "embedding")) { // Pass, do not convert embedding + } else if (contains(name, "scale_shift_table")) { + // Pass, do not convert. adaLN modulation tables (PixArt, LTXV) are sliced + // element-wise, which is invalid on quantized block layouts. } else if (ends_with(name, "_pad_token")) { // Pass, do not convert. LLaDA-Image stores its pad tokens far outside the f16 // range, so any format with an f16 scale or payload turns them into inf. diff --git a/src/model_loader.h b/src/model_loader.h index c5335b53a9..63ecfb63b8 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -13,6 +13,14 @@ TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules); +void convert_tensor(void* src, + ggml_type src_type, + void* dst, + ggml_type dst_type, + int nrows, + int n_per_row, + std::vector imatrix = {}); + class MmapWrapper; struct ModelFileData { @@ -20,6 +28,8 @@ struct ModelFileData { std::vector tensors; std::shared_ptr mmapped; std::shared_ptr mmbuffer; + // mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr) + std::map> device_mmbuffers; bool is_zip; }; @@ -120,7 +130,9 @@ class ModelLoader { void process_model_files(bool enable_mmap = false, bool writable_mmap = true); std::vector mmap_tensors(std::map& tensors, std::set ignore_tensors = {}, - bool writable = true); + bool writable = true, + ggml_backend_dev_t device = nullptr); + std::vector get_device_mmap_buffers() const; bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool use_mmap = false, const std::set* target_tensor_names = nullptr, diff --git a/src/model_manager.cpp b/src/model_manager.cpp index e02270940b..a47377b9cf 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -79,6 +79,10 @@ static bool device_supports_param_op(ggml_backend_dev_t device, if (op == GGML_OP_GET_ROWS) { ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); op_tensor = ggml_get_rows(ctx, weight, indices); + } else if (op == GGML_OP_CUSTOM) { + op_tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + op_tensor->op = op; + op_tensor->src[0] = weight; } if (op_tensor == nullptr) { ggml_free(ctx); @@ -246,17 +250,13 @@ bool ModelManager::register_param_tensors(ModelComponent component, } ggml_set_name(tensor, name.c_str()); - auto state = std::make_unique(); - state->name = name; - state->tensor = tensor; - state->component = component; - state->source_file = source_file; - state->source_version = source_version; - auto source = sources.find(name); - if (source != sources.end()) { - state->source = source->second; - state->has_source = true; - } + auto state = std::make_unique(); + state->name = name; + state->tensor = tensor; + state->component = component; + state->source_file = source_file; + state->source_version = source_version; + state->sources = find_tensor_sources(*state, sources); state->residency_mode = residency_mode; state->compute_backend = compute_backend; state->params_backend = params_backend; @@ -511,7 +511,9 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vectorname.c_str()); return false; } - if (state->compute_backend == state->params_backend || state->staged_to_compute_backend) { + // Custom CPU operators must retain host weights even when the runner uses a GPU. + if (state->usage_op == GGML_OP_CUSTOM || + state->compute_backend == state->params_backend || state->staged_to_compute_backend) { continue; } if (!state->loaded_to_params_backend || state->tensor == nullptr || state->tensor->data == nullptr) { @@ -757,12 +759,33 @@ bool ModelManager::validate_tensor(const TensorState& state) const { return true; } - if (!state.has_source) { + if (state.sources.empty()) { LOG_ERROR("%s tensor '%s' not in model metadata", model_component_name(state.component), state.name.c_str()); return false; } - const TensorStorage& tensor_storage = state.source; + TensorStorage tensor_storage = state.sources.front(); + if (state.sources.size() > 1) { + const int dim = tensor_storage.n_dims - 1; + if (dim < 0 || dim >= GGML_MAX_DIMS) { + return false; + } + tensor_storage.ne[dim] = 0; + for (const auto& part : state.sources) { + if (part.n_dims != tensor_storage.n_dims || part.ne[dim] < 0 || + part.ne[dim] > state.tensor->ne[dim] - tensor_storage.ne[dim]) { + LOG_ERROR("invalid tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str()); + return false; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (i != dim && part.ne[i] != tensor_storage.ne[i]) { + LOG_ERROR("incompatible tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str()); + return false; + } + } + tensor_storage.ne[dim] += part.ne[dim]; + } + } if (state.tensor->ne[0] != tensor_storage.ne[0] || state.tensor->ne[1] != tensor_storage.ne[1] || state.tensor->ne[2] != tensor_storage.ne[2] || @@ -780,44 +803,58 @@ bool ModelManager::validate_tensor(const TensorState& state) const { bool ModelManager::mmap_params(const std::vector& states, std::vector& created_storage_blocks) { - std::map mmap_candidates; - std::map mmap_states; + // A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing + // stages them for it, so they are mapped through a buffer of that GPU's device. + struct MmapGroup { + std::map candidates; + std::map states; + }; + std::map groups; for (TensorState* state : states) { if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr || state->tensor->data != nullptr || state->tensor->view_src != nullptr) { continue; } - mmap_candidates[state->name] = state->tensor; - mmap_states[state->name] = state; - } - if (mmap_candidates.empty()) { - return true; + ggml_backend_dev_t device = nullptr; + if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) { + device = ggml_backend_get_device(state->compute_backend); + } + MmapGroup& group = groups[device]; + group.candidates[state->name] = state->tensor; + group.states[state->name] = state; } - auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_); - if (mmap_store.empty()) { - return true; - } + for (auto& [device, group] : groups) { + // Device buffers wrap read-only mappings only; params that LoRAs are merged into in place + // 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(); - 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); + auto block = std::make_unique(); + 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); + if (!block->states.empty()) { + params_storage_blocks_.push_back(std::move(block)); + created_storage_blocks.push_back(raw); + } } return true; } bool ModelManager::can_mmap_storage(const TensorState& state) const { - if (state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) { + if (state.sources.size() > 1 || state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) { return false; } if (state.compute_backend == nullptr || state.params_backend == nullptr) { @@ -927,6 +964,62 @@ bool ModelManager::alloc_params_buffers(const std::vector& states, return true; } +bool ModelManager::load_tensor_parts(TensorState& state) { + auto ctx = std::unique_ptr( + ggml_init({state.sources.size() * ggml_tensor_overhead(), nullptr, true}), ggml_free); + if (!ctx) { + return false; + } + const size_t size = ggml_nbytes(state.tensor); + std::vector buffer; + void* data = state.tensor->data; + if (!ggml_backend_buffer_is_host(state.tensor->buffer)) { + buffer.resize(size); + data = buffer.data(); + } + std::map parts; + std::set names; + size_t offset = 0; + for (const auto& source : state.sources) { + auto part = ggml_new_tensor(ctx.get(), state.tensor->type, source.n_dims, source.ne); + const size_t part_size = ggml_nbytes(part); + if (part_size > size - offset) { + return false; + } + part->data = static_cast(data) + offset; + parts[source.name] = part; + names.insert(source.name); + offset += part_size; + } + if (offset != size) { + return false; + } + std::set loaded; + std::mutex mutex; + auto callback = [&](const TensorStorage& source, ggml_tensor** dst) { + *dst = nullptr; + auto part = parts.find(source.name); + if (part != parts.end()) { + *dst = part->second; + std::lock_guard lock(mutex); + loaded.insert(source.name); + } + return true; + }; + const bool success = state.source_file == 0 + ? model_loader_.load_tensors(callback, enable_mmap_, &names, false) + : model_loader_.load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_); + if (!success || loaded != names) { + return false; + } + if (!buffer.empty()) { + // Upload the assembled tensor once, including for row-split backend buffers. + ggml_backend_tensor_set(state.tensor, buffer.data(), 0, size); + } + state.loaded_to_params_backend = true; + return true; +} + bool ModelManager::load_tensors(const std::vector& states) { using ReadGroup = std::pair; using ReadBatch = std::map>; @@ -934,6 +1027,12 @@ bool ModelManager::load_tensors(const std::vector& states) { for (auto* state : states) { if (state == nullptr) continue; + if (state->sources.size() > 1) { + if (!load_tensor_parts(*state)) { + return false; + } + continue; + } auto& batches = groups[{state->source_file, state->source_version}]; // The loader supplies one destination per name; only conflicting types need another batch. auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) { @@ -1309,6 +1408,9 @@ size_t ModelManager::compute_backend_alloc_size(const std::vector& if (state == nullptr || state->tensor == nullptr) { continue; } + if (state->usage_op == GGML_OP_CUSTOM && !sd_backend_is_cpu(state->compute_backend)) { + continue; + } const bool compute_resident = state->compute_backend == state->params_backend ? state->loaded_to_params_backend @@ -1353,15 +1455,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe } size_t total_size = 0; - auto add_buffer = [&](ggml_backend_buffer_t buffer) { - if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) { + std::unordered_set seen; + auto add_buffer = [&](ggml_backend_buffer_t buffer) { + if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) { return; } ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer); auto split_devices = split_buffer_devices_.find(buffer_type); const bool on_device = split_devices == split_buffer_devices_.end() - ? 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) { + ? 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) { return ggml_backend_get_device(entry.first) == compute_device; }); if (!on_device) { @@ -1371,9 +1474,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; }; + // 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_) { if (block != nullptr) { add_buffer(block->buffer); + for (const auto& store : block->mmap_tensor_stores) { + add_buffer(store.mmbuffer.get()); + } } } for (const auto& block : compute_staging_blocks_) { @@ -1616,7 +1726,8 @@ void ModelManager::remove_runtime_owner(uintptr_t owner_id) { ModelManager::CapacityCheck ModelManager::check_capacity( const DeviceMemoryRequest& request, - const std::vector& states) const { + const std::vector& states, + bool log_details) const { CapacityCheck result; if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) { return result; @@ -1648,16 +1759,23 @@ ModelManager::CapacityCheck ModelManager::check_capacity( device_query_cache_.erase(cached); } } + const size_t weights_resident = compute_backend_resident_bytes(backend); + const size_t other_runtime = other_runtime_resident_bytes(request.owner_id, backend); + const size_t resident = add(weights_resident, add(other_runtime, request.runtime_resident_bytes)); + if (log_details) { + LOG_WARN("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB", + ggml_backend_name(backend), + free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0), + weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0), + request.runtime_resident_bytes / (1024.0 * 1024.0)); + } if (free_bytes == 0 && total_bytes == 0) { return SIZE_MAX; } // Vulkan's heap budget subtraction can underflow when usage exceeds the budget. - if (total_bytes > 0 && free_bytes > total_bytes) { + if (total_bytes > 0 && free_bytes > total_bytes && sd_backend_is(backend, "Vulkan")) { return size_t{0}; } - const size_t resident = add(compute_backend_resident_bytes(backend), - add(other_runtime_resident_bytes(request.owner_id, backend), - request.runtime_resident_bytes)); if (total_bytes > 0) { free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0); } @@ -1816,7 +1934,7 @@ bool ModelManager::ensure_compute_backend_capacity( } } - const auto capacity = check_capacity(request, required_states); + const auto capacity = check_capacity(request, required_states, true); const std::string available_device = capacity.available_device_bytes == SIZE_MAX ? "unknown" : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0)); diff --git a/src/model_manager.h b/src/model_manager.h index e45a329536..86a931d55a 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -38,8 +38,7 @@ class ModelManager : public DeviceResidencyManager { std::string name; ggml_tensor* tensor = nullptr; ModelComponent component = ModelComponent::Count; - TensorStorage source; - bool has_source = false; + std::vector sources; ModelLoader::FileId source_file = 0; SDVersion source_version = VERSION_COUNT; @@ -139,6 +138,8 @@ class ModelManager : public DeviceResidencyManager { bool apply_loras_to_params(const std::vector& states); bool mmap_params(const std::vector& states, std::vector& created_storage_blocks); + static std::vector find_tensor_sources(const TensorState& state, const String2TensorStorage& sources); + bool load_tensor_parts(TensorState& state); bool can_mmap_storage(const TensorState& state) const; bool alloc_params_buffers(const std::vector& states, std::vector& created_storage_blocks); @@ -158,7 +159,8 @@ class ModelManager : public DeviceResidencyManager { } }; CapacityCheck check_capacity(const DeviceMemoryRequest& request, - const std::vector& states) const; + const std::vector& states, + bool log_details = false) const; ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const; ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const; diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp index b5dfe7c7c5..10c7d2e68f 100644 --- a/src/model_manager_files.cpp +++ b/src/model_manager_files.cpp @@ -15,6 +15,26 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) { a.int8_convrot_group_size == b.int8_convrot_group_size; } +std::vector ModelManager::find_tensor_sources(const TensorState& state, const String2TensorStorage& sources) { + auto first = sources.find(state.name); + if (first == sources.end()) { + return {}; + } + std::vector result{first->second}; + if (state.component == ModelComponent::LoRA || + std::equal(first->second.ne, first->second.ne + GGML_MAX_DIMS, state.tensor->ne)) { + return result; + } + for (size_t i = 1;; ++i) { + auto part = sources.find(state.name + "." + std::to_string(i)); + if (part == sources.end()) { + break; + } + result.push_back(part->second); + } + return result; +} + void ModelManager::invalidate_sources(const std::unordered_set& states) { auto affected = states; for (const auto& block : params_storage_blocks_) { @@ -76,20 +96,16 @@ bool ModelManager::set_loader(ModelLoader loader) { } std::unordered_set changed; for (const auto& state : tensor_states_) { - const auto& sources = sources_for(*state); - auto source = sources.find(state->name); - const bool found = source != sources.end(); - if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) || + const auto sources = find_tensor_sources(*state, sources_for(*state)); + if (sources.size() != state->sources.size() || + !std::equal(sources.begin(), sources.end(), state->sources.begin(), same_tensor_source) || (lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) { changed.insert(state.get()); } } invalidate_sources(changed); for (auto* state : changed) { - const auto& sources = sources_for(*state); - auto source = sources.find(state->name); - state->has_source = source != sources.end(); - state->source = state->has_source ? source->second : TensorStorage{}; + state->sources = find_tensor_sources(*state, sources_for(*state)); } if (lora_changed) { ++current_lora_epoch_; @@ -134,9 +150,8 @@ ModelLoader::FileVersions ModelManager::source_versions(const std::setsource_file] = loader.file_revision(state->source_file); continue; } - auto source = sources.find(state->name); - if (source != sources.end()) { - versions[source->second.file_id] = source->second.file_revision; + for (const auto& source : find_tensor_sources(*state, sources)) { + versions[source.file_id] = source.file_revision; } } return versions; diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp index 1d3041938d..de96d5ae3b 100644 --- a/src/model_manager_prefetch.cpp +++ b/src/model_manager_prefetch.cpp @@ -210,7 +210,7 @@ WeightPrefetchResult ModelManager::prefetch_params( is_optional_missing_tensor(state->name)) { continue; } - if (state->compute_backend == state->params_backend) { + if (state->usage_op == GGML_OP_CUSTOM || state->compute_backend == state->params_backend) { needs_synchronous_load = needs_synchronous_load || !state->loaded_to_params_backend; continue; @@ -275,6 +275,7 @@ bool ModelManager::activate_prefetched_params( [&](TensorState* state) { return state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name) || + state->usage_op == GGML_OP_CUSTOM || state->compute_backend == state->params_backend || state->staged_to_compute_backend; }); diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 6f1b598e88..8b286e8131 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -105,7 +105,25 @@ std::string convert_open_clip_to_hf_clip_name(std::string name) { std::string convert_llada2_moe_te_name(std::string name); -std::string convert_cond_stage_model_name(std::string name, std::string prefix) { +static std::string convert_ming_image_te_name(std::string name) { + if (name == "llm.query_tokens") { + name = "llm.query_tokens_dict.16x16"; + } + static const std::vector> name_map = { + {"thinker.", "backbone."}, + {"attention.", "self_attn."}, + {"self_attn.dense.", "self_attn.o_proj."}, + {"gate.proj.", "gate."}, + }; + replace_with_name_map(name, name_map); + return name; +} + +std::string convert_cond_stage_model_name(std::string name, std::string prefix, SDVersion version) { + if (version == VERSION_MING_IMAGE && prefix == "text_encoders." && starts_with(name, "llm.")) { + name = convert_ming_image_te_name(name); + } + static const std::vector> clip_name_map{ {"transformer.text_projection.weight", "transformer.text_model.text_projection"}, {"model.text_projection.weight", "transformer.text_model.text_projection"}, @@ -806,6 +824,7 @@ std::string convert_diffusers_dit_to_original_lumina2(std::string name) { if (z_image_name_map.empty()) { z_image_name_map["all_x_embedder.2-1."] = "x_embedder."; z_image_name_map["all_final_layer.2-1."] = "final_layer."; + z_image_name_map["all_x_embedder.16-1."] = "x_embedder."; // --- transformer blocks --- auto add_attention_map = [&](const std::string& prefix, int num) { @@ -932,6 +951,79 @@ static bool is_diffusers_controlnet_name(const std::string& name) { return false; } +static std::string convert_diffusers_dit_to_original_pixart(std::string name) { + static const std::vector> prefix_map = { + {"pos_embed.proj.", "x_embedder.proj."}, + {"adaln_single.emb.timestep_embedder.linear_1.", "t_embedder.mlp.0."}, + {"adaln_single.emb.timestep_embedder.linear_2.", "t_embedder.mlp.2."}, + {"adaln_single.emb.resolution_embedder.linear_1.", "csize_embedder.mlp.0."}, + {"adaln_single.emb.resolution_embedder.linear_2.", "csize_embedder.mlp.2."}, + {"adaln_single.emb.aspect_ratio_embedder.linear_1.", "ar_embedder.mlp.0."}, + {"adaln_single.emb.aspect_ratio_embedder.linear_2.", "ar_embedder.mlp.2."}, + {"adaln_single.linear.", "t_block.1."}, + {"caption_projection.linear_1.", "y_embedder.y_proj.fc1."}, + {"caption_projection.linear_2.", "y_embedder.y_proj.fc2."}, + {"proj_out.", "final_layer.linear."}, + }; + for (const auto& entry : prefix_map) { + if (starts_with(name, entry.first)) { + return entry.second + name.substr(entry.first.size()); + } + } + if (name == "scale_shift_table") { + return "final_layer.scale_shift_table"; + } + const std::string block_prefix = "transformer_blocks."; + if (!starts_with(name, block_prefix)) { + return name; + } + const size_t block_end = name.find('.', block_prefix.size()); + if (block_end == std::string::npos) { + return name; + } + const std::string prefix = "blocks." + name.substr(block_prefix.size(), block_end - block_prefix.size()) + "."; + name = name.substr(block_end + 1); + static const std::vector> block_map = { + {"attn1.to_q.", "attn.qkv."}, + {"attn1.to_out.0.", "attn.proj."}, + {"attn2.to_q.", "cross_attn.q_linear."}, + {"attn2.to_k.", "cross_attn.kv_linear."}, + {"attn2.to_out.0.", "cross_attn.proj."}, + {"ff.net.0.proj.", "mlp.fc1."}, + {"ff.net.2.", "mlp.fc2."}, + }; + for (const auto& entry : block_map) { + if (starts_with(name, entry.first)) { + return prefix + entry.second + name.substr(entry.first.size()); + } + } + static const std::vector> part_map = { + {"attn1.to_k.weight", "attn.qkv.weight.1"}, + {"attn1.to_k.bias", "attn.qkv.bias.1"}, + {"attn1.to_v.weight", "attn.qkv.weight.2"}, + {"attn1.to_v.bias", "attn.qkv.bias.2"}, + {"attn2.to_v.weight", "cross_attn.kv_linear.weight.1"}, + {"attn2.to_v.bias", "cross_attn.kv_linear.bias.1"}, + }; + for (const auto& entry : part_map) { + if (name == entry.first || starts_with(name, entry.first + ".")) { + return prefix + entry.second + name.substr(entry.first.size()); + } + } + return prefix + name; +} + +static std::string convert_ming_image_dit_name(std::string name) { + static const std::vector> name_map = { + {"all_x_embedder.2-1.", "x_embedder."}, + {"all_final_layer.2-1.", "final_layer."}, + {"attention.norm_q.", "attention.q_norm."}, + {"attention.norm_k.", "attention.k_norm."}, + }; + replace_with_name_map(name, name_map); + return name; +} + std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) { if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) { name = convert_diffusers_unet_to_original_sd1(name); @@ -943,7 +1035,9 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_diffusers_dit_to_original_flux(name); } else if (sd_version_is_hunyuan_video(version)) { name = convert_hunyuan_video_to_original_flux(name); - } else if (sd_version_is_z_image(version)) { + } else if (version == VERSION_MING_IMAGE) { + name = convert_ming_image_dit_name(name); + } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) { name = convert_diffusers_dit_to_original_lumina2(name); } else if (sd_version_is_llada_image(version)) { name = convert_diffusers_dit_to_original_llada_image(name); @@ -951,6 +1045,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_other_dit_to_original_anima(name); } else if (sd_version_is_krea2(version)) { name = convert_diffusers_dit_to_original_krea2(name); + } else if (sd_version_is_pixart(version)) { + name = convert_diffusers_dit_to_original_pixart(name); } return name; } @@ -1349,9 +1445,8 @@ bool is_first_stage_model_name(const std::string& name) { } static std::string convert_esrgan_tensor_name(std::string name) { - static std::unordered_map esrgan_name_map; - - if (esrgan_name_map.empty()) { + static const auto esrgan_name_map = [] { + std::unordered_map esrgan_name_map; esrgan_name_map["model.0."] = "conv_first."; constexpr int max_num_blocks = 64; @@ -1377,7 +1472,8 @@ static std::string convert_esrgan_tensor_name(std::string name) { esrgan_name_map["model.7."] = "conv_last."; esrgan_name_map["model.8."] = "conv_hr."; esrgan_name_map["model.10."] = "conv_last."; - } + return esrgan_name_map; + }(); replace_with_prefix_map(name, esrgan_name_map); return name; @@ -1604,7 +1700,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { { for (const auto& prefix : cond_stage_model_prefix_vec) { if (starts_with(name, prefix)) { - name = convert_cond_stage_model_name(name.substr(prefix.size()), prefix); + name = convert_cond_stage_model_name(name.substr(prefix.size()), prefix, version); name = prefix + name; break; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 3735a1ffdc..a39abd936d 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -30,6 +31,7 @@ #include "stable-diffusion.h" #include "conditioning/conditioner.hpp" +#include "conditioning/conditioning_cache.h" #include "core/backend_fit.h" #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" @@ -104,6 +106,10 @@ const char* model_version_to_str[] = { "SenseNova U1.5", "LLaDA-Image", "ESRGAN", + "PixArt", + "Ming-Image", + "Z-Image L2P", + "Iris-3B", }; static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]), @@ -124,6 +130,18 @@ void calculate_alphas_cumprod(float* alphas_cumprod, } } +void calculate_alphas_cumprod_linear_beta(float* alphas_cumprod, + float beta_start, + float beta_end, + int timesteps = TIMESTEPS) { + float product = 1.0f; + for (int i = 0; i < timesteps; i++) { + float beta = beta_start + (beta_end - beta_start) * ((float)i / (timesteps - 1)); + product *= 1.0f - beta; + alphas_cumprod[i] = product; + } +} + template struct has_set_runtime_backends : std::false_type {}; template @@ -136,6 +154,7 @@ static_assert(std::atomic::is_always_lock_free, StableDiffusionGGML::StableDiffusionGGML() : rng(std::make_shared()), + conditioning_cache_(std::make_unique()), denoiser(std::make_shared()) {} StableDiffusionGGML::~StableDiffusionGGML() = default; @@ -205,6 +224,8 @@ void StableDiffusionGGML::end_runners() { } bool StableDiffusionGGML::reset_runners(const RunnerGroups& groups) { + conditioning_cache_->clear(); + conditioning_loras_.clear(); end_runners(); clear_lora_adapters(); runtime_lora_models.clear(); @@ -662,6 +683,10 @@ void StableDiffusionGGML::refresh_compvis_denoiser_sigmas() { std::vector alphas_cumprod(TIMESTEPS); if (file_alphas_cumprod.size() == TIMESTEPS) { alphas_cumprod = file_alphas_cumprod; + } else if (sd_version_is_pixart(version)) { + // PixArt checkpoints train with a linear beta schedule (0.0001 -> 0.02) + // instead of the scaled_linear schedule used by SD1.x/SDXL. + calculate_alphas_cumprod_linear_beta(alphas_cumprod.data(), 0.0001f, 0.02f); } else { calculate_alphas_cumprod(alphas_cumprod.data()); } @@ -916,6 +941,11 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { 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(sd_ctx_params->conditioning_cache_size)); auto configuration = std::make_unique(*sd_ctx_params); n_threads = sd_ctx_params->n_threads; tensor_executor = std::make_unique(n_threads > 0 ? n_threads : sd_get_num_physical_cores()); @@ -1197,6 +1227,11 @@ bool StableDiffusionGGML::validate_and_load_runners() { LOG_VERBOSE("validating model metadata"); std::set ignore_tensors; + if (version == VERSION_MING_IMAGE) { + ignore_tensors.insert("text_encoders.llm.vision."); + ignore_tensors.insert("text_encoders.llm.linear_proj."); + ignore_tensors.insert("text_encoders.llm.backbone.lm_head."); + } if (use_tae && !tae_preview_only) { ignore_tensors.insert("first_stage_model."); } @@ -1357,10 +1392,13 @@ bool StableDiffusionGGML::build_denoiser() { sd_version_is_anima(version) || sd_version_is_ernie_image(version) || sd_version_is_z_image(version) || + sd_version_is_z_image_l2p(version) || + version == VERSION_MING_IMAGE || sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_pid(version) || - sd_version_is_ideogram4(version)) { + sd_version_is_ideogram4(version) || + version == VERSION_IRIS) { pred_type = FLOW_PRED; if (sd_version_is_wan(version)) { default_flow_shift = 5.f; @@ -1368,7 +1406,7 @@ bool StableDiffusionGGML::build_denoiser() { default_flow_shift = 7.f; } else if (sd_version_is_minimax_h3(version)) { default_flow_shift = 12.f; - } else if (sd_version_is_ernie_image(version)) { + } else if (sd_version_is_ernie_image(version) || version == VERSION_IRIS) { default_flow_shift = 4.f; } else if (sd_version_is_pid(version)) { default_flow_shift = 1.5f; @@ -1378,6 +1416,8 @@ bool StableDiffusionGGML::build_denoiser() { default_flow_shift = 3.16f; } else if (sd_version_is_mage_flow(version)) { default_flow_shift = 6.f; + } else if (version == VERSION_MING_IMAGE) { + default_flow_shift = INFINITY; } else if (sd_version_is_llada_image(version)) { default_flow_shift = 1.0f; // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid } else { @@ -1438,6 +1478,8 @@ bool StableDiffusionGGML::build_denoiser() { } else if (sd_version_is_minimax_h3(version)) { LOG_INFO("running in MiniMax H3 AV FLOW mode"); denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); + } else if (version == VERSION_MING_IMAGE) { + denoiser = std::make_shared(); } else { LOG_INFO("running in FLOW mode"); denoiser = std::make_shared(); @@ -1779,8 +1821,19 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun int64_t t0 = ggml_time_ms(); end_runners(); clear_lora_adapters(); - if (!model_manager->prepare_lora_sources(all_loras)) + if (!model_manager->prepare_lora_sources(all_loras)) { + conditioning_cache_->clear(); 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) { return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) { return entry.matches(spec); @@ -1790,6 +1843,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) : apply_loras_at_runtime(all_loras); if (!success) { + conditioning_cache_->clear(); clear_lora_adapters(); runtime_lora_models.clear(); return false; @@ -1799,9 +1853,14 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun if (!all_loras.empty()) { LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); } + conditioning_loras_ = std::move(all_loras); 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() { for (auto& extension : generation_extensions) { extension->reset_runtime_condition(); @@ -1986,6 +2045,8 @@ void StableDiffusionGGML::preview_image(int step, int patch_sz = 1; const float(*latent_rgb_proj)[3] = nullptr; float* latent_rgb_bias = nullptr; + const float* latent_alpha_proj = nullptr; + float latent_alpha_bias = 1.f; if (channels == 128) { if (sd_version_uses_flux2_vae(version)) { @@ -1999,6 +2060,16 @@ void StableDiffusionGGML::preview_image(int step, LOG_WARN("No latent to RGB projection known for this model"); 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) { if (sd_version_is_wan(version)) { latent_rgb_proj = wan_22_latent_rgb_proj; @@ -2049,13 +2120,14 @@ void StableDiffusionGGML::preview_image(int step, uint32_t img_width = static_cast(_latents.shape()[0]) * patch_sz; uint32_t img_height = static_cast(_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); - 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)); GGML_ASSERT(images != nullptr); 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); free(data); @@ -2118,7 +2190,7 @@ std::vector StableDiffusionGGML::prepare_sample_timesteps(float sigma, if (version == VERSION_HIDREAM_O1) { return std::vector{1.0f - (t / static_cast(TIMESTEPS))}; } - if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version)) { + if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version) || sd_version_is_ideogram4(version) || version == VERSION_MING_IMAGE) { return std::vector{1000.f - t}; } return std::vector{t}; @@ -2230,6 +2302,15 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptris_enabled(); }); + using QwenPrefixInputs = std::tuple*, const sd::Tensor*, + const std::vector>*>; + std::vector qwen_prefix_inputs; + RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; const bool apply_denoise_mask = !denoise_mask.empty() && @@ -2308,7 +2389,8 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr(steps)); SamplePreviewContext preview = prepare_sample_preview_context(); sd::Tensor processed_init_latent = denoiser->process_latent_in(init_latent); @@ -2406,6 +2488,7 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr> empty_ref_latents; bool uncond_without_ref_latents = !img_uncond.empty() && @@ -2461,6 +2544,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::sample(const std::shared_ptrbefore_diffusion(diffusion_params, step); } + if (cache_qwen_prefix) { + auto* extra = std::get_if(&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(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); if (output_opt.empty()) { LOG_ERROR("diffusion model compute failed"); @@ -2522,41 +2631,69 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr* uncond_skip_layers = nullptr; - if (is_skiplayer_step && slg_uncond) { - LOG_VERBOSE("Skipping layers at uncond step %d\n", step); - uncond_skip_layers = &skip_layer_guidance.layers(); - } - uncond_out = run_condition(uncond, - uncond.c_concat.empty() ? nullptr : &uncond.c_concat, - uncond_skip_layers, - nullptr, - true); - if (uncond_out.empty()) { - return {}; + if (!skip_uncond) { + const std::vector* uncond_skip_layers = nullptr; + if (is_skiplayer_step && slg_uncond) { + LOG_VERBOSE("Skipping layers at uncond step %d\n", step); + uncond_skip_layers = &skip_layer_guidance.layers(); + } + uncond_out = run_condition(uncond, + uncond.c_concat.empty() ? nullptr : &uncond.c_concat, + uncond_skip_layers, + nullptr, + true); + if (uncond_out.empty()) { + return {}; + } + } else { + step_cache.invalidate_condition(&uncond); } } + if (!img_uncond.empty()) { - img_uncond_out = run_condition(img_uncond, - img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat, - nullptr, - uncond_without_ref_latents ? &empty_ref_latents : nullptr, - true); - if (img_uncond_out.empty()) { - return {}; + if (!skip_img_uncond) { + img_uncond_out = run_condition(img_uncond, + img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat, + nullptr, + uncond_without_ref_latents ? &empty_ref_latents : nullptr, + true); + if (img_uncond_out.empty()) { + return {}; + } + } else { + step_cache.invalidate_condition(&img_uncond); } } sd::guidance::GuidanceInput guidance_input; @@ -2566,7 +2703,7 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor& x) { - auto latents = first_stage_model->encode(n_threads, x, vae_tiling_params, circular_x, circular_y); - // Same fallback as decode_first_stage, on a copy: an encode that ran out of memory is retried with - // overlapping tiles for this call only. - sd_tiling_params_t retry_tiling_params = vae_tiling_params; - while (latents.empty() && - sd::backend_fit::prepare_vae_retry_tiling(retry_tiling_params, false, true)) { - latents = first_stage_model->encode(n_threads, x, retry_tiling_params, circular_x, circular_y); + // Same out-of-memory fallback as decode_first_stage, on a copy of the tiling parameters, so the + // smaller tiles apply to this call only. + auto tiling_params = first_stage_model->resolve_tiling_params(vae_tiling_params); + const int scale_factor = first_stage_model->get_scale_factor(); + const int64_t latent_w = x.shape()[0] / scale_factor; + const int64_t latent_h = x.shape()[1] / scale_factor; + sd::Tensor latents; + for (int last_w = 0, last_h = 0;;) { + int tile_size_w = static_cast(latent_w); + int tile_size_h = static_cast(latent_h); + float tile_overlap; + if (tiling_params.enabled && + !first_stage_model->get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, + latent_w, latent_h, circular_x, circular_y)) { + return {}; + } + // get_tile_sizes can widen a retry tile to keep the overlap, so stop once a retry no longer shrinks it. + if (tile_size_w == last_w && tile_size_h == last_h) { + return {}; + } + latents = first_stage_model->encode(n_threads, x, tiling_params, circular_x, circular_y); + if (!latents.empty() || + !sd::backend_fit::prepare_vae_retry_tiling(tiling_params, false, first_stage_model->last_compute_status(), + tile_size_w, tile_size_h, scale_factor, true)) { + break; + } + last_w = tile_size_w; + last_h = tile_size_h; } if (latents.empty()) { return {}; @@ -2780,17 +2944,37 @@ sd::Tensor StableDiffusionGGML::encode_first_stage(const sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensor& x, bool decode_video) { - if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) { + if (sd_version_is_pid(version) || sd_version_is_minit2i(version) || version == VERSION_IRIS) { return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f); } auto latents = first_stage_model->diffusion_to_vae_latents(x); - auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); - const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode(); - while (decoded.empty() && - sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { - decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + auto tiling_params = first_stage_model->resolve_tiling_params(vae_tiling_params); + const bool prefer_temporal_tiling = decode_video && latents.dim() == 5 && latents.shape()[2] > 1 && + first_stage_model->can_temporal_tile_decode(); + for (int last_w = 0, last_h = 0;;) { + int tile_size_w = static_cast(latents.shape()[0]); + int tile_size_h = static_cast(latents.shape()[1]); + float tile_overlap; + if (tiling_params.enabled && + !first_stage_model->get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, + latents.shape()[0], latents.shape()[1], circular_x, circular_y)) { + return {}; + } + // get_tile_sizes can widen a retry tile, so stop once a spatial retry no longer shrinks it. + if (tile_size_w == last_w && tile_size_h == last_h) { + return {}; + } + auto decoded = first_stage_model->decode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y); + const bool temporal_before = tiling_params.temporal_tiling; + if (!decoded.empty() || + !sd::backend_fit::prepare_vae_decode_retry_tiling(tiling_params, prefer_temporal_tiling, + first_stage_model->last_compute_status(), + tile_size_w, tile_size_h, first_stage_model->get_scale_factor())) { + return decoded; + } + last_w = tiling_params.temporal_tiling != temporal_before ? 0 : tile_size_w; + last_h = tiling_params.temporal_tiling != temporal_before ? 0 : tile_size_h; } - return decoded; } sd::Tensor StableDiffusionGGML::normalize_ltx_video_latents(const sd::Tensor& x) { diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index f2f8a8ced6..d1a8866be4 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -26,6 +26,7 @@ class RNG; struct Denoiser; struct LoraModel; struct ConditionerParams; +class ConditioningCache; struct SDCondition; struct RefImageParams; namespace Wav2Vec2 { @@ -178,6 +179,8 @@ class StableDiffusionGGML { std::recursive_mutex execution_mutex; std::unique_ptr config_; RunnerState runner_state_; + std::unique_ptr conditioning_cache_; + std::vector conditioning_loras_; bool executing_ = false; std::shared_ptr denoiser; @@ -361,6 +364,8 @@ class StableDiffusionGGML { bool apply_loras(const sd_lora_t* loras, uint32_t lora_count); + SDCondition get_learned_condition(const ConditionerParams& params); + void reset_generation_extensions(); void prepare_generation_extensions(const sd_pm_params_t& pm_params, diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 8c6ffffa4d..7de85b3595 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -10,6 +10,7 @@ #include "model/vae/vae.hpp" #include "request.h" #include "runtime/denoiser.hpp" +#include "runtime/image_preprocess.h" #include "upscaler.h" namespace sd::pipeline { @@ -34,22 +35,23 @@ namespace sd::pipeline { return original_axes; } - int tile_size_x, tile_size_y; + int tile_size_w, tile_size_h; float overlap; - int latent_size_x = request.width / request.vae_scale_factor; - int latent_size_y = request.height / request.vae_scale_factor; - sd->first_stage_model->get_tile_sizes(tile_size_x, - tile_size_y, - overlap, - sd_img_gen_params->vae_tiling_params, - latent_size_x, - latent_size_y, - 1.0f, - sd->circular_x, - sd->circular_y); - - sd->circular_x = sd->circular_x && (tile_size_x >= latent_size_x); - sd->circular_y = sd->circular_y && (tile_size_y >= latent_size_y); + int latent_size_w = request.width / request.vae_scale_factor; + int latent_size_h = request.height / request.vae_scale_factor; + if (!sd->first_stage_model->get_tile_sizes(tile_size_w, + tile_size_h, + overlap, + sd_img_gen_params->vae_tiling_params, + latent_size_w, + latent_size_h, + sd->circular_x, + sd->circular_y)) { + return original_axes; + } + + sd->circular_x = sd->circular_x && (tile_size_w >= latent_size_w); + sd->circular_y = sd->circular_y && (tile_size_h >= latent_size_h); if (sd->first_stage_model) { sd->first_stage_model->set_circular_axes(sd->circular_x, sd->circular_y); @@ -58,8 +60,8 @@ namespace sd::pipeline { sd->preview_vae->set_circular_axes(sd->circular_x, sd->circular_y); } - sd->circular_x = original_axes.circular_x && (tile_size_x < latent_size_x); - sd->circular_y = original_axes.circular_y && (tile_size_y < latent_size_y); + sd->circular_x = original_axes.circular_x && (tile_size_w < latent_size_w); + sd->circular_y = original_axes.circular_y && (tile_size_h < latent_size_h); return original_axes; } @@ -443,8 +445,7 @@ namespace sd::pipeline { 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(); condition_params.zero_out_masked = false; - auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + auto cond = sd->get_learned_condition(condition_params); if (cond.empty()) { LOG_ERROR("failed to encode prompt"); return std::nullopt; @@ -466,6 +467,11 @@ namespace sd::pipeline { // states with a zeroed prompt mask, so no extra text encode is needed. uncond.c_crossattn = cond.c_crossattn; uncond.c_vector = sd::Tensor::zeros_like(cond.c_vector); + } else if (sd->version == VERSION_MING_IMAGE) { + uncond.c_crossattn = sd::Tensor::zeros_like(cond.c_crossattn); + for (const auto& extra : cond.extra_c_crossattns) { + uncond.extra_c_crossattns.push_back(sd::Tensor::zeros_like(extra)); + } } else if (sd_version_is_sensenova_u1(sd->version)) { auto* sensenova_conditioner = static_cast(sd->cond_stage_model.get()); uncond = sensenova_conditioner->get_unconditional_condition(request->negative_prompt); @@ -482,8 +488,7 @@ namespace sd::pipeline { // LLaDA-Image CFG keeps the source latent but drops its SigVQ features. condition_params.ref_images = nullptr; } - uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + uncond = sd->get_learned_condition(condition_params); if (uncond.empty()) { LOG_ERROR("failed to encode negative prompt"); return std::nullopt; @@ -511,8 +516,7 @@ namespace sd::pipeline { if (use_ref_latent_img_cfg) { condition_params.ref_images = &empty_ref_images; } - img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + img_uncond = sd->get_learned_condition(condition_params); if (img_uncond.empty()) { LOG_ERROR("failed to encode image guidance prompt"); return std::nullopt; @@ -792,15 +796,9 @@ namespace sd::pipeline { return false; } - // MiniMax-H3 is video-only. Its denoiser always splits the packed latent into a video and an - // audio half, and only generate_video ever computes the audio length, so reaching this - // function with an H3 checkpoint is guaranteed to die on - // GGML_ASSERT(!audio_input_cache.empty()) with a core dump, after the several minutes it - // takes to load the weights, and with nothing in the output pointing at the missing --mode. - // (The AnimateDiff path below routes vid_gen back through here, but that is SD1.5 plus a - // motion module, never H3.) - if (sd_version_is_minimax_h3(sd->version)) { - LOG_ERROR("MiniMax-H3 is a video model and cannot be run in img_gen mode; use --mode vid_gen"); + if (!sd_version_supports_image_generation(sd->version)) { + LOG_ERROR("%s cannot be run with generate_image(); use generate_video() or --mode vid_gen in the CLI", + model_version_to_str[sd->version]); return false; } @@ -809,6 +807,12 @@ namespace sd::pipeline { int64_t t0 = ggml_time_ms(); sd->vae_tiling_params = sd_img_gen_params->vae_tiling_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); sd->rng->manual_seed(request.seed); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 626bdfee60..96b5a974c9 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -1,5 +1,6 @@ #include "model_builders.h" +#include #include #include @@ -17,23 +18,27 @@ #include "model/diffusion/hidream_o1.hpp" #include "model/diffusion/hunyuan.hpp" #include "model/diffusion/ideogram4.hpp" +#include "model/diffusion/iris.hpp" #include "model/diffusion/krea2.hpp" #include "model/diffusion/lens.hpp" #include "model/diffusion/lingbot_video.hpp" #include "model/diffusion/llada_image.hpp" #include "model/diffusion/ltxv.hpp" #include "model/diffusion/mage_flow.hpp" +#include "model/diffusion/ming_image.hpp" #include "model/diffusion/minimax_h3.hpp" #include "model/diffusion/minit2i.hpp" #include "model/diffusion/mmdit.hpp" #include "model/diffusion/model.hpp" #include "model/diffusion/pid.hpp" +#include "model/diffusion/pixart.hpp" #include "model/diffusion/qwen_image.hpp" #include "model/diffusion/qwen_image_2_1.hpp" #include "model/diffusion/sensenova_u1.h" #include "model/diffusion/unet.hpp" #include "model/diffusion/wan.hpp" #include "model/diffusion/z_image.hpp" +#include "model/diffusion/z_image_l2p.hpp" #include "model/vae/auto_encoder_kl.hpp" #include "model/vae/hunyuan_vae.hpp" #include "model/vae/ltx_audio_vae.hpp" @@ -291,7 +296,8 @@ namespace sd::model_builders { result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", - weight_manager); + weight_manager, + sd_ctx_params->model_args); } else { result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), tensor_storage_map, @@ -300,6 +306,19 @@ namespace sd::model_builders { weight_manager, sd_ctx_params->model_args); } + } else if (version == VERSION_PIXART) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + true, + 0, + false, + weight_manager, + sd_ctx_params->model_args); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager, + sd_ctx_params->model_args); } else if (sd_version_is_mage_flow(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, @@ -358,6 +377,11 @@ namespace sd::model_builders { tensor_storage_map, "model.diffusion_model", weight_manager); + } else if (version == VERSION_MING_IMAGE) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); } else if (sd_version_is_z_image(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, @@ -371,6 +395,18 @@ namespace sd::model_builders { "model.diffusion_model", version, weight_manager); + } else if (sd_version_is_z_image_l2p(version)) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); } else if (sd_version_is_llada_image(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, @@ -421,6 +457,11 @@ namespace sd::model_builders { tensor_storage_map, "model.diffusion_model", weight_manager); + } else if (version == VERSION_IRIS) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, version, "", false, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); } else { // SD1.x SD2.x SDXL std::map embbeding_map; for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) { @@ -561,6 +602,23 @@ namespace sd::model_builders { false, vae_version, weight_manager); + if (sd_version_is_pixart(version)) { + // Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override. + if (tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) { + model->scale_factor = 0.18215f; + } + for (const auto& [key, value] : parse_key_value_args(sd_ctx_params->model_args, "model arg")) { + if (key == "pixart_vae_scale_factor") { + float parsed = 0.f; + if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed > 0.f) { + model->scale_factor = parsed; + } else { + LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + LOG_VERBOSE("pixart: VAE scale factor = %.5f", model->scale_factor); + } if (sd_version_is_sdxl(version) && (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) { float vae_conv_2d_scale = 1.f / 32.f; @@ -574,7 +632,12 @@ namespace sd::model_builders { } }; - if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) { + if (version == VERSION_CHROMA_RADIANCE || + version == VERSION_HIDREAM_O1 || + sd_version_is_minit2i(version) || + sd_version_is_sensenova_u1(version) || + sd_version_is_z_image_l2p(version) || + version == VERSION_IRIS) { LOG_INFO("using FakeVAE"); result.vae = std::make_shared(version, ctx.backends.runtime_backend(SDBackendModule::VAE), diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index cfc54183af..4256fee4e2 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -384,17 +384,19 @@ namespace sd::pipeline { total_steps = sample_steps + std::max(0, high_noise_sample_steps); if (sample_params->custom_sigmas_count > 0) { - sigmas = std::vector(sample_params->custom_sigmas, + sigmas = std::vector(sample_params->custom_sigmas, sample_params->custom_sigmas + sample_params->custom_sigmas_count); - total_steps = static_cast(sigmas.size()) - 1; - LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps); - if (sample_steps >= total_steps) { - sample_steps = total_steps; - LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps); - } - if (high_noise_sample_steps > 0) { - high_noise_sample_steps = total_steps - sample_steps; - LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps); + if (total_steps != sample_params->custom_sigmas_count - 1) { + total_steps = sample_params->custom_sigmas_count - 1; + LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps); + if ((sample_steps > total_steps) || (high_noise_sample_steps <= 0 && sample_steps != total_steps)) { + sample_steps = total_steps; + LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps); + } + if (high_noise_sample_steps > 0) { + high_noise_sample_steps = total_steps - sample_steps; + LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps); + } } } else { scheduler_t scheduler = resolve_scheduler(sd, diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index 503afa530a..2fcff4c9d4 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -15,6 +15,7 @@ #include "model/vae/vae.hpp" #include "request.h" #include "runtime/denoiser.hpp" +#include "runtime/image_preprocess.h" namespace sd::pipeline { @@ -470,11 +471,15 @@ namespace sd::pipeline { sd::Tensor end_image; if (sd_vid_gen_params->init_image.data) { - start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height); + start_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height), + sd->get_image_channels()); } if (sd_vid_gen_params->end_image.data) { - end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height); + end_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height), + sd->get_image_channels()); } if (sd_version_is_minimax_h3(sd->version)) { @@ -1157,8 +1162,7 @@ namespace sd::pipeline { } int64_t prepare_start_ms = ggml_time_ms(); - embeds.cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + embeds.cond = sd->get_learned_condition(condition_params); if (embeds.cond.empty()) { LOG_ERROR("failed to encode video prompt"); return std::nullopt; @@ -1183,8 +1187,7 @@ namespace sd::pipeline { } if (request.use_uncond) { condition_params.text = request.negative_prompt; - embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + embeds.uncond = sd->get_learned_condition(condition_params); if (embeds.uncond.empty()) { LOG_ERROR("failed to encode negative video prompt"); return std::nullopt; @@ -1416,7 +1419,9 @@ namespace sd::pipeline { sd::Tensor video_mask = make_ltxav_video_denoise_mask(video_latent, 1.f); if (sd_vid_gen_params->init_image.data != nullptr) { - sd::Tensor start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height); + sd::Tensor start_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height), + sd->get_image_channels()); if (!apply_ltxav_condition_image_by_latent_index(sd, start_image, &video_latent, @@ -1429,7 +1434,9 @@ namespace sd::pipeline { } if (sd_vid_gen_params->end_image.data != nullptr) { - sd::Tensor end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height); + sd::Tensor end_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height), + sd->get_image_channels()); sd::Tensor end_image_latent = encode_ltxav_condition_image(sd, end_image, "end"); if (end_image_latent.empty()) { return false; @@ -1518,6 +1525,7 @@ namespace sd::pipeline { img_gen_params.qwen_image_layers = 0; 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.image_preprocess = sd_vid_gen_params->image_preprocess; sd->animatediff_num_frames = n_frames; bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); @@ -1548,6 +1556,11 @@ namespace sd::pipeline { sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); GenerationRequest request(sd, sd_vid_gen_params); + 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) { *fps_out = request.fps; } diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 416859e6da..02ebe57e0a 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -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 struct FluxScheduler : SigmaScheduler { - int image_seq_len = 0; - float base_shift = 0.5f; - float max_shift = 1.15f; - - explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr) + int image_seq_len = 0; + int base_image_seq_len = 256; + int max_image_seq_len = 4096; + float base_shift = 0.5f; + float max_shift = 1.15f; + float shift_terminal = 0.0f; + + FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr) : 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); } @@ -752,10 +760,8 @@ struct FluxScheduler : SigmaScheduler { } float compute_mu() const { - constexpr float base_shift_anchor = 256.0f; - constexpr float max_shift_anchor = 4096.0f; - float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor); - float b = base_shift - m * base_shift_anchor; + float m = (max_shift - base_shift) / static_cast(max_image_seq_len - base_image_seq_len); + float b = base_shift - m * static_cast(base_image_seq_len); return static_cast(image_seq_len) * m + b; } @@ -764,7 +770,7 @@ struct FluxScheduler : SigmaScheduler { sigmas.reserve(n + 1); 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) { 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; return sigmas; } @@ -1178,7 +1194,7 @@ struct Denoiser { } case FLUX_SCHEDULER: { LOG_INFO("get_sigmas with Flux scheduler"); - scheduler = std::make_shared(image_seq_len, extra_sample_args); + scheduler = std::make_shared(image_seq_len, version, extra_sample_args); break; } default: @@ -1353,6 +1369,42 @@ struct DiscreteFlowDenoiser : public Denoiser { } }; +struct MingImageFlowDenoiser : DiscreteFlowDenoiser { + float resolution_shift = std::exp(1.35f); + + MingImageFlowDenoiser() + : DiscreteFlowDenoiser(INFINITY) {} + + float sigma_min() override { return 0.f; } + float sigma_max() override { return 1.f; } + + float t_to_sigma(float t) override { + return time_snr_shift(std::isfinite(shift) ? shift : resolution_shift, (t + 1.f) / 1000.f); + } + + std::vector get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler, SDVersion version, const char* extra_sample_args = nullptr) override { + const float tokens = static_cast(image_seq_len) / 4.f; + const float mu = tokens >= 4096.f ? 1.35f : 0.5f + (tokens - 256.f) * (1.15f - 0.5f) / (4096.f - 256.f); + resolution_shift = std::exp(mu); + if (scheduler != DISCRETE_SCHEDULER) { + return DiscreteFlowDenoiser::get_sigmas(n, image_seq_len, scheduler, version, extra_sample_args); + } + if (n == 0) { + return {}; + } + const float factor = std::isfinite(shift) ? shift : resolution_shift; + std::vector sigmas; + sigmas.reserve(n + 1); + for (uint32_t i = 0; i < n; ++i) { + const float t = n == 1 ? 1.f : 1.f - static_cast(i) / static_cast(n - 1); + sigmas.push_back(time_snr_shift(factor, t)); + } + // Upstream sets sigma_min to zero before linspace, then appends the terminal zero. + sigmas.push_back(0.f); + return sigmas; + } +}; + struct H3AVFlowDenoiser : public DiscreteFlowDenoiser { int64_t video_channels; float audio_shift; @@ -1767,7 +1819,10 @@ static sd::Tensor sample_euler(denoise_cb_t model, const std::vector& sigmas) { int steps = static_cast(sigmas.size()) - 1; for (int i = 0; i < steps; i++) { - float sigma = sigmas[i]; + float sigma = sigmas[i]; + if (sigma == sigmas[i + 1]) { + continue; + } auto denoised_opt = model(x, sigma, i + 1); if (denoised_opt.pred.empty()) { return {}; diff --git a/src/runtime/image_preprocess.cpp b/src/runtime/image_preprocess.cpp new file mode 100644 index 0000000000..a8df7c3b56 --- /dev/null +++ b/src/runtime/image_preprocess.cpp @@ -0,0 +1,447 @@ +#include "image_preprocess.h" +#include "core/util.h" + +#include +#include + +namespace sd { + + static constexpr std::pair 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 image_resize_modes[] = { + {"auto", ImageResizeMode::Auto}, + {"none", ImageResizeMode::None}, + {"stretch", ImageResizeMode::Stretch}, + {"crop", ImageResizeMode::Crop}, + {"crop-resize", ImageResizeMode::CropResize}, + {"fit-pad", ImageResizeMode::FitPad}, + }; + + template + static bool parse_enum(const std::string& text, const std::pair (&names)[N], T& value) { + for (const auto& entry : names) { + if (text == entry.first) { + value = entry.second; + return true; + } + } + return false; + } + + template + static const char* enum_name(T value, const std::pair (&names)[N]) { + for (const auto& entry : names) { + if (value == entry.second) + return entry.first; + } + return "unknown"; + } + + template + static bool one_of(T value, std::initializer_list choices) { + return std::find(choices.begin(), choices.end(), value) != choices.end(); + } + + static bool one_of(const std::string& value, std::initializer_list choices) { + for (const char* choice : choices) { + if (value == choice) + return true; + } + return false; + } + + static ImageResizeMode resolve_mode(const std::map& 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 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(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 ImagePreprocessor::resolve_options(ImageTarget target, int index) const { + std::map 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 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 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 ImagePreprocessor::apply_transform(const Tensor& image, const std::map& 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(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 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(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(mode), BOOL_STR(antialias)); + return output; + } + + Tensor ImagePreprocessor::apply_geometry(const Tensor& 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(image.shape()[0]); + p.source_height = p.crop_height = static_cast(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(int64_t(p.source_height) * target_width / target_height)); + } else { + p.crop_height = std::max(1, static_cast(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(std::round(p.source_width * scale)))); + p.resize_height = std::max(1, std::min(target_height, static_cast(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(INT64_MAX, SIZE_MAX / sizeof(float)) / static_cast(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 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& 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 diff --git a/src/runtime/image_preprocess.h b/src/runtime/image_preprocess.h new file mode 100644 index 0000000000..b75c4e0e50 --- /dev/null +++ b/src/runtime/image_preprocess.h @@ -0,0 +1,88 @@ +#ifndef __SD_RUNTIME_IMAGE_PREPROCESS_H__ +#define __SD_RUNTIME_IMAGE_PREPROCESS_H__ + +#include +#include +#include + +#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 options; + }; + + class ImagePreprocessor { + std::vector rules_; + mutable bool valid_ = true; + ImageTransform init_transform_; + bool has_init_transform_ = false; + bool prepared_ = false; + std::vector owned_images_; + std::vector ref_images_; + std::vector id_images_; + std::vector control_frames_; + + bool fail(const std::string& message) const; + std::map resolve_options(ImageTarget target, int index) const; + Tensor apply_transform(const Tensor& image, const std::map& 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& 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 apply_geometry(const Tensor& 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 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__ diff --git a/src/runtime/latent-preview.h b/src/runtime/latent-preview.h index c19a0eebcf..5baf6f53cb 100644 --- a/src/runtime/latent-preview.h +++ b/src/runtime/latent-preview.h @@ -4,6 +4,86 @@ #include "core/tensor.hpp" #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] = { {0.19819857f, 0.11584999f, 0.07929777f}, {-0.16047224f, -0.10601170f, -0.15996324f}, @@ -324,7 +404,7 @@ const float sd_latent_rgb_proj[4][3] = { {-0.178022f, -0.200862f, -0.678514f}}; 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; uint32_t latent_width = static_cast(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_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 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 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) { 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]); r += value * latent_rgb_proj[d][0]; g += value * latent_rgb_proj[d][1]; b += value * latent_rgb_proj[d][2]; + if (latent_alpha_proj != nullptr) { + a += value * latent_alpha_proj[d]; + } } } else { // 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; b = b >= 0 ? b <= 1 ? b : 1 : 0; - buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255); - buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255); - buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255); + buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255); + buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 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& laten return latents.dim() == 5; } -void preview_latent_video(uint8_t* buffer, const sd::Tensor& 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& 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(latents.shape()[0]); uint32_t latent_height = static_cast(latents.shape()[1]); bool is_video = preview_latent_tensor_is_video(latents); uint32_t frames = is_video ? static_cast(latents.shape()[2]) : 1; uint32_t dim = is_video ? static_cast(latents.shape()[3]) : static_cast(latents.shape()[2]); - uint32_t rgb_width = latent_width * patch_size; - uint32_t rgb_height = latent_height * patch_size; - uint32_t unpatched_dim = dim / (patch_size * patch_size); + uint32_t rgb_width = latent_width * patch_size; + uint32_t rgb_height = latent_height * 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 rgb_x = 0; rgb_x < rgb_width; rgb_x++) { @@ -427,7 +516,7 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor& latents, con : 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) { for (uint32_t d = 0; d < unpatched_dim; d++) { 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& latents, con r += value * latent_rgb_proj[d][0]; g += value * latent_rgb_proj[d][1]; b += value * latent_rgb_proj[d][2]; + if (latent_alpha_proj != nullptr) { + a += value * latent_alpha_proj[d]; + } } } else { r = latent_value(0); @@ -450,9 +542,13 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor& latents, con g = std::min(1.0f, std::max(0.0f, g * .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 * 3 + 1] = (uint8_t)(g * 255); - buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255); + buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255); + buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 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); + } } } } diff --git a/src/runtime/preprocessing.hpp b/src/runtime/preprocessing.hpp index 1990d15f46..ab4ca29b64 100644 --- a/src/runtime/preprocessing.hpp +++ b/src/runtime/preprocessing.hpp @@ -165,16 +165,18 @@ static inline sd::Tensor convolve_tensor(const sd::Tensor& input, return output; } -static inline sd::Tensor grayscale_tensor(const sd::Tensor& rgb_img) { - GGML_ASSERT(rgb_img.dim() == 4); - GGML_ASSERT(rgb_img.shape()[2] >= 3); - sd::Tensor grayscale({rgb_img.shape()[0], rgb_img.shape()[1], 1, rgb_img.shape()[3]}); - for (int64_t iy = 0; iy < rgb_img.shape()[1]; ++iy) { - for (int64_t ix = 0; ix < rgb_img.shape()[0]; ++ix) { - float r = preprocessing_get_4d(rgb_img, ix, iy, 0, 0); - float g = preprocessing_get_4d(rgb_img, ix, iy, 1, 0); - float b = preprocessing_get_4d(rgb_img, ix, iy, 2, 0); - float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b; +static inline sd::Tensor grayscale_tensor(const sd::Tensor& image) { + GGML_ASSERT(image.dim() == 4); + GGML_ASSERT(image.shape()[2] >= 1); + sd::Tensor grayscale({image.shape()[0], image.shape()[1], 1, image.shape()[3]}); + for (int64_t iy = 0; iy < image.shape()[1]; ++iy) { + for (int64_t ix = 0; ix < image.shape()[0]; ++ix) { + float gray = preprocessing_get_4d(image, ix, iy, 0, 0); + if (image.shape()[2] >= 3) { + float g = preprocessing_get_4d(image, ix, iy, 1, 0); + 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); } } @@ -317,11 +319,12 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, image_gray = non_max_supression(G, theta); 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 ix = 0; ix < img.width; ++ix) { float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0); 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); } } diff --git a/src/runtime/sample-cache.cpp b/src/runtime/sample-cache.cpp index bfa3d4487f..0b92c73358 100644 --- a/src/runtime/sample-cache.cpp +++ b/src/runtime/sample-cache.cpp @@ -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 { switch (runtime.mode) { case SampleCacheMode::EASYCACHE: diff --git a/src/runtime/sample-cache.h b/src/runtime/sample-cache.h index cb0fe7bc58..30ff66e5fe 100644 --- a/src/runtime/sample-cache.h +++ b/src/runtime/sample-cache.h @@ -46,6 +46,7 @@ namespace sd_sample { bool before_condition(const void* condition, const sd::Tensor& input, sd::Tensor* output); void after_condition(const void* condition, const sd::Tensor& input, const sd::Tensor& output); + void invalidate_condition(const void* condition); bool is_step_skipped() const; }; diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index 2749b290e8..782a41a8b9 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -13,73 +13,70 @@ #include "core/util.h" #include "ggml.h" -static void sd_tiling_calc_tiles(int& num_tiles_dim, - float& tile_overlap_factor_dim, - int small_dim, - int tile_size, - const float tile_overlap_factor, - bool circular) { - int tile_overlap = static_cast(tile_size * tile_overlap_factor); - int non_tile_overlap = tile_size - tile_overlap; +struct TileSpan { + int offset; + int overlap_before; + int overlap_after; +}; - if (circular) { - // circular means the last and first tile are overlapping (wraping around) - num_tiles_dim = small_dim / non_tile_overlap; - - if (num_tiles_dim < 1) { - num_tiles_dim = 1; - } - - tile_overlap_factor_dim = (tile_size - small_dim / num_tiles_dim) / (float)tile_size; - - // if single tile and tile_overlap_factor is not 0, add one to ensure we have at least two overlapping tiles - if (num_tiles_dim == 1 && tile_overlap_factor_dim > 0) { - num_tiles_dim++; - tile_overlap_factor_dim = 0.5; - } - - return; +static std::vector sd_tiling_plan_axis(int dimension, + int tile_size, + float target_overlap, + bool circular) { + if (tile_size >= dimension) { + return {{0, 0, 0}}; } - // else, non-circular means the last and first tile are not overlapping - - num_tiles_dim = (small_dim - tile_overlap) / non_tile_overlap; - int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim; - if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) { - // if tiles don't fit perfectly using the desired overlap - // and there is enough room to squeeze an extra tile without overlap becoming >0.5 - num_tiles_dim++; + const int extent = circular ? dimension : dimension - tile_size; + int intervals = extent; + if (tile_size > 1) { + const float target = extent / ((1.0f - target_overlap) * tile_size); + // Adjacent tiles must overlap, but three tiles must not cover the same point. + const int min_intervals = 1 + (extent - 1) / (tile_size - 1); + const int max_intervals = std::max(1, static_cast(2LL * extent / tile_size)); + const int lower = std::clamp(static_cast(std::floor(target)), min_intervals, max_intervals); + const int upper = std::clamp(static_cast(std::ceil(target)), min_intervals, max_intervals); + auto error = [&](int count) { + return std::abs(1.0f / count - 1.0f / target); + }; + intervals = error(upper) < error(lower) ? upper : lower; } - tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1)); - if (num_tiles_dim <= 2) { - if (small_dim <= tile_size) { - num_tiles_dim = 1; - tile_overlap_factor_dim = 0; - } else { - num_tiles_dim = 2; - tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size; - } + const int num_tiles = circular ? intervals : intervals + 1; + auto position = [&](int index) -> int { + return static_cast(int64_t(index) * extent / intervals); + }; + std::vector tiles; + tiles.reserve(num_tiles); + for (int i = 0; i < num_tiles; ++i) { + const int offset = position(i); + // Keep wrapped neighbors in the same unwrapped coordinate space. + const int previous = i > 0 ? position(i - 1) : position(num_tiles - 1) - dimension; + const int next = i + 1 < num_tiles ? position(i + 1) : dimension; + const int before = i > 0 || circular ? previous + tile_size - offset : 0; + const int after = i + 1 < num_tiles || circular ? offset + tile_size - next : 0; + tiles.push_back({offset, before, after}); } + return tiles; } int sd_tiling_seam_safe_tile_size(int small_dim, int tile_size, float tile_overlap_factor) { - // sd_tiling_calc_tiles lets the real overlap fall well below the target for some sizes. The worst case is a - // tile a little over half of the axis (e.g. round(dim / 2)): it falls back to two tiles overlapping by - // 2 * tile_size - small_dim, which is 0 or 1 latent, so the cross-fade has nothing to blend over and the tile - // edge shows as a line. Pick the nearest tile size (smaller first; a larger one only up to twice the request, - // so an explicitly small tile cannot balloon) whose tiles overlap by at least half the target and by 2+ latents. + // A tile a little over half the axis (32 of 62 latents) gets two tiles overlapping by 1-2 latents, which shows + // as a seam. Pick the nearest size (smaller first, larger only up to 2x) whose overlaps are >= half the target + // and >= 2 latents. const int min_tile_size = 4; if (tile_overlap_factor <= 0.f || tile_size >= small_dim) { return tile_size; } auto overlaps_enough = [&](int size) { - int num_tiles; - float overlap_factor; - sd_tiling_calc_tiles(num_tiles, overlap_factor, small_dim, size, tile_overlap_factor, false); - const int overlap = static_cast(size * overlap_factor); const int min_overlap = std::max(2, static_cast(size * tile_overlap_factor * 0.5f)); - return num_tiles < 2 || overlap >= min_overlap; + const auto tiles = sd_tiling_plan_axis(small_dim, size, tile_overlap_factor, false); + for (size_t i = 0; i + 1 < tiles.size(); ++i) { + if (tiles[i].overlap_after < min_overlap) { + return false; + } + } + return true; }; for (int size = tile_size; size >= min_tile_size; --size) { if (overlaps_enough(size)) { @@ -152,70 +149,69 @@ static void sd_parallel_for_planes(int64_t plane_count, int64_t plane_elements, } static void sd_tensor_merge_2d(const float* input_data, - int64_t width, - int64_t height, + int64_t in_width, + int64_t in_height, int64_t plane_count, sd::Tensor* output, int x, int y, - int overlap_x, - int overlap_y, - bool circular_x, - bool circular_y, - int x_skip, - int y_skip) { + int overlap_left, + int overlap_right, + int overlap_top, + int overlap_bottom) { GGML_ASSERT(output != nullptr); - int64_t img_width = output->shape()[0]; - int64_t img_height = output->shape()[1]; - int64_t input_plane = width * height; - int64_t output_plane = sd_tensor_plane_size(*output); - GGML_ASSERT(output->numel() / output_plane == plane_count); - // unclamped -> expects x in the range [0-1] + int64_t out_width = output->shape()[0]; + int64_t out_height = output->shape()[1]; + int64_t in_size = in_width * in_height; + int64_t out_size = sd_tensor_plane_size(*output); + + GGML_ASSERT(output->numel() == plane_count * out_size); + GGML_ASSERT(x >= 0 && y >= 0); + GGML_ASSERT(x < out_width && in_width <= out_width); + GGML_ASSERT(y < out_height && in_height <= out_height); + GGML_ASSERT(overlap_left >= 0 && overlap_right >= 0); + GGML_ASSERT(overlap_top >= 0 && overlap_bottom >= 0); + auto smootherstep_f32 = [](const float x) -> float { - GGML_ASSERT(x >= 0.f && x <= 1.f); return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f); }; - - // Same per-element arithmetic as the per-pixel form, so the result stays bit-identical. - const bool blend = overlap_x > 0 || overlap_y > 0; - std::vector wx, wy; - std::vector dst_x(static_cast(width)); - for (int64_t ix = x_skip; ix < width; ix++) { - dst_x[ix] = (x + ix) % img_width; + // Weights depend only on the column or row; the per-element arithmetic order is unchanged. + std::vector x_weights(static_cast(in_width)); + std::vector y_weights(static_cast(in_height)); + std::vector dst_x(static_cast(in_width)); + for (int ix = 0; ix < in_width; ++ix) { + float x_f = 1.0f; + if (ix < overlap_left) { + x_f = static_cast(ix) / overlap_left; + } + if (ix >= in_width - overlap_right) { + x_f = static_cast(in_width - ix) / overlap_right; + } + x_weights[ix] = smootherstep_f32(std::clamp(x_f, 0.0f, 1.0f)); + dst_x[ix] = (x + ix) % out_width; } - if (blend) { - wx.resize(static_cast(width)); - wy.resize(static_cast(height)); - for (int64_t ix = x_skip; ix < width; ix++) { - const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f; - const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f; - wx[ix] = smootherstep_f32(std::min(std::min(x_f_0, x_f_1), 1.f)); + for (int iy = 0; iy < in_height; ++iy) { + float y_f = 1.0f; + if (iy < overlap_top) { + y_f = static_cast(iy) / overlap_top; } - for (int64_t iy = y_skip; iy < height; iy++) { - const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f; - const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f; - wy[iy] = smootherstep_f32(std::min(std::min(y_f_0, y_f_1), 1.f)); + if (iy >= in_height - overlap_bottom) { + y_f = static_cast(in_height - iy) / overlap_bottom; } + y_weights[iy] = smootherstep_f32(std::clamp(y_f, 0.0f, 1.0f)); } float* output_data = output->data(); - sd_parallel_for_planes(plane_count, input_plane, [&](int64_t first_plane, int64_t last_plane) { + sd_parallel_for_planes(plane_count, in_size, [&](int64_t first_plane, int64_t last_plane) { for (int64_t plane = first_plane; plane < last_plane; ++plane) { - const float* src = input_data + plane * input_plane; - float* dst = output_data + plane * output_plane; - for (int64_t iy = y_skip; iy < height; iy++) { - const float* src_row = src + width * iy; - float* dst_row = dst + img_width * ((y + iy) % img_height); - if (blend) { - const float sy = wy[iy]; - for (int64_t ix = x_skip; ix < width; ix++) { - float& out = dst_row[dst_x[ix]]; - out = out + src_row[ix] * sy * wx[ix]; - } - } else { - for (int64_t ix = x_skip; ix < width; ix++) { - dst_row[dst_x[ix]] = src_row[ix]; - } + const float* src = input_data + plane * in_size; + float* dst = output_data + plane * out_size; + for (int iy = 0; iy < in_height; ++iy) { + const float y_weight = y_weights[iy]; + const float* src_row = src + in_width * iy; + float* dst_row = dst + out_width * ((y + iy) % out_height); + for (int ix = 0; ix < in_width; ++ix) { + dst_row[dst_x[ix]] += x_weights[ix] * y_weight * src_row[ix]; } } } @@ -226,8 +222,8 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y, @@ -238,8 +234,8 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, output_width, output_height, scale, - p_tile_size_x, - p_tile_size_y, + p_tile_size_w, + p_tile_size_h, tile_overlap_factor, circular_x, circular_y, @@ -259,8 +255,8 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y, @@ -271,96 +267,40 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, int input_width = static_cast(input.shape()[0]); int input_height = static_cast(input.shape()[1]); - GGML_ASSERT(((input_width / output_width) == (input_height / output_height)) && - ((output_width / input_width) == (output_height / input_height))); GGML_ASSERT(((input_width / output_width) == scale) || ((output_width / input_width) == scale)); - int small_width = output_width; - int small_height = output_height; - bool decode = output_width > input_width; - if (decode) { - small_width = input_width; - small_height = input_height; - } - - int num_tiles_x; - float tile_overlap_factor_x; - sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x); + bool decode = output_width > input_width; // scale up + int small_width = decode ? input_width : output_width; + int small_height = decode ? input_height : output_height; + int scale_in = decode ? 1 : scale; + int scale_out = decode ? scale : 1; - int num_tiles_y; - float tile_overlap_factor_y; - sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y); + const auto tiles_x = sd_tiling_plan_axis(small_width, p_tile_size_w, tile_overlap_factor, circular_x); + const auto tiles_y = sd_tiling_plan_axis(small_height, p_tile_size_h, tile_overlap_factor, circular_y); - int tile_overlap_x = static_cast(p_tile_size_x * tile_overlap_factor_x); - int non_tile_overlap_x = p_tile_size_x - tile_overlap_x; - int tile_overlap_y = static_cast(p_tile_size_y * tile_overlap_factor_y); - int non_tile_overlap_y = p_tile_size_y - tile_overlap_y; - int tile_size_x = p_tile_size_x < small_width ? p_tile_size_x : small_width; - int tile_size_y = p_tile_size_y < small_height ? p_tile_size_y : small_height; - int input_tile_size_x = tile_size_x; - int input_tile_size_y = tile_size_y; - int output_tile_size_x = tile_size_x; - int output_tile_size_y = tile_size_y; - if (decode) { - output_tile_size_x *= scale; - output_tile_size_y *= scale; - } else { - input_tile_size_x *= scale; - input_tile_size_y *= scale; - } + int tile_width = std::min(p_tile_size_w, small_width); + int tile_height = std::min(p_tile_size_h, small_height); + int input_tile_width = tile_width * scale_in; + int input_tile_height = tile_height * scale_in; + int output_tile_width = tile_width * scale_out; + int output_tile_height = tile_height * scale_out; struct TilePlacement { - int x_in; - int y_in; - int x_out; - int y_out; - int dx; - int dy; + const TileSpan* x; + const TileSpan* y; }; std::vector placements; - bool last_y = false; - bool last_x = false; - for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) { - int dy = 0; - if (!circular_y && y + tile_size_y >= small_height) { - int original_y = y; - y = small_height - tile_size_y; - dy = original_y - y; - if (decode) { - dy *= scale; - } - last_y = true; + for (const auto& y : tiles_y) { + for (const auto& x : tiles_x) { + placements.push_back({&x, &y}); } - for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) { - int dx = 0; - if (!circular_x && x + tile_size_x >= small_width) { - int original_x = x; - x = small_width - tile_size_x; - dx = original_x - x; - if (decode) { - dx *= scale; - } - last_x = true; - } - placements.push_back({decode ? x : scale * x, - decode ? y : scale * y, - decode ? x * scale : x, - decode ? y * scale : y, - dx, - dy}); - } - last_x = false; } - int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x; - int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y; - int num_tiles = num_tiles_x * num_tiles_y; - int tile_count = 1; - float last_time = 0.0f; + const int num_tiles = static_cast(placements.size()); + int tile_count = 0; if (!silent) { - LOG_VERBOSE("num tiles : %d, %d ", num_tiles_x, num_tiles_y); - LOG_VERBOSE("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor); + LOG_VERBOSE("num tiles : %d, %d ", static_cast(tiles_x.size()), static_cast(tiles_y.size())); LOG_VERBOSE("processing %i tiles", num_tiles); pretty_progress(0, num_tiles, 0.0f); } @@ -392,7 +332,11 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, int64_t t1 = ggml_time_ms(); input_tiles.clear(); for (size_t i = first; i < last; ++i) { - input_tiles.push_back(sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, placements[i].x_in, placements[i].y_in)); + input_tiles.push_back(sd_tensor_split_2d(input, + input_tile_width, + input_tile_height, + placements[i].x->offset * scale_in, + placements[i].y->offset * scale_in)); } auto batch = on_processing(input_tiles); if (batch.data.empty()) { @@ -405,7 +349,7 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, } std::vector tile_shape = batch.data.shape(); tile_shape[batch.stack_dim] /= batch_tiles; - GGML_ASSERT(tile_shape[0] == output_tile_size_x && tile_shape[1] == output_tile_size_y); + GGML_ASSERT(tile_shape[0] == output_tile_width && tile_shape[1] == output_tile_height); auto merge_batch = [&, first, last, tile_shape](const sd::Tensor& data) { const int64_t tile_numel = data.numel() / static_cast(last - first); const int64_t plane_count = tile_numel / (tile_shape[0] * tile_shape[1]); @@ -416,20 +360,19 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, output = sd::Tensor::zeros(std::move(output_shape)); } for (size_t i = first; i < last; ++i) { - const auto& placement = placements[i]; + const auto& x = *placements[i].x; + const auto& y = *placements[i].y; sd_tensor_merge_2d(data.data() + static_cast(i - first) * tile_numel, tile_shape[0], tile_shape[1], plane_count, &output, - placement.x_out, - placement.y_out, - overlap_x_out, - overlap_y_out, - circular_x, - circular_y, - placement.dx, - placement.dy); + x.offset * scale_out, + y.offset * scale_out, + x.overlap_before * scale_out, + x.overlap_after * scale_out, + y.overlap_before * scale_out, + y.overlap_after * scale_out); } }; join_merger(); @@ -458,15 +401,12 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, merge_batch(batch.data); } - if (!silent) { - int64_t t2 = ggml_time_ms(); - last_time = (t2 - t1) / 1000.0f / static_cast(last - first); - for (size_t i = first; i < last; ++i) { - pretty_progress(tile_count, num_tiles, last_time); - tile_count++; + const float tile_time = (ggml_time_ms() - t1) / 1000.0f / static_cast(last - first); + for (size_t i = first; i < last; ++i) { + ++tile_count; + if (!silent) { + pretty_progress(tile_count, num_tiles, tile_time); } - } else { - tile_count += static_cast(last - first); } first = last; } @@ -474,11 +414,5 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, if (merge_failed) { return {}; } - if (!silent && tile_count < num_tiles) { - pretty_progress(num_tiles, num_tiles, last_time); - } - if (output.empty()) { - return {}; - } return output; } diff --git a/src/runtime/tiling.h b/src/runtime/tiling.h index 433ec08b93..2587dab283 100644 --- a/src/runtime/tiling.h +++ b/src/runtime/tiling.h @@ -21,8 +21,8 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y, @@ -35,8 +35,8 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y, diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index a8f83336d9..b13c055f3d 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -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) { *sd_ctx_params = {}; 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->rng_type = CUDA_RNG; 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" "tensor_type_rules: %s\n" "n_threads: %d\n" + "conditioning_cache_size: %d\n" "wtype: %s\n" "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->tensor_type_rules), sd_ctx_params->n_threads, + sd_ctx_params->conditioning_cache_size, sd_type_name(sd_ctx_params->wtype), sd_rng_type_name(sd_ctx_params->rng_type), sd_rng_type_name(sd_ctx_params->sampler_rng_type), @@ -630,14 +633,6 @@ struct sd_ctx_t { StableDiffusionGGML* sd = nullptr; }; -static bool sd_version_supports_video_generation(SDVersion version) { - return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); -} - -static bool sd_version_supports_image_generation(SDVersion version) { - return !sd_version_supports_video_generation(version); -} - sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t)); if (sd_ctx == nullptr) { @@ -759,28 +754,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, int* num_frames_out, sd_audio_t** audio_out, int* fps_out) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { - if (fps_out != nullptr) { - *fps_out = 0; - } - return false; - } - - if (frames_out != nullptr) { + if (frames_out != nullptr) *frames_out = nullptr; - } - if (audio_out != nullptr) { + if (audio_out != nullptr) *audio_out = nullptr; - } - if (num_frames_out != nullptr) { + if (num_frames_out != nullptr) *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); if (!execution.ready) { - if (fps_out != nullptr) { - *fps_out = 0; - } return false; } diff --git a/src/tokenizers/bpe_tokenizer.cpp b/src/tokenizers/bpe_tokenizer.cpp index 79b1029e01..ca86d65db1 100644 --- a/src/tokenizers/bpe_tokenizer.cpp +++ b/src/tokenizers/bpe_tokenizer.cpp @@ -242,7 +242,7 @@ bool BPETokenizer::encode(const std::string& text, std::vector& result, on_ ss << "\"" << token << "\", "; } ss << "]"; - LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str()); + LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", escape_newlines(text).c_str(), bpe_tokens.size(), escape_newlines(ss.str()).c_str()); result = std::move(bpe_tokens); return true; } diff --git a/src/tokenizers/hf_tokenizer.cpp b/src/tokenizers/hf_tokenizer.cpp index 821d41d797..f8daf6d2a0 100644 --- a/src/tokenizers/hf_tokenizer.cpp +++ b/src/tokenizers/hf_tokenizer.cpp @@ -735,7 +735,7 @@ bool HFTokenizer::encode(const std::string& text, std::vector& tokens, on_n } } ss << "]"; - LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), result.size(), ss.str().c_str()); + LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", escape_newlines(text).c_str(), result.size(), escape_newlines(ss.str()).c_str()); tokens = std::move(result); return true; } diff --git a/src/tokenizers/t5_unigram_tokenizer.cpp b/src/tokenizers/t5_unigram_tokenizer.cpp index 36526086e9..342dca1a2d 100644 --- a/src/tokenizers/t5_unigram_tokenizer.cpp +++ b/src/tokenizers/t5_unigram_tokenizer.cpp @@ -295,7 +295,13 @@ bool T5UniGramTokenizer::encode(const std::string& input, std::vector& resu std::vector tokens; std::vector token_strs; std::string normalized = normalize(input); - auto splited_texts = split_with_special_tokens(normalized, special_tokens); + if (normalized.empty()) { + // HF reference tokenizers emit no pieces for empty input; pad_tokens + // still appends EOS so the sequence becomes [EOS] + padding. + result = std::move(tokens); + return true; + } + auto splited_texts = split_with_special_tokens(normalized, special_tokens); if (splited_texts.empty()) { splited_texts.push_back(normalized); // for empty string } @@ -337,7 +343,7 @@ bool T5UniGramTokenizer::encode(const std::string& input, std::vector& resu ss << "\"" << token_str << "\", "; } ss << "]"; - LOG_VERBOSE("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str()); + LOG_VERBOSE("split prompt \"%s\" to tokens %s", escape_newlines(input).c_str(), escape_newlines(ss.str()).c_str()); result = std::move(tokens); return true; diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 341c76ecfc..dc3498e646 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -111,10 +111,22 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_tensor) { 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 rgb; + if (has_alpha) { + rgb = sd::ops::slice(input_tensor, 2, 0, 3); + } + const sd::Tensor& model_input = has_alpha ? rgb : input_tensor; sd::Tensor upscaled; const int scale = esrgan_upscaler->config.scale; 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 { auto on_processing = [&](const sd::Tensor& input_tile) -> sd::Tensor { auto output_tile = esrgan_upscaler->compute(n_threads, input_tile); @@ -125,7 +137,7 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te return output_tile; }; - upscaled = process_tiles_2d(input_tensor, + upscaled = process_tiles_2d(model_input, static_cast(input_tensor.shape()[0] * scale), static_cast(input_tensor.shape()[1] * scale), scale, @@ -141,6 +153,14 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te LOG_ERROR("esrgan compute failed"); 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; } @@ -242,6 +262,30 @@ int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) { 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) { if (upscaler_ctx->upscaler != nullptr) { delete upscaler_ctx->upscaler;