diff --git a/docs/performance.md b/docs/performance.md index ccfe4778a..e17616b53 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -73,6 +73,10 @@ sd-cli --diffusion-model flux1-dev.safetensors ... \ Ordered from fastest to smallest-VRAM: no flags → `--offload-to-cpu` → `--offload-to-cpu --max-vram ` → `--offload-to-cpu --max-vram --stream-layers`. Each step down costs a few percent of throughput to buy more room; combined they can run models roughly 3-4x larger than the raw VRAM would allow. +## Use SPEED to accelerate flow-model sampling. + +`--sampling-method speed_flow` runs the early denoising steps at a reduced spatial resolution and progressively expands to full resolution as the trajectory approaches the data manifold. Compatible with `--offload-to-cpu`, `--max-vram`, `--stream-layers`, and the cache-mode flags. Typical sampling-time speedups on flow models: ~1.6× on Flux, ~1.7× on Qwen Image, up to ~3.5× on Z-Image with a tuned transition point. Composition drifts slightly from the euler baseline (different trajectory), quality preserved for txt2img; edit / ref-guided models see reduced source fidelity — see [speed_sampler.md](./speed_sampler.md) for tuning knobs and compatibility notes. + ## Use quantization to reduce memory usage. [quantization](./quantization_and_gguf.md) diff --git a/docs/speed_sampler.md b/docs/speed_sampler.md new file mode 100644 index 000000000..eacbae496 --- /dev/null +++ b/docs/speed_sampler.md @@ -0,0 +1,84 @@ +# SPEED — Spectral Progressive Diffusion sampler + +SPEED (`--sampling-method speed_flow`) accelerates flow-model sampling by running the early denoising steps at a **reduced spatial resolution**, then progressively expanding to full resolution as the trajectory approaches the data manifold. Reference: [Xiao et al., arXiv:2605.18736](https://arxiv.org/abs/2605.18736); reference C++/Python: [howardhx/speed](https://github.com/howardhx/speed), [ruwwww/ComfyUI-SPEED](https://github.com/ruwwww/ComfyUI-SPEED). + +Between segments, the low-resolution latent is expanded via a spectral basis (DCT-II by default) — the low-frequency coefficients are preserved and the newly-exposed high-frequency bands are filled with Gaussian noise scaled by the current sigma. The flow-matching time at each transition is rescaled by `kappa` (paper Eq. 5-6) so the base sampler continues on a consistent schedule after resize. + +## Quick usage + +```sh +# Flux dev / schnell — default configuration +./bin/sd-cli --sampling-method speed_flow ... + +# Explicit scales; use ':' as the intra-list separator +./bin/sd-cli --sampling-method speed_flow \ + --extra-sample-args "speed_scales=0.5:1.0" + +# Manual transition point (overrides delta-optimal) +./bin/sd-cli --sampling-method speed_flow \ + --extra-sample-args "speed_scales=0.5:1.0,speed_manual_sigmas=0.7" +``` + +## Extra sampler args + +All parsed out of `--extra-sample-args`. The list is comma-separated at the top level; within a single value use `:` (or `/`, `|`) as the intra-value separator. + +| Arg | Default | Description | +|---|---|---| +| `speed_scales` | `0.5:1.0` | Colon-separated resolution fractions, strictly increasing, ending at `1.0`. E.g. `0.25:0.5:1.0` for a 3-level schedule. | +| `speed_levels` | *(unset)* | Shortcut: `speed_levels=2` → `{0.5, 1.0}`; `=3` → `{0.25, 0.5, 1.0}`. Ignored if `speed_scales` is set. | +| `speed_delta` | `0.01` | Noise-dominated tolerance (paper Eq. 9). Smaller values push transitions later in the trajectory. | +| `speed_manual_sigmas` | *(unset)* | Colon-separated sigma thresholds, one per transition (length `= scales.size() - 1`). Bypasses `speed_delta` and the power-spectrum formula. Values must be strictly decreasing and each in `(0, 1)`. | +| `speed_preset` | *auto* | Power-spectrum preset. Auto-selected from the loaded model: `wan21` (`A=219.485, β=2.423`) for WAN, `flux` (`A=203.615, β=1.915`) for every other flow model. Values fitted from the respective VAEs in the reference impl's `configs.yaml`. Set explicitly to override. | +| `speed_A`, `speed_beta` | *(preset)* | Power-spectrum override for custom models. Overrides `speed_preset`. | +| `speed_transform` | `dct` | Spectral basis for the resolution expansion. `dct` handles any ratio; `fft` also handles any ratio, provided both spatial dims are powers of 2. | +| `speed_seed` | `0` | Seed for the Gaussian noise placed into the new high-frequency bands at each transition. | + +## Typical speedups + +Measured locally on an RTX 3060 12 GiB. Numbers are sampling-time speedups over `--sampling-method euler` at the same seed and step count. + +| Model | Resolution | Steps | Config | Speedup | +|---|---|---|---|---| +| Flux schnell Q4_K | 1024×1024 | 8 | `speed_scales=0.5:1.0` (delta-optimal) | ~1.6× | +| Qwen Image Q2_K | 1024×1024 | 20 | `speed_scales=0.5:1.0` (delta-optimal) | ~1.7× | +| Z-Image turbo Q8 | 1024×688 | 8 | `speed_scales=0.5:1.0,speed_manual_sigmas=0.7` | ~3.5× | + +Compute savings scale with the number of low-resolution steps in the schedule. Models with front-loaded (shift-heavy) sigma schedules — Flux schnell, Z-Image turbo — often benefit from `speed_manual_sigmas` to push the transition into the middle of the trajectory rather than the very end. + +## Compatibility + +- `--offload-to-cpu`, `--max-vram`, `--stream-layers`: fully compatible. The graph-cut planner builds a second cached plan for the full-resolution shape at the transition (~5 ms overhead). +- `--cache-mode easycache` (and other caches): compatible. The cache runtime tolerates the resolution transition; the cache heuristic skips only steps where the shape context is stable. +- Non-flow denoisers (SD 1.5, SDXL): SPEED targets flow models specifically. If dispatched on an ε-prediction model it will run but the kappa/sigma rescale math is not valid — do not use. + +## Reference-latent models (Qwen Image Edit, Flux Kontext, etc.) + +**SPEED runs on ref-latent models but reduces fidelity to the source image.** A warning is logged when this combination is detected: + +``` +[WARN ] SPEED sampler with reference latents (edit / ref-guided model): the low-resolution + segments see a downscaled reference so fine detail from the source image (backgrounds, + textures, layout beyond broad structure) is likely to drift. +``` + +Why: the reference latents are the low-frequency structure the edit is supposed to preserve. During SPEED's low-resolution segments the reference is downscaled to match the working latent, so the model conditions on a coarser version of the reference — it retains the semantic subject and the requested edit direction, but loses fine detail (backgrounds, textures, exact framing). Composition remains plausible but is not faithful to the source. + +Practical guidance: + +- For **creative / style edits** where subject and edit intent matter more than reproducing the source scene: SPEED is fine. +- For **fidelity edits** (localized inpaint-style changes, preserving specific backgrounds or text): use `euler` or another full-resolution sampler. +- To trade off, use `speed_manual_sigmas` with a threshold in the middle of the trajectory so only 2–3 early steps run at low resolution. + +## Known limits (v1) + +- Naive O(N²) 2D DCT — ~1.8 s per direction on a 128×86×16 latent. Fine for image-sized latents; video and very large latents would need a proper FFT-based DCT. +- Only image latents (4-D `[W, H, C, N]`). Video (5-D) is not handled. +- Multi-resolution triggers 2 ggml graph builds per generation, cached thereafter. +- Inpainting masks with `!=` full-resolution shape are silently skipped at low-res steps (the mask blend guard checks the shape and falls through when it can't broadcast). + +## References + +- Paper: Xiao et al., *"SPEED: Spectral Progressive Diffusion for training-free acceleration"*, [arXiv:2605.18736](https://arxiv.org/abs/2605.18736) +- Official code: +- ComfyUI implementation: diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 1dfb6aa70..5d649bd3e 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -967,7 +967,7 @@ ArgOptions SDGenerationParams::get_options() { &hires_upscaler}, {"", "--extra-sample-args", - "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware", + "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; speed_flow supports speed_scales (use ':' between scales, e.g. speed_scales=0.25:0.5:1.0), speed_levels (2 or 3), speed_delta, speed_manual_sigmas (':'-separated, one per transition, overrides delta), speed_preset (flux|wan21), speed_transform (dct|fft), speed_A, speed_beta, speed_seed;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware", (int)',', &extra_sample_args}, {"", @@ -1488,12 +1488,12 @@ ArgOptions SDGenerationParams::get_options() { on_seed_arg}, {"", "--sampling-method", - "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]" + "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, speed_flow, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]" "(default: euler for Flux/SD3/Wan, euler_a otherwise)", on_sample_method_arg}, {"", "--high-noise-sampling-method", - "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]" + "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, speed_flow, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]" " default: euler for Flux/SD3/Wan, euler_a otherwise", on_high_noise_sample_method_arg}, {"", diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 93084fe29..840214590 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -56,6 +56,7 @@ enum sample_method_t { EULER_GE_SAMPLE_METHOD, DPMPP2M_SDE_SAMPLE_METHOD, DPMPP2M_SDE_BT_SAMPLE_METHOD, + SPEED_FLOW_SAMPLE_METHOD, SAMPLE_METHOD_COUNT }; diff --git a/src/core/spectral_ops.hpp b/src/core/spectral_ops.hpp new file mode 100644 index 000000000..23e724b82 --- /dev/null +++ b/src/core/spectral_ops.hpp @@ -0,0 +1,278 @@ +#ifndef __SD_CORE_SPECTRAL_OPS_HPP__ +#define __SD_CORE_SPECTRAL_OPS_HPP__ + +#include +#include +#include +#include +#include +#include + +#include "core/tensor.hpp" + +// CPU-side FFT / DCT / DWT primitives for use by spectral samplers (SPEED etc.). +// Operates on sd::Tensor laid out ggml-style [W, H, C, N] (W innermost). +// The 2D transforms run over the trailing two dims (W and H); leading dims are +// iterated as independent batches. +// +// Radix-2 only for FFT; DCT-II built on top via the mirror trick. Sizes must +// be powers of 2 — the SPEED transitions the paper uses (0.25/0.5/1.0 fractions +// of a power-of-two latent) always land on power-of-two spatial dims. + +namespace sd { +namespace spectral { + +using cplx = std::complex; + +// Locally defined so the header stays portable to compilers that don't +// expose M_PI in without _USE_MATH_DEFINES (MSVC, some clangs). +constexpr float k_pi = 3.14159265358979323846f; + +inline bool is_pow2(int64_t n) { + return n > 0 && (n & (n - 1)) == 0; +} + +// In-place radix-2 Cooley-Tukey FFT of length N (must be a power of 2). +// sign = -1: forward transform; sign = +1: inverse transform (unscaled). +inline void fft_1d(cplx* a, int64_t n, int sign) { + if (!is_pow2(n)) { + throw std::runtime_error("spectral::fft_1d: length " + std::to_string(n) + " is not a power of 2"); + } + for (int64_t i = 1, j = 0; i < n; ++i) { + int64_t bit = n >> 1; + for (; j & bit; bit >>= 1) { + j ^= bit; + } + j ^= bit; + if (i < j) { + std::swap(a[i], a[j]); + } + } + for (int64_t len = 2; len <= n; len <<= 1) { + float ang = sign * 2.0f * k_pi / static_cast(len); + cplx wlen(std::cos(ang), std::sin(ang)); + for (int64_t i = 0; i < n; i += len) { + cplx w(1.0f, 0.0f); + for (int64_t k = 0; k < len / 2; ++k) { + cplx u = a[i + k]; + cplx v = a[i + k + len / 2] * w; + a[i + k] = u + v; + a[i + k + len / 2] = u - v; + w *= wlen; + } + } + } +} + +// 2D FFT over the trailing two dims [..., H, W] of a real 4D tensor in ggml +// layout [W, H, C, N]. Returns a complex-valued grid of the same [W, H] +// dimensions for each (C, N). Ortho normalization (divides by sqrt(H*W)). +inline std::vector fft_2d(const sd::Tensor& x) { + const auto& s = x.shape(); + if (s.size() != 4) { + throw std::runtime_error("spectral::fft_2d: expected 4D tensor [W,H,C,N]"); + } + int64_t W = s[0]; + int64_t H = s[1]; + int64_t C = s[2]; + int64_t N = s[3]; + int64_t HW = H * W; + if (!is_pow2(H) || !is_pow2(W)) { + throw std::runtime_error("spectral::fft_2d: spatial dims must be powers of 2, got W=" + + std::to_string(W) + " H=" + std::to_string(H)); + } + std::vector out(static_cast(N * C * H * W)); + float norm = 1.0f / std::sqrt(static_cast(HW)); + const float* xp = x.data(); + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + cplx* plane = out.data() + (n * C + c) * HW; + for (int64_t i = 0; i < HW; ++i) { + plane[i] = cplx(xp[(n * C + c) * HW + i], 0.0f); + } + std::vector row(W); + for (int64_t h = 0; h < H; ++h) { + for (int64_t w = 0; w < W; ++w) { + row[w] = plane[h * W + w]; + } + fft_1d(row.data(), W, -1); + for (int64_t w = 0; w < W; ++w) { + plane[h * W + w] = row[w]; + } + } + std::vector col(H); + for (int64_t w = 0; w < W; ++w) { + for (int64_t h = 0; h < H; ++h) { + col[h] = plane[h * W + w]; + } + fft_1d(col.data(), H, -1); + for (int64_t h = 0; h < H; ++h) { + plane[h * W + w] = col[h]; + } + } + for (int64_t i = 0; i < HW; ++i) { + plane[i] *= norm; + } + } + } + return out; +} + +// Inverse of fft_2d: takes a complex [N, C, H, W] grid and returns the real +// tensor of the given shape. Ortho normalization. +inline sd::Tensor ifft_2d(const std::vector& in, const std::vector& shape) { + if (shape.size() != 4) { + throw std::runtime_error("spectral::ifft_2d: expected 4D output shape"); + } + int64_t W = shape[0]; + int64_t H = shape[1]; + int64_t C = shape[2]; + int64_t N = shape[3]; + int64_t HW = H * W; + if (!is_pow2(H) || !is_pow2(W)) { + throw std::runtime_error("spectral::ifft_2d: spatial dims must be powers of 2"); + } + if (static_cast(in.size()) != N * C * HW) { + throw std::runtime_error("spectral::ifft_2d: input buffer size does not match shape"); + } + sd::Tensor out(shape); + float norm = 1.0f / std::sqrt(static_cast(HW)); + float* op = out.data(); + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + std::vector plane(HW); + for (int64_t i = 0; i < HW; ++i) { + plane[i] = in[(n * C + c) * HW + i]; + } + std::vector row(W); + for (int64_t h = 0; h < H; ++h) { + for (int64_t w = 0; w < W; ++w) { + row[w] = plane[h * W + w]; + } + fft_1d(row.data(), W, +1); + for (int64_t w = 0; w < W; ++w) { + plane[h * W + w] = row[w] / static_cast(W); + } + } + std::vector col(H); + for (int64_t w = 0; w < W; ++w) { + for (int64_t h = 0; h < H; ++h) { + col[h] = plane[h * W + w]; + } + fft_1d(col.data(), H, +1); + for (int64_t h = 0; h < H; ++h) { + plane[h * W + w] = col[h] / static_cast(H); + } + } + for (int64_t i = 0; i < HW; ++i) { + op[(n * C + c) * HW + i] = plane[i].real() / norm; + } + } + } + return out; +} + +// 1D orthonormal DCT-II by direct summation (O(N^2)). +// X[k] = f_k * sum_n x[n] * cos(pi * (2n+1) * k / (2N)), +// where f_0 = sqrt(1/N) and f_k = sqrt(2/N) for k > 0. +inline void dct1d_direct(const float* in, float* out_row, int64_t N) { + float scale0 = 1.0f / std::sqrt(static_cast(N)); + float scale = std::sqrt(2.0f / static_cast(N)); + float inv_2N = 1.0f / static_cast(2 * N); + for (int64_t k = 0; k < N; ++k) { + float sum = 0.0f; + for (int64_t n = 0; n < N; ++n) { + sum += in[n] * std::cos(k_pi * (2 * n + 1) * k * inv_2N); + } + out_row[k] = sum * ((k == 0) ? scale0 : scale); + } +} + +// 1D orthonormal DCT-III (inverse of DCT-II by direct summation). +inline void idct1d_direct(const float* in, float* out_row, int64_t N) { + float scale0 = 1.0f / std::sqrt(static_cast(N)); + float scale = std::sqrt(2.0f / static_cast(N)); + float inv_2N = 1.0f / static_cast(2 * N); + for (int64_t n = 0; n < N; ++n) { + float sum = in[0] * scale0; + for (int64_t k = 1; k < N; ++k) { + sum += in[k] * scale * std::cos(k_pi * (2 * n + 1) * k * inv_2N); + } + out_row[n] = sum; + } +} + +// 2D orthonormal DCT-II. Row-column decomposition using dct1d_direct. +inline sd::Tensor dct2_2d(const sd::Tensor& x) { + const auto& s = x.shape(); + if (s.size() != 4) { + throw std::runtime_error("spectral::dct2_2d: expected 4D tensor"); + } + int64_t W = s[0]; + int64_t H = s[1]; + int64_t C = s[2]; + int64_t N = s[3]; + int64_t HW = H * W; + sd::Tensor out(s); + const float* xp = x.data(); + float* op = out.data(); + std::vector col_in(H), col_out(H); + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + float* plane = op + (n * C + c) * HW; + for (int64_t h = 0; h < H; ++h) { + dct1d_direct(xp + (n * C + c) * HW + h * W, plane + h * W, W); + } + for (int64_t w = 0; w < W; ++w) { + for (int64_t h = 0; h < H; ++h) { + col_in[h] = plane[h * W + w]; + } + dct1d_direct(col_in.data(), col_out.data(), H); + for (int64_t h = 0; h < H; ++h) { + plane[h * W + w] = col_out[h]; + } + } + } + } + return out; +} + +// 2D orthonormal inverse DCT (DCT-III). +inline sd::Tensor idct2_2d(const sd::Tensor& c_in) { + const auto& s = c_in.shape(); + if (s.size() != 4) { + throw std::runtime_error("spectral::idct2_2d: expected 4D tensor"); + } + int64_t W = s[0]; + int64_t H = s[1]; + int64_t C = s[2]; + int64_t N = s[3]; + int64_t HW = H * W; + sd::Tensor out(s); + const float* cp = c_in.data(); + float* op = out.data(); + std::vector col_in(H), col_out(H); + for (int64_t n = 0; n < N; ++n) { + for (int64_t ch = 0; ch < C; ++ch) { + float* plane = op + (n * C + ch) * HW; + for (int64_t h = 0; h < H; ++h) { + idct1d_direct(cp + (n * C + ch) * HW + h * W, plane + h * W, W); + } + for (int64_t w = 0; w < W; ++w) { + for (int64_t h = 0; h < H; ++h) { + col_in[h] = plane[h * W + w]; + } + idct1d_direct(col_in.data(), col_out.data(), H); + for (int64_t h = 0; h < H; ++h) { + plane[h * W + w] = col_out[h]; + } + } + } + } + return out; +} + +} // namespace spectral +} // namespace sd + +#endif // __SD_CORE_SPECTRAL_OPS_HPP__ diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 115175b14..c096c4882 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -15,6 +15,7 @@ #include "core/tensor.hpp" #include "runtime/gits_noise.h" #include "runtime/guidance.h" +#include "runtime/speed_sampler.hpp" /*================================================= CompVisDenoiser ==================================================*/ @@ -2668,7 +2669,8 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, float eta, bool is_flow_denoiser, const char* extra_sample_args, - std::shared_ptr denoiser_for_dispatch = nullptr) { + std::shared_ptr denoiser_for_dispatch = nullptr, + SDVersion version = VERSION_COUNT) { if (denoiser_for_dispatch) { if (auto sefi = std::dynamic_pointer_cast(denoiser_for_dispatch)) { return sample_sefi_euler(sefi.get(), model, std::move(x)); @@ -2720,6 +2722,11 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, rng, eta); case EULER_GE_SAMPLE_METHOD: return sample_gradient_estimation(model, std::move(x), sigmas, rng, is_flow_denoiser, eta, extra_args); + case SPEED_FLOW_SAMPLE_METHOD: { + sd::speed::Config cfg = sd::speed::default_config_for(version); + sd::speed::parse_config_from_args(cfg, extra_args); + return sd::speed::sample_speed_flow(model, std::move(x), sigmas, cfg); + } default: return {}; } diff --git a/src/runtime/speed_sampler.hpp b/src/runtime/speed_sampler.hpp new file mode 100644 index 000000000..765b1830b --- /dev/null +++ b/src/runtime/speed_sampler.hpp @@ -0,0 +1,456 @@ +#ifndef __SD_RUNTIME_SPEED_SAMPLER_HPP__ +#define __SD_RUNTIME_SPEED_SAMPLER_HPP__ + +#include +#include +#include +#include +#include +#include +#include + +#include "core/rng.hpp" +#include "core/spectral_ops.hpp" +#include "core/tensor.hpp" + +// SPEED — Spectral Progressive Diffusion sampler. +// Ref: Xiao et al. arXiv:2605.18736; https://github.com/howardhx/speed; +// ComfyUI-SPEED (https://github.com/ruwwww/ComfyUI-SPEED). +// +// Runs the denoising trajectory of a flow model at progressively higher +// spatial resolutions, saving the majority of the DiT compute during the +// low-frequency-dominated early steps. Between resolution levels, the latent +// is expanded via a spectral basis (DCT or FFT) — high-frequency bands are +// filled with fresh Gaussian noise scaled by the current sigma. The +// flow-matching time at each transition is rescaled by kappa (paper Eq. 5-6) +// so the base sampler operates on a consistent schedule after resize. +// +// Kept lean in this initial implementation: image latents only, DCT + FFT +// transforms, base sampler is deterministic euler, delta-optimal transitions +// only. DWT / manual sigmas / video latents can be added incrementally. + +namespace sd { +namespace speed { + +enum class Transform { + DCT, + FFT, +}; + +struct Config { + std::vector scales = {0.5f, 1.0f}; // strictly increasing, ends at 1.0 + float delta = 0.01f; // noise-dominated tolerance (paper Eq. 9) + float spectrum_A = 203.615097f; // Flux dev preset + float spectrum_beta = 1.915461f; + Transform transform = Transform::DCT; + uint64_t seed = 0; // for the spectral noise padding + // If non-empty, overrides the delta-optimal schedule: use these sigma + // thresholds directly, one per transition (length == scales.size() - 1). + // Values must be strictly decreasing and each in (0, 1). + std::vector manual_sigmas; +}; + +// Power-spectrum presets fitted per VAE in the reference impl's configs.yaml. +struct Preset { + float A; + float beta; +}; +inline Preset preset_flux() { return {203.615097f, 1.915461f}; } +inline Preset preset_wan21() { return {219.484718f, 2.422687f}; } + +// Auto-select a preset based on which VAE the model uses. All Flux-family +// image models (Flux, Flux2, SeFi, Z-Image, Qwen Image, ...) share the Flux +// VAE calibration in the reference impl; only WAN gets a distinct preset. +inline Preset preset_for(SDVersion version) { + if (sd_version_is_wan(version)) { + return preset_wan21(); + } + return preset_flux(); +} + +inline Config default_config_for(SDVersion version) { + Config cfg; + Preset p = preset_for(version); + cfg.spectrum_A = p.A; + cfg.spectrum_beta = p.beta; + return cfg; +} + +// Power-law radial spectrum P(omega) = A * |omega|^(-beta) (Eq. 8). +inline float power_spectrum(float omega, float A, float beta) { + return A * std::pow(std::abs(omega), -beta); +} + +// Activation time for one radial frequency (Eq. 9). +inline float activation_time(float P_omega, float delta) { + return 1.0f / (1.0f + std::sqrt(delta / (P_omega * (1.0f + P_omega - delta)))); +} + +// Rescale factor at a transition (paper Eq. 5). +inline float kappa(float t, float r) { + return r / (1.0f + (r - 1.0f) * t); +} + +// Aligned flow-matching time after spectral noise expansion (Eq. 6). +inline float align_timestep(float t, float r) { + return t * kappa(t, r); +} + +// Delta-optimal transition times between consecutive scales (Eq. 10). +inline std::vector delta_optimal_transitions(const std::vector& scales, + float delta, + float A, + float beta, + int64_t H, + int64_t W) { + float omega_max = 0.5f * static_cast(std::min(H, W)); + std::vector out; + out.reserve(scales.size() > 0 ? scales.size() - 1 : 0); + for (size_t i = 0; i + 1 < scales.size(); ++i) { + float omega_i = scales[i] * omega_max; + out.push_back(activation_time(power_spectrum(omega_i, A, beta), delta)); + } + return out; +} + +struct Transition { + int step_idx; + float s_from; + float s_to; +}; + +// Resolve which sampler step index each threshold in `t_stars` first crosses. +inline std::vector resolve_transitions(const std::vector& sigmas, + const std::vector& scales, + const std::vector& t_stars) { + std::vector out; + if (scales.size() < 2) { + return out; + } + int n_steps = static_cast(sigmas.size()) - 1; + for (size_t i = 0; i + 1 < scales.size(); ++i) { + float thr = t_stars[i]; + int step_idx = n_steps; + for (int j = 0; j < n_steps; ++j) { + if (sigmas[j] <= thr) { + step_idx = j; + break; + } + } + if (step_idx >= n_steps) { + break; + } + out.push_back({step_idx, scales[i], scales[i + 1]}); + } + return out; +} + +// Round a scale fraction to an even integer size so the DiT patchify stays valid. +inline int64_t scaled_dim(float s, int64_t full) { + int64_t v = static_cast(std::llround(static_cast(s) * static_cast(full))); + if (v < 2) { + v = 2; + } + if ((v & 1) != 0) { + v++; // keep even so patch_size=2 grids align + } + return v; +} + +// DCT-truncation downscale of a 4D [W, H, C, N] tensor to (H_lo, W_lo). +// This is the "initial coarse latent" step of SPEED. +inline sd::Tensor dct_downscale(const sd::Tensor& x, int64_t H_lo, int64_t W_lo) { + const auto& s = x.shape(); + if (s.size() != 4) { + throw std::runtime_error("speed::dct_downscale: expected 4D tensor"); + } + int64_t W_hi = s[0]; + int64_t H_hi = s[1]; + int64_t C = s[2]; + int64_t N = s[3]; + if (H_lo >= H_hi && W_lo >= W_hi) { + return x; + } + sd::Tensor coeffs = sd::spectral::dct2_2d(x); + sd::Tensor trunc({W_lo, H_lo, C, N}); + const float* cp = coeffs.data(); + float* tp = trunc.data(); + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + for (int64_t h = 0; h < H_lo; ++h) { + for (int64_t w = 0; w < W_lo; ++w) { + tp[((n * C + c) * H_lo + h) * W_lo + w] = + cp[((n * C + c) * H_hi + h) * W_hi + w]; + } + } + } + } + return sd::spectral::idct2_2d(trunc); +} + +// Spectral noise expansion (DCT variant): promote a low-res latent to a +// higher-res one by placing its DCT coefficients in the low-frequency block of +// a larger coefficient plane and filling the rest with scaled Gaussian noise. +inline sd::Tensor dct_expand(const sd::Tensor& x, + int64_t H_tgt, + int64_t W_tgt, + float t_noise, + uint64_t seed) { + const auto& s = x.shape(); + if (s.size() != 4) { + throw std::runtime_error("speed::dct_expand: expected 4D tensor"); + } + int64_t W_src = s[0]; + int64_t H_src = s[1]; + int64_t C = s[2]; + int64_t N = s[3]; + if (H_tgt < H_src || W_tgt < W_src) { + throw std::runtime_error("speed::dct_expand: target smaller than source"); + } + sd::Tensor src_coeffs = sd::spectral::dct2_2d(x); + sd::Tensor big({W_tgt, H_tgt, C, N}); + auto rng = std::make_shared(); + rng->manual_seed(seed); + auto noise_flat = rng->randn(static_cast(big.numel())); + std::memcpy(big.data(), noise_flat.data(), noise_flat.size() * sizeof(float)); + for (int64_t i = 0; i < big.numel(); ++i) { + big.data()[i] *= t_noise; + } + // Overwrite the low-freq block with the source coefficients. + const float* sp = src_coeffs.data(); + float* bp = big.data(); + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + for (int64_t h = 0; h < H_src; ++h) { + for (int64_t w = 0; w < W_src; ++w) { + bp[((n * C + c) * H_tgt + h) * W_tgt + w] = + sp[((n * C + c) * H_src + h) * W_src + w]; + } + } + } + } + return sd::spectral::idct2_2d(big); +} + +// FFT variant of the noise expansion: shifted-spectrum padding with the source +// spectrum centered in the higher-frequency plane. Matches the reference +// _fft_expand_np. Currently unused by the default DCT transform path but kept +// for completeness so we can switch via config. +inline sd::Tensor fft_expand(const sd::Tensor& x, + int64_t H_tgt, + int64_t W_tgt, + float t_noise, + uint64_t seed) { + const auto& s = x.shape(); + int64_t W_src = s[0]; + int64_t H_src = s[1]; + int64_t C = s[2]; + int64_t N = s[3]; + auto src_fft = sd::spectral::fft_2d(x); + std::vector big(static_cast(N * C * H_tgt * W_tgt)); + auto rng = std::make_shared(); + rng->manual_seed(seed); + auto noise_r = rng->randn(static_cast(N * C * H_tgt * W_tgt)); + auto noise_i = rng->randn(static_cast(N * C * H_tgt * W_tgt)); + float inv_sqrt2 = 1.0f / std::sqrt(2.0f); + for (size_t i = 0; i < big.size(); ++i) { + big[i] = sd::spectral::cplx(t_noise * noise_r[i] * inv_sqrt2, + t_noise * noise_i[i] * inv_sqrt2); + } + int64_t pad_h = (H_tgt - H_src) / 2; + int64_t pad_w = (W_tgt - W_src) / 2; + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + for (int64_t h = 0; h < H_src; ++h) { + for (int64_t w = 0; w < W_src; ++w) { + big[((n * C + c) * H_tgt + pad_h + h) * W_tgt + pad_w + w] = + src_fft[((n * C + c) * H_src + h) * W_src + w]; + } + } + } + } + return sd::spectral::ifft_2d(big, {W_tgt, H_tgt, C, N}); +} + +// Signature matches the other flow samplers: (model, x, sigmas, ...). +// The extra_sample_args pointer is parsed for speed_* overrides. +template +inline sd::Tensor sample_speed_flow(ModelCb model, + sd::Tensor x, + std::vector sigmas, + const Config& cfg) { + const auto& s = x.shape(); + if (s.size() != 4) { + throw std::runtime_error("speed::sample_speed_flow: expected 4D latent [W,H,C,N]"); + } + int64_t W_full = s[0]; + int64_t H_full = s[1]; + if (cfg.scales.size() < 2) { + // Nothing to do; run plain euler over the full schedule. + int steps = static_cast(sigmas.size()) - 1; + for (int i = 0; i < steps; ++i) { + auto denoised_opt = model(x, sigmas[i], i + 1); + if (denoised_opt.pred.empty()) { + return {}; + } + sd::Tensor denoised = std::move(denoised_opt.pred); + if (sigmas[i + 1] == 0.f) { + x = denoised; + } else { + float r = sigmas[i + 1] / sigmas[i]; + x = r * x + (1.0f - r) * denoised; + } + } + return x; + } + + std::vector t_stars; + if (!cfg.manual_sigmas.empty()) { + if (cfg.manual_sigmas.size() != cfg.scales.size() - 1) { + throw std::runtime_error("speed::sample_speed_flow: speed_manual_sigmas length must be scales.size() - 1"); + } + t_stars = cfg.manual_sigmas; + } else { + t_stars = delta_optimal_transitions(cfg.scales, cfg.delta, cfg.spectrum_A, + cfg.spectrum_beta, H_full, W_full); + } + auto transitions = resolve_transitions(sigmas, cfg.scales, t_stars); + // Only downscale if there's at least one transition that will bring the + // latent back up; otherwise the loop would return a low-res tensor to the + // VAE and crash. If no transitions fire, run at full resolution. + if (!transitions.empty() && cfg.scales.front() < 1.0f) { + int64_t H_lo = scaled_dim(cfg.scales.front(), H_full); + int64_t W_lo = scaled_dim(cfg.scales.front(), W_full); + x = dct_downscale(x, H_lo, W_lo); + } + int total_steps = static_cast(sigmas.size()) - 1; + int seg_start = 0; + for (size_t seg = 0; seg <= transitions.size(); ++seg) { + int seg_end = (seg < transitions.size()) ? transitions[seg].step_idx : total_steps; + for (int i = seg_start; i < seg_end; ++i) { + auto denoised_opt = model(x, sigmas[i], i + 1); + if (denoised_opt.pred.empty()) { + return {}; + } + sd::Tensor denoised = std::move(denoised_opt.pred); + if (sigmas[i + 1] == 0.f) { + x = denoised; + } else { + float r = sigmas[i + 1] / sigmas[i]; + x = r * x + (1.0f - r) * denoised; + } + } + if (seg == transitions.size()) { + break; + } + const Transition& tr = transitions[seg]; + int64_t H_tgt = scaled_dim(tr.s_to, H_full); + int64_t W_tgt = scaled_dim(tr.s_to, W_full); + float r = tr.s_to / tr.s_from; + float t_at = sigmas[tr.step_idx]; + uint64_t seed = cfg.seed + static_cast(seg + 1) * 10000ULL; + if (cfg.transform == Transform::FFT) { + x = fft_expand(x, H_tgt, W_tgt, t_at, seed); + } else { + x = dct_expand(x, H_tgt, W_tgt, t_at, seed); + } + float k = kappa(t_at, r); + for (int64_t i = 0; i < x.numel(); ++i) { + x.data()[i] *= k; + } + // Patch the transition sigma so the base solver continues on the aligned schedule. + sigmas[tr.step_idx] = align_timestep(t_at, r); + seg_start = tr.step_idx; + } + return x; +} + +// Parse "speed_scales=0.5,1.0", "speed_delta=0.01", "speed_transform=dct|fft", +// "speed_A=...", "speed_beta=...", "speed_seed=..." out of extra_sample_args. +// Missing keys inherit from the caller-supplied defaults. +inline void parse_config_from_args(Config& cfg, + const std::vector>& kv) { + for (const auto& [key, value] : kv) { + if (key == "speed_scales") { + // Accept ':' as the scales separator so the value survives the + // comma-splitting parse_key_value_args does at the top level. + // e.g. --extra-sample-args speed_scales=0.25:0.5:1.0 + std::vector parsed; + std::string cur; + for (char ch : value) { + if (ch == ':' || ch == '/' || ch == '|') { + if (!cur.empty()) { + parsed.push_back(std::stof(cur)); + cur.clear(); + } + } else if (ch != ' ') { + cur.push_back(ch); + } + } + if (!cur.empty()) { + parsed.push_back(std::stof(cur)); + } + if (!parsed.empty()) { + cfg.scales = std::move(parsed); + } + } else if (key == "speed_levels") { + // Convenience shortcut: speed_levels=2 -> {0.5, 1.0}; =3 -> {0.25, 0.5, 1.0}; etc. + // scales[i] = 2^(i - (n-1)) for i in [0, n-1]. + int n = std::stoi(value); + if (n >= 2) { + std::vector parsed; + parsed.reserve(static_cast(n)); + for (int i = 0; i < n; ++i) { + parsed.push_back(std::ldexp(1.0f, i - (n - 1))); + } + cfg.scales = std::move(parsed); + } + } else if (key == "speed_manual_sigmas") { + std::vector parsed; + std::string cur; + for (char ch : value) { + if (ch == ':' || ch == '/' || ch == '|') { + if (!cur.empty()) { + parsed.push_back(std::stof(cur)); + cur.clear(); + } + } else if (ch != ' ') { + cur.push_back(ch); + } + } + if (!cur.empty()) { + parsed.push_back(std::stof(cur)); + } + cfg.manual_sigmas = std::move(parsed); + } else if (key == "speed_preset") { + if (value == "wan21" || value == "wan" || value == "wan2.1") { + Preset p = preset_wan21(); + cfg.spectrum_A = p.A; + cfg.spectrum_beta = p.beta; + } else { + Preset p = preset_flux(); + cfg.spectrum_A = p.A; + cfg.spectrum_beta = p.beta; + } + } else if (key == "speed_delta") { + cfg.delta = std::stof(value); + } else if (key == "speed_A") { + cfg.spectrum_A = std::stof(value); + } else if (key == "speed_beta") { + cfg.spectrum_beta = std::stof(value); + } else if (key == "speed_seed") { + cfg.seed = static_cast(std::stoull(value)); + } else if (key == "speed_transform") { + if (value == "fft") { + cfg.transform = Transform::FFT; + } else { + cfg.transform = Transform::DCT; + } + } + } +} + +} // namespace speed +} // namespace sd + +#endif // __SD_RUNTIME_SPEED_SAMPLER_HPP__ diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 28628be7c..9169f6aab 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -127,6 +127,9 @@ const char* sampling_methods_str[] = { "Euler CFG++", "Euler A CFG++", "Euler GE", + "DPM++ (2M) SDE", + "DPM++ (2M) SDE (Brownian tree)", + "SPEED (Spectral Progressive)", }; /*================================================== Helper Functions ================================================*/ @@ -2298,7 +2301,37 @@ class StableDiffusionGGML { : init_latent; sd::Tensor denoised = x_t; + // Lazy resize cache for ref_latents so multi-resolution samplers (SPEED + // etc.) that downscale x can still condition on shape-matched refs. + std::map, std::vector>> resized_ref_latents_cache; + auto ref_latents_for = [&](int64_t target_w, int64_t target_h) -> const std::vector>* { + if (ref_latents.empty()) { + return &ref_latents; + } + const auto& first_shape = ref_latents[0].shape(); + if (first_shape[0] == target_w && first_shape[1] == target_h) { + return &ref_latents; + } + auto key = std::make_pair(target_w, target_h); + auto it = resized_ref_latents_cache.find(key); + if (it != resized_ref_latents_cache.end()) { + return &it->second; + } + std::vector> resized; + resized.reserve(ref_latents.size()); + for (const auto& rl : ref_latents) { + std::vector new_shape = rl.shape(); + new_shape[0] = target_w; + new_shape[1] = target_h; + resized.push_back(sd::ops::interpolate(rl, new_shape)); + } + auto ins = resized_ref_latents_cache.emplace(std::move(key), std::move(resized)); + return &ins.first->second; + }; + auto denoise = [&](const sd::Tensor& x, float sigma, int step) -> sd::guidance::GuiderOutput { + const std::vector>* effective_ref_latents = + ref_latents_for(x.shape()[0], x.shape()[1]); if (get_cancel_flag() == SD_CANCEL_ALL) { LOG_DEBUG("cancelling generation"); return {}; @@ -2454,7 +2487,7 @@ class StableDiffusionGGML { } } - cond_out = run_condition(*positive_condition, c_concat_override); + cond_out = run_condition(*positive_condition, c_concat_override, nullptr, effective_ref_latents); if (cond_out.empty()) { return {}; } @@ -2474,7 +2507,8 @@ class StableDiffusionGGML { } uncond_out = run_condition(uncond, uncond.c_concat.empty() ? nullptr : &uncond.c_concat, - uncond_skip_layers); + uncond_skip_layers, + effective_ref_latents); if (uncond_out.empty()) { return {}; } @@ -2483,7 +2517,7 @@ class StableDiffusionGGML { 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); + uncond_without_ref_latents ? &empty_ref_latents : effective_ref_latents); if (img_uncond_out.empty()) { return {}; } @@ -2506,7 +2540,8 @@ class StableDiffusionGGML { guidance_input.predict_skip_layer = [&]() -> sd::Tensor { return run_condition(cond, cond.c_concat.empty() ? nullptr : &cond.c_concat, - &skip_layer_guidance.layers()); + &skip_layer_guidance.layers(), + effective_ref_latents); }; } } @@ -2528,7 +2563,7 @@ class StableDiffusionGGML { if (cache_runtime.spectrum_enabled) { cache_runtime.spectrum.update(denoised); } - if (!denoise_mask.empty()) { + if (!denoise_mask.empty() && denoise_mask.shape() == denoised.shape()) { denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask); } if (sd_should_preview_denoised() && preview.callback != nullptr) { @@ -2539,7 +2574,10 @@ class StableDiffusionGGML { return output; }; - auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args, denoiser); + if (method == SPEED_FLOW_SAMPLE_METHOD && !ref_latents.empty()) { + LOG_WARN("SPEED sampler with reference latents (edit / ref-guided model): the low-resolution segments see a downscaled reference so fine detail from the source image (backgrounds, textures, layout beyond broad structure) is likely to drift. Composition will remain plausible but fidelity to the reference is reduced. See docs/speed_sampler.md."); + } + auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args, denoiser, version); if (x0_opt.empty()) { LOG_ERROR("Diffusion model sampling failed"); if (control_net) { @@ -2808,6 +2846,7 @@ const char* sample_method_to_str[] = { "euler_ge", "dpm++2m_sde", "dpm++2m_sde_bt", + "speed_flow", }; const char* sd_sample_method_name(enum sample_method_t sample_method) {