mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-10-26 07:04:10 +08:00
582 lines
22 KiB
C++
582 lines
22 KiB
C++
#include "runtime.h"
|
|
|
|
#include <algorithm>
|
|
#include <cctype>
|
|
#include <chrono>
|
|
#include <cmath>
|
|
#include <cstdlib>
|
|
#include <filesystem>
|
|
#include <mutex>
|
|
#include <regex>
|
|
#include <sstream>
|
|
|
|
#include "common/common.h"
|
|
#include "common/log.h"
|
|
#include "model_catalog.h"
|
|
|
|
#include <map>
|
|
#include <set>
|
|
|
|
namespace fs = std::filesystem;
|
|
|
|
bool validate_server_img_gen_request(const ServerRuntime& runtime,
|
|
const ImgGenJobRequest& request,
|
|
std::string& error_message) {
|
|
auto dimensions_valid = [&](double width, double height) {
|
|
if (!std::isfinite(width) || !std::isfinite(height) ||
|
|
width < 1 || height < 1 || width > k_max_generation_dimension || height > k_max_generation_dimension) {
|
|
error_message = "generation width and height must be between 1 and " + std::to_string(k_max_generation_dimension);
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
const auto& params = request.gen_params;
|
|
const int width = params.get_resolved_width();
|
|
const int height = params.get_resolved_height();
|
|
if (!dimensions_valid(width, height)) {
|
|
return false;
|
|
}
|
|
if (params.hires_enabled) {
|
|
const double hires_width = params.hires_width > 0 ? params.hires_width : params.hires_height > 0 ? params.hires_height
|
|
: std::round(width * params.hires_scale);
|
|
const double hires_height = params.hires_height > 0 ? params.hires_height : params.hires_width > 0 ? params.hires_width
|
|
: std::round(height * params.hires_scale);
|
|
if (!dimensions_valid(hires_width, hires_height)) {
|
|
return false;
|
|
}
|
|
}
|
|
if (!request.adetailer_enabled || runtime.default_gen_params->ad_model_path.empty()) {
|
|
return true;
|
|
}
|
|
auto pass_valid = [&](const ServerADetailerPass& pass) {
|
|
if (!dimensions_valid(pass.width, pass.height)) {
|
|
return false;
|
|
}
|
|
// Match the native key normalization, but leave full option parsing to ADetailer.
|
|
auto trim = [](const std::string& text) {
|
|
const auto first = text.find_first_not_of(" \t\r\n");
|
|
return first == std::string::npos ? std::string() : text.substr(first, text.find_last_not_of(" \t\r\n") - first + 1);
|
|
};
|
|
std::stringstream stream(pass.extra_args);
|
|
std::string item;
|
|
while (std::getline(stream, item, ',')) {
|
|
const auto equals = item.find('=');
|
|
if (equals == std::string::npos) {
|
|
continue;
|
|
}
|
|
auto key = trim(item.substr(0, equals));
|
|
std::transform(key.begin(), key.end(), key.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
|
if (key != "inpaint_width" && key != "inpaint_height" && key != "steps") {
|
|
continue;
|
|
}
|
|
const auto value = trim(item.substr(equals + 1));
|
|
const int minimum = key == "steps" ? 0 : 1;
|
|
const int maximum = key == "steps" ? k_max_adetailer_steps : k_max_generation_dimension;
|
|
try {
|
|
size_t used = 0;
|
|
const int parsed = std::stoi(value, &used);
|
|
if (used == value.size() && parsed >= minimum && parsed <= maximum) {
|
|
continue;
|
|
}
|
|
} catch (const std::exception&) {
|
|
}
|
|
error_message = "ADetailer " + key + " must be between " + std::to_string(minimum) + " and " + std::to_string(maximum);
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
if (request.adetailer_passes.empty()) {
|
|
ServerADetailerPass pass;
|
|
pass.extra_args = params.extra_ad_args;
|
|
return pass_valid(pass);
|
|
}
|
|
return std::all_of(request.adetailer_passes.begin(), request.adetailer_passes.end(), pass_valid);
|
|
}
|
|
|
|
bool apply_server_adetailer(ServerRuntime& runtime,
|
|
const ImgGenJobRequest& request,
|
|
const sd_img_gen_params_t& params,
|
|
SDImageVec& results) {
|
|
const auto& gen_params = request.gen_params;
|
|
const auto& ad_model_path = runtime.default_gen_params->ad_model_path;
|
|
if (!request.adetailer_enabled || ad_model_path.empty()) {
|
|
return true;
|
|
}
|
|
std::string error_message;
|
|
if (!validate_server_img_gen_request(runtime, request, error_message)) {
|
|
LOG_ERROR("%s", error_message.c_str());
|
|
return false;
|
|
}
|
|
|
|
std::vector<ServerADetailerPass> passes = request.adetailer_passes;
|
|
if (passes.empty()) {
|
|
ServerADetailerPass pass;
|
|
pass.prompt = gen_params.ad_prompt;
|
|
pass.negative_prompt = gen_params.ad_negative_prompt;
|
|
pass.extra_args = gen_params.extra_ad_args;
|
|
passes.push_back(std::move(pass));
|
|
}
|
|
|
|
ADetailerCtxPtr ad_ctx(new_adetailer_ctx(ad_model_path.c_str(),
|
|
runtime.ctx_params->n_threads,
|
|
runtime.ctx_params->effective_backend.c_str(),
|
|
runtime.ctx_params->effective_params_backend.c_str()));
|
|
if (ad_ctx == nullptr) {
|
|
LOG_ERROR("new_adetailer_ctx failed");
|
|
return false;
|
|
}
|
|
|
|
for (const auto& pass : passes) {
|
|
sd_adetailer_params_t ad_params{};
|
|
ad_params.prompt = pass.prompt.empty() ? nullptr : pass.prompt.c_str();
|
|
ad_params.negative_prompt = pass.negative_prompt.empty() ? nullptr : pass.negative_prompt.c_str();
|
|
ad_params.extra_ad_args = pass.extra_args.c_str();
|
|
for (int i = 0; i < results.count(); ++i) {
|
|
if (results[i].data == nullptr) {
|
|
continue;
|
|
}
|
|
sd_img_gen_params_t ad_generation_params = params;
|
|
ad_generation_params.seed = generation_seed(results.generation_info(i));
|
|
ad_generation_params.width = pass.width;
|
|
ad_generation_params.height = pass.height;
|
|
ad_generation_params.strength = pass.strength;
|
|
if (pass.clip_skip > 0) {
|
|
ad_generation_params.clip_skip = pass.clip_skip;
|
|
}
|
|
if (pass.automatic_scheduler) {
|
|
ad_generation_params.sample_params.scheduler = SCHEDULER_COUNT;
|
|
}
|
|
|
|
sd_image_generation_result_t* detailed = nullptr;
|
|
if (!adetail_image(ad_ctx.get(),
|
|
runtime.sd_ctx,
|
|
results[i],
|
|
&ad_params,
|
|
&ad_generation_params,
|
|
&detailed) ||
|
|
detailed == nullptr || detailed->count <= 0 || detailed->items[0].image.data == nullptr) {
|
|
free_sd_image_generation_result(detailed);
|
|
LOG_ERROR("ADetailer failed for image %d", i + 1);
|
|
return false;
|
|
}
|
|
results.append_generation_info(i, detailed->items[0].generation_info, runtime.ctx_params->lora_model_dir);
|
|
free(results[i].data);
|
|
results[i] = detailed->items[0].image;
|
|
detailed->items[0].image = {};
|
|
free_sd_image_generation_result(detailed);
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static std::string lower_ascii(std::string value) {
|
|
std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
|
|
return static_cast<char>(std::tolower(c));
|
|
});
|
|
return value;
|
|
}
|
|
|
|
static bool is_supported_model_ext(const fs::path& p) {
|
|
auto ext = lower_ascii(p.extension().u8string());
|
|
return ext == ".gguf" || ext == ".pt" || ext == ".pth" || ext == ".safetensors";
|
|
}
|
|
|
|
static const std::string k_base64_chars =
|
|
"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
|
|
"abcdefghijklmnopqrstuvwxyz"
|
|
"0123456789+/";
|
|
|
|
std::string base64_encode(const std::vector<uint8_t>& bytes) {
|
|
std::string ret;
|
|
int val = 0;
|
|
int valb = -6;
|
|
for (uint8_t c : bytes) {
|
|
val = (val << 8) + c;
|
|
valb += 8;
|
|
while (valb >= 0) {
|
|
ret.push_back(k_base64_chars[(val >> valb) & 0x3F]);
|
|
valb -= 6;
|
|
}
|
|
}
|
|
if (valb > -6) {
|
|
ret.push_back(k_base64_chars[((val << 8) >> (valb + 8)) & 0x3F]);
|
|
}
|
|
while (ret.size() % 4) {
|
|
ret.push_back('=');
|
|
}
|
|
return ret;
|
|
}
|
|
|
|
std::string normalize_output_format(std::string output_format) {
|
|
std::transform(output_format.begin(), output_format.end(), output_format.begin(),
|
|
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
|
return output_format;
|
|
}
|
|
|
|
std::vector<std::string> supported_img_output_formats(bool allow_webp) {
|
|
std::vector<std::string> formats = {"png", "jpeg"};
|
|
#ifdef SD_USE_WEBP
|
|
if (allow_webp) {
|
|
formats.push_back("webp");
|
|
}
|
|
#else
|
|
(void)allow_webp;
|
|
#endif
|
|
return formats;
|
|
}
|
|
|
|
std::vector<std::string> supported_vid_output_formats() {
|
|
std::vector<std::string> formats;
|
|
#ifdef SD_USE_WEBM
|
|
formats.push_back("webm");
|
|
#endif
|
|
#ifdef SD_USE_WEBP
|
|
formats.push_back("webp");
|
|
#endif
|
|
formats.push_back("avi");
|
|
return formats;
|
|
}
|
|
|
|
static std::string valid_vid_output_formats_message() {
|
|
const std::vector<std::string> formats = supported_vid_output_formats();
|
|
|
|
std::string message = "invalid output_format, must be one of [";
|
|
for (size_t i = 0; i < formats.size(); ++i) {
|
|
if (i > 0) {
|
|
message += ", ";
|
|
}
|
|
message += formats[i];
|
|
}
|
|
message += "]";
|
|
return message;
|
|
}
|
|
|
|
bool assign_output_options(ImgGenJobRequest& request,
|
|
std::string output_format,
|
|
int output_compression,
|
|
bool allow_webp,
|
|
std::string& error_message) {
|
|
request.output_format = normalize_output_format(std::move(output_format));
|
|
request.output_compression = std::clamp(output_compression, 0, 100);
|
|
|
|
const std::vector<std::string> valid_formats = supported_img_output_formats(allow_webp);
|
|
const bool valid_format = std::find(valid_formats.begin(),
|
|
valid_formats.end(),
|
|
request.output_format) != valid_formats.end();
|
|
if (!valid_format) {
|
|
error_message = "invalid output_format, must be one of [";
|
|
for (size_t i = 0; i < valid_formats.size(); ++i) {
|
|
if (i > 0) {
|
|
error_message += ", ";
|
|
}
|
|
error_message += valid_formats[i];
|
|
}
|
|
error_message += "]";
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
bool assign_output_options(VidGenJobRequest& request,
|
|
std::string output_format,
|
|
int output_compression,
|
|
std::string& error_message) {
|
|
request.output_format = normalize_output_format(std::move(output_format));
|
|
request.output_compression = std::clamp(output_compression, 0, 100);
|
|
|
|
if (request.output_format == "avi") {
|
|
return true;
|
|
}
|
|
|
|
if (request.output_format == "webm") {
|
|
#ifdef SD_USE_WEBM
|
|
return true;
|
|
#else
|
|
error_message = valid_vid_output_formats_message();
|
|
return false;
|
|
#endif
|
|
}
|
|
|
|
if (request.output_format == "webp") {
|
|
#ifdef SD_USE_WEBP
|
|
return true;
|
|
#else
|
|
error_message = valid_vid_output_formats_message();
|
|
return false;
|
|
#endif
|
|
}
|
|
|
|
error_message = valid_vid_output_formats_message();
|
|
return false;
|
|
}
|
|
|
|
std::string video_mime_type(const std::string& output_format) {
|
|
if (output_format == "webm") {
|
|
return "video/webm";
|
|
}
|
|
if (output_format == "webp") {
|
|
return "image/webp";
|
|
}
|
|
return "video/x-msvideo";
|
|
}
|
|
|
|
bool runtime_supports_generation_mode(const ServerRuntime& runtime, SDMode mode) {
|
|
if (mode == VID_GEN) {
|
|
return sd_ctx_supports_video_generation(runtime.sd_ctx);
|
|
}
|
|
if (mode == IMG_GEN) {
|
|
return sd_ctx_supports_image_generation(runtime.sd_ctx);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::string unsupported_generation_mode_error(SDMode mode) {
|
|
if (mode == VID_GEN) {
|
|
return "loaded model does not support vid_gen";
|
|
}
|
|
if (mode == IMG_GEN) {
|
|
return "loaded model does not support img_gen";
|
|
}
|
|
return "loaded model does not support requested mode";
|
|
}
|
|
|
|
ArgOptions SDSvrParams::get_options() {
|
|
ArgOptions options;
|
|
|
|
options.string_options = {
|
|
{"-l", "--listen-ip", "server listen ip (default: 127.0.0.1)", 0, &listen_ip},
|
|
{"", "--serve-html-path", "path to HTML file to serve at root (optional)", 0, &serve_html_path},
|
|
};
|
|
|
|
options.int_options = {
|
|
{"", "--listen-port", "server listen port (default: 1234)", &listen_port},
|
|
};
|
|
|
|
options.bool_options = {
|
|
{"", "--color", "colors the logging tags according to level", true, &color},
|
|
};
|
|
|
|
auto on_help_arg = [&](int, const char**, int, bool& valid) {
|
|
normal_exit = true;
|
|
valid = true;
|
|
return -1;
|
|
};
|
|
|
|
options.manual_options = {
|
|
{"-h", "--help", "show this help message and exit", on_help_arg},
|
|
{"", "--model-dir", "model root; repeat to search multiple directories (default: models)",
|
|
[this](int argc, const char** argv, int index) {
|
|
if (++index >= argc || argv[index][0] == '\0')
|
|
return -1;
|
|
model_dirs.push_back(argv_to_utf8(index, argv));
|
|
return 1;
|
|
}},
|
|
};
|
|
add_log_options(options, log_level);
|
|
return options;
|
|
}
|
|
|
|
bool SDSvrParams::validate() {
|
|
if (listen_ip.empty()) {
|
|
LOG_ERROR("error: the following arguments are required: listen_ip");
|
|
return false;
|
|
}
|
|
|
|
if (listen_port < 0 || listen_port > 65535) {
|
|
LOG_ERROR("error: listen_port should be in the range [0, 65535]");
|
|
return false;
|
|
}
|
|
|
|
if (!serve_html_path.empty() && !fs::exists(fs::u8path(serve_html_path))) {
|
|
LOG_ERROR("error: serve_html_path file does not exist: %s", serve_html_path.c_str());
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool SDSvrParams::resolve_and_validate() {
|
|
if (!validate()) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::string SDSvrParams::to_string() const {
|
|
std::ostringstream oss;
|
|
oss << "SDSvrParams {\n"
|
|
<< " log_level: " << log_level_name(log_level) << ",\n"
|
|
<< " listen_ip: " << listen_ip << ",\n"
|
|
<< " listen_port: \"" << listen_port << "\",\n"
|
|
<< " serve_html_path: \"" << serve_html_path << "\",\n"
|
|
<< "}";
|
|
return oss.str();
|
|
}
|
|
|
|
void refresh_lora_cache(ServerRuntime& rt) {
|
|
std::vector<LoraEntry> new_cache;
|
|
std::vector<ModelRoot> directories;
|
|
if (!rt.ctx_params->lora_model_dir.empty()) {
|
|
directories.push_back(make_model_root(fs::u8path(rt.ctx_params->lora_model_dir), true));
|
|
} else {
|
|
for (const auto& root : get_model_roots(*rt.svr_params)) {
|
|
try {
|
|
const auto paths = model_category_directories(root, "loras");
|
|
for (size_t i = 0; i < paths.size(); ++i) {
|
|
if (!model_path_inside(root.path, fs::canonical(paths[i]))) {
|
|
LOG_WARN("LoRA directory escapes model directory: %s", paths[i].u8string().c_str());
|
|
continue;
|
|
}
|
|
directories.push_back(make_model_root(paths[i], root.primary && lower_ascii(paths[i].filename().u8string()) == "loras"));
|
|
}
|
|
} catch (const std::exception& e) {
|
|
LOG_WARN("error while finding LoRA directories in %s: %s", root.label.c_str(), e.what());
|
|
}
|
|
}
|
|
}
|
|
std::set<std::string> seen;
|
|
std::map<std::string, size_t> names;
|
|
for (const auto& directory : directories) {
|
|
try {
|
|
if (!fs::exists(directory.path))
|
|
continue;
|
|
for (fs::recursive_directory_iterator it(directory.path), end; it != end; ++it) {
|
|
std::error_code ec;
|
|
const auto canonical = fs::canonical(it->path(), ec);
|
|
if (ec || !model_path_inside(directory.path, canonical)) {
|
|
it.disable_recursion_pending();
|
|
LOG_WARN("skipping inaccessible LoRA or path outside LoRA directory: %s", it->path().u8string().c_str());
|
|
continue;
|
|
}
|
|
const auto ext = lower_ascii(it->path().extension().u8string());
|
|
if (!it->is_regular_file() || (!is_supported_model_ext(it->path()) && ext != ".ckpt" && ext != ".sft") ||
|
|
!seen.insert(model_path_key(canonical)).second)
|
|
continue;
|
|
const auto relative = it->path().lexically_relative(directory.path).generic_u8string();
|
|
LoraEntry entry;
|
|
entry.name = it->path().stem().u8string();
|
|
entry.path = directory.primary ? relative : directory.id + "/" + relative;
|
|
entry.fullpath = it->path().u8string();
|
|
entry.display_path = relative;
|
|
entry.root = directory.path;
|
|
++names[relative];
|
|
new_cache.push_back(std::move(entry));
|
|
}
|
|
} catch (const std::exception& e) {
|
|
LOG_WARN("error while scanning LoRA directory '%s': %s", directory.path.u8string().c_str(), e.what());
|
|
}
|
|
}
|
|
for (auto& entry : new_cache) {
|
|
if (names[entry.display_path] > 1)
|
|
entry.display_path = fs::u8path(entry.fullpath).generic_u8string();
|
|
}
|
|
std::sort(new_cache.begin(), new_cache.end(), [](const LoraEntry& a, const LoraEntry& b) {
|
|
return a.path < b.path;
|
|
});
|
|
std::lock_guard<std::mutex> lock(*rt.lora_mutex);
|
|
*rt.lora_cache = std::move(new_cache);
|
|
}
|
|
|
|
std::string get_lora_full_path(ServerRuntime& rt, const std::string& path) {
|
|
std::lock_guard<std::mutex> lock(*rt.lora_mutex);
|
|
auto it = std::find_if(rt.lora_cache->begin(), rt.lora_cache->end(),
|
|
[&](const LoraEntry& entry) { return entry.path == path; });
|
|
if (it == rt.lora_cache->end())
|
|
return "";
|
|
std::error_code ec;
|
|
const auto canonical = fs::canonical(fs::u8path(it->fullpath), ec);
|
|
if (ec || !model_path_inside(it->root, canonical) || !fs::is_regular_file(canonical, ec) || ec)
|
|
return "";
|
|
return canonical.u8string();
|
|
}
|
|
|
|
void refresh_upscaler_cache(ServerRuntime& rt) {
|
|
std::vector<UpscalerEntry> new_cache;
|
|
std::vector<UpscalerEntry> previous_cache;
|
|
{
|
|
std::lock_guard<std::mutex> lock(*rt.upscaler_mutex);
|
|
previous_cache = *rt.upscaler_cache;
|
|
}
|
|
|
|
fs::path upscaler_dir = fs::u8path(rt.ctx_params->hires_upscalers_dir);
|
|
std::error_code ec;
|
|
if (fs::exists(upscaler_dir, ec) && !ec && fs::is_directory(upscaler_dir, ec) && !ec) {
|
|
try {
|
|
auto it = fs::directory_iterator(
|
|
upscaler_dir,
|
|
fs::directory_options::skip_permission_denied,
|
|
ec);
|
|
auto end = fs::directory_iterator();
|
|
while (!ec && it != end) {
|
|
std::error_code entry_ec;
|
|
bool is_reg = it->is_regular_file(entry_ec);
|
|
if (entry_ec || !is_reg) {
|
|
it.increment(ec);
|
|
continue;
|
|
}
|
|
const fs::path& p = it->path();
|
|
if (!is_supported_model_ext(p)) {
|
|
it.increment(ec);
|
|
continue;
|
|
}
|
|
|
|
UpscalerEntry upscaler_entry;
|
|
upscaler_entry.name = p.stem().u8string();
|
|
std::error_code abs_ec;
|
|
fs::path abs_path = fs::absolute(p, abs_ec);
|
|
upscaler_entry.fullpath = (abs_ec ? p : abs_path.lexically_normal()).u8string();
|
|
upscaler_entry.model_name = "ESRGAN_4x";
|
|
upscaler_entry.path = p.filename().u8string();
|
|
|
|
std::error_code size_ec;
|
|
upscaler_entry.file_size = it->file_size(size_ec);
|
|
if (size_ec) {
|
|
it.increment(ec);
|
|
continue;
|
|
}
|
|
|
|
std::error_code time_ec;
|
|
upscaler_entry.last_modified = it->last_write_time(time_ec);
|
|
auto previous = std::find_if(previous_cache.begin(), previous_cache.end(), [&](const UpscalerEntry& cached) {
|
|
return cached.fullpath == upscaler_entry.fullpath &&
|
|
cached.file_size == upscaler_entry.file_size &&
|
|
(!time_ec && cached.last_modified == upscaler_entry.last_modified);
|
|
});
|
|
upscaler_entry.image_upscale_factor = previous != previous_cache.end()
|
|
? previous->image_upscale_factor
|
|
: get_upscaler_model_scale(upscaler_entry.fullpath.c_str());
|
|
if (upscaler_entry.image_upscale_factor > 0) {
|
|
upscaler_entry.scale = upscaler_entry.image_upscale_factor;
|
|
upscaler_entry.model_name = "ESRGAN_" + std::to_string(upscaler_entry.scale) + "x";
|
|
}
|
|
|
|
new_cache.push_back(std::move(upscaler_entry));
|
|
it.increment(ec);
|
|
}
|
|
} catch (const std::exception& e) {
|
|
LOG_WARN("error while scanning upscalers directory '%s': %s", upscaler_dir.u8string().c_str(), e.what());
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (ec) {
|
|
LOG_WARN("error while scanning upscalers directory '%s': %s", upscaler_dir.u8string().c_str(), ec.message().c_str());
|
|
return;
|
|
}
|
|
|
|
std::sort(new_cache.begin(), new_cache.end(), [](const UpscalerEntry& a, const UpscalerEntry& b) {
|
|
return a.name < b.name;
|
|
});
|
|
|
|
{
|
|
std::lock_guard<std::mutex> lock(*rt.upscaler_mutex);
|
|
*rt.upscaler_cache = std::move(new_cache);
|
|
}
|
|
}
|
|
|
|
int64_t unix_timestamp_now_ms() {
|
|
return std::chrono::duration_cast<std::chrono::milliseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count();
|
|
}
|