mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-08-04 17:20:41 -05:00
refactor: reorganize src model layout (#1615)
This commit is contained in:
leejet
4148
src/core/ggml_extend.hpp
Normal file
4148
src/core/ggml_extend.hpp
Normal file
File diff suppressed because it is too large
Load Diff
656
src/core/ggml_extend_backend.cpp
Normal file
656
src/core/ggml_extend_backend.cpp
Normal file
@@ -0,0 +1,656 @@
|
||||
#include "core/ggml_extend_backend.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <vector>
|
||||
|
||||
#include "core/util.h"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
static std::string trim_copy(const std::string& value) {
|
||||
size_t begin = 0;
|
||||
while (begin < value.size() && std::isspace(static_cast<unsigned char>(value[begin]))) {
|
||||
++begin;
|
||||
}
|
||||
size_t end = value.size();
|
||||
while (end > begin && std::isspace(static_cast<unsigned char>(value[end - 1]))) {
|
||||
--end;
|
||||
}
|
||||
return value.substr(begin, end - begin);
|
||||
}
|
||||
|
||||
static std::string lower_copy(std::string value) {
|
||||
std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
return value;
|
||||
}
|
||||
|
||||
static std::vector<std::string> split_copy(const std::string& value, char delimiter) {
|
||||
std::vector<std::string> parts;
|
||||
std::string part;
|
||||
std::istringstream stream(value);
|
||||
while (std::getline(stream, part, delimiter)) {
|
||||
parts.push_back(part);
|
||||
}
|
||||
return parts;
|
||||
}
|
||||
|
||||
static bool is_default_backend_token(const std::string& name) {
|
||||
const std::string lower = lower_copy(trim_copy(name));
|
||||
return lower.empty() || lower == "default" || lower == "auto";
|
||||
}
|
||||
|
||||
static bool parse_backend_module(const std::string& raw_name, SDBackendModule* module) {
|
||||
std::string name = lower_copy(trim_copy(raw_name));
|
||||
name.erase(std::remove(name.begin(), name.end(), '-'), name.end());
|
||||
name.erase(std::remove(name.begin(), name.end(), '_'), name.end());
|
||||
|
||||
if (name == "diffusion" || name == "model" || name == "unet" || name == "dit") {
|
||||
*module = SDBackendModule::DIFFUSION;
|
||||
return true;
|
||||
}
|
||||
if (name == "te" || name == "clip" || name == "text" || name == "textencoder" || name == "textencoders" || name == "conditioner" || name == "cond" || name == "llm" || name == "t5" || name == "t5xxl") {
|
||||
*module = SDBackendModule::TE;
|
||||
return true;
|
||||
}
|
||||
if (name == "clipvision" || name == "vision") {
|
||||
*module = SDBackendModule::CLIP_VISION;
|
||||
return true;
|
||||
}
|
||||
if (name == "vae" || name == "firststage" || name == "autoencoder" || name == "tae") {
|
||||
*module = SDBackendModule::VAE;
|
||||
return true;
|
||||
}
|
||||
if (name == "controlnet" || name == "control") {
|
||||
*module = SDBackendModule::CONTROL_NET;
|
||||
return true;
|
||||
}
|
||||
if (name == "photomaker" || name == "photomakerid" || name == "pmid" || name == "photo") {
|
||||
*module = SDBackendModule::PHOTOMAKER;
|
||||
return true;
|
||||
}
|
||||
if (name == "upscaler" || name == "esrgan" || name == "hires") {
|
||||
*module = SDBackendModule::UPSCALER;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static std::string module_assignment_name(const SDBackendAssignment& assignment, SDBackendModule module) {
|
||||
auto it = assignment.module_names.find(module);
|
||||
if (it != assignment.module_names.end()) {
|
||||
return it->second;
|
||||
}
|
||||
return assignment.default_name;
|
||||
}
|
||||
|
||||
static std::string backend_cache_key(ggml_backend_t backend) {
|
||||
if (backend == nullptr) {
|
||||
return "";
|
||||
}
|
||||
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
|
||||
if (dev != nullptr) {
|
||||
return lower_copy(ggml_backend_dev_name(dev));
|
||||
}
|
||||
const char* backend_name = ggml_backend_name(backend);
|
||||
return backend_name != nullptr ? lower_copy(backend_name) : "";
|
||||
}
|
||||
|
||||
static std::string resolve_first_device_by_type(enum ggml_backend_dev_type type) {
|
||||
ggml_backend_dev_t dev = ggml_backend_dev_by_type(type);
|
||||
if (dev == nullptr) {
|
||||
return "";
|
||||
}
|
||||
return ggml_backend_dev_name(dev);
|
||||
}
|
||||
|
||||
static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
return tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
|
||||
}
|
||||
|
||||
static bool ggml_backend_tensor_is_host_accessible(const struct ggml_tensor* tensor) {
|
||||
if (tensor == nullptr || tensor->data == nullptr) {
|
||||
return false;
|
||||
}
|
||||
|
||||
ggml_backend_buffer_t buffer = ggml_backend_tensor_buffer(tensor);
|
||||
return buffer == nullptr || ggml_backend_buffer_is_host(buffer);
|
||||
}
|
||||
|
||||
static size_t ggml_backend_tensor_offset(const struct ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
|
||||
return static_cast<size_t>(i0 * tensor->nb[0] + i1 * tensor->nb[1] + i2 * tensor->nb[2] + i3 * tensor->nb[3]);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void ggml_backend_tensor_write_scalar(const struct ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3, T value) {
|
||||
const size_t offset = ggml_backend_tensor_offset(tensor, i0, i1, i2, i3);
|
||||
|
||||
if (ggml_backend_tensor_is_host_accessible(tensor)) {
|
||||
auto* dst = reinterpret_cast<T*>(reinterpret_cast<char*>(tensor->data) + offset);
|
||||
*dst = value;
|
||||
return;
|
||||
}
|
||||
|
||||
ggml_backend_tensor_set(const_cast<struct ggml_tensor*>(tensor), &value, offset, sizeof(T));
|
||||
}
|
||||
|
||||
static void ggml_set_f32_nd(const struct ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3, float value) {
|
||||
switch (tensor->type) {
|
||||
case GGML_TYPE_I8:
|
||||
ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, static_cast<int8_t>(value));
|
||||
break;
|
||||
case GGML_TYPE_I16:
|
||||
ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, static_cast<int16_t>(value));
|
||||
break;
|
||||
case GGML_TYPE_I32:
|
||||
ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, static_cast<int32_t>(value));
|
||||
break;
|
||||
case GGML_TYPE_F16:
|
||||
ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, ggml_fp32_to_fp16(value));
|
||||
break;
|
||||
case GGML_TYPE_BF16:
|
||||
ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, ggml_fp32_to_bf16(value));
|
||||
break;
|
||||
case GGML_TYPE_F32:
|
||||
ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, value);
|
||||
break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value) {
|
||||
if (!ggml_is_contiguous(tensor)) {
|
||||
int64_t id[4] = {0, 0, 0, 0};
|
||||
ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]);
|
||||
ggml_set_f32_nd(tensor, id[0], id[1], id[2], id[3], value);
|
||||
return;
|
||||
}
|
||||
|
||||
switch (tensor->type) {
|
||||
case GGML_TYPE_I8:
|
||||
ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, static_cast<int8_t>(value));
|
||||
break;
|
||||
case GGML_TYPE_I16:
|
||||
ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, static_cast<int16_t>(value));
|
||||
break;
|
||||
case GGML_TYPE_I32:
|
||||
ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, static_cast<int32_t>(value));
|
||||
break;
|
||||
case GGML_TYPE_F16:
|
||||
ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, ggml_fp32_to_fp16(value));
|
||||
break;
|
||||
case GGML_TYPE_BF16:
|
||||
ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, ggml_fp32_to_bf16(value));
|
||||
break;
|
||||
case GGML_TYPE_F32:
|
||||
ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, value);
|
||||
break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
}
|
||||
|
||||
static void ggml_backend_load_all_once() {
|
||||
// If the registry already has devices and the CPU backend is present,
|
||||
// assume either static registration or explicit host-side preloading has
|
||||
// completed and avoid rescanning the default paths.
|
||||
if (ggml_backend_dev_count() > 0 && ggml_backend_reg_by_name("CPU") != nullptr) {
|
||||
return;
|
||||
}
|
||||
// In dynamic-backend mode the backend modules are discovered at runtime,
|
||||
// so we must load them before asking for the CPU backend or its proc table.
|
||||
// If the host preloaded only a subset of backends, allow one default-path
|
||||
// scan so missing modules can still be discovered.
|
||||
static std::once_flag once;
|
||||
std::call_once(once, []() {
|
||||
if (ggml_backend_dev_count() > 0 && ggml_backend_reg_by_name("CPU") != nullptr) {
|
||||
return;
|
||||
}
|
||||
ggml_backend_load_all();
|
||||
});
|
||||
}
|
||||
|
||||
bool sd_backend_is(ggml_backend_t backend, const std::string& name) {
|
||||
if (!backend) {
|
||||
return false;
|
||||
}
|
||||
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
|
||||
if (!dev) {
|
||||
return false;
|
||||
}
|
||||
std::string dev_name = ggml_backend_dev_name(dev);
|
||||
return lower_copy(dev_name).find(lower_copy(name)) != std::string::npos;
|
||||
}
|
||||
|
||||
static std::string get_default_backend_name() {
|
||||
ggml_backend_load_all_once();
|
||||
// should pick the same backend preference as ggml_backend_init_best
|
||||
std::string name = resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_GPU);
|
||||
if (!name.empty()) {
|
||||
return name;
|
||||
}
|
||||
name = resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
|
||||
if (!name.empty()) {
|
||||
return name;
|
||||
}
|
||||
return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
|
||||
}
|
||||
|
||||
static std::string sd_resolve_backend_name(const std::string& name) {
|
||||
ggml_backend_load_all_once();
|
||||
std::string requested = trim_copy(name);
|
||||
std::string lower = lower_copy(requested);
|
||||
|
||||
if (is_default_backend_token(lower)) {
|
||||
return get_default_backend_name();
|
||||
}
|
||||
if (lower == "gpu") {
|
||||
std::string result = resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_GPU);
|
||||
if (!result.empty()) {
|
||||
return result;
|
||||
}
|
||||
return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
|
||||
}
|
||||
|
||||
const size_t device_count = ggml_backend_dev_count();
|
||||
for (size_t i = 0; i < device_count; ++i) {
|
||||
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
|
||||
std::string dev_name = ggml_backend_dev_name(dev);
|
||||
if (lower_copy(dev_name) == lower) {
|
||||
return dev_name;
|
||||
}
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < device_count; ++i) {
|
||||
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
|
||||
std::string dev_name = ggml_backend_dev_name(dev);
|
||||
std::string dev_lower = lower_copy(dev_name);
|
||||
if (dev_lower.rfind(lower, 0) == 0) {
|
||||
return dev_name;
|
||||
}
|
||||
}
|
||||
|
||||
return "";
|
||||
}
|
||||
|
||||
static bool backend_name_exists(const std::string& name) {
|
||||
return !sd_resolve_backend_name(name).empty();
|
||||
}
|
||||
|
||||
static ggml_backend_t init_named_backend(const std::string& name) {
|
||||
ggml_backend_load_all_once();
|
||||
LOG_DEBUG("Initializing backend: %s", name.c_str());
|
||||
if (trim_copy(name).empty()) {
|
||||
return ggml_backend_init_best();
|
||||
}
|
||||
|
||||
std::string resolved = sd_resolve_backend_name(name);
|
||||
if (resolved.empty()) {
|
||||
return nullptr;
|
||||
}
|
||||
return ggml_backend_init_by_name(resolved.c_str(), nullptr);
|
||||
}
|
||||
|
||||
bool sd_backend_is_cpu(ggml_backend_t backend) {
|
||||
if (backend == nullptr) {
|
||||
return false;
|
||||
}
|
||||
auto dev = ggml_backend_get_device(backend);
|
||||
return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
|
||||
}
|
||||
|
||||
ggml_backend_t sd_backend_cpu_init() {
|
||||
ggml_backend_load_all_once();
|
||||
return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
|
||||
}
|
||||
|
||||
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
|
||||
if (backend == nullptr) {
|
||||
return false;
|
||||
}
|
||||
auto dev = ggml_backend_get_device(backend);
|
||||
if (dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
|
||||
auto reg = ggml_backend_dev_backend_reg(dev);
|
||||
auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");
|
||||
if (ggml_backend_set_n_threads_fn != nullptr) {
|
||||
ggml_backend_set_n_threads_fn(backend, n_threads);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
const char* sd_get_system_info() {
|
||||
static std::string cache_info = []() -> std::string {
|
||||
ggml_backend_load_all_once();
|
||||
std::stringstream ss;
|
||||
ss << "System Info: \n";
|
||||
auto dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
|
||||
if (dev != nullptr) {
|
||||
auto reg = ggml_backend_dev_backend_reg(dev);
|
||||
auto ggml_backend_get_features_fn = (ggml_backend_get_features_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features");
|
||||
if (ggml_backend_get_features_fn != nullptr) {
|
||||
ggml_backend_feature* feat = ggml_backend_get_features_fn(reg);
|
||||
while (feat->name && feat->value) {
|
||||
ss << " " << feat->name << " = " << feat->value << " | ";
|
||||
feat++;
|
||||
}
|
||||
} else {
|
||||
LOG_WARN("unable to get CPU features");
|
||||
}
|
||||
} else {
|
||||
LOG_WARN("unable to get CPU features");
|
||||
}
|
||||
return ss.str();
|
||||
}();
|
||||
return cache_info.c_str();
|
||||
}
|
||||
|
||||
static ggml_backend_t sd_get_default_backend() {
|
||||
ggml_backend_load_all_once();
|
||||
static std::once_flag once;
|
||||
std::call_once(once, []() {
|
||||
size_t dev_count = ggml_backend_dev_count();
|
||||
if (dev_count == 0) {
|
||||
LOG_ERROR("No devices found!");
|
||||
} else {
|
||||
LOG_DEBUG("Found %zu backend devices:", dev_count);
|
||||
for (size_t i = 0; i < dev_count; ++i) {
|
||||
auto dev = ggml_backend_dev_get(i);
|
||||
LOG_DEBUG("#%zu: %s", i, ggml_backend_dev_name(dev));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
ggml_backend_t backend = nullptr;
|
||||
const char* SD_VK_DEVICE = getenv("SD_VK_DEVICE");
|
||||
if (SD_VK_DEVICE != nullptr) {
|
||||
std::string sd_vk_device_str = SD_VK_DEVICE;
|
||||
try {
|
||||
unsigned long long device = std::stoull(sd_vk_device_str);
|
||||
std::string vk_device_name = "Vulkan" + std::to_string(device);
|
||||
if (backend_name_exists(vk_device_name)) {
|
||||
LOG_INFO("Selecting %s as main device by env var SD_VK_DEVICE", vk_device_name.c_str());
|
||||
backend = init_named_backend(vk_device_name);
|
||||
if (!backend) {
|
||||
LOG_WARN("Device %s requested by SD_VK_DEVICE failed to init. Falling back to the default device.", vk_device_name.c_str());
|
||||
}
|
||||
} else {
|
||||
LOG_WARN("Device %s requested by SD_VK_DEVICE was not found. Falling back to the default device.", vk_device_name.c_str());
|
||||
}
|
||||
} catch (const std::invalid_argument&) {
|
||||
LOG_WARN("SD_VK_DEVICE environment variable is not a valid integer (%s). Falling back to the default device.", SD_VK_DEVICE);
|
||||
} catch (const std::out_of_range&) {
|
||||
LOG_WARN("SD_VK_DEVICE environment variable value is out of range for `unsigned long long` type (%s). Falling back to the default device.", SD_VK_DEVICE);
|
||||
}
|
||||
}
|
||||
|
||||
if (!backend) {
|
||||
std::string dev_name = get_default_backend_name();
|
||||
backend = init_named_backend(dev_name);
|
||||
if (!backend && !dev_name.empty()) {
|
||||
LOG_WARN("device %s failed to init", dev_name.c_str());
|
||||
}
|
||||
}
|
||||
|
||||
if (!backend) {
|
||||
LOG_WARN("loading CPU backend");
|
||||
backend = sd_backend_cpu_init();
|
||||
}
|
||||
|
||||
if (sd_backend_is_cpu(backend)) {
|
||||
LOG_DEBUG("Using CPU backend");
|
||||
}
|
||||
|
||||
return backend;
|
||||
}
|
||||
|
||||
static bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
|
||||
if (assignment == nullptr) {
|
||||
return false;
|
||||
}
|
||||
|
||||
*assignment = {};
|
||||
const std::string in = trim_copy(spec);
|
||||
if (in.empty()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
for (const std::string& raw_part : split_copy(in, ',')) {
|
||||
const std::string part = trim_copy(raw_part);
|
||||
if (part.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const size_t eq = part.find('=');
|
||||
if (eq == std::string::npos) {
|
||||
assignment->set_default(part);
|
||||
continue;
|
||||
}
|
||||
|
||||
const std::string key = trim_copy(part.substr(0, eq));
|
||||
const std::string value = trim_copy(part.substr(eq + 1));
|
||||
if (key.empty() || value.empty()) {
|
||||
if (error != nullptr) {
|
||||
*error = "invalid backend assignment '" + part + "'";
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
const std::string key_lower = lower_copy(key);
|
||||
if (key_lower == "all" || key_lower == "default" || key_lower == "*") {
|
||||
assignment->set_default(value);
|
||||
continue;
|
||||
}
|
||||
|
||||
SDBackendModule module = SDBackendModule::DIFFUSION;
|
||||
if (!parse_backend_module(key, &module)) {
|
||||
if (error != nullptr) {
|
||||
*error = "unknown backend module '" + key + "'";
|
||||
}
|
||||
return false;
|
||||
}
|
||||
assignment->set_module(module, value);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SDBackendAssignment::empty() const {
|
||||
return default_name.empty() && module_names.empty();
|
||||
}
|
||||
|
||||
std::string SDBackendAssignment::get(SDBackendModule module) const {
|
||||
return module_assignment_name(*this, module);
|
||||
}
|
||||
|
||||
void SDBackendAssignment::set_default(const std::string& name) {
|
||||
default_name = trim_copy(name);
|
||||
}
|
||||
|
||||
void SDBackendAssignment::set_module(SDBackendModule module, const std::string& name) {
|
||||
module_names[module] = trim_copy(name);
|
||||
}
|
||||
|
||||
void SDBackendHandleDeleter::operator()(ggml_backend_t backend) const {
|
||||
ggml_backend_free(backend);
|
||||
}
|
||||
|
||||
SDBackendManager::~SDBackendManager() {
|
||||
reset();
|
||||
}
|
||||
|
||||
void SDBackendManager::reset() {
|
||||
backends_.clear();
|
||||
runtime_assignment_ = {};
|
||||
params_assignment_ = {};
|
||||
}
|
||||
|
||||
ggml_backend_t SDBackendManager::runtime_backend(SDBackendModule module) {
|
||||
return init_cached_backend(runtime_assignment_.get(module));
|
||||
}
|
||||
|
||||
ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
|
||||
std::string name = params_assignment_.get(module);
|
||||
if (name.empty()) {
|
||||
return runtime_backend(module);
|
||||
}
|
||||
return init_cached_backend(name);
|
||||
}
|
||||
|
||||
bool SDBackendManager::runtime_backend_is_cpu(SDBackendModule module) {
|
||||
return sd_backend_is_cpu(runtime_backend(module));
|
||||
}
|
||||
|
||||
bool SDBackendManager::params_backend_is_cpu(SDBackendModule module) {
|
||||
return sd_backend_is_cpu(params_backend(module));
|
||||
}
|
||||
|
||||
bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
|
||||
ggml_backend_t backend = runtime_backend(module);
|
||||
if (backend == nullptr) {
|
||||
return false;
|
||||
}
|
||||
if (sd_backend_is_cpu(backend)) {
|
||||
return true;
|
||||
}
|
||||
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
|
||||
if (dev == nullptr) {
|
||||
return false;
|
||||
}
|
||||
ggml_backend_dev_props props;
|
||||
ggml_backend_dev_get_props(dev, &props);
|
||||
return props.caps.buffer_from_host_ptr;
|
||||
}
|
||||
|
||||
bool SDBackendManager::init(const char* backend_spec,
|
||||
const char* params_backend_spec,
|
||||
bool offload_params_to_cpu,
|
||||
bool keep_clip_on_cpu,
|
||||
bool keep_vae_on_cpu,
|
||||
bool keep_control_net_on_cpu,
|
||||
std::string* error) {
|
||||
reset();
|
||||
|
||||
if (!sd_parse_backend_assignment(SAFE_STR(backend_spec), &runtime_assignment_, error)) {
|
||||
return false;
|
||||
}
|
||||
if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), ¶ms_assignment_, error)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (runtime_assignment_.empty()) {
|
||||
if (keep_clip_on_cpu) {
|
||||
runtime_assignment_.set_module(SDBackendModule::TE, "cpu");
|
||||
}
|
||||
if (keep_vae_on_cpu) {
|
||||
runtime_assignment_.set_module(SDBackendModule::VAE, "cpu");
|
||||
}
|
||||
if (keep_control_net_on_cpu) {
|
||||
runtime_assignment_.set_module(SDBackendModule::CONTROL_NET, "cpu");
|
||||
}
|
||||
}
|
||||
|
||||
if (params_assignment_.empty() && offload_params_to_cpu) {
|
||||
params_assignment_.set_default("cpu");
|
||||
}
|
||||
|
||||
return validate(error);
|
||||
}
|
||||
|
||||
bool SDBackendManager::validate(std::string* error) const {
|
||||
auto validate_name = [&](const std::string& name) -> bool {
|
||||
if (is_default_backend_token(name)) {
|
||||
return true;
|
||||
}
|
||||
if (!sd_resolve_backend_name(name).empty()) {
|
||||
return true;
|
||||
}
|
||||
if (error != nullptr) {
|
||||
*error = "backend '" + name + "' was not found";
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
if (!validate_name(runtime_assignment_.default_name) ||
|
||||
!validate_name(params_assignment_.default_name)) {
|
||||
return false;
|
||||
}
|
||||
for (const auto& kv : runtime_assignment_.module_names) {
|
||||
if (!validate_name(kv.second)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
for (const auto& kv : params_assignment_.module_names) {
|
||||
if (!validate_name(kv.second)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_backend_t SDBackendManager::init_cached_backend(const std::string& name) {
|
||||
std::string resolved = sd_resolve_backend_name(name);
|
||||
std::string key = lower_copy(resolved);
|
||||
ggml_backend_t backend = nullptr;
|
||||
|
||||
if (!key.empty()) {
|
||||
auto it = backends_.find(key);
|
||||
if (it != backends_.end()) {
|
||||
return it->second.get();
|
||||
}
|
||||
} else if (!is_default_backend_token(name)) {
|
||||
LOG_ERROR("backend '%s' was not found", name.c_str());
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
backend = is_default_backend_token(name) ? sd_get_default_backend() : init_named_backend(resolved);
|
||||
if (backend == nullptr) {
|
||||
LOG_ERROR("failed to initialize backend '%s'", name.c_str());
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
std::string actual_key = backend_cache_key(backend);
|
||||
if (actual_key.empty()) {
|
||||
actual_key = !key.empty() ? key : lower_copy(trim_copy(name));
|
||||
}
|
||||
|
||||
auto it = backends_.find(actual_key);
|
||||
if (it != backends_.end()) {
|
||||
ggml_backend_free(backend);
|
||||
return it->second.get();
|
||||
}
|
||||
|
||||
SDBackendHandle handle(backend);
|
||||
backends_.emplace(actual_key, std::move(handle));
|
||||
return backend;
|
||||
}
|
||||
|
||||
const char* sd_backend_module_name(SDBackendModule module) {
|
||||
switch (module) {
|
||||
case SDBackendModule::DIFFUSION:
|
||||
return "diffusion";
|
||||
case SDBackendModule::TE:
|
||||
return "te";
|
||||
case SDBackendModule::CLIP_VISION:
|
||||
return "clip_vision";
|
||||
case SDBackendModule::VAE:
|
||||
return "vae";
|
||||
case SDBackendModule::CONTROL_NET:
|
||||
return "controlnet";
|
||||
case SDBackendModule::PHOTOMAKER:
|
||||
return "photomaker";
|
||||
case SDBackendModule::UPSCALER:
|
||||
return "upscaler";
|
||||
}
|
||||
return "unknown";
|
||||
}
|
||||
79
src/core/ggml_extend_backend.h
Normal file
79
src/core/ggml_extend_backend.h
Normal file
@@ -0,0 +1,79 @@
|
||||
#ifndef __SD_CORE_GGML_EXTEND_BACKEND_H__
|
||||
#define __SD_CORE_GGML_EXTEND_BACKEND_H__
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
|
||||
enum class SDBackendModule {
|
||||
DIFFUSION,
|
||||
TE,
|
||||
CLIP_VISION,
|
||||
VAE,
|
||||
CONTROL_NET,
|
||||
PHOTOMAKER,
|
||||
UPSCALER,
|
||||
};
|
||||
|
||||
struct SDBackendAssignment {
|
||||
std::string default_name;
|
||||
std::unordered_map<SDBackendModule, std::string> module_names;
|
||||
|
||||
bool empty() const;
|
||||
std::string get(SDBackendModule module) const;
|
||||
void set_default(const std::string& name);
|
||||
void set_module(SDBackendModule module, const std::string& name);
|
||||
};
|
||||
|
||||
struct SDBackendHandleDeleter {
|
||||
void operator()(ggml_backend_t backend) const;
|
||||
};
|
||||
|
||||
using SDBackendHandle = std::unique_ptr<struct ggml_backend, SDBackendHandleDeleter>;
|
||||
|
||||
class SDBackendManager {
|
||||
private:
|
||||
SDBackendAssignment runtime_assignment_;
|
||||
SDBackendAssignment params_assignment_;
|
||||
std::unordered_map<std::string, SDBackendHandle> backends_;
|
||||
|
||||
public:
|
||||
SDBackendManager() = default;
|
||||
~SDBackendManager();
|
||||
|
||||
SDBackendManager(const SDBackendManager&) = delete;
|
||||
SDBackendManager& operator=(const SDBackendManager&) = delete;
|
||||
|
||||
bool init(const char* backend_spec,
|
||||
const char* params_backend_spec,
|
||||
bool offload_params_to_cpu,
|
||||
bool keep_clip_on_cpu,
|
||||
bool keep_vae_on_cpu,
|
||||
bool keep_control_net_on_cpu,
|
||||
std::string* error);
|
||||
void reset();
|
||||
|
||||
ggml_backend_t runtime_backend(SDBackendModule module);
|
||||
ggml_backend_t params_backend(SDBackendModule module);
|
||||
|
||||
bool runtime_backend_is_cpu(SDBackendModule module);
|
||||
bool params_backend_is_cpu(SDBackendModule module);
|
||||
bool runtime_backend_supports_host_buffer(SDBackendModule module);
|
||||
|
||||
private:
|
||||
bool validate(std::string* error) const;
|
||||
ggml_backend_t init_cached_backend(const std::string& name);
|
||||
};
|
||||
|
||||
bool sd_backend_is(ggml_backend_t backend, const std::string& name);
|
||||
bool sd_backend_is_cpu(ggml_backend_t backend);
|
||||
ggml_backend_t sd_backend_cpu_init();
|
||||
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
|
||||
const char* sd_backend_module_name(SDBackendModule module);
|
||||
void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
|
||||
#endif // __SD_CORE_GGML_EXTEND_BACKEND_H__
|
||||
806
src/core/ggml_graph_cut.cpp
Normal file
806
src/core/ggml_graph_cut.cpp
Normal file
@@ -0,0 +1,806 @@
|
||||
#include "core/ggml_graph_cut.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <sstream>
|
||||
#include <stack>
|
||||
#include <unordered_map>
|
||||
|
||||
#include "core/util.h"
|
||||
#include "ggml-alloc.h"
|
||||
#include "ggml-backend.h"
|
||||
|
||||
#include "ggml/src/ggml-impl.h"
|
||||
|
||||
namespace sd::ggml_graph_cut {
|
||||
|
||||
static constexpr double MAX_VRAM_BYTES_PER_GIB = 1024.0 * 1024.0 * 1024.0;
|
||||
|
||||
static std::string graph_cut_tensor_display_name(const ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return "<null>";
|
||||
}
|
||||
if (tensor->name[0] != '\0') {
|
||||
return tensor->name;
|
||||
}
|
||||
return sd_format("<tensor@%p>", (const void*)tensor);
|
||||
}
|
||||
|
||||
static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
GGML_ASSERT(tensor != nullptr);
|
||||
for (int i = 0; i < gf->n_leafs; ++i) {
|
||||
if (gf->leafs[i] == tensor) {
|
||||
return i;
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
static bool is_params_tensor(const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return params_tensor_set.find(tensor) != params_tensor_set.end();
|
||||
}
|
||||
|
||||
static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
if (name == nullptr || name[0] == '\0') {
|
||||
return -1;
|
||||
}
|
||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
||||
for (int i = 0; i < n_nodes; ++i) {
|
||||
ggml_tensor* node = ggml_graph_node(gf, i);
|
||||
if (node != nullptr && std::strcmp(node->name, name) == 0) {
|
||||
return i;
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
static Plan::InputShape input_shape(const ggml_tensor* tensor) {
|
||||
Plan::InputShape shape;
|
||||
if (tensor == nullptr) {
|
||||
return shape;
|
||||
}
|
||||
shape.type = tensor->type;
|
||||
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
|
||||
shape.ne[static_cast<size_t>(i)] = tensor->ne[i];
|
||||
}
|
||||
return shape;
|
||||
}
|
||||
|
||||
static size_t graph_cut_segment_vram_bytes(const Segment& segment) {
|
||||
return segment.compute_buffer_size +
|
||||
segment.input_param_bytes +
|
||||
segment.input_previous_cut_bytes +
|
||||
segment.output_bytes;
|
||||
}
|
||||
|
||||
size_t max_vram_gib_to_bytes(float max_vram) {
|
||||
if (max_vram <= 0.f) {
|
||||
return 0;
|
||||
}
|
||||
return static_cast<size_t>(static_cast<double>(max_vram) * MAX_VRAM_BYTES_PER_GIB);
|
||||
}
|
||||
|
||||
static float max_vram_bytes_to_gib(size_t max_vram_bytes) {
|
||||
return static_cast<float>(static_cast<double>(max_vram_bytes) / MAX_VRAM_BYTES_PER_GIB);
|
||||
}
|
||||
|
||||
static size_t resolve_auto_max_vram_bytes(float spare_vram, ggml_backend_t backend) {
|
||||
if (backend == nullptr) {
|
||||
LOG_WARN("--max-vram < 0 requested, but no backend is available; disabling graph splitting");
|
||||
return 0;
|
||||
}
|
||||
|
||||
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
|
||||
if (dev == nullptr) {
|
||||
LOG_WARN("--max-vram < 0 requested, but no backend device is available; disabling graph splitting");
|
||||
return 0;
|
||||
}
|
||||
if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
|
||||
LOG_WARN("--max-vram < 0 requested, but the main backend is CPU; disabling graph splitting");
|
||||
return 0;
|
||||
}
|
||||
|
||||
size_t free_vram = 0;
|
||||
size_t total_vram = 0;
|
||||
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
|
||||
size_t spare_bytes = static_cast<size_t>(MAX_VRAM_BYTES_PER_GIB * spare_vram);
|
||||
|
||||
if (free_vram <= spare_bytes) {
|
||||
LOG_WARN("--max-vram < 0 requested, but free VRAM is %.2f GiB; reserving %.2f GiB leaves no graph budget",
|
||||
free_vram / MAX_VRAM_BYTES_PER_GIB, spare_vram);
|
||||
return 0;
|
||||
}
|
||||
|
||||
const size_t max_vram_bytes = free_vram - spare_bytes;
|
||||
LOG_INFO("--max-vram < 0 auto-detected %.2f GiB free VRAM (%.2f GiB total), reserving %.2f GiB; using %.2f GiB",
|
||||
free_vram / MAX_VRAM_BYTES_PER_GIB,
|
||||
total_vram / MAX_VRAM_BYTES_PER_GIB,
|
||||
spare_vram,
|
||||
max_vram_bytes / MAX_VRAM_BYTES_PER_GIB);
|
||||
return max_vram_bytes;
|
||||
}
|
||||
|
||||
float resolve_max_vram_gib(float max_vram, ggml_backend_t backend) {
|
||||
if (max_vram >= 0.f) {
|
||||
return max_vram;
|
||||
}
|
||||
return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
|
||||
}
|
||||
|
||||
static Segment make_segment_seed(const Plan& plan,
|
||||
size_t start_segment_index,
|
||||
size_t end_segment_index) {
|
||||
GGML_ASSERT(start_segment_index < plan.segments.size());
|
||||
GGML_ASSERT(end_segment_index < plan.segments.size());
|
||||
GGML_ASSERT(start_segment_index <= end_segment_index);
|
||||
|
||||
Segment seed;
|
||||
const auto& start_segment = plan.segments[start_segment_index];
|
||||
const auto& target_segment = plan.segments[end_segment_index];
|
||||
std::unordered_set<int> seen_output_node_indices;
|
||||
for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
|
||||
for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
|
||||
if (seen_output_node_indices.insert(output_node_index).second) {
|
||||
seed.output_node_indices.push_back(output_node_index);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (start_segment_index == end_segment_index) {
|
||||
seed.group_name = target_segment.group_name;
|
||||
} else {
|
||||
seed.group_name = sd_format("%s..%s",
|
||||
start_segment.group_name.c_str(),
|
||||
target_segment.group_name.c_str());
|
||||
}
|
||||
return seed;
|
||||
}
|
||||
|
||||
static void build_segment(ggml_cgraph* gf,
|
||||
Plan& plan,
|
||||
Segment& segment,
|
||||
const std::unordered_map<const ggml_tensor*, int>& producer_index,
|
||||
std::unordered_set<int>& available_cut_output_node_indices,
|
||||
ggml_backend_t backend,
|
||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const char* log_desc) {
|
||||
std::set<int> internal_nodes;
|
||||
std::unordered_set<const ggml_tensor*> input_seen;
|
||||
std::vector<Segment::InputRef> input_refs;
|
||||
|
||||
std::stack<ggml_tensor*> work_stack;
|
||||
for (int output_node_index : segment.output_node_indices) {
|
||||
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
|
||||
if (output != nullptr) {
|
||||
work_stack.push(output);
|
||||
}
|
||||
}
|
||||
|
||||
while (!work_stack.empty()) {
|
||||
ggml_tensor* tensor = work_stack.top();
|
||||
work_stack.pop();
|
||||
|
||||
if (tensor == nullptr) {
|
||||
continue;
|
||||
}
|
||||
|
||||
auto producer_it = producer_index.find(tensor);
|
||||
if (producer_it == producer_index.end()) {
|
||||
if (input_seen.insert(tensor).second) {
|
||||
Segment::InputRef input_ref;
|
||||
input_ref.type = is_params_tensor(params_tensor_set, tensor) ? Segment::INPUT_PARAM : Segment::INPUT_EXTERNAL;
|
||||
input_ref.display_name = graph_cut_tensor_display_name(tensor);
|
||||
input_ref.leaf_index = graph_leaf_index(gf, tensor);
|
||||
input_refs.push_back(std::move(input_ref));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
int node_idx = producer_it->second;
|
||||
if (available_cut_output_node_indices.find(node_idx) != available_cut_output_node_indices.end()) {
|
||||
if (input_seen.insert(tensor).second) {
|
||||
Segment::InputRef input_ref;
|
||||
input_ref.type = Segment::INPUT_PREVIOUS_CUT;
|
||||
input_ref.display_name = graph_cut_tensor_display_name(tensor);
|
||||
input_ref.node_index = node_idx;
|
||||
input_refs.push_back(std::move(input_ref));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!internal_nodes.insert(node_idx).second) {
|
||||
continue;
|
||||
}
|
||||
|
||||
ggml_tensor* node = ggml_graph_node(gf, node_idx);
|
||||
for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
|
||||
if (node->src[src_idx] != nullptr) {
|
||||
work_stack.push(node->src[src_idx]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!internal_nodes.empty()) {
|
||||
segment.internal_node_indices.assign(internal_nodes.begin(), internal_nodes.end());
|
||||
}
|
||||
|
||||
std::sort(input_refs.begin(),
|
||||
input_refs.end(),
|
||||
[](const Segment::InputRef& a, const Segment::InputRef& b) {
|
||||
if (a.type != b.type) {
|
||||
return a.type < b.type;
|
||||
}
|
||||
return a.display_name < b.display_name;
|
||||
});
|
||||
segment.input_refs = input_refs;
|
||||
for (const auto& input : input_refs) {
|
||||
ggml_tensor* current_input = input_tensor(gf, input);
|
||||
size_t tensor_bytes = current_input == nullptr
|
||||
? 0
|
||||
: (input.type == Segment::INPUT_PREVIOUS_CUT
|
||||
? cache_tensor_bytes(current_input)
|
||||
: ggml_nbytes(current_input));
|
||||
switch (input.type) {
|
||||
case Segment::INPUT_PREVIOUS_CUT:
|
||||
segment.input_previous_cut_bytes += tensor_bytes;
|
||||
break;
|
||||
case Segment::INPUT_PARAM:
|
||||
segment.input_param_bytes += tensor_bytes;
|
||||
break;
|
||||
case Segment::INPUT_EXTERNAL:
|
||||
default:
|
||||
segment.input_external_bytes += tensor_bytes;
|
||||
break;
|
||||
}
|
||||
}
|
||||
for (int output_node_index : segment.output_node_indices) {
|
||||
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
|
||||
segment.output_bytes += cache_tensor_bytes(output);
|
||||
}
|
||||
segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc);
|
||||
|
||||
for (int output_node_index : segment.output_node_indices) {
|
||||
available_cut_output_node_indices.insert(output_node_index);
|
||||
}
|
||||
plan.segments.push_back(std::move(segment));
|
||||
}
|
||||
|
||||
bool is_graph_cut_tensor(const ggml_tensor* tensor) {
|
||||
if (tensor == nullptr || tensor->name[0] == '\0') {
|
||||
return false;
|
||||
}
|
||||
return std::strncmp(tensor->name, GGML_RUNNER_CUT_PREFIX, std::strlen(GGML_RUNNER_CUT_PREFIX)) == 0;
|
||||
}
|
||||
|
||||
std::string make_graph_cut_name(const std::string& group, const std::string& output) {
|
||||
return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output;
|
||||
}
|
||||
|
||||
void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output) {
|
||||
if (tensor == nullptr) {
|
||||
return;
|
||||
}
|
||||
auto name = make_graph_cut_name(group, output);
|
||||
ggml_set_name(tensor, name.c_str());
|
||||
}
|
||||
|
||||
int leaf_count(ggml_cgraph* gf) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
return gf->n_leafs;
|
||||
}
|
||||
|
||||
ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
if (leaf_index < 0 || leaf_index >= gf->n_leafs) {
|
||||
return nullptr;
|
||||
}
|
||||
return gf->leafs[leaf_index];
|
||||
}
|
||||
|
||||
ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
return tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
|
||||
}
|
||||
|
||||
ggml_tensor* cache_source_tensor(ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
if (tensor_buffer(tensor) == nullptr && tensor->src[0] != nullptr &&
|
||||
ggml_nelements(tensor->src[0]) == ggml_nelements(tensor) &&
|
||||
ggml_nbytes(tensor->src[0]) == ggml_nbytes(tensor)) {
|
||||
return cache_source_tensor(tensor->src[0]);
|
||||
}
|
||||
return tensor->view_src ? tensor->view_src : tensor;
|
||||
}
|
||||
|
||||
size_t cache_tensor_bytes(const ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return 0;
|
||||
}
|
||||
const ggml_tensor* cache_src = tensor->view_src ? tensor->view_src : tensor;
|
||||
return ggml_nbytes(cache_src);
|
||||
}
|
||||
|
||||
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) {
|
||||
return false;
|
||||
}
|
||||
for (const auto& input_shape_ref : plan.input_shapes) {
|
||||
if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) {
|
||||
return false;
|
||||
}
|
||||
ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index];
|
||||
if (leaf == nullptr || input_shape_ref.type != leaf->type) {
|
||||
return false;
|
||||
}
|
||||
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
|
||||
if (input_shape_ref.ne[static_cast<size_t>(d)] != leaf->ne[d]) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
if (output_index >= segment.output_node_indices.size()) {
|
||||
return nullptr;
|
||||
}
|
||||
int node_index = segment.output_node_indices[output_index];
|
||||
if (node_index < 0 || node_index >= ggml_graph_n_nodes(gf)) {
|
||||
return nullptr;
|
||||
}
|
||||
return ggml_graph_node(gf, node_index);
|
||||
}
|
||||
|
||||
ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
if (input_ref.type == Segment::INPUT_PREVIOUS_CUT) {
|
||||
if (input_ref.node_index < 0 || input_ref.node_index >= ggml_graph_n_nodes(gf)) {
|
||||
return nullptr;
|
||||
}
|
||||
return ggml_graph_node(gf, input_ref.node_index);
|
||||
}
|
||||
if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) {
|
||||
return nullptr;
|
||||
}
|
||||
return leaf_tensor(gf, input_ref.leaf_index);
|
||||
}
|
||||
|
||||
std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
std::vector<ggml_tensor*> tensors;
|
||||
std::unordered_set<ggml_tensor*> seen_tensors;
|
||||
tensors.reserve(segment.input_refs.size());
|
||||
seen_tensors.reserve(segment.input_refs.size());
|
||||
for (const auto& input_ref : segment.input_refs) {
|
||||
if (input_ref.type != Segment::INPUT_PARAM) {
|
||||
continue;
|
||||
}
|
||||
ggml_tensor* tensor = input_tensor(gf, input_ref);
|
||||
if (tensor == nullptr) {
|
||||
continue;
|
||||
}
|
||||
if (seen_tensors.insert(tensor).second) {
|
||||
tensors.push_back(tensor);
|
||||
}
|
||||
}
|
||||
return tensors;
|
||||
}
|
||||
|
||||
std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc) {
|
||||
std::vector<ggml_tensor*> tensors = param_tensors(gf, segment);
|
||||
std::vector<ggml_tensor*> filtered_tensors;
|
||||
filtered_tensors.reserve(tensors.size());
|
||||
for (ggml_tensor* tensor : tensors) {
|
||||
if (tensor_buffer(tensor) == nullptr) {
|
||||
LOG_WARN("%s graph cut skipping param input without buffer: segment=%s tensor=%s",
|
||||
log_desc == nullptr ? "unknown" : log_desc,
|
||||
segment.group_name.c_str(),
|
||||
tensor->name);
|
||||
continue;
|
||||
}
|
||||
filtered_tensors.push_back(tensor);
|
||||
}
|
||||
return filtered_tensors;
|
||||
}
|
||||
|
||||
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
|
||||
const Plan& plan,
|
||||
size_t current_segment_index) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
std::unordered_set<std::string> future_input_names;
|
||||
for (size_t seg_idx = current_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
|
||||
const auto& segment = plan.segments[seg_idx];
|
||||
for (const auto& input_ref : segment.input_refs) {
|
||||
if (input_ref.type != Segment::INPUT_PREVIOUS_CUT) {
|
||||
continue;
|
||||
}
|
||||
ggml_tensor* current_input = input_tensor(gf, input_ref);
|
||||
if (current_input != nullptr && current_input->name[0] != '\0') {
|
||||
future_input_names.insert(current_input->name);
|
||||
}
|
||||
}
|
||||
}
|
||||
return future_input_names;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
|
||||
const Segment& segment,
|
||||
ggml_context** graph_ctx_out) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
GGML_ASSERT(graph_ctx_out != nullptr);
|
||||
|
||||
const size_t graph_size = segment.internal_node_indices.size() + segment.input_refs.size() + 8;
|
||||
ggml_init_params params = {
|
||||
/*.mem_size =*/ggml_graph_overhead_custom(graph_size, false) + 1024,
|
||||
/*.mem_buffer =*/nullptr,
|
||||
/*.no_alloc =*/true,
|
||||
};
|
||||
ggml_context* graph_ctx = ggml_init(params);
|
||||
GGML_ASSERT(graph_ctx != nullptr);
|
||||
ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false);
|
||||
GGML_ASSERT(segment_graph != nullptr);
|
||||
|
||||
for (const auto& input : segment.input_refs) {
|
||||
ggml_tensor* current_input = input_tensor(gf, input);
|
||||
if (current_input == nullptr) {
|
||||
continue;
|
||||
}
|
||||
GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
|
||||
segment_graph->leafs[segment_graph->n_leafs++] = current_input;
|
||||
}
|
||||
|
||||
for (int output_node_index : segment.output_node_indices) {
|
||||
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
|
||||
if (output == nullptr) {
|
||||
continue;
|
||||
}
|
||||
ggml_set_output(output);
|
||||
}
|
||||
for (int node_idx : segment.internal_node_indices) {
|
||||
ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
|
||||
}
|
||||
*graph_ctx_out = graph_ctx;
|
||||
return segment_graph;
|
||||
}
|
||||
|
||||
size_t measure_segment_compute_buffer(ggml_backend_t backend,
|
||||
ggml_cgraph* gf,
|
||||
const Segment& segment,
|
||||
const char* log_desc) {
|
||||
GGML_ASSERT(backend != nullptr);
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
if (segment.internal_node_indices.empty()) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
ggml_context* graph_ctx = nullptr;
|
||||
ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx);
|
||||
ggml_gallocr_t allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
|
||||
|
||||
size_t sizes[1] = {0};
|
||||
ggml_gallocr_reserve_n_size(
|
||||
allocr,
|
||||
segment_graph,
|
||||
nullptr,
|
||||
nullptr,
|
||||
sizes);
|
||||
size_t buffer_size = sizes[0];
|
||||
|
||||
ggml_gallocr_free(allocr);
|
||||
ggml_free(graph_ctx);
|
||||
return buffer_size;
|
||||
}
|
||||
|
||||
Plan build_plan(ggml_backend_t backend,
|
||||
ggml_cgraph* gf,
|
||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const char* log_desc) {
|
||||
GGML_ASSERT(backend != nullptr);
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
Plan plan;
|
||||
plan.available = true;
|
||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
||||
if (n_nodes <= 0) {
|
||||
return plan;
|
||||
}
|
||||
plan.n_nodes = n_nodes;
|
||||
plan.n_leafs = gf->n_leafs;
|
||||
for (int i = 0; i < gf->n_leafs; ++i) {
|
||||
ggml_tensor* leaf = gf->leafs[i];
|
||||
if (is_params_tensor(params_tensor_set, leaf)) {
|
||||
continue;
|
||||
}
|
||||
auto shape = input_shape(leaf);
|
||||
shape.leaf_index = i;
|
||||
plan.input_shapes.push_back(shape);
|
||||
}
|
||||
|
||||
std::unordered_map<const ggml_tensor*, int> producer_index;
|
||||
producer_index.reserve(static_cast<size_t>(n_nodes));
|
||||
for (int i = 0; i < n_nodes; ++i) {
|
||||
producer_index[ggml_graph_node(gf, i)] = i;
|
||||
}
|
||||
|
||||
std::vector<Segment> grouped_segments;
|
||||
std::unordered_map<std::string, size_t> group_to_segment;
|
||||
for (int i = 0; i < n_nodes; ++i) {
|
||||
ggml_tensor* node = ggml_graph_node(gf, i);
|
||||
if (!is_graph_cut_tensor(node)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
plan.has_cuts = true;
|
||||
std::string full_name(node->name);
|
||||
std::string payload = full_name.substr(std::strlen(GGML_RUNNER_CUT_PREFIX));
|
||||
size_t sep = payload.find('|');
|
||||
std::string group = sep == std::string::npos ? payload : payload.substr(0, sep);
|
||||
|
||||
auto it = group_to_segment.find(group);
|
||||
if (it == group_to_segment.end()) {
|
||||
Segment segment;
|
||||
segment.group_name = group;
|
||||
segment.output_node_indices.push_back(i);
|
||||
group_to_segment[group] = grouped_segments.size();
|
||||
grouped_segments.push_back(std::move(segment));
|
||||
} else {
|
||||
auto& segment = grouped_segments[it->second];
|
||||
segment.output_node_indices.push_back(i);
|
||||
}
|
||||
}
|
||||
|
||||
if (!plan.has_cuts) {
|
||||
return plan;
|
||||
}
|
||||
|
||||
std::unordered_set<int> available_cut_output_node_indices;
|
||||
available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
|
||||
for (auto& segment : grouped_segments) {
|
||||
build_segment(gf,
|
||||
plan,
|
||||
segment,
|
||||
producer_index,
|
||||
available_cut_output_node_indices,
|
||||
backend,
|
||||
params_tensor_set,
|
||||
log_desc);
|
||||
}
|
||||
|
||||
int final_output_index = graph_node_index_by_name(gf, "ggml_runner_final_result_tensor");
|
||||
if (final_output_index < 0) {
|
||||
final_output_index = n_nodes - 1;
|
||||
}
|
||||
ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr;
|
||||
if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) {
|
||||
Segment final_segment;
|
||||
final_segment.group_name = "ggml_runner.final";
|
||||
final_segment.output_node_indices.push_back(final_output_index);
|
||||
build_segment(gf,
|
||||
plan,
|
||||
final_segment,
|
||||
producer_index,
|
||||
available_cut_output_node_indices,
|
||||
backend,
|
||||
params_tensor_set,
|
||||
log_desc);
|
||||
}
|
||||
|
||||
return plan;
|
||||
}
|
||||
|
||||
Plan apply_max_vram_budget(ggml_cgraph* gf,
|
||||
const Plan& base_plan,
|
||||
size_t max_graph_vram_bytes,
|
||||
ggml_backend_t backend,
|
||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const char* log_desc) {
|
||||
GGML_ASSERT(backend != nullptr);
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
int64_t t_budget_begin = ggml_time_ms();
|
||||
if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) {
|
||||
return base_plan;
|
||||
}
|
||||
|
||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
||||
std::unordered_map<const ggml_tensor*, int> producer_index;
|
||||
producer_index.reserve(static_cast<size_t>(n_nodes));
|
||||
for (int i = 0; i < n_nodes; ++i) {
|
||||
producer_index[ggml_graph_node(gf, i)] = i;
|
||||
}
|
||||
|
||||
Plan merged_plan;
|
||||
merged_plan.available = true;
|
||||
merged_plan.has_cuts = base_plan.has_cuts;
|
||||
merged_plan.valid = base_plan.valid;
|
||||
merged_plan.n_nodes = base_plan.n_nodes;
|
||||
merged_plan.n_leafs = base_plan.n_leafs;
|
||||
|
||||
std::unordered_set<int> available_cut_output_node_indices;
|
||||
available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
|
||||
|
||||
size_t start_segment_index = 0;
|
||||
while (start_segment_index < base_plan.segments.size()) {
|
||||
Plan single_plan;
|
||||
auto single_available_cut_output_node_indices = available_cut_output_node_indices;
|
||||
auto single_seed = make_segment_seed(base_plan,
|
||||
start_segment_index,
|
||||
start_segment_index);
|
||||
build_segment(gf,
|
||||
single_plan,
|
||||
single_seed,
|
||||
producer_index,
|
||||
single_available_cut_output_node_indices,
|
||||
backend,
|
||||
params_tensor_set,
|
||||
log_desc);
|
||||
GGML_ASSERT(!single_plan.segments.empty());
|
||||
|
||||
size_t best_end_segment_index = start_segment_index;
|
||||
bool can_merge_next_segment = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes;
|
||||
|
||||
while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) {
|
||||
const size_t next_end_segment_index = best_end_segment_index + 1;
|
||||
Plan candidate_plan;
|
||||
auto candidate_available_cut_output_node_indices = available_cut_output_node_indices;
|
||||
auto candidate_seed = make_segment_seed(base_plan,
|
||||
start_segment_index,
|
||||
next_end_segment_index);
|
||||
build_segment(gf,
|
||||
candidate_plan,
|
||||
candidate_seed,
|
||||
producer_index,
|
||||
candidate_available_cut_output_node_indices,
|
||||
backend,
|
||||
params_tensor_set,
|
||||
log_desc);
|
||||
GGML_ASSERT(!candidate_plan.segments.empty());
|
||||
|
||||
const auto& candidate_segment = candidate_plan.segments.back();
|
||||
if (graph_cut_segment_vram_bytes(candidate_segment) > max_graph_vram_bytes) {
|
||||
break;
|
||||
}
|
||||
|
||||
best_end_segment_index = next_end_segment_index;
|
||||
}
|
||||
|
||||
auto best_seed = make_segment_seed(base_plan,
|
||||
start_segment_index,
|
||||
best_end_segment_index);
|
||||
build_segment(gf,
|
||||
merged_plan,
|
||||
best_seed,
|
||||
producer_index,
|
||||
available_cut_output_node_indices,
|
||||
backend,
|
||||
params_tensor_set,
|
||||
log_desc);
|
||||
start_segment_index = best_end_segment_index + 1;
|
||||
}
|
||||
|
||||
if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) {
|
||||
LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments",
|
||||
log_desc,
|
||||
max_graph_vram_bytes / 1024.0 / 1024.0,
|
||||
base_plan.segments.size(),
|
||||
merged_plan.segments.size());
|
||||
}
|
||||
|
||||
if (log_desc != nullptr) {
|
||||
LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms",
|
||||
log_desc,
|
||||
ggml_time_ms() - t_budget_begin);
|
||||
}
|
||||
|
||||
return merged_plan;
|
||||
}
|
||||
|
||||
Plan resolve_plan(ggml_backend_t backend,
|
||||
ggml_cgraph* gf,
|
||||
PlanCache* cache,
|
||||
size_t max_graph_vram_bytes,
|
||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const char* log_desc) {
|
||||
GGML_ASSERT(backend != nullptr);
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
GGML_ASSERT(cache != nullptr);
|
||||
|
||||
int64_t t_prepare_begin = ggml_time_ms();
|
||||
Plan base_plan;
|
||||
int64_t t_plan_begin = ggml_time_ms();
|
||||
if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) {
|
||||
base_plan = cache->graph_cut_plan;
|
||||
} else {
|
||||
base_plan = build_plan(backend, gf, params_tensor_set, log_desc);
|
||||
cache->graph_cut_plan = base_plan;
|
||||
cache->graph_cut_plan.available = true;
|
||||
cache->budgeted_graph_cut_plan.available = false;
|
||||
if (log_desc != nullptr) {
|
||||
LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin);
|
||||
}
|
||||
}
|
||||
|
||||
Plan resolved_plan = base_plan;
|
||||
if (max_graph_vram_bytes > 0 && base_plan.has_cuts) {
|
||||
if (cache->budgeted_graph_cut_plan.available &&
|
||||
cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes &&
|
||||
plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) {
|
||||
resolved_plan = cache->budgeted_graph_cut_plan;
|
||||
} else {
|
||||
resolved_plan = apply_max_vram_budget(gf,
|
||||
base_plan,
|
||||
max_graph_vram_bytes,
|
||||
backend,
|
||||
params_tensor_set,
|
||||
log_desc);
|
||||
cache->budgeted_graph_cut_plan = resolved_plan;
|
||||
cache->budgeted_graph_cut_plan.available = true;
|
||||
cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes;
|
||||
}
|
||||
}
|
||||
return resolved_plan;
|
||||
}
|
||||
|
||||
void annotate_residency(Plan& plan, size_t max_graph_vram_bytes) {
|
||||
// Cached plans may be reused with a smaller live budget.
|
||||
for (auto& seg : plan.segments) {
|
||||
seg.residency = SegmentResidency::STREAMED;
|
||||
}
|
||||
if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) {
|
||||
return;
|
||||
}
|
||||
|
||||
bool any_param_bearing = false;
|
||||
for (const auto& seg : plan.segments) {
|
||||
if (seg.input_param_bytes > 0) {
|
||||
any_param_bearing = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!any_param_bearing) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Leave room for the largest active streamed segment.
|
||||
size_t worst_streamed_footprint = 0;
|
||||
for (const auto& seg : plan.segments) {
|
||||
const size_t seg_footprint = seg.input_param_bytes +
|
||||
seg.compute_buffer_size +
|
||||
seg.output_bytes +
|
||||
seg.input_previous_cut_bytes +
|
||||
seg.input_external_bytes;
|
||||
if (seg_footprint > worst_streamed_footprint) {
|
||||
worst_streamed_footprint = seg_footprint;
|
||||
}
|
||||
}
|
||||
constexpr size_t safety = 512ull * 1024 * 1024;
|
||||
const size_t reserved = safety + worst_streamed_footprint;
|
||||
|
||||
if (max_graph_vram_bytes <= reserved) {
|
||||
return;
|
||||
}
|
||||
const size_t available = max_graph_vram_bytes - reserved;
|
||||
|
||||
size_t cumulative = 0;
|
||||
for (auto& seg : plan.segments) {
|
||||
if (cumulative + seg.input_param_bytes > available) {
|
||||
break;
|
||||
}
|
||||
seg.residency = SegmentResidency::RESIDENT;
|
||||
cumulative += seg.input_param_bytes;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace sd::ggml_graph_cut
|
||||
117
src/core/ggml_graph_cut.h
Normal file
117
src/core/ggml_graph_cut.h
Normal file
@@ -0,0 +1,117 @@
|
||||
#ifndef __SD_CORE_GGML_GRAPH_CUT_H__
|
||||
#define __SD_CORE_GGML_GRAPH_CUT_H__
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
|
||||
namespace sd::ggml_graph_cut {
|
||||
|
||||
// Streaming residency for a segment's params.
|
||||
enum class SegmentResidency : uint8_t {
|
||||
STREAMED = 0,
|
||||
RESIDENT = 1,
|
||||
};
|
||||
|
||||
struct Segment {
|
||||
enum InputType {
|
||||
INPUT_EXTERNAL = 0,
|
||||
INPUT_PREVIOUS_CUT,
|
||||
INPUT_PARAM,
|
||||
};
|
||||
|
||||
struct InputRef {
|
||||
InputType type = INPUT_EXTERNAL;
|
||||
std::string display_name;
|
||||
int leaf_index = -1;
|
||||
int node_index = -1;
|
||||
};
|
||||
|
||||
size_t compute_buffer_size = 0;
|
||||
size_t output_bytes = 0;
|
||||
size_t input_external_bytes = 0;
|
||||
size_t input_previous_cut_bytes = 0;
|
||||
size_t input_param_bytes = 0;
|
||||
std::string group_name;
|
||||
std::vector<int> internal_node_indices;
|
||||
std::vector<int> output_node_indices;
|
||||
std::vector<InputRef> input_refs;
|
||||
SegmentResidency residency = SegmentResidency::STREAMED;
|
||||
};
|
||||
|
||||
struct Plan {
|
||||
struct InputShape {
|
||||
int leaf_index = -1;
|
||||
ggml_type type = GGML_TYPE_COUNT;
|
||||
std::array<int64_t, GGML_MAX_DIMS> ne = {0, 0, 0, 0};
|
||||
};
|
||||
|
||||
bool available = false;
|
||||
bool has_cuts = false;
|
||||
bool valid = true;
|
||||
int n_nodes = 0;
|
||||
int n_leafs = 0;
|
||||
std::vector<InputShape> input_shapes;
|
||||
std::vector<Segment> segments;
|
||||
};
|
||||
|
||||
struct PlanCache {
|
||||
Plan graph_cut_plan;
|
||||
Plan budgeted_graph_cut_plan;
|
||||
size_t budgeted_graph_cut_plan_max_vram_bytes = 0;
|
||||
};
|
||||
|
||||
static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";
|
||||
|
||||
bool is_graph_cut_tensor(const ggml_tensor* tensor);
|
||||
std::string make_graph_cut_name(const std::string& group, const std::string& output);
|
||||
void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output);
|
||||
int leaf_count(ggml_cgraph* gf);
|
||||
ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index);
|
||||
ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor);
|
||||
ggml_tensor* cache_source_tensor(ggml_tensor* tensor);
|
||||
size_t cache_tensor_bytes(const ggml_tensor* tensor);
|
||||
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan);
|
||||
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
|
||||
ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
|
||||
std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
|
||||
std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc);
|
||||
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
|
||||
const Plan& plan,
|
||||
size_t current_segment_index);
|
||||
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
|
||||
const Segment& segment,
|
||||
ggml_context** graph_ctx_out);
|
||||
size_t measure_segment_compute_buffer(ggml_backend_t backend,
|
||||
ggml_cgraph* gf,
|
||||
const Segment& segment,
|
||||
const char* log_desc);
|
||||
size_t max_vram_gib_to_bytes(float max_vram);
|
||||
float resolve_max_vram_gib(float max_vram, ggml_backend_t backend);
|
||||
Plan build_plan(ggml_backend_t backend,
|
||||
ggml_cgraph* gf,
|
||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const char* log_desc);
|
||||
Plan apply_max_vram_budget(ggml_cgraph* gf,
|
||||
const Plan& base_plan,
|
||||
size_t max_graph_vram_bytes,
|
||||
ggml_backend_t backend,
|
||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const char* log_desc);
|
||||
Plan resolve_plan(ggml_backend_t backend,
|
||||
ggml_cgraph* gf,
|
||||
PlanCache* cache,
|
||||
size_t max_graph_vram_bytes,
|
||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
||||
const char* log_desc);
|
||||
|
||||
// Mark leading segments resident when they fit after streamed-segment headroom.
|
||||
void annotate_residency(Plan& plan, size_t max_graph_vram_bytes);
|
||||
} // namespace sd::ggml_graph_cut
|
||||
|
||||
#endif // __SD_CORE_GGML_GRAPH_CUT_H__
|
||||
132
src/core/layer_registry.cpp
Normal file
132
src/core/layer_registry.cpp
Normal file
@@ -0,0 +1,132 @@
|
||||
#include "core/layer_registry.h"
|
||||
|
||||
#include <utility>
|
||||
|
||||
#include "core/util.h"
|
||||
|
||||
namespace sd::layer_registry {
|
||||
|
||||
void LayerRegistry::register_layer(const std::string& name, ggml_tensor* tensor) {
|
||||
auto& info = layers_[name];
|
||||
info.tensors.push_back(tensor);
|
||||
info.bytes += ggml_nbytes(tensor);
|
||||
}
|
||||
|
||||
bool LayerRegistry::move_layer_to_gpu(const std::string& name) {
|
||||
auto it = layers_.find(name);
|
||||
if (it == layers_.end())
|
||||
return false;
|
||||
|
||||
LayerInfo& info = it->second;
|
||||
if (info.on_gpu)
|
||||
return true;
|
||||
if (gpu_backend_ == nullptr || cpu_backend_ == nullptr) {
|
||||
LOG_ERROR("layer_registry: backends not set; cannot move '%s' to GPU",
|
||||
name.c_str());
|
||||
return false;
|
||||
}
|
||||
if (info.tensors.empty()) {
|
||||
info.on_gpu = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
// 1. Build a no_alloc context big enough to hold one twin tensor per CPU
|
||||
// tensor, plus a little overhead.
|
||||
const size_t ctx_size = info.tensors.size() * ggml_tensor_overhead() + 1024;
|
||||
ggml_init_params ctx_params{ctx_size, /*mem_buffer=*/nullptr, /*no_alloc=*/true};
|
||||
ggml_context* twin_ctx = ggml_init(ctx_params);
|
||||
if (twin_ctx == nullptr) {
|
||||
LOG_ERROR("layer_registry: failed to allocate twin context for '%s'",
|
||||
name.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
// 2. Create one GPU twin per CPU tensor. The twin shares the original
|
||||
// name so any name-based lookup keeps working.
|
||||
std::vector<ggml_tensor*> gpu_twins;
|
||||
gpu_twins.reserve(info.tensors.size());
|
||||
for (ggml_tensor* cpu_t : info.tensors) {
|
||||
ggml_tensor* twin = ggml_dup_tensor(twin_ctx, cpu_t);
|
||||
if (cpu_t->name[0] != '\0') {
|
||||
ggml_set_name(twin, cpu_t->name);
|
||||
}
|
||||
gpu_twins.push_back(twin);
|
||||
}
|
||||
|
||||
// 3. Back the twins with a GPU buffer in one alloc call.
|
||||
ggml_backend_buffer_t gpu_buffer = ggml_backend_alloc_ctx_tensors(twin_ctx, gpu_backend_);
|
||||
if (gpu_buffer == nullptr) {
|
||||
LOG_ERROR("layer_registry: failed to allocate GPU buffer for '%s'",
|
||||
name.c_str());
|
||||
ggml_free(twin_ctx);
|
||||
return false;
|
||||
}
|
||||
|
||||
// 4. H2D copy + sync.
|
||||
for (size_t i = 0; i < info.tensors.size(); ++i) {
|
||||
ggml_backend_tensor_copy(info.tensors[i], gpu_twins[i]);
|
||||
}
|
||||
ggml_backend_synchronize(gpu_backend_);
|
||||
|
||||
// 5. Swap buffer/data/extra so the originals now point at GPU memory.
|
||||
for (size_t i = 0; i < info.tensors.size(); ++i) {
|
||||
std::swap(info.tensors[i]->buffer, gpu_twins[i]->buffer);
|
||||
std::swap(info.tensors[i]->data, gpu_twins[i]->data);
|
||||
std::swap(info.tensors[i]->extra, gpu_twins[i]->extra);
|
||||
}
|
||||
|
||||
info.gpu_twins = std::move(gpu_twins);
|
||||
info.twin_ctx = twin_ctx;
|
||||
info.gpu_buffer = gpu_buffer;
|
||||
info.on_gpu = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool LayerRegistry::move_layer_to_cpu(const std::string& name) {
|
||||
auto it = layers_.find(name);
|
||||
if (it == layers_.end())
|
||||
return false;
|
||||
|
||||
LayerInfo& info = it->second;
|
||||
if (!info.on_gpu)
|
||||
return true;
|
||||
if (info.tensors.size() != info.gpu_twins.size()) {
|
||||
LOG_ERROR("layer_registry: twin/tensor count mismatch for '%s'",
|
||||
name.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
// 1. Swap back: originals point at CPU memory again.
|
||||
for (size_t i = 0; i < info.tensors.size(); ++i) {
|
||||
if (info.gpu_twins[i] == nullptr)
|
||||
continue;
|
||||
std::swap(info.tensors[i]->buffer, info.gpu_twins[i]->buffer);
|
||||
std::swap(info.tensors[i]->data, info.gpu_twins[i]->data);
|
||||
std::swap(info.tensors[i]->extra, info.gpu_twins[i]->extra);
|
||||
}
|
||||
|
||||
// 2. Free the GPU buffer + twin context.
|
||||
if (info.gpu_buffer != nullptr) {
|
||||
ggml_backend_buffer_free(info.gpu_buffer);
|
||||
info.gpu_buffer = nullptr;
|
||||
}
|
||||
if (info.twin_ctx != nullptr) {
|
||||
ggml_free(info.twin_ctx);
|
||||
info.twin_ctx = nullptr;
|
||||
}
|
||||
info.gpu_twins.clear();
|
||||
info.on_gpu = false;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool LayerRegistry::is_layer_on_gpu(const std::string& name) const {
|
||||
auto it = layers_.find(name);
|
||||
return it != layers_.end() && it->second.on_gpu;
|
||||
}
|
||||
|
||||
size_t LayerRegistry::get_layer_size(const std::string& name) const {
|
||||
auto it = layers_.find(name);
|
||||
return it != layers_.end() ? it->second.bytes : 0;
|
||||
}
|
||||
|
||||
} // namespace sd::layer_registry
|
||||
50
src/core/layer_registry.h
Normal file
50
src/core/layer_registry.h
Normal file
@@ -0,0 +1,50 @@
|
||||
#ifndef __SD_CORE_LAYER_REGISTRY_H__
|
||||
#define __SD_CORE_LAYER_REGISTRY_H__
|
||||
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
|
||||
namespace sd::layer_registry {
|
||||
|
||||
struct LayerInfo {
|
||||
std::vector<ggml_tensor*> tensors;
|
||||
std::vector<ggml_tensor*> gpu_twins;
|
||||
ggml_context* twin_ctx = nullptr;
|
||||
ggml_backend_buffer_t gpu_buffer = nullptr;
|
||||
bool on_gpu = false;
|
||||
size_t bytes = 0;
|
||||
};
|
||||
|
||||
class LayerRegistry {
|
||||
public:
|
||||
LayerRegistry() = default;
|
||||
LayerRegistry(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend)
|
||||
: gpu_backend_(gpu_backend), cpu_backend_(cpu_backend) {}
|
||||
|
||||
void set_backends(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend) {
|
||||
gpu_backend_ = gpu_backend;
|
||||
cpu_backend_ = cpu_backend;
|
||||
}
|
||||
void register_layer(const std::string& name, ggml_tensor* tensor);
|
||||
bool move_layer_to_gpu(const std::string& name);
|
||||
bool move_layer_to_cpu(const std::string& name);
|
||||
bool is_layer_on_gpu(const std::string& name) const;
|
||||
size_t get_layer_size(const std::string& name) const;
|
||||
size_t get_layer_count() const { return layers_.size(); }
|
||||
|
||||
const std::map<std::string, LayerInfo>& layers() const { return layers_; }
|
||||
|
||||
private:
|
||||
ggml_backend_t gpu_backend_ = nullptr;
|
||||
ggml_backend_t cpu_backend_ = nullptr;
|
||||
std::map<std::string, LayerInfo> layers_;
|
||||
};
|
||||
|
||||
} // namespace sd::layer_registry
|
||||
|
||||
#endif // __SD_CORE_LAYER_REGISTRY_H__
|
||||
177
src/core/ordered_map.hpp
Normal file
177
src/core/ordered_map.hpp
Normal file
@@ -0,0 +1,177 @@
|
||||
#ifndef __SD_CORE_ORDERED_MAP_HPP__
|
||||
#define __SD_CORE_ORDERED_MAP_HPP__
|
||||
|
||||
#include <iostream>
|
||||
#include <list>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
|
||||
#include <initializer_list>
|
||||
#include <iterator>
|
||||
#include <list>
|
||||
#include <stdexcept>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
|
||||
template <typename Key, typename T>
|
||||
class OrderedMap {
|
||||
public:
|
||||
using key_type = Key;
|
||||
using mapped_type = T;
|
||||
using value_type = std::pair<const Key, T>;
|
||||
using list_type = std::list<value_type>;
|
||||
using size_type = typename list_type::size_type;
|
||||
using difference_type = typename list_type::difference_type;
|
||||
using iterator = typename list_type::iterator;
|
||||
using const_iterator = typename list_type::const_iterator;
|
||||
|
||||
private:
|
||||
list_type data_;
|
||||
std::unordered_map<Key, iterator> index_;
|
||||
|
||||
public:
|
||||
// --- constructors ---
|
||||
OrderedMap() = default;
|
||||
|
||||
OrderedMap(std::initializer_list<value_type> init) {
|
||||
for (const auto& kv : init)
|
||||
insert(kv);
|
||||
}
|
||||
|
||||
OrderedMap(const OrderedMap&) = default;
|
||||
OrderedMap(OrderedMap&&) noexcept = default;
|
||||
OrderedMap& operator=(const OrderedMap&) = default;
|
||||
OrderedMap& operator=(OrderedMap&&) noexcept = default;
|
||||
|
||||
// --- element access ---
|
||||
T& at(const Key& key) {
|
||||
auto it = index_.find(key);
|
||||
if (it == index_.end())
|
||||
throw std::out_of_range("OrderedMap::at: key not found");
|
||||
return it->second->second;
|
||||
}
|
||||
|
||||
const T& at(const Key& key) const {
|
||||
auto it = index_.find(key);
|
||||
if (it == index_.end())
|
||||
throw std::out_of_range("OrderedMap::at: key not found");
|
||||
return it->second->second;
|
||||
}
|
||||
|
||||
T& operator[](const Key& key) {
|
||||
auto it = index_.find(key);
|
||||
if (it == index_.end()) {
|
||||
data_.emplace_back(key, T{});
|
||||
auto iter = std::prev(data_.end());
|
||||
index_[key] = iter;
|
||||
return iter->second;
|
||||
}
|
||||
return it->second->second;
|
||||
}
|
||||
|
||||
// --- iterators ---
|
||||
iterator begin() noexcept { return data_.begin(); }
|
||||
const_iterator begin() const noexcept { return data_.begin(); }
|
||||
const_iterator cbegin() const noexcept { return data_.cbegin(); }
|
||||
|
||||
iterator end() noexcept { return data_.end(); }
|
||||
const_iterator end() const noexcept { return data_.end(); }
|
||||
const_iterator cend() const noexcept { return data_.cend(); }
|
||||
|
||||
// --- capacity ---
|
||||
bool empty() const noexcept { return data_.empty(); }
|
||||
size_type size() const noexcept { return data_.size(); }
|
||||
|
||||
// --- modifiers ---
|
||||
void clear() noexcept {
|
||||
data_.clear();
|
||||
index_.clear();
|
||||
}
|
||||
|
||||
std::pair<iterator, bool> insert(const value_type& value) {
|
||||
auto it = index_.find(value.first);
|
||||
if (it != index_.end()) {
|
||||
return {it->second, false};
|
||||
}
|
||||
data_.push_back(value);
|
||||
auto iter = std::prev(data_.end());
|
||||
index_[value.first] = iter;
|
||||
return {iter, true};
|
||||
}
|
||||
|
||||
std::pair<iterator, bool> insert(value_type&& value) {
|
||||
auto it = index_.find(value.first);
|
||||
if (it != index_.end()) {
|
||||
return {it->second, false};
|
||||
}
|
||||
data_.push_back(std::move(value));
|
||||
auto iter = std::prev(data_.end());
|
||||
index_[iter->first] = iter;
|
||||
return {iter, true};
|
||||
}
|
||||
|
||||
void erase(const Key& key) {
|
||||
auto it = index_.find(key);
|
||||
if (it != index_.end()) {
|
||||
data_.erase(it->second);
|
||||
index_.erase(it);
|
||||
}
|
||||
}
|
||||
|
||||
iterator erase(iterator pos) {
|
||||
index_.erase(pos->first);
|
||||
return data_.erase(pos);
|
||||
}
|
||||
|
||||
// --- lookup ---
|
||||
size_type count(const Key& key) const {
|
||||
return index_.count(key);
|
||||
}
|
||||
|
||||
iterator find(const Key& key) {
|
||||
auto it = index_.find(key);
|
||||
if (it == index_.end())
|
||||
return data_.end();
|
||||
return it->second;
|
||||
}
|
||||
|
||||
const_iterator find(const Key& key) const {
|
||||
auto it = index_.find(key);
|
||||
if (it == index_.end())
|
||||
return data_.end();
|
||||
return it->second;
|
||||
}
|
||||
|
||||
bool contains(const Key& key) const {
|
||||
return index_.find(key) != index_.end();
|
||||
}
|
||||
|
||||
// --- comparison ---
|
||||
bool operator==(const OrderedMap& other) const {
|
||||
return data_ == other.data_;
|
||||
}
|
||||
|
||||
bool operator!=(const OrderedMap& other) const {
|
||||
return !(*this == other);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
std::pair<iterator, bool> emplace(Args&&... args) {
|
||||
value_type value(std::forward<Args>(args)...);
|
||||
auto it = index_.find(value.first);
|
||||
if (it != index_.end()) {
|
||||
return {it->second, false};
|
||||
}
|
||||
data_.push_back(std::move(value));
|
||||
auto iter = std::prev(data_.end());
|
||||
index_[iter->first] = iter;
|
||||
return {iter, true};
|
||||
}
|
||||
|
||||
void swap(OrderedMap& other) noexcept {
|
||||
data_.swap(other.data_);
|
||||
index_.swap(other.index_);
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_CORE_ORDERED_MAP_HPP__
|
||||
35
src/core/rng.hpp
Normal file
35
src/core/rng.hpp
Normal file
@@ -0,0 +1,35 @@
|
||||
#ifndef __SD_CORE_RNG_HPP__
|
||||
#define __SD_CORE_RNG_HPP__
|
||||
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
class RNG {
|
||||
public:
|
||||
virtual void manual_seed(uint64_t seed) = 0;
|
||||
virtual std::vector<float> randn(uint32_t n) = 0;
|
||||
};
|
||||
|
||||
class STDDefaultRNG : public RNG {
|
||||
private:
|
||||
std::default_random_engine generator;
|
||||
|
||||
public:
|
||||
void manual_seed(uint64_t seed) override {
|
||||
generator.seed((unsigned int)seed);
|
||||
}
|
||||
|
||||
std::vector<float> randn(uint32_t n) override {
|
||||
std::vector<float> result;
|
||||
float mean = 0.0;
|
||||
float stddev = 1.0;
|
||||
std::normal_distribution<float> distribution(mean, stddev);
|
||||
for (uint32_t i = 0; i < n; i++) {
|
||||
float random_number = distribution(generator);
|
||||
result.push_back(random_number);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_CORE_RNG_HPP__
|
||||
147
src/core/rng_mt19937.hpp
Normal file
147
src/core/rng_mt19937.hpp
Normal file
@@ -0,0 +1,147 @@
|
||||
#ifndef __SD_CORE_RNG_MT19937_HPP__
|
||||
#define __SD_CORE_RNG_MT19937_HPP__
|
||||
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
#include "core/rng.hpp"
|
||||
|
||||
// RNG imitiating torch cpu randn on CPU.
|
||||
// Port from pytorch, original license: https://github.com/pytorch/pytorch/blob/d01a7b0241ed1c4cded7e7ca097249feb343f072/LICENSE
|
||||
// Ref: https://github.com/pytorch/pytorch/blob/d01a7b0241ed1c4cded7e7ca097249feb343f072/aten/src/ATen/core/TransformationHelper.h, for uniform_real
|
||||
// Ref: https://github.com/pytorch/pytorch/blob/d01a7b0241ed1c4cded7e7ca097249feb343f072/aten/src/ATen/native/cpu/DistributionTemplates.h, for normal_kernel/normal_fill/normal_fill_16
|
||||
// Ref: https://github.com/pytorch/pytorch/blob/d01a7b0241ed1c4cded7e7ca097249feb343f072/aten/src/ATen/core/MT19937RNGEngine.h, for mt19937_engine
|
||||
// Ref: https://github.com/pytorch/pytorch/blob/d01a7b0241ed1c4cded7e7ca097249feb343f072/aten/src/ATen/core/DistributionsHelper.h, for uniform_real_distribution/normal_distribution
|
||||
class MT19937RNG : public RNG {
|
||||
static const int N = 624;
|
||||
static const int M = 397;
|
||||
static const uint32_t MATRIX_A = 0x9908b0dfU;
|
||||
static const uint32_t UMASK = 0x80000000U;
|
||||
static const uint32_t LMASK = 0x7fffffffU;
|
||||
|
||||
struct State {
|
||||
uint64_t seed_;
|
||||
int left_;
|
||||
bool seeded_;
|
||||
uint32_t next_;
|
||||
std::array<uint32_t, N> state_;
|
||||
bool has_next_gauss = false;
|
||||
double next_gauss = 0.0f;
|
||||
};
|
||||
|
||||
State s;
|
||||
|
||||
uint32_t mix_bits(uint32_t u, uint32_t v) { return (u & UMASK) | (v & LMASK); }
|
||||
uint32_t twist(uint32_t u, uint32_t v) { return (mix_bits(u, v) >> 1) ^ ((v & 1) ? MATRIX_A : 0); }
|
||||
void next_state() {
|
||||
uint32_t* p = s.state_.data();
|
||||
s.left_ = N;
|
||||
s.next_ = 0;
|
||||
for (int j = N - M + 1; --j; p++)
|
||||
p[0] = p[M] ^ twist(p[0], p[1]);
|
||||
for (int j = M; --j; p++)
|
||||
p[0] = p[M - N] ^ twist(p[0], p[1]);
|
||||
p[0] = p[M - N] ^ twist(p[0], s.state_[0]);
|
||||
}
|
||||
|
||||
uint32_t rand_uint32() {
|
||||
if (--s.left_ == 0)
|
||||
next_state();
|
||||
uint32_t y = s.state_[s.next_++];
|
||||
y ^= (y >> 11);
|
||||
y ^= (y << 7) & 0x9d2c5680U;
|
||||
y ^= (y << 15) & 0xefc60000U;
|
||||
y ^= (y >> 18);
|
||||
return y;
|
||||
}
|
||||
|
||||
uint64_t rand_uint64() {
|
||||
uint64_t high = (uint64_t)rand_uint32();
|
||||
uint64_t low = (uint64_t)rand_uint32();
|
||||
return (high << 32) | low;
|
||||
}
|
||||
|
||||
template <typename T, typename V>
|
||||
T uniform_real(V val, T from, T to) {
|
||||
constexpr auto MASK = static_cast<V>((static_cast<uint64_t>(1) << std::numeric_limits<T>::digits) - 1);
|
||||
constexpr auto DIVISOR = static_cast<T>(1) / (static_cast<uint64_t>(1) << std::numeric_limits<T>::digits);
|
||||
T x = (val & MASK) * DIVISOR;
|
||||
return (x * (to - from) + from);
|
||||
}
|
||||
|
||||
double normal_double_value(double mean, double std) {
|
||||
if (s.has_next_gauss) {
|
||||
s.has_next_gauss = false;
|
||||
return s.next_gauss;
|
||||
}
|
||||
double u1 = uniform_real(rand_uint64(), 0., 1.); // double
|
||||
double u2 = uniform_real(rand_uint64(), 0., 1.); // double
|
||||
|
||||
double r = std::sqrt(-2.0 * std::log1p(-u2));
|
||||
double theta = 2.0 * 3.14159265358979323846 * u1;
|
||||
double value = r * std::cos(theta) * std + mean;
|
||||
s.next_gauss = r * std::sin(theta) * std + mean;
|
||||
s.has_next_gauss = true;
|
||||
return value;
|
||||
}
|
||||
|
||||
void normal_fill_16(float* data, float mean, float std) {
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
float u1 = 1.0f - data[j];
|
||||
float u2 = data[j + 8];
|
||||
float r = std::sqrt(-2.0f * std::log(u1));
|
||||
float theta = 2.0f * 3.14159265358979323846f * u2;
|
||||
data[j] = r * std::cos(theta) * std + mean;
|
||||
data[j + 8] = r * std::sin(theta) * std + mean;
|
||||
}
|
||||
}
|
||||
|
||||
void randn(float* data, int64_t size, float mean = 0.0f, float std = 1.0f) {
|
||||
if (size >= 16) {
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
data[i] = uniform_real(rand_uint32(), 0.f, 1.f);
|
||||
}
|
||||
for (int64_t i = 0; i < size - 15; i += 16) {
|
||||
normal_fill_16(data + i, mean, std);
|
||||
}
|
||||
if (size % 16 != 0) {
|
||||
// Recompute the last 16 values.
|
||||
data = data + size - 16;
|
||||
for (int64_t i = 0; i < 16; i++) {
|
||||
data[i] = uniform_real(rand_uint32(), 0.f, 1.f);
|
||||
}
|
||||
normal_fill_16(data, mean, std);
|
||||
}
|
||||
} else {
|
||||
// Strange handling, hard to understand, but keeping it consistent with PyTorch.
|
||||
for (int64_t i = 0; i < size; i++) {
|
||||
data[i] = (float)normal_double_value(mean, std);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
MT19937RNG(uint64_t seed = 0) { manual_seed(seed); }
|
||||
|
||||
void manual_seed(uint64_t seed) override {
|
||||
s.seed_ = seed;
|
||||
s.seeded_ = true;
|
||||
s.state_[0] = (uint32_t)(seed & 0xffffffffU);
|
||||
for (int j = 1; j < N; j++) {
|
||||
uint32_t prev = s.state_[j - 1];
|
||||
s.state_[j] = 1812433253U * (prev ^ (prev >> 30)) + j;
|
||||
}
|
||||
s.left_ = 1;
|
||||
s.next_ = 0;
|
||||
s.has_next_gauss = false;
|
||||
}
|
||||
|
||||
std::vector<float> randn(uint32_t n) override {
|
||||
std::vector<float> out;
|
||||
out.resize(n);
|
||||
randn((float*)out.data(), out.size());
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_CORE_RNG_MT19937_HPP__
|
||||
125
src/core/rng_philox.hpp
Normal file
125
src/core/rng_philox.hpp
Normal file
@@ -0,0 +1,125 @@
|
||||
#ifndef __SD_CORE_RNG_PHILOX_HPP__
|
||||
#define __SD_CORE_RNG_PHILOX_HPP__
|
||||
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
#include "core/rng.hpp"
|
||||
|
||||
// RNG imitiating torch cuda randn on CPU.
|
||||
// Port from: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/5ef669de080814067961f28357256e8fe27544f4/modules/rng_philox.py
|
||||
class PhiloxRNG : public RNG {
|
||||
private:
|
||||
uint64_t seed;
|
||||
uint32_t offset;
|
||||
|
||||
private:
|
||||
std::vector<uint32_t> philox_m = {0xD2511F53, 0xCD9E8D57};
|
||||
std::vector<uint32_t> philox_w = {0x9E3779B9, 0xBB67AE85};
|
||||
float two_pow32_inv = 2.3283064e-10f;
|
||||
float two_pow32_inv_2pi = 2.3283064e-10f * 6.2831855f;
|
||||
|
||||
std::vector<uint32_t> uint32(uint64_t x) {
|
||||
std::vector<uint32_t> result(2);
|
||||
result[0] = static_cast<uint32_t>(x & 0xFFFFFFFF);
|
||||
result[1] = static_cast<uint32_t>(x >> 32);
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<std::vector<uint32_t>> uint32(const std::vector<uint64_t>& x) {
|
||||
uint32_t N = (uint32_t)x.size();
|
||||
std::vector<std::vector<uint32_t>> result(2, std::vector<uint32_t>(N));
|
||||
|
||||
for (uint32_t i = 0; i < N; ++i) {
|
||||
result[0][i] = static_cast<uint32_t>(x[i] & 0xFFFFFFFF);
|
||||
result[1][i] = static_cast<uint32_t>(x[i] >> 32);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// A single round of the Philox 4x32 random number generator.
|
||||
void philox4_round(std::vector<std::vector<uint32_t>>& counter,
|
||||
const std::vector<std::vector<uint32_t>>& key) {
|
||||
uint32_t N = (uint32_t)counter[0].size();
|
||||
for (uint32_t i = 0; i < N; i++) {
|
||||
std::vector<uint32_t> v1 = uint32(static_cast<uint64_t>(counter[0][i]) * static_cast<uint64_t>(philox_m[0]));
|
||||
std::vector<uint32_t> v2 = uint32(static_cast<uint64_t>(counter[2][i]) * static_cast<uint64_t>(philox_m[1]));
|
||||
|
||||
counter[0][i] = v2[1] ^ counter[1][i] ^ key[0][i];
|
||||
counter[1][i] = v2[0];
|
||||
counter[2][i] = v1[1] ^ counter[3][i] ^ key[1][i];
|
||||
counter[3][i] = v1[0];
|
||||
}
|
||||
}
|
||||
|
||||
// Generates 32-bit random numbers using the Philox 4x32 random number generator.
|
||||
// Parameters:
|
||||
// counter : A 4xN array of 32-bit integers representing the counter values (offset into generation).
|
||||
// key : A 2xN array of 32-bit integers representing the key values (seed).
|
||||
// rounds : The number of rounds to perform.
|
||||
// Returns:
|
||||
// std::vector<std::vector<uint32_t>>: A 4xN array of 32-bit integers containing the generated random numbers.
|
||||
std::vector<std::vector<uint32_t>> philox4_32(std::vector<std::vector<uint32_t>>& counter,
|
||||
std::vector<std::vector<uint32_t>>& key,
|
||||
int rounds = 10) {
|
||||
uint32_t N = (uint32_t)counter[0].size();
|
||||
for (int i = 0; i < rounds - 1; ++i) {
|
||||
philox4_round(counter, key);
|
||||
|
||||
for (uint32_t j = 0; j < N; ++j) {
|
||||
key[0][j] += philox_w[0];
|
||||
key[1][j] += philox_w[1];
|
||||
}
|
||||
}
|
||||
|
||||
philox4_round(counter, key);
|
||||
return counter;
|
||||
}
|
||||
|
||||
float box_muller(float x, float y) {
|
||||
float u = x * two_pow32_inv + two_pow32_inv / 2;
|
||||
float v = y * two_pow32_inv_2pi + two_pow32_inv_2pi / 2;
|
||||
|
||||
float s = sqrt(-2.0f * log(u));
|
||||
|
||||
float r1 = s * sin(v);
|
||||
return r1;
|
||||
}
|
||||
|
||||
public:
|
||||
PhiloxRNG(uint64_t seed = 0) {
|
||||
this->seed = seed;
|
||||
this->offset = 0;
|
||||
}
|
||||
|
||||
void manual_seed(uint64_t seed) override {
|
||||
this->seed = seed;
|
||||
this->offset = 0;
|
||||
}
|
||||
|
||||
std::vector<float> randn(uint32_t n) override {
|
||||
std::vector<std::vector<uint32_t>> counter(4, std::vector<uint32_t>(n, 0));
|
||||
for (uint32_t i = 0; i < n; i++) {
|
||||
counter[0][i] = this->offset;
|
||||
}
|
||||
|
||||
for (uint32_t i = 0; i < n; i++) {
|
||||
counter[2][i] = i;
|
||||
}
|
||||
this->offset += 1;
|
||||
|
||||
std::vector<uint64_t> key(n, this->seed);
|
||||
std::vector<std::vector<uint32_t>> key_uint32 = uint32(key);
|
||||
|
||||
std::vector<std::vector<uint32_t>> g = philox4_32(counter, key_uint32);
|
||||
|
||||
std::vector<float> result;
|
||||
for (uint32_t i = 0; i < n; ++i) {
|
||||
result.push_back(box_muller((float)g[0][i], (float)g[1][i]));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_CORE_RNG_PHILOX_HPP__
|
||||
1664
src/core/tensor.hpp
Normal file
1664
src/core/tensor.hpp
Normal file
File diff suppressed because it is too large
Load Diff
127
src/core/tensor_ggml.hpp
Normal file
127
src/core/tensor_ggml.hpp
Normal file
@@ -0,0 +1,127 @@
|
||||
#ifndef __SD_CORE_TENSOR_GGML_HPP__
|
||||
#define __SD_CORE_TENSOR_GGML_HPP__
|
||||
|
||||
#include <array>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
|
||||
#include "core/tensor.hpp"
|
||||
#include "ggml.h"
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
struct GGMLTypeTraits;
|
||||
|
||||
template <>
|
||||
struct GGMLTypeTraits<float> {
|
||||
static constexpr ggml_type type = GGML_TYPE_F32;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GGMLTypeTraits<ggml_fp16_t> {
|
||||
static constexpr ggml_type type = GGML_TYPE_F16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GGMLTypeTraits<int32_t> {
|
||||
static constexpr ggml_type type = GGML_TYPE_I32;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GGMLTypeTraits<int64_t> {
|
||||
static constexpr ggml_type type = GGML_TYPE_I64;
|
||||
};
|
||||
|
||||
inline std::vector<int64_t> shape_from_ggml(const ggml_tensor* tensor) {
|
||||
std::vector<int64_t> shape;
|
||||
shape.reserve(static_cast<size_t>(ggml_n_dims(tensor)));
|
||||
for (int i = 0; i < ggml_n_dims(tensor); ++i) {
|
||||
shape.push_back(tensor->ne[i]);
|
||||
}
|
||||
return shape;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline Tensor<T> make_sd_tensor_from_ggml(const ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return {};
|
||||
}
|
||||
if (tensor->type != GGMLTypeTraits<T>::type) {
|
||||
GGML_ABORT("ggml tensor type does not match sd::Tensor type");
|
||||
}
|
||||
Tensor<T> result(shape_from_ggml(tensor));
|
||||
if (tensor->buffer != nullptr) {
|
||||
ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor));
|
||||
} else {
|
||||
std::memcpy(result.data(), tensor->data, ggml_nbytes(tensor));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline ggml_tensor* make_ggml_tensor(ggml_context* ctx, const Tensor<T>& tensor, bool copy_data = true) {
|
||||
GGML_ASSERT(tensor.dim() > 0 && tensor.dim() <= 5);
|
||||
|
||||
int n_dims = std::min(static_cast<int>(tensor.dim()), GGML_MAX_DIMS);
|
||||
|
||||
std::array<int64_t, GGML_MAX_DIMS> ne = {1, 1, 1, 1};
|
||||
for (int64_t i = 0; i < n_dims; ++i) {
|
||||
ne[static_cast<size_t>(i)] = tensor.shape()[static_cast<size_t>(i)];
|
||||
}
|
||||
|
||||
if (tensor.dim() == 5) {
|
||||
ne[3] *= tensor.shape()[4];
|
||||
}
|
||||
|
||||
ggml_tensor* result = ggml_new_tensor(ctx, GGMLTypeTraits<T>::type, n_dims, ne.data());
|
||||
if (copy_data && tensor.numel() > 0) {
|
||||
std::memcpy(result->data, tensor.data(), static_cast<size_t>(ggml_nbytes(result)));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline Tensor<T> load_tensor_from_file_as_tensor(const std::string& file_path) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
throw std::runtime_error("failed to open tensor file: " + file_path);
|
||||
}
|
||||
|
||||
int32_t n_dims = 0;
|
||||
int32_t length = 0;
|
||||
int32_t ttype = 0;
|
||||
file.read(reinterpret_cast<char*>(&n_dims), sizeof(n_dims));
|
||||
file.read(reinterpret_cast<char*>(&length), sizeof(length));
|
||||
file.read(reinterpret_cast<char*>(&ttype), sizeof(ttype));
|
||||
if (!file.good()) {
|
||||
throw std::runtime_error("incomplete tensor file header: " + file_path);
|
||||
}
|
||||
if (static_cast<ggml_type>(ttype) != GGMLTypeTraits<T>::type) {
|
||||
throw std::invalid_argument("tensor file type does not match requested sd::Tensor type");
|
||||
}
|
||||
|
||||
std::vector<int64_t> shape(n_dims, 1);
|
||||
for (int i = 0; i < n_dims; ++i) {
|
||||
int32_t dim = 1;
|
||||
file.read(reinterpret_cast<char*>(&dim), sizeof(dim));
|
||||
shape[static_cast<size_t>(i)] = dim;
|
||||
}
|
||||
std::string name(static_cast<size_t>(length), '\0');
|
||||
file.read(name.data(), length);
|
||||
|
||||
shape.resize(static_cast<size_t>(n_dims));
|
||||
Tensor<T> tensor(shape);
|
||||
file.read(reinterpret_cast<char*>(tensor.data()), static_cast<std::streamsize>(tensor.numel() * sizeof(T)));
|
||||
if (!file.good()) {
|
||||
throw std::runtime_error("incomplete tensor file data: " + file_path);
|
||||
}
|
||||
return tensor;
|
||||
}
|
||||
|
||||
} // namespace sd
|
||||
|
||||
#endif // __SD_CORE_TENSOR_GGML_HPP__
|
||||
974
src/core/util.cpp
Normal file
974
src/core/util.cpp
Normal file
@@ -0,0 +1,974 @@
|
||||
#include "core/util.h"
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cmath>
|
||||
#include <codecvt>
|
||||
#include <cstdarg>
|
||||
#include <exception>
|
||||
#include <fstream>
|
||||
#include <locale>
|
||||
#include <regex>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
#include "runtime/preprocessing.hpp"
|
||||
|
||||
#if defined(__APPLE__) && defined(__MACH__)
|
||||
#include <sys/sysctl.h>
|
||||
#include <sys/types.h>
|
||||
#endif
|
||||
|
||||
#if !defined(_WIN32)
|
||||
#include <sys/ioctl.h>
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include "ggml.h"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
bool ends_with(const std::string& str, const std::string& ending) {
|
||||
if (str.length() >= ending.length()) {
|
||||
return (str.compare(str.length() - ending.length(), ending.length(), ending) == 0);
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool starts_with(const std::string& str, const std::string& start) {
|
||||
if (str.find(start) == 0) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool contains(const std::string& str, const std::string& substr) {
|
||||
if (str.find(substr) != std::string::npos) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
void replace_all_chars(std::string& str, char target, char replacement) {
|
||||
for (size_t i = 0; i < str.length(); ++i) {
|
||||
if (str[i] == target) {
|
||||
str[i] = replacement;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::string sd_format(const char* fmt, ...) {
|
||||
va_list ap;
|
||||
va_list ap2;
|
||||
va_start(ap, fmt);
|
||||
va_copy(ap2, ap);
|
||||
int size = vsnprintf(nullptr, 0, fmt, ap);
|
||||
std::vector<char> buf(size + 1);
|
||||
int size2 = vsnprintf(buf.data(), size + 1, fmt, ap2);
|
||||
va_end(ap2);
|
||||
va_end(ap);
|
||||
return std::string(buf.data(), size);
|
||||
}
|
||||
|
||||
int round_up_to(int value, int base) {
|
||||
if (base <= 0) {
|
||||
return value;
|
||||
}
|
||||
if (value % base == 0) {
|
||||
return value;
|
||||
} else {
|
||||
return ((value / base) + 1) * base;
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef _WIN32 // code for windows
|
||||
#define NOMINMAX
|
||||
#include <windows.h>
|
||||
|
||||
bool file_exists(const std::string& filename) {
|
||||
DWORD attributes = GetFileAttributesA(filename.c_str());
|
||||
return (attributes != INVALID_FILE_ATTRIBUTES && !(attributes & FILE_ATTRIBUTE_DIRECTORY));
|
||||
}
|
||||
|
||||
bool is_directory(const std::string& path) {
|
||||
DWORD attributes = GetFileAttributesA(path.c_str());
|
||||
return (attributes != INVALID_FILE_ATTRIBUTES && (attributes & FILE_ATTRIBUTE_DIRECTORY));
|
||||
}
|
||||
|
||||
class MmapWrapperImpl : public MmapWrapper {
|
||||
public:
|
||||
MmapWrapperImpl(void* data, size_t size, HANDLE hfile, HANDLE hmapping)
|
||||
: MmapWrapper(data, size), hfile_(hfile), hmapping_(hmapping) {}
|
||||
|
||||
~MmapWrapperImpl() override {
|
||||
UnmapViewOfFile(data_);
|
||||
CloseHandle(hmapping_);
|
||||
CloseHandle(hfile_);
|
||||
}
|
||||
|
||||
private:
|
||||
HANDLE hfile_;
|
||||
HANDLE hmapping_;
|
||||
};
|
||||
|
||||
std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename, bool writable) {
|
||||
void* mapped_data = nullptr;
|
||||
size_t file_size = 0;
|
||||
|
||||
HANDLE file_handle = CreateFileA(
|
||||
filename.c_str(),
|
||||
GENERIC_READ,
|
||||
FILE_SHARE_READ,
|
||||
nullptr,
|
||||
OPEN_EXISTING,
|
||||
FILE_ATTRIBUTE_NORMAL,
|
||||
nullptr);
|
||||
|
||||
if (file_handle == INVALID_HANDLE_VALUE) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
LARGE_INTEGER size;
|
||||
if (!GetFileSizeEx(file_handle, &size)) {
|
||||
CloseHandle(file_handle);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
file_size = static_cast<size_t>(size.QuadPart);
|
||||
|
||||
DWORD page_prot = writable ? PAGE_WRITECOPY : PAGE_READONLY;
|
||||
|
||||
HANDLE mapping_handle = CreateFileMapping(file_handle, nullptr, page_prot, 0, 0, nullptr);
|
||||
|
||||
if (mapping_handle == nullptr) {
|
||||
CloseHandle(file_handle);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
DWORD view_access = writable ? FILE_MAP_COPY : FILE_MAP_READ;
|
||||
|
||||
mapped_data = MapViewOfFile(mapping_handle, view_access, 0, 0, file_size);
|
||||
|
||||
if (mapped_data == nullptr) {
|
||||
CloseHandle(mapping_handle);
|
||||
CloseHandle(file_handle);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
return std::make_unique<MmapWrapperImpl>(mapped_data, file_size, file_handle, mapping_handle);
|
||||
}
|
||||
|
||||
#else // Unix
|
||||
#include <dirent.h>
|
||||
#include <fcntl.h>
|
||||
#include <sys/mman.h>
|
||||
#include <sys/stat.h>
|
||||
#include <unistd.h>
|
||||
|
||||
bool file_exists(const std::string& filename) {
|
||||
struct stat buffer;
|
||||
return (stat(filename.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
|
||||
}
|
||||
|
||||
bool is_directory(const std::string& path) {
|
||||
struct stat buffer;
|
||||
return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
|
||||
}
|
||||
|
||||
struct MmapFlags {
|
||||
bool sequential;
|
||||
bool populate;
|
||||
bool willneed;
|
||||
bool dontneed;
|
||||
};
|
||||
|
||||
static MmapFlags get_mmap_flags() {
|
||||
MmapFlags result = {};
|
||||
const char* SD_MMAP_FLAGS = std::getenv("SD_MMAP_FLAGS");
|
||||
if (SD_MMAP_FLAGS && *SD_MMAP_FLAGS) {
|
||||
std::stringstream ss(SD_MMAP_FLAGS);
|
||||
std::string token;
|
||||
while (std::getline(ss, token, ',')) {
|
||||
std::string ntoken = trim(token);
|
||||
std::transform(ntoken.begin(), ntoken.end(), ntoken.begin(), ::tolower);
|
||||
if (ntoken == "sequential") {
|
||||
result.sequential = true;
|
||||
} else if (ntoken == "populate") {
|
||||
result.populate = true;
|
||||
} else if (ntoken == "willneed") {
|
||||
result.willneed = true;
|
||||
} else if (ntoken == "dontneed") {
|
||||
result.dontneed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
class MmapWrapperImpl : public MmapWrapper {
|
||||
public:
|
||||
MmapWrapperImpl(void* data, size_t size, int fd)
|
||||
: MmapWrapper(data, size), fd_(fd) {}
|
||||
|
||||
~MmapWrapperImpl() override {
|
||||
#ifdef __linux__
|
||||
auto cfg_flags = get_mmap_flags();
|
||||
|
||||
// Drop the kernel pagecache pages for this file. madvise(DONTNEED)
|
||||
// alone only unmaps from the process address space; pagecache
|
||||
// entries persist (`free` reports them as buff/cache and the OOM
|
||||
// killer doesn't touch them, but they ARE counted against
|
||||
// overcommit and can starve other allocations on tight-RAM
|
||||
// systems). posix_fadvise(POSIX_FADV_DONTNEED) is the documented
|
||||
// way to evict pagecache for a specific fd's pages.
|
||||
if (cfg_flags.dontneed) {
|
||||
madvise(data_, size_, MADV_DONTNEED);
|
||||
posix_fadvise(fd_, 0, 0, POSIX_FADV_DONTNEED);
|
||||
}
|
||||
#endif
|
||||
munmap(data_, size_);
|
||||
close(fd_);
|
||||
}
|
||||
|
||||
private:
|
||||
int fd_;
|
||||
};
|
||||
|
||||
std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename, bool writable) {
|
||||
int file_descriptor = open(filename.c_str(), O_RDONLY);
|
||||
if (file_descriptor == -1) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
auto cfg_flags = get_mmap_flags();
|
||||
|
||||
int mmap_flags = MAP_PRIVATE;
|
||||
|
||||
#ifdef __linux__
|
||||
// Sequential access hint helps the kernel read-ahead efficiently and
|
||||
// also encourages eviction of already-read pages (the kernel keeps
|
||||
// a smaller working set when this is set).
|
||||
if (cfg_flags.sequential) {
|
||||
posix_fadvise(file_descriptor, 0, 0, POSIX_FADV_SEQUENTIAL);
|
||||
}
|
||||
if (cfg_flags.populate) {
|
||||
mmap_flags |= MAP_POPULATE;
|
||||
}
|
||||
#endif
|
||||
|
||||
struct stat sb;
|
||||
if (fstat(file_descriptor, &sb) == -1) {
|
||||
close(file_descriptor);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
size_t file_size = sb.st_size;
|
||||
|
||||
if (file_size == 0) {
|
||||
close(file_descriptor);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
int mmap_prot = PROT_READ | (writable ? PROT_WRITE : 0);
|
||||
|
||||
void* mapped_data = mmap(nullptr, file_size, mmap_prot, mmap_flags, file_descriptor, 0);
|
||||
|
||||
if (mapped_data == MAP_FAILED) {
|
||||
close(file_descriptor);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
#ifdef __linux__
|
||||
if (cfg_flags.willneed) {
|
||||
posix_madvise(mapped_data, file_size, POSIX_MADV_WILLNEED);
|
||||
}
|
||||
#endif
|
||||
|
||||
return std::make_unique<MmapWrapperImpl>(mapped_data, file_size, file_descriptor);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
bool MmapWrapper::copy_data(void* buf, size_t n, size_t offset) const {
|
||||
if (offset >= size_ || n > (size_ - offset)) {
|
||||
return false;
|
||||
}
|
||||
std::memcpy(buf, data() + offset, n);
|
||||
return true;
|
||||
}
|
||||
|
||||
// get_num_physical_cores is copy from
|
||||
// https://github.com/ggerganov/llama.cpp/blob/master/examples/common.cpp
|
||||
// LICENSE: https://github.com/ggerganov/llama.cpp/blob/master/LICENSE
|
||||
int32_t sd_get_num_physical_cores() {
|
||||
#ifdef __linux__
|
||||
// enumerate the set of thread siblings, num entries is num cores
|
||||
std::unordered_set<std::string> siblings;
|
||||
for (uint32_t cpu = 0; cpu < UINT32_MAX; ++cpu) {
|
||||
std::ifstream thread_siblings("/sys/devices/system/cpu" + std::to_string(cpu) + "/topology/thread_siblings");
|
||||
if (!thread_siblings.is_open()) {
|
||||
break; // no more cpus
|
||||
}
|
||||
std::string line;
|
||||
if (std::getline(thread_siblings, line)) {
|
||||
siblings.insert(line);
|
||||
}
|
||||
}
|
||||
if (siblings.size() > 0) {
|
||||
return static_cast<int32_t>(siblings.size());
|
||||
}
|
||||
#elif defined(__APPLE__) && defined(__MACH__)
|
||||
int32_t num_physical_cores;
|
||||
size_t len = sizeof(num_physical_cores);
|
||||
int result = sysctlbyname("hw.perflevel0.physicalcpu", &num_physical_cores, &len, nullptr, 0);
|
||||
if (result == 0) {
|
||||
return num_physical_cores;
|
||||
}
|
||||
result = sysctlbyname("hw.physicalcpu", &num_physical_cores, &len, nullptr, 0);
|
||||
if (result == 0) {
|
||||
return num_physical_cores;
|
||||
}
|
||||
#elif defined(_WIN32)
|
||||
// TODO: Implement
|
||||
#endif
|
||||
unsigned int n_threads = std::thread::hardware_concurrency();
|
||||
return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
|
||||
}
|
||||
|
||||
static sd_progress_cb_t sd_progress_cb = nullptr;
|
||||
void* sd_progress_cb_data = nullptr;
|
||||
|
||||
static sd_preview_cb_t sd_preview_cb = nullptr;
|
||||
static void* sd_preview_cb_data = nullptr;
|
||||
preview_t sd_preview_mode = PREVIEW_NONE;
|
||||
int sd_preview_interval = 1;
|
||||
bool sd_preview_denoised = true;
|
||||
bool sd_preview_noisy = false;
|
||||
|
||||
std::u32string utf8_to_utf32(const std::string& utf8_str) {
|
||||
std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
|
||||
return converter.from_bytes(utf8_str);
|
||||
}
|
||||
|
||||
std::string utf32_to_utf8(const std::u32string& utf32_str) {
|
||||
std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
|
||||
return converter.to_bytes(utf32_str);
|
||||
}
|
||||
|
||||
std::u32string unicode_value_to_utf32(int unicode_value) {
|
||||
std::u32string utf32_string = {static_cast<char32_t>(unicode_value)};
|
||||
return utf32_string;
|
||||
}
|
||||
|
||||
static std::string sd_basename(const std::string& path) {
|
||||
size_t pos = path.find_last_of('/');
|
||||
if (pos != std::string::npos) {
|
||||
return path.substr(pos + 1);
|
||||
}
|
||||
pos = path.find_last_of('\\');
|
||||
if (pos != std::string::npos) {
|
||||
return path.substr(pos + 1);
|
||||
}
|
||||
return path;
|
||||
}
|
||||
|
||||
std::string path_join(const std::string& p1, const std::string& p2) {
|
||||
if (p1.empty()) {
|
||||
return p2;
|
||||
}
|
||||
|
||||
if (p2.empty()) {
|
||||
return p1;
|
||||
}
|
||||
|
||||
if (p1[p1.length() - 1] == '/' || p1[p1.length() - 1] == '\\') {
|
||||
return p1 + p2;
|
||||
}
|
||||
|
||||
return p1 + "/" + p2;
|
||||
}
|
||||
|
||||
std::vector<std::string> split_string(const std::string& str, char delimiter) {
|
||||
std::vector<std::string> result;
|
||||
size_t start = 0;
|
||||
size_t end = str.find(delimiter);
|
||||
|
||||
while (end != std::string::npos) {
|
||||
result.push_back(str.substr(start, end - start));
|
||||
start = end + 1;
|
||||
end = str.find(delimiter, start);
|
||||
}
|
||||
|
||||
// Add the last segment after the last delimiter
|
||||
result.push_back(str.substr(start));
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
KeyValueArgs parse_key_value_args(const char* args, const char* context) {
|
||||
KeyValueArgs pairs;
|
||||
|
||||
if (args == nullptr || args[0] == '\0') {
|
||||
return pairs;
|
||||
}
|
||||
|
||||
std::string raw(args);
|
||||
size_t start = 0;
|
||||
for (size_t pos = 0; pos <= raw.size(); ++pos) {
|
||||
if (pos != raw.size() && raw[pos] != ',' && raw[pos] != ';') {
|
||||
continue;
|
||||
}
|
||||
|
||||
std::string token = trim(raw.substr(start, pos - start));
|
||||
if (!token.empty()) {
|
||||
size_t eq = token.find('=');
|
||||
if (eq == std::string::npos) {
|
||||
const char* log_context = context ? context : "key=value arg";
|
||||
LOG_WARN("ignoring malformed %s '%s'", log_context, token.c_str());
|
||||
} else {
|
||||
std::string key = trim(token.substr(0, eq));
|
||||
std::string value = trim(token.substr(eq + 1));
|
||||
pairs.emplace_back(std::move(key), std::move(value));
|
||||
}
|
||||
}
|
||||
|
||||
start = pos + 1;
|
||||
}
|
||||
|
||||
return pairs;
|
||||
}
|
||||
|
||||
KeyValueArgs parse_key_value_args(const std::string& args, const char* context) {
|
||||
return parse_key_value_args(args.c_str(), context);
|
||||
}
|
||||
|
||||
bool parse_strict_float(const std::string& text, float& value) {
|
||||
try {
|
||||
size_t consumed = 0;
|
||||
float parsed = std::stof(text, &consumed);
|
||||
if (!trim(text.substr(consumed)).empty()) {
|
||||
return false;
|
||||
}
|
||||
value = parsed;
|
||||
return true;
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool parse_strict_int(const std::string& text, int& value) {
|
||||
try {
|
||||
size_t consumed = 0;
|
||||
int parsed = std::stoi(text, &consumed);
|
||||
if (!trim(text.substr(consumed)).empty()) {
|
||||
return false;
|
||||
}
|
||||
value = parsed;
|
||||
return true;
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool parse_strict_bool(const std::string& text, bool& value) {
|
||||
std::string lowered = trim(text);
|
||||
std::transform(lowered.begin(), lowered.end(), lowered.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
|
||||
if (lowered == "1" || lowered == "true" || lowered == "yes" || lowered == "on") {
|
||||
value = true;
|
||||
return true;
|
||||
}
|
||||
if (lowered == "0" || lowered == "false" || lowered == "no" || lowered == "off") {
|
||||
value = false;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static std::string build_progress_bar(int step, int steps) {
|
||||
std::string progress = " |";
|
||||
int max_progress = 50;
|
||||
int32_t current = 0;
|
||||
if (steps > 0) {
|
||||
current = (int32_t)(step * 1.f * max_progress / steps);
|
||||
}
|
||||
for (int i = 0; i < 50; i++) {
|
||||
if (i > current) {
|
||||
progress += " ";
|
||||
} else if (i == current && i != max_progress - 1) {
|
||||
progress += ">";
|
||||
} else {
|
||||
progress += "=";
|
||||
}
|
||||
}
|
||||
progress += "|";
|
||||
return progress;
|
||||
}
|
||||
|
||||
static void print_progress_line(int step, int steps, const std::string& speed_text) {
|
||||
if (step == 0) {
|
||||
return;
|
||||
}
|
||||
std::string progress = build_progress_bar(step, steps);
|
||||
const char* lf = (step == steps ? "\n" : "");
|
||||
printf("\r%s %i/%i - %s\033[K%s", progress.c_str(), step, steps, speed_text.c_str(), lf);
|
||||
fflush(stdout); // for linux
|
||||
}
|
||||
|
||||
void pretty_progress(int step, int steps, float time) {
|
||||
if (sd_progress_cb) {
|
||||
sd_progress_cb(step, steps, time, sd_progress_cb_data);
|
||||
return;
|
||||
}
|
||||
if (step == 0) {
|
||||
return;
|
||||
}
|
||||
const char* unit = "s/it";
|
||||
float speed = time;
|
||||
if (speed < 1.0f && speed > 0.f) {
|
||||
speed = 1.0f / speed;
|
||||
unit = "it/s";
|
||||
}
|
||||
print_progress_line(step, steps, sd_format("%.2f%s", speed, unit));
|
||||
}
|
||||
|
||||
void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds) {
|
||||
if (sd_progress_cb) {
|
||||
float time = elapsed_seconds / (step + 1e-6f);
|
||||
sd_progress_cb(step, steps, time, sd_progress_cb_data);
|
||||
return;
|
||||
}
|
||||
if (step == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
double bytes_per_second = 0.0;
|
||||
if (elapsed_seconds > 0.0f) {
|
||||
bytes_per_second = bytes_processed / (double)elapsed_seconds;
|
||||
}
|
||||
|
||||
double speed_mb = bytes_per_second / (1024.0 * 1024.0);
|
||||
if (speed_mb >= 1024.0) {
|
||||
print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0));
|
||||
} else {
|
||||
print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb));
|
||||
}
|
||||
}
|
||||
|
||||
std::string ltrim(const std::string& s) {
|
||||
auto it = std::find_if(s.begin(), s.end(), [](int ch) {
|
||||
return !std::isspace(ch);
|
||||
});
|
||||
return std::string(it, s.end());
|
||||
}
|
||||
|
||||
std::string rtrim(const std::string& s) {
|
||||
auto it = std::find_if(s.rbegin(), s.rend(), [](int ch) {
|
||||
return !std::isspace(ch);
|
||||
});
|
||||
return std::string(s.begin(), it.base());
|
||||
}
|
||||
|
||||
std::string trim(const std::string& s) {
|
||||
return rtrim(ltrim(s));
|
||||
}
|
||||
|
||||
static sd_log_cb_t sd_log_cb = nullptr;
|
||||
void* sd_log_cb_data = nullptr;
|
||||
|
||||
#define LOG_BUFFER_SIZE 4096
|
||||
|
||||
void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) {
|
||||
va_list args;
|
||||
va_start(args, format);
|
||||
|
||||
static char log_buffer[LOG_BUFFER_SIZE + 1];
|
||||
int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", sd_basename(file).c_str(), line);
|
||||
|
||||
if (written >= 0 && written < LOG_BUFFER_SIZE) {
|
||||
vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args);
|
||||
}
|
||||
size_t len = strlen(log_buffer);
|
||||
if (log_buffer[len - 1] != '\n') {
|
||||
strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len);
|
||||
}
|
||||
|
||||
if (sd_log_cb) {
|
||||
sd_log_cb(level, log_buffer, sd_log_cb_data);
|
||||
}
|
||||
|
||||
va_end(args);
|
||||
}
|
||||
|
||||
void sd_set_log_callback(sd_log_cb_t cb, void* data) {
|
||||
sd_log_cb = cb;
|
||||
sd_log_cb_data = data;
|
||||
}
|
||||
void sd_set_progress_callback(sd_progress_cb_t cb, void* data) {
|
||||
sd_progress_cb = cb;
|
||||
sd_progress_cb_data = data;
|
||||
}
|
||||
void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, bool denoised, bool noisy, void* data) {
|
||||
sd_preview_cb = cb;
|
||||
sd_preview_cb_data = data;
|
||||
sd_preview_mode = mode;
|
||||
sd_preview_interval = interval;
|
||||
sd_preview_denoised = denoised;
|
||||
sd_preview_noisy = noisy;
|
||||
}
|
||||
|
||||
sd_preview_cb_t sd_get_preview_callback() {
|
||||
return sd_preview_cb;
|
||||
}
|
||||
void* sd_get_preview_callback_data() {
|
||||
return sd_preview_cb_data;
|
||||
}
|
||||
|
||||
preview_t sd_get_preview_mode() {
|
||||
return sd_preview_mode;
|
||||
}
|
||||
int sd_get_preview_interval() {
|
||||
return sd_preview_interval;
|
||||
}
|
||||
bool sd_should_preview_denoised() {
|
||||
return sd_preview_denoised;
|
||||
}
|
||||
bool sd_should_preview_noisy() {
|
||||
return sd_preview_noisy;
|
||||
}
|
||||
|
||||
sd_progress_cb_t sd_get_progress_callback() {
|
||||
return sd_progress_cb;
|
||||
}
|
||||
void* sd_get_progress_callback_data() {
|
||||
return sd_progress_cb_data;
|
||||
}
|
||||
|
||||
sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index) {
|
||||
const auto& shape = tensor.shape();
|
||||
GGML_ASSERT(shape.size() == 4 || shape.size() == 5);
|
||||
int width = static_cast<int>(shape[0]);
|
||||
int height = static_cast<int>(shape[1]);
|
||||
int channel = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
|
||||
uint8_t* data = (uint8_t*)malloc(static_cast<size_t>(width * height * channel));
|
||||
GGML_ASSERT(data != nullptr);
|
||||
preprocessing_tensor_frame_to_sd_image(tensor, frame_index, data);
|
||||
return {
|
||||
static_cast<uint32_t>(width),
|
||||
static_cast<uint32_t>(height),
|
||||
static_cast<uint32_t>(channel),
|
||||
data,
|
||||
};
|
||||
}
|
||||
|
||||
sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
|
||||
int target_width,
|
||||
int target_height,
|
||||
bool scale) {
|
||||
sd::Tensor<float> tensor = sd::zeros<float>({static_cast<int64_t>(image.width),
|
||||
static_cast<int64_t>(image.height),
|
||||
static_cast<int64_t>(image.channel),
|
||||
1});
|
||||
for (uint32_t iw = 0; iw < image.width; ++iw) {
|
||||
for (uint32_t ih = 0; ih < image.height; ++ih) {
|
||||
for (uint32_t ic = 0; ic < image.channel; ++ic) {
|
||||
tensor.index(iw, ih, ic, 0) = sd_image_get_f32(image, iw, ih, ic, scale);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (target_width >= 0 && target_height >= 0 &&
|
||||
(tensor.shape()[0] != target_width || tensor.shape()[1] != target_height)) {
|
||||
tensor = sd::ops::interpolate(tensor,
|
||||
{target_width,
|
||||
target_height,
|
||||
tensor.shape()[2],
|
||||
tensor.shape()[3]});
|
||||
}
|
||||
return tensor;
|
||||
}
|
||||
|
||||
// Constants for means and std
|
||||
float means[3] = {0.48145466f, 0.4578275f, 0.40821073f};
|
||||
float stds[3] = {0.26862954f, 0.26130258f, 0.27577711f};
|
||||
|
||||
sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_width, int target_height) {
|
||||
GGML_ASSERT(image.dim() == 4);
|
||||
GGML_ASSERT(image.shape()[2] == 3);
|
||||
GGML_ASSERT(image.shape()[3] == 1);
|
||||
GGML_ASSERT(target_width > 0 && target_height > 0);
|
||||
|
||||
float width_scale = static_cast<float>(target_width) / static_cast<float>(image.shape()[0]);
|
||||
float height_scale = static_cast<float>(target_height) / static_cast<float>(image.shape()[1]);
|
||||
float scale = std::fmax(width_scale, height_scale);
|
||||
|
||||
int64_t resized_width = static_cast<int64_t>(scale * static_cast<float>(image.shape()[0]));
|
||||
int64_t resized_height = static_cast<int64_t>(scale * static_cast<float>(image.shape()[1]));
|
||||
|
||||
sd::Tensor<float> resized = sd::ops::interpolate(
|
||||
image,
|
||||
{resized_width, resized_height, image.shape()[2], image.shape()[3]});
|
||||
|
||||
int64_t h_offset = std::max<int64_t>((resized_height - target_height) / 2, 0);
|
||||
int64_t w_offset = std::max<int64_t>((resized_width - target_width) / 2, 0);
|
||||
|
||||
sd::Tensor<float> cropped({target_width, target_height, image.shape()[2], image.shape()[3]});
|
||||
for (int64_t y = 0; y < target_height; ++y) {
|
||||
for (int64_t x = 0; x < target_width; ++x) {
|
||||
for (int64_t c = 0; c < image.shape()[2]; ++c) {
|
||||
cropped.index(x, y, c, 0) = resized.index(x + w_offset, y + h_offset, c, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sd::Tensor<float> normalized = sd::ops::clamp(cropped, 0.0f, 1.0f);
|
||||
sd::Tensor<float> mean({1, 1, 3, 1}, {means[0], means[1], means[2]});
|
||||
sd::Tensor<float> std({1, 1, 3, 1}, {stds[0], stds[1], stds[2]});
|
||||
return (normalized - mean) / std;
|
||||
}
|
||||
|
||||
// Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/prompt_parser.py#L345
|
||||
//
|
||||
// Parses a string with attention tokens and returns a list of pairs: text and its associated weight.
|
||||
// Accepted tokens are:
|
||||
// (abc) - increases attention to abc by a multiplier of 1.1
|
||||
// (abc:3.12) - increases attention to abc by a multiplier of 3.12
|
||||
// [abc] - decreases attention to abc by a multiplier of 1.1
|
||||
// BREAK - separates the prompt into conceptually distinct parts for sequential processing
|
||||
// B - internal helper pattern; prevents 'B' in 'BREAK' from being consumed as normal text
|
||||
// \( - literal character '('
|
||||
// \[ - literal character '['
|
||||
// \) - literal character ')'
|
||||
// \] - literal character ']'
|
||||
// \\ - literal character '\'
|
||||
// anything else - just text
|
||||
//
|
||||
// >>> parse_prompt_attention('normal text')
|
||||
// [['normal text', 1.0]]
|
||||
// >>> parse_prompt_attention('an (important) word')
|
||||
// [['an ', 1.0], ['important', 1.1], [' word', 1.0]]
|
||||
// >>> parse_prompt_attention('(unbalanced')
|
||||
// [['unbalanced', 1.1]]
|
||||
// >>> parse_prompt_attention('\(literal\]')
|
||||
// [['(literal]', 1.0]]
|
||||
// >>> parse_prompt_attention('(unnecessary)(parens)')
|
||||
// [['unnecessaryparens', 1.1]]
|
||||
// >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')
|
||||
// [['a ', 1.0],
|
||||
// ['house', 1.5730000000000004],
|
||||
// [' ', 1.1],
|
||||
// ['on', 1.0],
|
||||
// [' a ', 1.1],
|
||||
// ['hill', 0.55],
|
||||
// [', sun, ', 1.1],
|
||||
// ['sky', 1.4641000000000006],
|
||||
// ['.', 1.1]]
|
||||
std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::string& text) {
|
||||
std::vector<std::pair<std::string, float>> res;
|
||||
std::vector<int> round_brackets;
|
||||
std::vector<int> square_brackets;
|
||||
|
||||
float round_bracket_multiplier = 1.1f;
|
||||
float square_bracket_multiplier = 1 / 1.1f;
|
||||
|
||||
std::regex re_attention(R"(\\\(|\\\)|\\\[|\\\]|\\\\|\\|\(|\[|:([+-]?[.\d]+)\)|\)|\]|\bBREAK\b|[^\\()\[\]:B]+|:|\bB)");
|
||||
std::regex re_break(R"(\s*\bBREAK\b\s*)");
|
||||
|
||||
auto multiply_range = [&](int start_position, float multiplier) {
|
||||
for (int p = start_position; p < res.size(); ++p) {
|
||||
res[p].second *= multiplier;
|
||||
}
|
||||
};
|
||||
|
||||
std::smatch m, m2;
|
||||
std::string remaining_text = text;
|
||||
|
||||
while (std::regex_search(remaining_text, m, re_attention)) {
|
||||
std::string text = m[0];
|
||||
std::string weight = m[1];
|
||||
|
||||
if (text == "(") {
|
||||
round_brackets.push_back((int)res.size());
|
||||
} else if (text == "[") {
|
||||
square_brackets.push_back((int)res.size());
|
||||
} else if (!weight.empty()) {
|
||||
if (!round_brackets.empty()) {
|
||||
multiply_range(round_brackets.back(), std::stof(weight));
|
||||
round_brackets.pop_back();
|
||||
}
|
||||
} else if (text == ")" && !round_brackets.empty()) {
|
||||
multiply_range(round_brackets.back(), round_bracket_multiplier);
|
||||
round_brackets.pop_back();
|
||||
} else if (text == "]" && !square_brackets.empty()) {
|
||||
multiply_range(square_brackets.back(), square_bracket_multiplier);
|
||||
square_brackets.pop_back();
|
||||
} else if (text == "\\(") {
|
||||
res.push_back({text.substr(1), 1.0f});
|
||||
} else if (std::regex_search(text, m2, re_break)) {
|
||||
res.push_back({"BREAK", -1.0f});
|
||||
} else {
|
||||
res.push_back({text, 1.0f});
|
||||
}
|
||||
|
||||
remaining_text = m.suffix();
|
||||
}
|
||||
|
||||
for (int pos : round_brackets) {
|
||||
multiply_range(pos, round_bracket_multiplier);
|
||||
}
|
||||
|
||||
for (int pos : square_brackets) {
|
||||
multiply_range(pos, square_bracket_multiplier);
|
||||
}
|
||||
|
||||
if (res.empty()) {
|
||||
res.push_back({"", 1.0f});
|
||||
}
|
||||
|
||||
int i = 0;
|
||||
while (i + 1 < res.size()) {
|
||||
if (res[i].second == res[i + 1].second) {
|
||||
res[i].first += res[i + 1].first;
|
||||
res.erase(res.begin() + i + 1);
|
||||
} else {
|
||||
++i;
|
||||
}
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
static size_t get_utf8_char_len(char c) {
|
||||
unsigned char uc = static_cast<unsigned char>(c);
|
||||
if ((uc & 0x80) == 0) {
|
||||
return 1;
|
||||
}
|
||||
if ((uc & 0xE0) == 0xC0) {
|
||||
return 2;
|
||||
}
|
||||
if ((uc & 0xF0) == 0xE0) {
|
||||
return 3;
|
||||
}
|
||||
if ((uc & 0xF8) == 0xF0) {
|
||||
return 4;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
static bool is_ascii_alpha(char c) {
|
||||
return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z');
|
||||
}
|
||||
|
||||
static bool starts_with_at(const std::string& text, size_t pos, const std::string& needle) {
|
||||
return pos + needle.size() <= text.size() && text.compare(pos, needle.size(), needle) == 0;
|
||||
}
|
||||
|
||||
static bool is_word_internal_apostrophe(const std::string& text, size_t pos) {
|
||||
return pos > 0 && pos + 1 < text.size() &&
|
||||
is_ascii_alpha(text[pos - 1]) && is_ascii_alpha(text[pos + 1]);
|
||||
}
|
||||
|
||||
static std::vector<std::pair<std::string, bool>> split_quotation(const std::string& text) {
|
||||
static const std::vector<std::pair<std::string, std::string>> quote_pairs = {
|
||||
{"'", "'"},
|
||||
{"\"", "\""},
|
||||
{"\xE2\x80\x98", "\xE2\x80\x99"},
|
||||
{"\xE2\x80\x9C", "\xE2\x80\x9D"},
|
||||
};
|
||||
|
||||
std::vector<std::pair<std::string, bool>> result;
|
||||
size_t segment_start = 0;
|
||||
size_t i = 0;
|
||||
|
||||
auto push_segment = [&](size_t begin, size_t end, bool matched) {
|
||||
if (end > begin) {
|
||||
result.emplace_back(text.substr(begin, end - begin), matched);
|
||||
}
|
||||
};
|
||||
|
||||
while (i < text.size()) {
|
||||
bool matched_quote = false;
|
||||
for (const auto& quote_pair : quote_pairs) {
|
||||
const std::string& open_quote = quote_pair.first;
|
||||
const std::string& close_quote = quote_pair.second;
|
||||
if (!starts_with_at(text, i, open_quote)) {
|
||||
continue;
|
||||
}
|
||||
if (open_quote == "'" && is_word_internal_apostrophe(text, i)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t search_pos = i + open_quote.size();
|
||||
size_t close_pos = std::string::npos;
|
||||
bool invalid = false;
|
||||
while (search_pos < text.size()) {
|
||||
if (open_quote != close_quote && starts_with_at(text, search_pos, open_quote)) {
|
||||
invalid = true;
|
||||
break;
|
||||
}
|
||||
if (starts_with_at(text, search_pos, close_quote)) {
|
||||
if (close_quote == "'" && is_word_internal_apostrophe(text, search_pos)) {
|
||||
search_pos += close_quote.size();
|
||||
continue;
|
||||
}
|
||||
close_pos = search_pos;
|
||||
break;
|
||||
}
|
||||
|
||||
size_t char_len = get_utf8_char_len(text[search_pos]);
|
||||
if (search_pos + char_len > text.size()) {
|
||||
char_len = 1;
|
||||
}
|
||||
search_pos += char_len;
|
||||
}
|
||||
if (invalid || close_pos == std::string::npos) {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t quote_start = i;
|
||||
push_segment(segment_start, quote_start, false);
|
||||
i = close_pos + close_quote.size();
|
||||
push_segment(quote_start, i, true);
|
||||
segment_start = i;
|
||||
matched_quote = true;
|
||||
break;
|
||||
}
|
||||
if (!matched_quote) {
|
||||
size_t char_len = get_utf8_char_len(text[i]);
|
||||
if (i + char_len > text.size()) {
|
||||
char_len = 1;
|
||||
}
|
||||
i += char_len;
|
||||
}
|
||||
}
|
||||
|
||||
push_segment(segment_start, text.size(), false);
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<std::pair<std::string, float>> split_quotation_attention(
|
||||
const std::vector<std::pair<std::string, float>>& parsed_attention) {
|
||||
std::vector<std::pair<std::string, float>> result;
|
||||
for (const auto& item : parsed_attention) {
|
||||
const std::string& text = item.first;
|
||||
float weight = item.second;
|
||||
for (const auto& part : split_quotation(text)) {
|
||||
if (part.second) {
|
||||
size_t i = 0;
|
||||
while (i < part.first.size()) {
|
||||
size_t char_len = get_utf8_char_len(part.first[i]);
|
||||
if (i + char_len > part.first.size()) {
|
||||
char_len = 1;
|
||||
}
|
||||
result.emplace_back(part.first.substr(i, char_len), weight);
|
||||
i += char_len;
|
||||
}
|
||||
} else {
|
||||
result.emplace_back(part.first, weight);
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
106
src/core/util.h
Normal file
106
src/core/util.h
Normal file
@@ -0,0 +1,106 @@
|
||||
#ifndef __SD_CORE_UTIL_H__
|
||||
#define __SD_CORE_UTIL_H__
|
||||
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "core/tensor.hpp"
|
||||
#include "ggml-backend.h"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
#define SAFE_STR(s) ((s) ? (s) : "")
|
||||
#define BOOL_STR(b) ((b) ? "true" : "false")
|
||||
|
||||
bool ends_with(const std::string& str, const std::string& ending);
|
||||
bool starts_with(const std::string& str, const std::string& start);
|
||||
bool contains(const std::string& str, const std::string& substr);
|
||||
|
||||
std::string sd_format(const char* fmt, ...);
|
||||
|
||||
void replace_all_chars(std::string& str, char target, char replacement);
|
||||
|
||||
int round_up_to(int value, int base);
|
||||
|
||||
bool file_exists(const std::string& filename);
|
||||
bool is_directory(const std::string& path);
|
||||
|
||||
std::u32string utf8_to_utf32(const std::string& utf8_str);
|
||||
std::string utf32_to_utf8(const std::u32string& utf32_str);
|
||||
std::u32string unicode_value_to_utf32(int unicode_value);
|
||||
// std::string sd_basename(const std::string& path);
|
||||
|
||||
sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index = 0);
|
||||
|
||||
sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
|
||||
int target_width = -1,
|
||||
int target_height = -1,
|
||||
bool scale = true);
|
||||
|
||||
sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_width, int target_height);
|
||||
|
||||
class MmapWrapper {
|
||||
public:
|
||||
static std::unique_ptr<MmapWrapper> create(const std::string& filename, bool writable = false);
|
||||
|
||||
virtual ~MmapWrapper() = default;
|
||||
|
||||
MmapWrapper(const MmapWrapper&) = delete;
|
||||
MmapWrapper& operator=(const MmapWrapper&) = delete;
|
||||
MmapWrapper(MmapWrapper&&) = delete;
|
||||
MmapWrapper& operator=(MmapWrapper&&) = delete;
|
||||
|
||||
const uint8_t* data() const { return static_cast<uint8_t*>(data_); }
|
||||
uint8_t* writable_data() { return static_cast<uint8_t*>(data_); }
|
||||
size_t size() const { return size_; }
|
||||
bool copy_data(void* buf, size_t n, size_t offset) const;
|
||||
|
||||
protected:
|
||||
MmapWrapper(void* data, size_t size)
|
||||
: data_(data), size_(size) {}
|
||||
void* data_ = nullptr;
|
||||
size_t size_ = 0;
|
||||
};
|
||||
|
||||
std::string path_join(const std::string& p1, const std::string& p2);
|
||||
std::vector<std::string> split_string(const std::string& str, char delimiter);
|
||||
|
||||
using KeyValueArgs = std::vector<std::pair<std::string, std::string>>;
|
||||
|
||||
KeyValueArgs parse_key_value_args(const char* args, const char* context = "key=value arg");
|
||||
KeyValueArgs parse_key_value_args(const std::string& args, const char* context = "key=value arg");
|
||||
bool parse_strict_float(const std::string& text, float& value);
|
||||
bool parse_strict_int(const std::string& text, int& value);
|
||||
bool parse_strict_bool(const std::string& text, bool& value);
|
||||
|
||||
void pretty_progress(int step, int steps, float time);
|
||||
void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds);
|
||||
|
||||
void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
|
||||
|
||||
std::string trim(const std::string& s);
|
||||
|
||||
std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::string& text);
|
||||
std::vector<std::pair<std::string, float>> split_quotation_attention(
|
||||
const std::vector<std::pair<std::string, float>>& parsed_attention);
|
||||
|
||||
sd_progress_cb_t sd_get_progress_callback();
|
||||
void* sd_get_progress_callback_data();
|
||||
|
||||
sd_preview_cb_t sd_get_preview_callback();
|
||||
void* sd_get_preview_callback_data();
|
||||
preview_t sd_get_preview_mode();
|
||||
int sd_get_preview_interval();
|
||||
bool sd_should_preview_denoised();
|
||||
bool sd_should_preview_noisy();
|
||||
|
||||
// test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc.
|
||||
bool sd_backend_is(ggml_backend_t backend, const std::string& name);
|
||||
|
||||
#define LOG_DEBUG(format, ...) log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
|
||||
#define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
|
||||
#define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
|
||||
#define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
|
||||
#endif // __SD_CORE_UTIL_H__
|
||||
Reference in New Issue
Block a user