mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-25 04:32:29 +00:00
Compare commits
6 Commits
master-907
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
b167b942f7 | ||
|
|
1a2330de68 | ||
|
|
740c7ae193 | ||
|
|
4dfe8f5d45 | ||
|
|
caa111adf3 | ||
|
|
88411ef1e0 |
@ -63,6 +63,12 @@ See [backend selection](./backend.md) for full syntax.
|
||||
|
||||
When a graph has cut markers and its missing weights plus incremental compute workspace exceed the available device headroom, it runs its fixed segment list in order. A reusable monolithic compute buffer is not counted as a new allocation. An explicit `--max-vram` budget deducts already-resident managed weights and compute/cache buffers registered by every runner sharing the device, so later graph runs remain segmented when the full graph exceeds the budget. The current segment's weights are pinned during compute, and the next parameter-bearing segment is prefetched when the device supports asynchronous transfer. No opt-in streaming flag is required.
|
||||
|
||||
When choosing between monolithic and segmented execution, the runner requires
|
||||
an additional 128 MiB of headroom in both available device memory and any explicit
|
||||
managed budget. This planning headroom absorbs small allocation estimate changes;
|
||||
subsequent capacity checks can consume it while still preserving the 512 MiB device
|
||||
scratch reserve and respecting the managed budget.
|
||||
|
||||
- `--max-vram <GiB>` optionally lowers the live-memory limit. A positive value is a managed per-device budget, `0` uses the device's current free memory without an explicit budget, and a negative value snapshots free memory at startup while reserving that many GiB (`--max-vram -1` reserves about 1 GiB). Driver contexts and unrelated external allocations remain outside the managed budget.
|
||||
- `--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.
|
||||
|
||||
@ -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
|
||||
```
|
||||
|
||||
<img alt="Qwen Image 2.1 example" src="../assets/qwen/qwen_image_2.1.png" />
|
||||
@ -35,7 +35,7 @@ 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`.
|
||||
@ -44,7 +44,27 @@ For multiple reference images, repeat `-r` in the desired order, for example `-r
|
||||
|
||||
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.
|
||||
|
||||
The cache uses FP32 on all attention backends. For the default 32-layer model, a prefix of 4096 tokens takes about 4 GiB per condition, in addition to weights and working buffers. 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.
|
||||
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.
|
||||
|
||||
|
||||
@ -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, qwen_image_2_1_prefix_cache",
|
||||
"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},
|
||||
{"",
|
||||
|
||||
@ -56,6 +56,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`
|
||||
|
||||
@ -434,7 +435,8 @@ Top-level fields:
|
||||
| `samplers` | `array<string>` | Available sampling methods |
|
||||
| `schedulers` | `array<string>` | Available schedulers |
|
||||
| `loras` | `array<object>` | Available LoRA entries |
|
||||
| `upscalers` | `array<object>` | Available model-backed highres upscalers |
|
||||
| `upscalers` | `array<object>` | Available highres upscalers, built-in and model-backed |
|
||||
| `upscale` | `boolean` | Whether a compatible RGB ESRGAN model is available for `POST /sdcpp/v1/upscale` |
|
||||
| `limits` | `object` | Shared queue and size limits |
|
||||
|
||||
`model`
|
||||
@ -476,6 +478,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.
|
||||
|
||||
@ -489,6 +493,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`:
|
||||
|
||||
@ -641,6 +647,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<object>` | One image |
|
||||
| `images[].index` | `integer` | |
|
||||
| `images[].b64_json` | `string` | Base64-encoded image bytes |
|
||||
| `upscaler` | `string` | The upscaler actually used |
|
||||
| `scale` | `integer` | The model's scale factor |
|
||||
| `repeats` | `integer` | How many times it was run |
|
||||
| `width` | `integer` | Result width |
|
||||
| `height` | `integer` | Result height |
|
||||
| `output_format` | `string` | Final encoded image format |
|
||||
|
||||
Typical status codes:
|
||||
|
||||
- `200 OK`
|
||||
- `400 Bad Request` (invalid request, unsupported output format, unreadable image, incompatible upscaler, or output dimensions exceeding the limit)
|
||||
- `500 Internal Server Error`
|
||||
|
||||
Notes:
|
||||
|
||||
- Final output dimensions, including all repeats, must not exceed 8192 pixels on either axis (`limits.max_upscale_width` and `limits.max_upscale_height`). Requests exceeding this bound are rejected before upscaling.
|
||||
- The upscaler models are three-channel; alpha is not preserved.
|
||||
- The request holds the generation context lock, so an upscale and a
|
||||
generation never run on the device at the same time.
|
||||
|
||||
#### `POST /sdcpp/v1/jobs/{id}/cancel`
|
||||
|
||||
Attempts to cancel an accepted job.
|
||||
|
||||
@ -3,12 +3,33 @@
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <filesystem>
|
||||
#include <limits>
|
||||
|
||||
#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;
|
||||
@ -241,37 +262,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<std::mutex> 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;
|
||||
}
|
||||
}
|
||||
|
||||
@ -341,6 +384,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;
|
||||
@ -350,6 +395,7 @@ 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;
|
||||
return result;
|
||||
}
|
||||
|
||||
@ -415,6 +461,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<uint64_t>() <= static_cast<uint64_t>(std::numeric_limits<int>::max())
|
||||
: value.is_number_integer() && value.get<int64_t>() >= std::numeric_limits<int>::min() &&
|
||||
value.get<int64_t>() <= std::numeric_limits<int>::max();
|
||||
if (!valid) {
|
||||
res.status = 400;
|
||||
res.set_content(json({{"error", std::string(key) + " must be a 32-bit integer"}}).dump(), "application/json");
|
||||
return;
|
||||
}
|
||||
}
|
||||
ImgGenJobRequest output_options;
|
||||
std::string error_message;
|
||||
if (!assign_output_options(output_options,
|
||||
body.value("output_format", std::string("png")),
|
||||
body.value("output_compression", 100),
|
||||
true,
|
||||
error_message)) {
|
||||
res.status = 400;
|
||||
res.set_content(json({{"error", error_message}}).dump(), "application/json");
|
||||
return;
|
||||
}
|
||||
const int tile_size = std::max(32, body.value("tile_size", runtime->default_gen_params->upscale_tile_size));
|
||||
const int repeats = std::clamp(body.value("repeats", 1), 1, 4);
|
||||
const std::string wanted = body.value("upscaler", std::string());
|
||||
|
||||
const std::string encoded = body.value("image", std::string());
|
||||
if (encoded.empty()) {
|
||||
res.status = 400;
|
||||
res.set_content(R"({"error":"image is required"})", "application/json");
|
||||
return;
|
||||
}
|
||||
SDImageOwner input;
|
||||
if (!decode_base64_image(encoded, 3, 0, 0, input) || input.get().data == nullptr) {
|
||||
res.status = 400;
|
||||
res.set_content(R"({"error":"image could not be read"})", "application/json");
|
||||
return;
|
||||
}
|
||||
|
||||
refresh_upscaler_cache(*runtime);
|
||||
int model_scale = 0;
|
||||
std::string model_path;
|
||||
std::string used_name;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(*runtime->upscaler_mutex);
|
||||
for (const auto& entry : *runtime->upscaler_cache) {
|
||||
if (entry.image_upscale_factor > 0 && (wanted.empty() || entry.name == wanted)) {
|
||||
model_path = entry.fullpath;
|
||||
used_name = entry.name;
|
||||
model_scale = entry.image_upscale_factor;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (model_path.empty()) {
|
||||
res.status = 400;
|
||||
res.set_content(json({{"error", wanted.empty()
|
||||
? std::string("no RGB ESRGAN upscaler models are available; "
|
||||
"start the server with --hires-upscalers-dir")
|
||||
: "no compatible image upscaler called " + wanted}})
|
||||
.dump(),
|
||||
"application/json");
|
||||
return;
|
||||
}
|
||||
|
||||
if (!valid_upscale_dimensions(input.get(), model_scale, repeats)) {
|
||||
res.status = 400;
|
||||
res.set_content(R"({"error":"upscaled dimensions must not exceed 8192 x 8192"})", "application/json");
|
||||
return;
|
||||
}
|
||||
|
||||
// One GPU: an upscale must not run while a generation is using it.
|
||||
std::lock_guard<std::mutex> ctx_lock(*runtime->sd_ctx_mutex);
|
||||
UpscalerCtxPtr upscaler_ctx(new_upscaler_ctx(model_path.c_str(),
|
||||
runtime->ctx_params->diffusion_conv_direct,
|
||||
runtime->ctx_params->n_threads,
|
||||
tile_size,
|
||||
runtime->ctx_params->backend.c_str(),
|
||||
runtime->ctx_params->params_backend.c_str()));
|
||||
if (upscaler_ctx == nullptr) {
|
||||
res.status = 500;
|
||||
res.set_content(R"({"error":"the upscaler model could not be loaded"})", "application/json");
|
||||
return;
|
||||
}
|
||||
const int factor = get_upscale_factor(upscaler_ctx.get());
|
||||
// The model file may have changed since its metadata was cached.
|
||||
if (!valid_upscale_dimensions(input.get(), factor, repeats)) {
|
||||
res.status = 400;
|
||||
res.set_content(R"({"error":"upscaled dimensions must not exceed 8192 x 8192"})", "application/json");
|
||||
return;
|
||||
}
|
||||
|
||||
SDImageOwner current(input.release());
|
||||
for (int i = 0; i < repeats; ++i) {
|
||||
sd_image_t* out_images = nullptr;
|
||||
int out_count = 0;
|
||||
if (!upscale(upscaler_ctx.get(), current.get(), (uint32_t)factor, &out_images, &out_count) ||
|
||||
out_count <= 0 || out_images[0].data == nullptr) {
|
||||
free_sd_images(out_images, out_count);
|
||||
res.status = 500;
|
||||
res.set_content(R"({"error":"upscale failed"})", "application/json");
|
||||
return;
|
||||
}
|
||||
sd_image_t produced = out_images[0];
|
||||
out_images[0] = {0, 0, 0, nullptr};
|
||||
free_sd_images(out_images, out_count);
|
||||
current.reset(produced);
|
||||
}
|
||||
|
||||
const std::string& format = output_options.output_format;
|
||||
const int compression = output_options.output_compression;
|
||||
const sd_image_t result = current.get();
|
||||
auto image_bytes = encode_image_to_vector(format == "jpeg" ? EncodedImageFormat::JPEG
|
||||
: format == "webp" ? EncodedImageFormat::WEBP
|
||||
: EncodedImageFormat::PNG,
|
||||
result.data,
|
||||
result.width,
|
||||
result.height,
|
||||
result.channel,
|
||||
"",
|
||||
compression);
|
||||
if (image_bytes.empty()) {
|
||||
res.status = 500;
|
||||
res.set_content(R"({"error":"the result could not be encoded"})", "application/json");
|
||||
return;
|
||||
}
|
||||
|
||||
json out;
|
||||
out["upscaler"] = used_name;
|
||||
out["scale"] = factor;
|
||||
out["repeats"] = repeats;
|
||||
out["width"] = result.width;
|
||||
out["height"] = result.height;
|
||||
out["output_format"] = format;
|
||||
json images = json::array();
|
||||
images.push_back({{"index", 0}, {"b64_json", base64_encode(image_bytes)}});
|
||||
out["images"] = std::move(images);
|
||||
res.set_content(out.dump(), "application/json");
|
||||
res.status = 200;
|
||||
} catch (const json::exception& e) {
|
||||
res.status = 400;
|
||||
res.set_content(json({{"error", "invalid request"}, {"message", e.what()}}).dump(), "application/json");
|
||||
} catch (const std::exception& e) {
|
||||
res.status = 500;
|
||||
res.set_content(json({{"error", std::string("server_error: ") + e.what()}}).dump(), "application/json");
|
||||
}
|
||||
});
|
||||
|
||||
svr.Post("/sdcpp/v1/img_gen", [runtime](const httplib::Request& req, httplib::Response& res) {
|
||||
try {
|
||||
if (req.body.empty()) {
|
||||
|
||||
@ -295,6 +295,11 @@ std::string get_lora_full_path(ServerRuntime& rt, const std::string& path) {
|
||||
|
||||
void refresh_upscaler_cache(ServerRuntime& rt) {
|
||||
std::vector<UpscalerEntry> new_cache;
|
||||
std::vector<UpscalerEntry> previous_cache;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(*rt.upscaler_mutex);
|
||||
previous_cache = *rt.upscaler_cache;
|
||||
}
|
||||
|
||||
fs::path upscaler_dir = rt.ctx_params->hires_upscalers_dir;
|
||||
if (fs::exists(upscaler_dir) && fs::is_directory(upscaler_dir)) {
|
||||
@ -312,6 +317,20 @@ void refresh_upscaler_cache(ServerRuntime& rt) {
|
||||
upscaler_entry.fullpath = fs::absolute(p).lexically_normal().u8string();
|
||||
upscaler_entry.model_name = "ESRGAN_4x";
|
||||
upscaler_entry.path = p.filename().u8string();
|
||||
upscaler_entry.file_size = entry.file_size();
|
||||
upscaler_entry.last_modified = entry.last_write_time();
|
||||
auto previous = std::find_if(previous_cache.begin(), previous_cache.end(), [&](const UpscalerEntry& cached) {
|
||||
return cached.fullpath == upscaler_entry.fullpath &&
|
||||
cached.file_size == upscaler_entry.file_size &&
|
||||
cached.last_modified == upscaler_entry.last_modified;
|
||||
});
|
||||
upscaler_entry.image_upscale_factor = previous != previous_cache.end()
|
||||
? previous->image_upscale_factor
|
||||
: get_upscaler_model_scale(upscaler_entry.fullpath.c_str());
|
||||
if (upscaler_entry.image_upscale_factor > 0) {
|
||||
upscaler_entry.scale = upscaler_entry.image_upscale_factor;
|
||||
upscaler_entry.model_name = "ESRGAN_" + std::to_string(upscaler_entry.scale) + "x";
|
||||
}
|
||||
|
||||
new_cache.push_back(std::move(upscaler_entry));
|
||||
}
|
||||
|
||||
@ -2,6 +2,7 @@
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <filesystem>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
@ -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 {
|
||||
|
||||
@ -558,6 +558,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;
|
||||
|
||||
|
||||
@ -623,7 +623,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;
|
||||
@ -755,6 +759,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,
|
||||
|
||||
@ -221,7 +221,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
|
||||
bool skip_reshape = false,
|
||||
bool flash_attn = false,
|
||||
float kv_scale = 1.0f,
|
||||
bool sage_attn = false);
|
||||
bool sage_attn = false,
|
||||
bool* used_flash_attn = nullptr);
|
||||
|
||||
ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
|
||||
@ -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<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
|
||||
@ -831,11 +837,20 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
|
||||
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();
|
||||
|
||||
@ -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;
|
||||
@ -125,7 +132,8 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* mask = nullptr,
|
||||
bool skip_reshape = false,
|
||||
bool flash_attn = false,
|
||||
float kv_scale = 1.f);
|
||||
float kv_scale = 1.f,
|
||||
bool* used_flash_attn = nullptr);
|
||||
|
||||
struct GGMLRunner {
|
||||
private:
|
||||
@ -289,7 +297,8 @@ public:
|
||||
|
||||
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 +333,7 @@ public:
|
||||
|
||||
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);
|
||||
|
||||
@ -5,6 +5,7 @@
|
||||
#include <cassert>
|
||||
#include <cmath>
|
||||
#include <set>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
#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<SpatialRegion> images;
|
||||
size_t token_count = 0;
|
||||
|
||||
void append_tokens(size_t count) {
|
||||
token_count += count;
|
||||
}
|
||||
|
||||
void append_image(int height, int width, int frames = 1, float height_step = 1.f, float width_step = 1.f) {
|
||||
size_t count = static_cast<size_t>(height) * width * frames;
|
||||
images.push_back({token_count, count, height * height_step, width * width_step});
|
||||
append_tokens(count);
|
||||
}
|
||||
};
|
||||
|
||||
struct Frequency {
|
||||
size_t axis;
|
||||
float omega;
|
||||
};
|
||||
|
||||
struct Embedding {
|
||||
std::vector<float> values;
|
||||
std::vector<std::vector<float>> ids;
|
||||
PositionLayout positions;
|
||||
std::vector<Frequency> frequencies;
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix;
|
||||
int batch_size = 1;
|
||||
};
|
||||
|
||||
enum class RefIndexMode {
|
||||
FIXED,
|
||||
INCREASE,
|
||||
@ -56,40 +96,25 @@ namespace Rope {
|
||||
return flat_vec;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
|
||||
int dim,
|
||||
float theta,
|
||||
const std::vector<int>& axis_wrap_dims = {}) {
|
||||
__STATIC_INLINE__ std::vector<float> rope_frequencies(int dim, float theta) {
|
||||
assert(dim % 2 == 0);
|
||||
int half_dim = dim / 2;
|
||||
|
||||
std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
|
||||
|
||||
std::vector<float> 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<std::vector<float>> rope(const std::vector<float>& pos,
|
||||
const std::vector<float>& omega) {
|
||||
int half_dim = static_cast<int>(omega.size());
|
||||
size_t pos_size = pos.size();
|
||||
std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
|
||||
for (size_t i = 0; i < pos_size; ++i) {
|
||||
for (size_t 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<std::vector<float>> rope(const std::vector<float>& pos,
|
||||
int dim,
|
||||
float theta) {
|
||||
return rope(pos, rope_frequencies(dim, theta));
|
||||
}
|
||||
|
||||
// Generate IDs for image patches and text
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) {
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
|
||||
@ -139,9 +170,13 @@ namespace Rope {
|
||||
int index = 0,
|
||||
int h_offset = 0,
|
||||
int w_offset = 0,
|
||||
bool scale_rope = false) {
|
||||
bool scale_rope = false,
|
||||
PositionLayout* layout = nullptr) {
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
if (layout) {
|
||||
layout->append_image(h_len, w_len);
|
||||
}
|
||||
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
|
||||
|
||||
int h_start = h_offset;
|
||||
@ -192,8 +227,8 @@ namespace Rope {
|
||||
int bs,
|
||||
const std::vector<float>& axis_thetas,
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {},
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix,
|
||||
std::vector<Frequency>* frequencies = nullptr) {
|
||||
std::vector<std::vector<float>> trans_ids = transpose(ids);
|
||||
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<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
|
||||
size_t offset = 0;
|
||||
for (size_t i = 0; i < num_axes; ++i) {
|
||||
std::vector<int> axis_wrap_dims;
|
||||
if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
|
||||
axis_wrap_dims = wrap_dims[i];
|
||||
}
|
||||
float axis_theta = 10000.0f;
|
||||
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<std::vector<float>> rope_emb =
|
||||
rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
|
||||
rope(trans_ids[i], omega); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
|
||||
for (int b = 0; b < bs; ++b) {
|
||||
for (int 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<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {},
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix,
|
||||
std::vector<Frequency>* frequencies = nullptr) {
|
||||
std::vector<float> axis_thetas(axes_dim.size(), theta);
|
||||
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
|
||||
return embed_nd(ids, bs, axis_thetas, axes_dim, layout, frequencies);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
|
||||
@ -264,7 +305,7 @@ namespace Rope {
|
||||
float theta,
|
||||
int head_dim,
|
||||
const std::vector<int>& mrope_section,
|
||||
const std::vector<std::vector<int>>& axis_wrap_dims = {}) {
|
||||
std::vector<Frequency>* frequencies = nullptr) {
|
||||
GGML_ASSERT(bs > 0);
|
||||
GGML_ASSERT(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<std::vector<std::vector<float>>> axis_embs;
|
||||
axis_embs.reserve(3);
|
||||
for (int axis = 0; axis < 3; ++axis) {
|
||||
std::vector<int> axis_wrap;
|
||||
if (axis < static_cast<int>(axis_wrap_dims.size())) {
|
||||
axis_wrap = axis_wrap_dims[axis];
|
||||
}
|
||||
axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
|
||||
axis_embs.push_back(rope(trans_ids[axis], omega));
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> emb = axis_embs[0];
|
||||
for (int axis = 1; axis < 3; ++axis) {
|
||||
int length = std::min<int>(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,7 +345,7 @@ namespace Rope {
|
||||
return flatten(emb);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height,
|
||||
__STATIC_INLINE__ Embedding embed_2d_interleaved(int height,
|
||||
int width,
|
||||
int dim,
|
||||
float theta = 10000.f,
|
||||
@ -318,6 +365,10 @@ namespace Rope {
|
||||
w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
|
||||
}
|
||||
|
||||
Embedding result;
|
||||
result.positions.append_image(height, width, 1,
|
||||
height > 1 ? scale / (height - 1) : 1.f,
|
||||
width > 1 ? scale / (width - 1) : 1.f);
|
||||
std::vector<float> x_pos;
|
||||
std::vector<float> y_pos;
|
||||
x_pos.reserve(static_cast<size_t>(height) * width);
|
||||
@ -326,13 +377,20 @@ namespace Rope {
|
||||
float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
|
||||
for (int ix = 0; ix < width; ++ix) {
|
||||
float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
|
||||
result.ids.push_back({0.f, y, x});
|
||||
x_pos.push_back(x);
|
||||
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<float> out(static_cast<size_t>(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<std::vector<float>> 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<std::vector<float>> ids;
|
||||
int curr_h_offset = 0;
|
||||
int curr_w_offset = 0;
|
||||
@ -386,7 +446,8 @@ namespace Rope {
|
||||
static_cast<int>(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,23 +470,27 @@ namespace Rope {
|
||||
const std::vector<ggml_tensor*>& 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<float> gen_flux_pe(int h,
|
||||
__STATIC_INLINE__ Embedding gen_flux_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
@ -435,11 +500,11 @@ namespace Rope {
|
||||
RefIndexMode ref_index_mode,
|
||||
float ref_index_scale,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim,
|
||||
bool is_longcat) {
|
||||
std::vector<std::vector<float>> ids = gen_flux_ids(h,
|
||||
Embedding result;
|
||||
result.batch_size = bs;
|
||||
result.ids = gen_flux_ids(h,
|
||||
w,
|
||||
patch_size,
|
||||
bs,
|
||||
@ -449,48 +514,9 @@ namespace Rope {
|
||||
ref_latents,
|
||||
ref_index_mode,
|
||||
ref_index_scale,
|
||||
is_longcat);
|
||||
std::vector<std::vector<int>> 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<int>(pos_len, 0));
|
||||
size_t cursor = context_len; // text first
|
||||
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
cursor += img_tokens;
|
||||
// reference latents
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
if (ref == nullptr) {
|
||||
continue;
|
||||
}
|
||||
int ref_h = static_cast<int>(ref->ne[1]);
|
||||
int ref_w = static_cast<int>(ref->ne[0]);
|
||||
int ref_h_l = (ref_h + (patch_size / 2)) / patch_size;
|
||||
int ref_w_l = (ref_w + (patch_size / 2)) / patch_size;
|
||||
size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
|
||||
for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = ref_h_l;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = ref_w_l;
|
||||
}
|
||||
}
|
||||
cursor += ref_tokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
is_longcat, &result.positions);
|
||||
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
|
||||
@ -503,11 +529,15 @@ namespace Rope {
|
||||
int t_offset = 0,
|
||||
int h_offset = 0,
|
||||
int w_offset = 0,
|
||||
bool scale_rope = false) {
|
||||
bool scale_rope = false,
|
||||
PositionLayout* layout = nullptr) {
|
||||
int t_len = (t + (pt / 2)) / pt;
|
||||
int 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<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
|
||||
|
||||
if (scale_rope) {
|
||||
@ -573,7 +603,11 @@ namespace Rope {
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode) {
|
||||
RefIndexMode ref_index_mode,
|
||||
PositionLayout* layout = nullptr) {
|
||||
if (layout) {
|
||||
layout->append_tokens(context_len);
|
||||
}
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
int txt_id_start = std::max(h_len, w_len) / 2;
|
||||
@ -585,18 +619,18 @@ 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<float> gen_qwen_image_pe(int t,
|
||||
__STATIC_INLINE__ Embedding gen_qwen_image_pe(int t,
|
||||
int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
@ -605,70 +639,29 @@ namespace Rope {
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
|
||||
std::vector<std::vector<int>> 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<int>(total_tokens / bs, 0));
|
||||
size_t cursor = context_len; // ignore text tokens
|
||||
const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
cursor += img_tokens;
|
||||
// For each reference image, store wrap sizes as well
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
if (ref == nullptr) {
|
||||
continue;
|
||||
}
|
||||
int ref_h = static_cast<int>(ref->ne[1]);
|
||||
int ref_w = static_cast<int>(ref->ne[0]);
|
||||
int ref_pad_h = (patch_size - (ref_h % patch_size)) % patch_size;
|
||||
int ref_pad_w = (patch_size - (ref_w % patch_size)) % patch_size;
|
||||
int ref_h_len = (ref_h + ref_pad_h) / patch_size;
|
||||
int ref_w_len = (ref_w + ref_pad_w) / patch_size;
|
||||
size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
|
||||
for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = ref_h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = ref_w_len;
|
||||
}
|
||||
}
|
||||
cursor += ref_n_tokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
|
||||
__STATIC_INLINE__ Embedding gen_mage_flow_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
Embedding result;
|
||||
result.batch_size = bs;
|
||||
result.positions.append_tokens(context_len);
|
||||
const int axes_dim_num = static_cast<int>(axes_dim.size());
|
||||
auto make_image_ids = [=](int image_h, int image_w, int image_index) {
|
||||
auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) {
|
||||
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
|
||||
std::vector<float>(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<int>(i + 1));
|
||||
ids = concat_ids(ids, ref_ids, bs);
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
result.ids = std::move(ids);
|
||||
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
|
||||
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<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
|
||||
@ -711,44 +707,37 @@ namespace Rope {
|
||||
}
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(context_len);
|
||||
}
|
||||
return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
|
||||
__STATIC_INLINE__ Embedding gen_lens_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
|
||||
std::vector<std::vector<int>> 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<int>(pos_len, 0));
|
||||
const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][token_i] = h;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][token_i] = w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> 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<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
|
||||
std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
|
||||
std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
|
||||
@ -774,39 +763,25 @@ namespace Rope {
|
||||
}
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(context_len);
|
||||
}
|
||||
return concat_ids(img_ids_repeated, txt_ids, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
|
||||
__STATIC_INLINE__ Embedding 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<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
|
||||
std::vector<std::vector<int>> 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<int>(pos_len, 0));
|
||||
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
|
||||
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<float>(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<ggml_tensor*>& 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<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
|
||||
@ -913,11 +889,17 @@ namespace Rope {
|
||||
txt_ids[i][0] = (i % padded_context_len) + 1.f;
|
||||
}
|
||||
|
||||
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<int>(img_ids.size() / bs), seq_multi_of);
|
||||
if (layout) {
|
||||
layout->append_tokens(img_pad_len);
|
||||
}
|
||||
if (img_pad_len > 0) {
|
||||
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
|
||||
img_ids = concat_ids(img_ids, img_pad_ids, bs);
|
||||
@ -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<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
|
||||
@ -947,11 +930,17 @@ namespace Rope {
|
||||
}
|
||||
}
|
||||
|
||||
if (layout) {
|
||||
layout->append_tokens(padded_context_len);
|
||||
}
|
||||
int axes_dim_num = 3;
|
||||
int index = padded_context_len + 1;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
|
||||
|
||||
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
|
||||
if (layout) {
|
||||
layout->append_tokens(img_pad_len);
|
||||
}
|
||||
if (img_pad_len > 0) {
|
||||
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
|
||||
img_ids = concat_ids(img_ids, img_pad_ids, bs);
|
||||
@ -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<std::vector<float>> 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<float>(3, 0.f));
|
||||
if (layout) {
|
||||
layout->append_tokens(image_pad);
|
||||
}
|
||||
}
|
||||
|
||||
const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
|
||||
@ -1016,11 +1012,14 @@ 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<float> gen_llada_image_edit_pe(int h,
|
||||
__STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int context_len,
|
||||
@ -1028,48 +1027,30 @@ namespace Rope {
|
||||
int seq_multi_of,
|
||||
int theta,
|
||||
const std::vector<int>& 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<float>(theta), axes_dim, {});
|
||||
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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_llada_image_pe(int h,
|
||||
__STATIC_INLINE__ Embedding gen_llada_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int seq_multi_of,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of);
|
||||
std::vector<std::vector<int>> 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<int>(pos_len, 0));
|
||||
size_t cursor = context_len + bound_mod(context_len, seq_multi_of);
|
||||
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Generate z_image positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
|
||||
__STATIC_INLINE__ Embedding gen_z_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
@ -1078,33 +1059,12 @@ namespace Rope {
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
RefIndexMode ref_index_mode,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
|
||||
std::vector<std::vector<int>> 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<int>(pos_len, 0));
|
||||
size_t cursor = context_len + bound_mod(context_len, seq_multi_of); // skip text (and its padding)
|
||||
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
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<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
|
||||
|
||||
65
src/model/common/rope_circular.hpp
Normal file
65
src/model/common/rope_circular.hpp
Normal file
@ -0,0 +1,65 @@
|
||||
#ifndef __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
|
||||
#define __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
|
||||
|
||||
#include "model/common/rope.hpp"
|
||||
|
||||
namespace Rope {
|
||||
__STATIC_INLINE__ void apply_circular(Embedding& embedding, bool circular_x, bool circular_y) {
|
||||
if (!circular_x && !circular_y) {
|
||||
return;
|
||||
}
|
||||
|
||||
GGML_ASSERT(embedding.batch_size > 0);
|
||||
GGML_ASSERT(embedding.ids.size() % embedding.batch_size == 0);
|
||||
size_t pos_len = embedding.ids.size() / embedding.batch_size;
|
||||
size_t half_dim = embedding.frequencies.size();
|
||||
GGML_ASSERT(embedding.positions.token_count == pos_len);
|
||||
GGML_ASSERT(embedding.values.size() == embedding.ids.size() * half_dim * 4);
|
||||
|
||||
constexpr float TWO_PI = 6.28318530717958647692f;
|
||||
for (const auto& region : embedding.positions.images) {
|
||||
GGML_ASSERT(region.begin <= pos_len && region.count <= pos_len - region.begin);
|
||||
for (size_t j = 0; j < half_dim; ++j) {
|
||||
const auto& frequency = embedding.frequencies[j];
|
||||
float period = 0.f;
|
||||
if (circular_y && frequency.axis == static_cast<size_t>(region.height_axis)) {
|
||||
period = region.height_period;
|
||||
} else if (circular_x && frequency.axis == static_cast<size_t>(region.width_axis)) {
|
||||
period = region.width_period;
|
||||
}
|
||||
if (period <= 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Quantize to periodic harmonics while preserving the original coordinate offsets.
|
||||
float rounded = std::round(frequency.omega * period / TWO_PI);
|
||||
for (int b = 0; b < embedding.batch_size; ++b) {
|
||||
size_t begin = b * pos_len + region.begin;
|
||||
for (size_t i = begin; i < begin + region.count; ++i) {
|
||||
GGML_ASSERT(frequency.axis < embedding.ids[i].size());
|
||||
float angle = embedding.ids[i][frequency.axis] * TWO_PI * rounded / period;
|
||||
float cos_val = std::cos(angle);
|
||||
float sin_val = std::sin(angle);
|
||||
if (embedding.layout == EmbedNDLayout::ErnieImage) {
|
||||
size_t cos_offset = (i * half_dim + j) * 2;
|
||||
size_t sin_offset = embedding.ids.size() * half_dim * 2 + cos_offset;
|
||||
embedding.values[cos_offset] = cos_val;
|
||||
embedding.values[cos_offset + 1] = cos_val;
|
||||
embedding.values[sin_offset] = sin_val;
|
||||
embedding.values[sin_offset + 1] = sin_val;
|
||||
} else {
|
||||
size_t offset = (i * half_dim + j) * 4;
|
||||
embedding.values[offset] = cos_val;
|
||||
embedding.values[offset + 1] = -sin_val;
|
||||
embedding.values[offset + 2] = sin_val;
|
||||
embedding.values[offset + 3] = cos_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace Rope
|
||||
|
||||
#endif // __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
|
||||
@ -603,7 +603,7 @@ namespace Anima {
|
||||
return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2));
|
||||
}
|
||||
|
||||
static std::vector<float> gen_anima_image_pe_vec(int bs,
|
||||
static Rope::Embedding gen_anima_image_pe_vec(int bs,
|
||||
int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
@ -613,7 +613,9 @@ namespace Anima {
|
||||
float w_extrapolation_ratio,
|
||||
float t_extrapolation_ratio,
|
||||
const std::vector<ggml_tensor*>& ref_latents) {
|
||||
auto ids = Rope::gen_flux_ids(h,
|
||||
Rope::Embedding result;
|
||||
result.batch_size = bs;
|
||||
result.ids = Rope::gen_flux_ids(h,
|
||||
w,
|
||||
patch_size,
|
||||
bs,
|
||||
@ -623,14 +625,15 @@ namespace Anima {
|
||||
ref_latents,
|
||||
Rope::RefIndexMode::FIXED,
|
||||
1.0f,
|
||||
false);
|
||||
false, &result.positions);
|
||||
|
||||
std::vector<float> axis_thetas = {
|
||||
static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
|
||||
static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]),
|
||||
static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]),
|
||||
};
|
||||
return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
|
||||
result.values = Rope::embed_nd(result.ids, bs, axis_thetas, axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
@ -657,7 +660,7 @@ 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,
|
||||
image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1,
|
||||
static_cast<int>(h_pad),
|
||||
static_cast<int>(w_pad),
|
||||
static_cast<int>(config.patch_size),
|
||||
@ -666,7 +669,7 @@ namespace Anima {
|
||||
4.0f,
|
||||
4.0f,
|
||||
1.0f,
|
||||
ref_latents);
|
||||
ref_latents));
|
||||
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
|
||||
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());
|
||||
|
||||
@ -720,7 +720,7 @@ namespace Boogu {
|
||||
}
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h,
|
||||
__STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
@ -728,7 +728,10 @@ namespace Boogu {
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids;
|
||||
Rope::Embedding result;
|
||||
result.batch_size = bs;
|
||||
result.positions.append_tokens(context_len);
|
||||
auto& ids = result.ids;
|
||||
ids.reserve(static_cast<size_t>(bs) * context_len);
|
||||
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<float>(theta), axes_dim);
|
||||
result.values = Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
|
||||
return result;
|
||||
}
|
||||
|
||||
struct BooguImageRunner : public DiffusionModelRunner {
|
||||
@ -793,14 +799,14 @@ namespace Boogu {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
|
||||
pe_vec = gen_boogu_pe(static_cast<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
ref_latents,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
@ -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);
|
||||
|
||||
@ -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<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
@ -1548,7 +1548,7 @@ namespace Flux {
|
||||
} else if (version == VERSION_OVIS_IMAGE) {
|
||||
txt_arange_dims = {1, 2};
|
||||
}
|
||||
pe_vec = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
@ -1558,10 +1558,8 @@ namespace Flux {
|
||||
ref_index_mode,
|
||||
config.ref_index_scale,
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim,
|
||||
sd_version_is_longcat(version));
|
||||
sd_version_is_longcat(version)));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
// LOG_VERBOSE("pos_len %d", pos_len);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
|
||||
@ -149,18 +149,21 @@ namespace Ideogram4 {
|
||||
return std::make_shared<Linear>(in_features, out_features, bias);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
|
||||
__STATIC_INLINE__ Rope::Embedding gen_ideogram4_pe(int grid_h,
|
||||
int grid_w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int head_dim,
|
||||
int rope_theta,
|
||||
const std::vector<int>& mrope_section,
|
||||
bool circular_x = false,
|
||||
bool circular_y = false) {
|
||||
const std::vector<int>& mrope_section) {
|
||||
GGML_ASSERT(bs == 1);
|
||||
std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
|
||||
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<size_t>(bs) * (context_len + grid_h * grid_w),
|
||||
std::vector<float>(3, 0.f));
|
||||
auto& ids = result.ids;
|
||||
|
||||
for (int i = 0; i < context_len; ++i) {
|
||||
ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
|
||||
@ -175,29 +178,13 @@ namespace Ideogram4 {
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<int>> axis_wrap_dims(3);
|
||||
if (circular_y || circular_x) {
|
||||
size_t total_len = static_cast<size_t>(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<size_t>(context_len); idx < total_len; ++idx) {
|
||||
axis_wrap_dims[1][idx] = grid_h;
|
||||
}
|
||||
}
|
||||
if (circular_x) {
|
||||
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
|
||||
axis_wrap_dims[2][idx] = grid_w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return Rope::embed_interleaved_mrope(ids,
|
||||
result.values = Rope::embed_interleaved_mrope(ids,
|
||||
bs,
|
||||
static_cast<float>(rope_theta),
|
||||
head_dim,
|
||||
mrope_section,
|
||||
axis_wrap_dims);
|
||||
&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<int>(grid_h),
|
||||
pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast<int>(grid_h),
|
||||
static_cast<int>(grid_w),
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context_len),
|
||||
static_cast<int>(head_dim),
|
||||
static_cast<int>(config.rope_theta),
|
||||
config.mrope_section,
|
||||
runner_ctx.circular_x_enabled,
|
||||
runner_ctx.circular_y_enabled);
|
||||
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());
|
||||
|
||||
|
||||
@ -689,7 +689,7 @@ namespace Krea2 {
|
||||
}
|
||||
};
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h,
|
||||
__STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
@ -698,14 +698,19 @@ namespace Krea2 {
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<ggml_tensor*>& 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,7 +754,7 @@ namespace Krea2 {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
|
||||
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
@ -757,7 +762,7 @@ namespace Krea2 {
|
||||
config.theta,
|
||||
config.axes_dim,
|
||||
ref_latents,
|
||||
ref_image_params.ref_index_mode);
|
||||
ref_image_params.ref_index_mode));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
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());
|
||||
|
||||
@ -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<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
@ -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<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
ZImage::SEQ_MULTI_OF,
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
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<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(context->ne[1]),
|
||||
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
|
||||
ZImage::SEQ_MULTI_OF,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
@ -110,13 +110,13 @@ namespace MageFlow {
|
||||
}
|
||||
|
||||
int batch_size = static_cast<int>(x->ne[3]);
|
||||
pe_vec = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
batch_size,
|
||||
static_cast<int>(context->ne[1]),
|
||||
ref_latents,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
@ -154,18 +154,26 @@ namespace MiniT2I {
|
||||
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
|
||||
}
|
||||
|
||||
inline std::vector<float> make_vision_rope(int side, int head_dim) {
|
||||
inline Rope::Embedding make_vision_rope(int side, int head_dim) {
|
||||
GGML_ASSERT(head_dim % 4 == 0);
|
||||
int dim = head_dim / 2;
|
||||
int quarter = dim / 2;
|
||||
int length = side * side;
|
||||
Rope::Embedding result;
|
||||
result.positions.append_image(side, side);
|
||||
std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
|
||||
std::vector<float> freqs(quarter);
|
||||
for (int i = 0; i < quarter; ++i) {
|
||||
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
|
||||
}
|
||||
for (int axis : {1, 2}) {
|
||||
for (float frequency : freqs) {
|
||||
result.frequencies.push_back({static_cast<size_t>(axis), frequency});
|
||||
}
|
||||
}
|
||||
for (int y = 0; y < side; ++y) {
|
||||
for (int x = 0; x < side; ++x) {
|
||||
result.ids.push_back({0.f, static_cast<float>(y), static_cast<float>(x)});
|
||||
int pos = y * side + x;
|
||||
size_t base = static_cast<size_t>(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<int>(img_side), static_cast<int>(config.hidden_size));
|
||||
auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
|
||||
auto img_pe_vec = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
|
||||
auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
|
||||
auto joint_pe_vec = txt_pe_vec;
|
||||
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<float>& x_tensor,
|
||||
|
||||
@ -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 {
|
||||
@ -184,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner {
|
||||
protected:
|
||||
std::string prefix;
|
||||
|
||||
std::vector<float> finish_rope_pe(Rope::Embedding embedding) {
|
||||
Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled);
|
||||
return std::move(embedding.values);
|
||||
}
|
||||
|
||||
public:
|
||||
DiffusionModelRunner(ggml_backend_t backend,
|
||||
const std::string& prefix,
|
||||
|
||||
@ -135,7 +135,7 @@ namespace Pid {
|
||||
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
|
||||
}
|
||||
|
||||
inline std::vector<float> make_rope_2d(int height,
|
||||
inline Rope::Embedding make_rope_2d(int height,
|
||||
int width,
|
||||
int dim,
|
||||
float theta = 10000.f,
|
||||
@ -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<int>(Hs),
|
||||
pos_img_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
|
||||
static_cast<int>(Ws),
|
||||
static_cast<int>(config.hidden_size / config.num_groups),
|
||||
10000.f,
|
||||
16.f,
|
||||
static_cast<int>(config.rope_ref_grid_h),
|
||||
static_cast<int>(config.rope_ref_grid_w));
|
||||
static_cast<int>(config.rope_ref_grid_w)));
|
||||
auto pos_img = ggml_new_tensor_4d(compute_ctx,
|
||||
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<int>(Hs),
|
||||
pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
|
||||
static_cast<int>(Ws),
|
||||
static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
|
||||
10000.f,
|
||||
16.f,
|
||||
static_cast<int>(config.rope_ref_grid_h),
|
||||
static_cast<int>(config.rope_ref_grid_w));
|
||||
static_cast<int>(config.rope_ref_grid_w)));
|
||||
auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
|
||||
GGML_TYPE_F32,
|
||||
2,
|
||||
|
||||
@ -635,7 +635,7 @@ namespace Qwen {
|
||||
ref_index_mode = Rope::RefIndexMode::DECREASE;
|
||||
}
|
||||
|
||||
pe_vec = Rope::gen_qwen_image_pe(time_len,
|
||||
pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len,
|
||||
static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
@ -644,9 +644,7 @@ namespace Qwen {
|
||||
ref_latents,
|
||||
ref_index_mode,
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
// LOG_VERBOSE("pos_len %d", pos_len);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
|
||||
@ -68,6 +68,7 @@ namespace Qwen {
|
||||
std::vector<QwenImage21Segment> segments;
|
||||
std::vector<std::vector<float>> positions;
|
||||
int64_t prefix_length = 0;
|
||||
Rope::PositionLayout rope_layout;
|
||||
|
||||
static QwenImage21Layout build(int64_t text_length,
|
||||
const sd::Tensor<int32_t>& image_slots,
|
||||
@ -82,6 +83,7 @@ namespace Qwen {
|
||||
auto [height, width] = image_shapes[index];
|
||||
int64_t start = static_cast<int64_t>(layout.positions.size());
|
||||
layout.segments.push_back({start, start + height * width, context_start, index});
|
||||
layout.rope_layout.append_image(static_cast<int>(height), static_cast<int>(width));
|
||||
for (int64_t h = 0; h < height; ++h) {
|
||||
for (int64_t w = 0; w < width; ++w) {
|
||||
layout.positions.push_back({static_cast<float>(position),
|
||||
@ -106,6 +108,7 @@ namespace Qwen {
|
||||
} else {
|
||||
int64_t start = static_cast<int64_t>(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<float>(position);
|
||||
layout.positions.push_back({p, p, p});
|
||||
@ -131,6 +134,8 @@ namespace Qwen {
|
||||
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 {
|
||||
@ -186,10 +191,15 @@ namespace Qwen {
|
||||
q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
|
||||
k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
|
||||
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
|
||||
// Preserve query-first attention evaluation while writing each layer's
|
||||
// prefix before its full-sequence K/V can accumulate across layers.
|
||||
ctx->expand_graph(q);
|
||||
auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
|
||||
auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
|
||||
auto copy = ggml_new_tensor(ctx->ggml_ctx, GGML_TYPE_F32, 4, part->ne);
|
||||
copy = ggml_cpy(ctx->ggml_ctx, part, copy);
|
||||
// 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);
|
||||
@ -198,21 +208,37 @@ namespace Qwen {
|
||||
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;
|
||||
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 = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, true, ctx->flash_attn_enabled);
|
||||
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 = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled);
|
||||
auto out = attend(sq, sk, sv, masks[i]);
|
||||
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
|
||||
}
|
||||
}
|
||||
@ -348,8 +374,35 @@ namespace Qwen {
|
||||
QwenImage21Model model;
|
||||
std::vector<float> pe_data;
|
||||
std::vector<sd::Tensor<float>> mask_data;
|
||||
ggml_type prefix_cache_type = GGML_TYPE_COUNT;
|
||||
bool prefix_cache_enabled = true;
|
||||
bool prefix_cache_disabled = false;
|
||||
bool prefix_cache_auto_f32 = false;
|
||||
|
||||
static bool supports_prefix_cache_type(ggml_type type) {
|
||||
if (type == GGML_TYPE_F32) {
|
||||
return true;
|
||||
}
|
||||
const auto* traits = ggml_get_type_traits(type);
|
||||
if (traits->from_float_ref == nullptr || traits->to_float == nullptr) {
|
||||
return false;
|
||||
}
|
||||
auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
|
||||
if (cpu == nullptr) {
|
||||
return false;
|
||||
}
|
||||
auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
|
||||
ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free);
|
||||
if (ctx == nullptr) {
|
||||
return false;
|
||||
}
|
||||
// Some reference quantizers have no runtime copy support. Query the
|
||||
// device through the registry so dynamically loaded CPU backends work.
|
||||
auto source = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type));
|
||||
auto encoded = ggml_cast(ctx.get(), source, type);
|
||||
auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32);
|
||||
return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded);
|
||||
}
|
||||
|
||||
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
|
||||
: DiffusionModelRunner(backend, prefix, weight_manager),
|
||||
@ -358,6 +411,22 @@ namespace Qwen {
|
||||
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);
|
||||
@ -374,11 +443,9 @@ namespace Qwen {
|
||||
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 != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 ||
|
||||
k->ne[0] != config.head_dim || k->ne[1] != cache.prefix_length ||
|
||||
k->ne[2] != config.hidden_size / config.head_dim || k->ne[3] != 1 ||
|
||||
v->ne[0] != config.head_dim || v->ne[1] != config.hidden_size / config.head_dim ||
|
||||
v->ne[2] != cache.prefix_length || v->ne[3] != 1) {
|
||||
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;
|
||||
}
|
||||
}
|
||||
@ -415,17 +482,42 @@ namespace Qwen {
|
||||
}
|
||||
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);
|
||||
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;
|
||||
}
|
||||
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
|
||||
}
|
||||
bool flash_attn_used = true;
|
||||
auto run = [&](const QwenImage21PrefixCache& active_cache) {
|
||||
flash_attn_used = true;
|
||||
auto checked_cache = active_cache;
|
||||
if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
|
||||
checked_cache.flash_attn_used = &flash_attn_used;
|
||||
}
|
||||
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
|
||||
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
|
||||
pe_data = Rope::embed_nd(std::vector<std::vector<float>>(first_position, layout.positions.end()), 1, 10000.f, config.axes_dim);
|
||||
Rope::Embedding embedding;
|
||||
embedding.ids.assign(first_position, layout.positions.end());
|
||||
const size_t offset = cached ? static_cast<size_t>(layout.prefix_length) : 0;
|
||||
embedding.positions.token_count = embedding.ids.size();
|
||||
for (auto region : layout.rope_layout.images) {
|
||||
if (region.begin >= offset) {
|
||||
region.begin -= offset;
|
||||
embedding.positions.images.push_back(region);
|
||||
}
|
||||
}
|
||||
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) {
|
||||
@ -455,15 +547,28 @@ namespace Qwen {
|
||||
ref_inputs.push_back(make_input(ref));
|
||||
}
|
||||
}
|
||||
auto ctx = get_context();
|
||||
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, active_cache);
|
||||
ref_inputs, pe, layout, masks, checked_cache);
|
||||
if (!flash_attn_used) {
|
||||
return static_cast<ggml_cgraph*>(nullptr);
|
||||
}
|
||||
ggml_build_forward_expand(graph, out);
|
||||
return graph;
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
|
||||
};
|
||||
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.
|
||||
@ -478,7 +583,7 @@ namespace Qwen {
|
||||
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)", extra->prefix_id, layout.prefix_length);
|
||||
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;
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -642,7 +642,7 @@ namespace ZImage {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
|
||||
pe_vec = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
|
||||
pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
@ -651,9 +651,7 @@ namespace ZImage {
|
||||
ref_latents,
|
||||
ref_index_mode,
|
||||
config.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
config.axes_dim);
|
||||
config.axes_dim));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
// LOG_VERBOSE("pos_len %d", pos_len);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
|
||||
@ -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<int>(feat_map.size()); ++feat_idx) {
|
||||
ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
|
||||
if (feat_cache != nullptr) {
|
||||
cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
|
||||
ggml_build_forward_expand(gf, feat_cache);
|
||||
runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -1443,15 +1443,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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -1349,9 +1349,8 @@ bool is_first_stage_model_name(const std::string& name) {
|
||||
}
|
||||
|
||||
static std::string convert_esrgan_tensor_name(std::string name) {
|
||||
static std::unordered_map<std::string, std::string> esrgan_name_map;
|
||||
|
||||
if (esrgan_name_map.empty()) {
|
||||
static const auto esrgan_name_map = [] {
|
||||
std::unordered_map<std::string, std::string> esrgan_name_map;
|
||||
esrgan_name_map["model.0."] = "conv_first.";
|
||||
|
||||
constexpr int max_num_blocks = 64;
|
||||
@ -1377,7 +1376,8 @@ static std::string convert_esrgan_tensor_name(std::string name) {
|
||||
esrgan_name_map["model.7."] = "conv_last.";
|
||||
esrgan_name_map["model.8."] = "conv_hr.";
|
||||
esrgan_name_map["model.10."] = "conv_last.";
|
||||
}
|
||||
return esrgan_name_map;
|
||||
}();
|
||||
|
||||
replace_with_prefix_map(name, esrgan_name_map);
|
||||
return name;
|
||||
|
||||
@ -2441,6 +2441,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
timesteps_tensor,
|
||||
cond,
|
||||
&controls);
|
||||
bool uncond_controls_ready = false;
|
||||
|
||||
static const std::vector<sd::Tensor<float>> empty_ref_latents;
|
||||
bool uncond_without_ref_latents = !img_uncond.empty() &&
|
||||
@ -2530,6 +2531,17 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
return std::move(cached_output);
|
||||
}
|
||||
|
||||
// A re-enabled condition can miss the cache even when the positive pass was reused.
|
||||
if (!uncond_controls_ready && !uncond.empty() &&
|
||||
(&condition == &uncond || &condition == &img_uncond)) {
|
||||
compute_sample_controls(control_image,
|
||||
noised_input,
|
||||
timesteps_tensor,
|
||||
uncond,
|
||||
&controls);
|
||||
uncond_controls_ready = true;
|
||||
}
|
||||
|
||||
for (const auto& extension : generation_extensions) {
|
||||
extension->before_diffusion(diffusion_params, step);
|
||||
}
|
||||
@ -2569,19 +2581,39 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
}
|
||||
}
|
||||
|
||||
float effective_guidance_scale = guidance_schedule.empty()
|
||||
? cfg_scale
|
||||
: guidance_schedule[guidance_schedule.size() - 1 - step];
|
||||
|
||||
float image_guidance_scale = img_cfg_scale;
|
||||
|
||||
constexpr float kEpsilon = 1e-5f;
|
||||
|
||||
bool skip_uncond = false;
|
||||
if (!uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
|
||||
if (!img_uncond.empty()) {
|
||||
skip_uncond = std::abs(image_guidance_scale - effective_guidance_scale) < kEpsilon;
|
||||
} else {
|
||||
skip_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
|
||||
}
|
||||
}
|
||||
|
||||
bool skip_img_uncond = false;
|
||||
if (!img_uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
|
||||
if (!uncond.empty()) {
|
||||
skip_img_uncond = std::abs(image_guidance_scale - 1.0f) < kEpsilon;
|
||||
} else {
|
||||
skip_img_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
|
||||
}
|
||||
}
|
||||
|
||||
cond_out = run_condition(*positive_condition, c_concat_override);
|
||||
if (cond_out.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
if (!uncond.empty()) {
|
||||
if (!step_cache.is_step_skipped()) {
|
||||
compute_sample_controls(control_image,
|
||||
noised_input,
|
||||
timesteps_tensor,
|
||||
uncond,
|
||||
&controls);
|
||||
}
|
||||
if (!skip_uncond) {
|
||||
const std::vector<int>* uncond_skip_layers = nullptr;
|
||||
if (is_skiplayer_step && slg_uncond) {
|
||||
LOG_VERBOSE("Skipping layers at uncond step %d\n", step);
|
||||
@ -2595,8 +2627,13 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
if (uncond_out.empty()) {
|
||||
return {};
|
||||
}
|
||||
} else {
|
||||
step_cache.invalidate_condition(&uncond);
|
||||
}
|
||||
}
|
||||
|
||||
if (!img_uncond.empty()) {
|
||||
if (!skip_img_uncond) {
|
||||
img_uncond_out = run_condition(img_uncond,
|
||||
img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat,
|
||||
nullptr,
|
||||
@ -2605,6 +2642,9 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
if (img_uncond_out.empty()) {
|
||||
return {};
|
||||
}
|
||||
} else {
|
||||
step_cache.invalidate_condition(&img_uncond);
|
||||
}
|
||||
}
|
||||
sd::guidance::GuidanceInput guidance_input;
|
||||
guidance_input.step = step;
|
||||
@ -2613,7 +2653,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
|
||||
guidance_input.pred_uncond = uncond_out.empty() ? nullptr : &uncond_out;
|
||||
guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out;
|
||||
|
||||
sd::guidance::GuiderOutput guided = guidance_schedule.empty() ? primary_guidance.forward(guidance_input, {}) : primary_guidance.forward(guidance_input, {}, guidance_schedule[guidance_schedule.size() - 1 - step]);
|
||||
sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {}, effective_guidance_scale);
|
||||
if (guided.pred.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
@ -729,11 +729,19 @@ 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;
|
||||
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;
|
||||
|
||||
explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
|
||||
FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr)
|
||||
: image_seq_len(image_seq_len) {
|
||||
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<float>(max_image_seq_len - base_image_seq_len);
|
||||
float b = base_shift - m * static_cast<float>(base_image_seq_len);
|
||||
return static_cast<float>(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<FluxScheduler>(image_seq_len, extra_sample_args);
|
||||
scheduler = std::make_shared<FluxScheduler>(image_seq_len, version, extra_sample_args);
|
||||
break;
|
||||
}
|
||||
default:
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -46,6 +46,7 @@ namespace sd_sample {
|
||||
|
||||
bool before_condition(const void* condition, const sd::Tensor<float>& input, sd::Tensor<float>* output);
|
||||
void after_condition(const void* condition, const sd::Tensor<float>& input, const sd::Tensor<float>& output);
|
||||
void invalidate_condition(const void* condition);
|
||||
bool is_step_skipped() const;
|
||||
};
|
||||
|
||||
|
||||
@ -262,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;
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user