mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-10-26 07:04:10 +08:00
671 lines
31 KiB
C++
671 lines
31 KiB
C++
#include "model_catalog.h"
|
|
|
|
#include <algorithm>
|
|
#include <cctype>
|
|
#include <cmath>
|
|
#include <fstream>
|
|
#include <iomanip>
|
|
#include <set>
|
|
#include <sstream>
|
|
#include <stdexcept>
|
|
#include <vector>
|
|
|
|
#include "common/log.h"
|
|
#include "recipes_builtin.h"
|
|
|
|
namespace fs = std::filesystem;
|
|
|
|
struct ModelSlot {
|
|
const char* id;
|
|
const char* label;
|
|
std::vector<std::string> categories;
|
|
std::string SDContextParams::* field;
|
|
};
|
|
|
|
static const ModelSlot k_model_slots[] = {
|
|
{"model", "Checkpoint", {"checkpoints"}, &SDContextParams::model_path},
|
|
{"diffusion_model", "Diffusion model", {"diffusion_models"}, &SDContextParams::diffusion_model_path},
|
|
{"high_noise_diffusion_model", "High-noise diffusion model", {"diffusion_models"}, &SDContextParams::high_noise_diffusion_model_path},
|
|
{"uncond_diffusion_model", "Unconditional diffusion model", {"diffusion_models"}, &SDContextParams::uncond_diffusion_model_path},
|
|
{"clip_l", "CLIP-L", {"text_encoders"}, &SDContextParams::clip_l_path},
|
|
{"clip_g", "CLIP-G", {"text_encoders"}, &SDContextParams::clip_g_path},
|
|
{"clip_vision", "CLIP Vision", {"clip_vision"}, &SDContextParams::clip_vision_path},
|
|
{"t5xxl", "T5", {"text_encoders"}, &SDContextParams::t5xxl_path},
|
|
{"llm", "LLM Text Encoder", {"text_encoders"}, &SDContextParams::llm_path},
|
|
{"llm_vision", "LLM Vision Encoder", {"text_encoders", "clip_vision"}, &SDContextParams::llm_vision_path},
|
|
{"tokenizer", "Tokenizer", {"tokenizers"}, &SDContextParams::tokenizer},
|
|
{"vae", "VAE", {"vae"}, &SDContextParams::vae_path},
|
|
{"audio_vae", "Audio VAE", {"vae"}, &SDContextParams::audio_vae_path},
|
|
{"audio_encoder", "Audio Encoder", {"text_encoders"}, &SDContextParams::audio_encoder_path},
|
|
{"embeddings_connectors", "Embeddings Connectors", {"text_encoders"}, &SDContextParams::embeddings_connectors_path},
|
|
{"tae", "TAE", {"vae", "vae_approx"}, &SDContextParams::taesd_path},
|
|
{"taesd", "TAE", {"vae_approx", "vae"}, &SDContextParams::taesd_path},
|
|
{"control_net", "ControlNet", {"controlnet"}, &SDContextParams::control_net_path},
|
|
{"ip_adapter", "IP-Adapter", {"ipadapter"}, &SDContextParams::ip_adapter_path},
|
|
{"motion_module", "Motion Module", {"motion_modules"}, &SDContextParams::motion_module_path},
|
|
{"photo_maker", "PhotoMaker", {"photomaker"}, &SDContextParams::photo_maker_path},
|
|
{"pulid_weights", "PuLID", {"pulid"}, &SDContextParams::pulid_weights_path},
|
|
};
|
|
|
|
static const ModelSlot* find_model_slot(const std::string& id) {
|
|
for (const auto& slot : k_model_slots) {
|
|
if (id == slot.id) {
|
|
return &slot;
|
|
}
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
static bool matches_model_category(const ModelSlot& slot, const std::string& category) {
|
|
return std::find(slot.categories.begin(), slot.categories.end(), category) != slot.categories.end();
|
|
}
|
|
|
|
static void clear_model_paths(SDContextParams& params) {
|
|
for (const auto& slot : k_model_slots) {
|
|
(params.*slot.field).clear();
|
|
}
|
|
}
|
|
|
|
std::string model_path_key(const fs::path& path) {
|
|
#ifdef _WIN32
|
|
const auto value = path.generic_wstring();
|
|
const int length = LCMapStringEx(LOCALE_NAME_INVARIANT, LCMAP_LOWERCASE, value.data(),
|
|
static_cast<int>(value.size()), nullptr, 0, nullptr, nullptr, 0);
|
|
if (length > 0) {
|
|
std::wstring folded(length, L'\0');
|
|
LCMapStringEx(LOCALE_NAME_INVARIANT, LCMAP_LOWERCASE, value.data(),
|
|
static_cast<int>(value.size()), &folded[0], length, nullptr, nullptr, 0);
|
|
return fs::path(folded).generic_u8string();
|
|
}
|
|
#endif
|
|
return path.generic_u8string();
|
|
}
|
|
|
|
bool model_path_inside(const fs::path& root, const fs::path& path) {
|
|
auto parent = root.begin();
|
|
auto child = path.begin();
|
|
for (; parent != root.end(); ++parent, ++child) {
|
|
if (child == path.end() || model_path_key(*parent) != model_path_key(*child)) {
|
|
return false;
|
|
}
|
|
}
|
|
return child != path.end();
|
|
}
|
|
|
|
ModelRoot make_model_root(const fs::path& path, bool primary) {
|
|
ModelRoot root;
|
|
const auto absolute = fs::absolute(path);
|
|
std::error_code ec;
|
|
root.path = fs::weakly_canonical(absolute, ec);
|
|
if (ec)
|
|
root.path = absolute.lexically_normal();
|
|
root.primary = primary;
|
|
uint64_t hash = 14695981039346656037ULL;
|
|
for (unsigned char c : model_path_key(root.path)) {
|
|
hash ^= c;
|
|
hash *= 1099511628211ULL;
|
|
}
|
|
std::ostringstream id;
|
|
id << '@' << std::hex << std::setfill('0') << std::setw(16) << hash;
|
|
root.id = id.str();
|
|
root.label = root.path.filename().u8string() + " [" + root.id.substr(1, 8) + "]";
|
|
return root;
|
|
}
|
|
|
|
std::vector<ModelRoot> get_model_roots(const SDSvrParams& params) {
|
|
const auto directories = params.model_dirs.empty() ? std::vector<std::string>{"models"} : params.model_dirs;
|
|
std::vector<ModelRoot> roots;
|
|
std::set<std::string> seen;
|
|
for (size_t i = 0; i < directories.size(); ++i) {
|
|
auto root = make_model_root(fs::u8path(directories[i]), i == 0);
|
|
if (seen.insert(model_path_key(root.path)).second)
|
|
roots.push_back(std::move(root));
|
|
}
|
|
return roots;
|
|
}
|
|
|
|
std::vector<fs::path> model_category_directories(const ModelRoot& root, const std::string& category) {
|
|
std::vector<fs::path> directories;
|
|
if (!fs::exists(root.path))
|
|
return directories;
|
|
for (const auto& entry : fs::directory_iterator(root.path)) {
|
|
auto name = entry.path().filename().u8string();
|
|
std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return std::tolower(c); });
|
|
if (name == category || (category == "checkpoints" && name == "stable-diffusion") ||
|
|
(category == "loras" && name == "lora")) {
|
|
directories.push_back(entry.path());
|
|
}
|
|
}
|
|
std::sort(directories.begin(), directories.end(), [&](const fs::path& a, const fs::path& b) {
|
|
const bool a_standard = a.filename().u8string() == category;
|
|
const bool b_standard = b.filename().u8string() == category;
|
|
return a_standard != b_standard ? a_standard : a < b;
|
|
});
|
|
return directories;
|
|
}
|
|
|
|
static bool valid_identifier(const std::string& id) {
|
|
return !id.empty() && id.size() <= 128 &&
|
|
std::all_of(id.begin(), id.end(), [](unsigned char c) {
|
|
return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
|
|
(c >= '0' && c <= '9') || c == '-' || c == '_' || c == '.';
|
|
});
|
|
}
|
|
|
|
static void check_keys(const json& object, const std::set<std::string>& keys) {
|
|
if (!object.is_object()) {
|
|
throw std::runtime_error("expected an object");
|
|
}
|
|
for (auto it = object.begin(); it != object.end(); ++it) {
|
|
if (!keys.count(it.key())) {
|
|
throw std::runtime_error("unknown field: " + it.key());
|
|
}
|
|
}
|
|
}
|
|
|
|
static void apply_recipe_defaults(const json& defaults, SDGenerationParams& params) {
|
|
check_keys(defaults, {"width", "height", "steps", "cfg_scale", "strength", "distilled_guidance", "sampling_method", "scheduler", "custom_sigmas",
|
|
"video_frames", "fps", "flow_shift", "moe_boundary", "high_noise_steps", "high_noise_cfg_scale",
|
|
"high_noise_flow_shift", "high_noise_sampling_method", "high_noise_scheduler", "img_cfg_scale", "shifted_timestep"});
|
|
for (const char* key : {"width", "height", "steps", "video_frames", "fps", "high_noise_steps", "shifted_timestep"}) {
|
|
if (!defaults.contains(key)) {
|
|
continue;
|
|
}
|
|
if (!defaults[key].is_number_integer()) {
|
|
throw std::runtime_error(std::string(key) + " must be an integer");
|
|
}
|
|
const int64_t value = defaults[key].get<int64_t>();
|
|
const std::string field = key;
|
|
const bool dimension = field == "width" || field == "height";
|
|
const int64_t maximum = dimension || field == "video_frames" ? 4096 : field == "fps" ? 240
|
|
: 1000;
|
|
const int64_t minimum = dimension ? 64 : field == "shifted_timestep" ? 0
|
|
: 1;
|
|
if ((value < minimum || value > maximum) && !(field == "high_noise_steps" && value == -1)) {
|
|
throw std::runtime_error(std::string(key) + " is out of range");
|
|
}
|
|
if (std::string(key) == "width")
|
|
params.width = static_cast<int>(value);
|
|
if (std::string(key) == "height")
|
|
params.height = static_cast<int>(value);
|
|
if (std::string(key) == "steps")
|
|
params.sample_params.sample_steps = static_cast<int>(value);
|
|
if (field == "shifted_timestep")
|
|
params.sample_params.shifted_timestep = static_cast<int>(value);
|
|
if (field == "video_frames")
|
|
params.video_frames = static_cast<int>(value);
|
|
if (field == "fps")
|
|
params.fps = static_cast<int>(value);
|
|
if (field == "high_noise_steps")
|
|
params.high_noise_sample_params.sample_steps = static_cast<int>(value);
|
|
}
|
|
for (const char* key : {"cfg_scale", "strength", "moe_boundary", "img_cfg_scale", "distilled_guidance", "flow_shift", "high_noise_cfg_scale", "high_noise_flow_shift"}) {
|
|
if (!defaults.contains(key)) {
|
|
continue;
|
|
}
|
|
if (!defaults[key].is_number()) {
|
|
throw std::runtime_error(std::string(key) + " must be a number");
|
|
}
|
|
const float value = defaults[key].get<float>();
|
|
if (!std::isfinite(value) || value < 0 || value > ((std::string(key) == "strength" || std::string(key) == "moe_boundary") ? 1.f : 1000.f)) {
|
|
throw std::runtime_error(std::string(key) + " is out of range");
|
|
}
|
|
if (std::string(key) == "strength")
|
|
params.strength = value;
|
|
if (std::string(key) == "moe_boundary")
|
|
params.moe_boundary = value;
|
|
if (std::string(key) == "cfg_scale")
|
|
params.sample_params.guidance.txt_cfg = value;
|
|
if (std::string(key) == "img_cfg_scale")
|
|
params.sample_params.guidance.img_cfg = value;
|
|
if (std::string(key) == "distilled_guidance")
|
|
params.sample_params.guidance.distilled_guidance = value;
|
|
if (std::string(key) == "flow_shift")
|
|
params.sample_params.flow_shift = value;
|
|
if (std::string(key) == "high_noise_cfg_scale")
|
|
params.high_noise_sample_params.guidance.txt_cfg = value;
|
|
if (std::string(key) == "high_noise_flow_shift")
|
|
params.high_noise_sample_params.flow_shift = value;
|
|
}
|
|
if (defaults.contains("custom_sigmas")) {
|
|
const auto& sigmas = defaults.at("custom_sigmas");
|
|
if (!sigmas.is_array() || sigmas.size() == 1 || sigmas.size() > 1001)
|
|
throw std::runtime_error("custom_sigmas must be empty or contain 2-1001 values");
|
|
for (const auto& sigma : sigmas) {
|
|
if (!sigma.is_number() || !std::isfinite(sigma.get<float>()) || sigma.get<float>() < 0)
|
|
throw std::runtime_error("custom_sigmas must contain finite nonnegative numbers");
|
|
}
|
|
params.custom_sigmas = sigmas.get<std::vector<float>>();
|
|
}
|
|
for (const char* key : {"sampling_method", "high_noise_sampling_method"}) {
|
|
if (!defaults.contains(key))
|
|
continue;
|
|
auto& sample = std::string(key) == "sampling_method" ? params.sample_params : params.high_noise_sample_params;
|
|
const auto name = defaults.at(key).get<std::string>();
|
|
bool found = false;
|
|
for (int i = 0; i < SAMPLE_METHOD_COUNT; ++i) {
|
|
if (name == sd_sample_method_name(static_cast<sample_method_t>(i))) {
|
|
sample.sample_method = static_cast<sample_method_t>(i);
|
|
found = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!found)
|
|
throw std::runtime_error(std::string("unknown ") + key);
|
|
}
|
|
for (const char* key : {"scheduler", "high_noise_scheduler"}) {
|
|
if (!defaults.contains(key))
|
|
continue;
|
|
auto& sample = std::string(key) == "scheduler" ? params.sample_params : params.high_noise_sample_params;
|
|
const auto name = defaults.at(key).get<std::string>();
|
|
bool found = false;
|
|
for (int i = 0; i < SCHEDULER_COUNT; ++i) {
|
|
if (name == sd_scheduler_name(static_cast<scheduler_t>(i)) ||
|
|
(name == "normal" && i == DISCRETE_SCHEDULER)) {
|
|
sample.scheduler = static_cast<scheduler_t>(i);
|
|
found = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!found)
|
|
throw std::runtime_error(std::string("unknown ") + key);
|
|
}
|
|
}
|
|
|
|
static void apply_recipe_context(const json& settings, SDContextParams& params) {
|
|
check_keys(settings, {"vae_format", "model_args", "backend", "params_backend"});
|
|
if (settings.contains("vae_format")) {
|
|
const auto format = settings.at("vae_format").get<std::string>();
|
|
if (format != "auto" && format != "flux" && format != "sd3" && format != "flux2" && format != "wan")
|
|
throw std::runtime_error("unknown vae_format");
|
|
params.vae_format = format;
|
|
}
|
|
if (settings.contains("model_args")) {
|
|
params.model_args = settings.at("model_args").get<std::string>();
|
|
}
|
|
if (settings.contains("backend")) {
|
|
params.backend = settings.at("backend").get<std::string>();
|
|
}
|
|
if (settings.contains("params_backend")) {
|
|
params.params_backend = settings.at("params_backend").get<std::string>();
|
|
}
|
|
}
|
|
|
|
static const std::map<std::string, bool> k_recipe_features = {
|
|
{"img2img", true},
|
|
{"mask_image", true},
|
|
{"lora", true},
|
|
{"vae_tiling", true},
|
|
{"cache", true},
|
|
{"cache_dit", true},
|
|
{"preview", true},
|
|
{"flow_shift", false},
|
|
{"distilled_guidance", false},
|
|
{"shifted_timestep", false},
|
|
{"img_cfg", false},
|
|
{"slg", false},
|
|
{"clip_skip", false},
|
|
{"ref_images", false},
|
|
{"controlnet", false},
|
|
{"control_image", false},
|
|
{"ip_adapter", false},
|
|
{"init_image", false},
|
|
{"end_image", false},
|
|
{"control_frames", false},
|
|
{"vace", false},
|
|
{"high_noise", false},
|
|
{"cache_unet", false},
|
|
{"preview_tae", false},
|
|
};
|
|
|
|
static void validate_recipe_features(const json& recipe) {
|
|
for (const char* key : {"features", "disabled_features"}) {
|
|
if (!recipe.contains(key))
|
|
continue;
|
|
const auto& values = recipe.at(key);
|
|
if (!values.is_array() || values.size() > k_recipe_features.size())
|
|
throw std::runtime_error(std::string(key) + " must be an array of feature names");
|
|
std::set<std::string> seen;
|
|
for (const auto& value : values) {
|
|
const auto name = value.get<std::string>();
|
|
if (!k_recipe_features.count(name) || !seen.insert(name).second)
|
|
throw std::runtime_error("unknown or duplicate feature: " + name);
|
|
}
|
|
}
|
|
}
|
|
|
|
json resolve_recipe_features(const json& recipe) {
|
|
json features = k_recipe_features;
|
|
for (const auto& slot : recipe.at("slots")) {
|
|
const auto id = slot.at("id").get<std::string>();
|
|
if (id == "ip_adapter")
|
|
features["ip_adapter"] = true;
|
|
if (id == "control_net")
|
|
features["controlnet"] = true;
|
|
if (id == "high_noise_diffusion_model")
|
|
features["high_noise"] = true;
|
|
if (id == "tae" || id == "taesd")
|
|
features["preview_tae"] = true;
|
|
}
|
|
for (const auto& name : recipe.value("features", json::array()))
|
|
features[name.get<std::string>()] = true;
|
|
for (const auto& name : recipe.value("disabled_features", json::array()))
|
|
features[name.get<std::string>()] = false;
|
|
return features;
|
|
}
|
|
|
|
json model_generation_features(const ServerRuntime& runtime) {
|
|
if (!runtime.sd_ctx || runtime.models->active_recipe.is_null()) {
|
|
json features = k_recipe_features;
|
|
for (auto& feature : features)
|
|
feature = false;
|
|
return features;
|
|
}
|
|
json features = resolve_recipe_features(runtime.models->active_recipe);
|
|
if (runtime.models->active_recipe.value("source", "") == "command_line")
|
|
return features;
|
|
// Optional components only become usable once selected and loaded.
|
|
features["ip_adapter"] = features.value("ip_adapter", false) && !runtime.ctx_params->ip_adapter_path.empty();
|
|
features["controlnet"] = features.value("controlnet", false) && !runtime.ctx_params->control_net_path.empty();
|
|
features["high_noise"] = features.value("high_noise", false) && !runtime.ctx_params->high_noise_diffusion_model_path.empty();
|
|
features["preview_tae"] = features.value("preview_tae", false) && !runtime.ctx_params->taesd_path.empty();
|
|
if (!runtime.ctx_params->pulid_weights_path.empty())
|
|
features["slg"] = false;
|
|
return features;
|
|
}
|
|
|
|
static void validate_recipe(const json& recipe) {
|
|
check_keys(recipe, {"schema_version", "id", "name", "description", "model_url", "slots", "required_any", "defaults", "context", "features", "disabled_features"});
|
|
if (!recipe.at("schema_version").is_number_integer() || recipe.at("schema_version") != 1 ||
|
|
!valid_identifier(recipe.at("id").get<std::string>()) || recipe.at("name").get<std::string>().empty()) {
|
|
throw std::runtime_error("invalid schema_version, id or name");
|
|
}
|
|
if (recipe.contains("description") && !recipe["description"].is_string()) {
|
|
throw std::runtime_error("description must be a string");
|
|
}
|
|
if (recipe.contains("model_url")) {
|
|
const auto url = recipe.at("model_url").get<std::string>();
|
|
const std::string prefix = "https://huggingface.co/";
|
|
if (url.size() <= prefix.size() || url.size() > 2048 || url.rfind(prefix, 0) != 0 ||
|
|
std::any_of(url.begin(), url.end(), [](unsigned char c) { return std::isspace(c) || std::iscntrl(c); })) {
|
|
throw std::runtime_error("model_url must be a Hugging Face HTTPS URL");
|
|
}
|
|
}
|
|
const auto& slots = recipe.at("slots");
|
|
if (!slots.is_array() || slots.empty() || slots.size() > 32) {
|
|
throw std::runtime_error("slots must be a nonempty array of at most 32 entries");
|
|
}
|
|
std::set<std::string> ids;
|
|
bool has_required_model = false;
|
|
for (const auto& slot : slots) {
|
|
check_keys(slot, {"id", "label", "category", "required"});
|
|
const auto id = slot.at("id").get<std::string>();
|
|
const auto* spec = find_model_slot(id);
|
|
if (!spec || !ids.insert(id == "taesd" ? "tae" : id).second ||
|
|
(slot.contains("category") && !matches_model_category(*spec, slot.at("category").get<std::string>())) ||
|
|
(slot.contains("label") && slot.at("label").get<std::string>().empty()) || !slot.at("required").is_boolean()) {
|
|
throw std::runtime_error("invalid or duplicate model slot: " + id);
|
|
}
|
|
has_required_model |= (id == "model" || id == "diffusion_model") && slot["required"].get<bool>();
|
|
}
|
|
if (!has_required_model) {
|
|
throw std::runtime_error("a required model or diffusion_model slot is needed");
|
|
}
|
|
if (recipe.contains("required_any")) {
|
|
const auto& groups = recipe.at("required_any");
|
|
if (!groups.is_array() || groups.size() > 32)
|
|
throw std::runtime_error("required_any must be an array of at most 32 slot groups");
|
|
for (const auto& group : groups) {
|
|
if (!group.is_array() || group.size() < 2 || group.size() > slots.size())
|
|
throw std::runtime_error("required_any groups need at least two declared slots");
|
|
std::set<std::string> seen;
|
|
for (const auto& value : group) {
|
|
const auto id = value.get<std::string>();
|
|
if (!seen.insert(id).second || std::none_of(slots.begin(), slots.end(), [&](const json& slot) { return slot.at("id") == id; }))
|
|
throw std::runtime_error("unknown or duplicate required_any slot: " + id);
|
|
}
|
|
}
|
|
}
|
|
validate_recipe_features(recipe);
|
|
SDContextParams context;
|
|
apply_recipe_context(recipe.value("context", json::object()), context);
|
|
SDGenerationParams defaults;
|
|
apply_recipe_defaults(recipe.value("defaults", json::object()), defaults);
|
|
}
|
|
|
|
static json normalize_recipe(json recipe) {
|
|
validate_recipe(recipe);
|
|
recipe["effective_features"] = resolve_recipe_features(recipe);
|
|
for (auto& slot : recipe.at("slots")) {
|
|
const auto& spec = *find_model_slot(slot.at("id").get<std::string>());
|
|
slot["label"] = slot.value("label", std::string(spec.label));
|
|
slot["category"] = slot.value("category", spec.categories.front());
|
|
slot["categories"] = spec.categories;
|
|
}
|
|
return recipe;
|
|
}
|
|
|
|
static bool supported_model_file(const fs::path& path, const std::string& category) {
|
|
std::string ext = path.extension().u8string();
|
|
std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
|
|
if (category == "tokenizers") {
|
|
return ext == ".json";
|
|
}
|
|
return ext == ".gguf" || ext == ".safetensors" || ext == ".sft" ||
|
|
ext == ".ckpt" || ext == ".pt" || ext == ".pth" || ext == ".bin";
|
|
}
|
|
|
|
void refresh_model_catalog(ServerRuntime& runtime) {
|
|
auto& models = *runtime.models;
|
|
std::lock_guard<std::mutex> lock(models.catalog_mutex);
|
|
std::map<std::string, json> recipes;
|
|
std::map<std::string, ModelFile> files;
|
|
json errors = json::array();
|
|
auto report = [&](const std::string& file, const std::string& message) {
|
|
errors.push_back({{"file", file}, {"message", message}});
|
|
};
|
|
for (const char* builtin : k_builtin_recipes) {
|
|
auto recipe = normalize_recipe(json::parse(builtin));
|
|
const auto id = recipe.at("id").get<std::string>();
|
|
recipe["source"] = "builtin";
|
|
recipes[id] = std::move(recipe);
|
|
}
|
|
|
|
std::set<std::string> seen_files;
|
|
for (const auto& source : get_model_roots(*runtime.svr_params)) {
|
|
const auto& root = source.path;
|
|
auto display_path = [&](const fs::path& path) {
|
|
return source.id + "/" + path.lexically_relative(root).generic_u8string();
|
|
};
|
|
const fs::path recipe_dir = root / "recipes";
|
|
std::map<std::string, std::vector<json>> user_recipes;
|
|
try {
|
|
if (fs::exists(recipe_dir)) {
|
|
if (!model_path_inside(root, fs::canonical(recipe_dir)))
|
|
throw std::runtime_error("recipes directory escapes model directory");
|
|
for (const auto& entry : fs::directory_iterator(recipe_dir)) {
|
|
const std::string display = display_path(entry.path());
|
|
if (entry.path().extension() != ".json")
|
|
continue;
|
|
try {
|
|
if (!entry.is_regular_file() || !model_path_inside(root, fs::canonical(entry.path())))
|
|
throw std::runtime_error("recipe must be a file inside the model directory");
|
|
if (entry.file_size() > 1024 * 1024)
|
|
throw std::runtime_error("recipe exceeds 1 MiB");
|
|
std::ifstream input(entry.path(), std::ios::binary);
|
|
json recipe;
|
|
input >> recipe;
|
|
recipe = normalize_recipe(std::move(recipe));
|
|
recipe["source"] = "user";
|
|
recipe["file"] = entry.path().filename().u8string();
|
|
recipe["root"] = source.id;
|
|
const auto id = recipe.at("id").get<std::string>();
|
|
user_recipes[id].push_back(std::move(recipe));
|
|
} catch (const std::exception& e) {
|
|
report(display, e.what());
|
|
}
|
|
}
|
|
}
|
|
} catch (const std::exception& e) {
|
|
report(display_path(recipe_dir), e.what());
|
|
}
|
|
for (auto& item : user_recipes) {
|
|
if (item.second.size() != 1) {
|
|
for (const auto& recipe : item.second)
|
|
report(source.id + "/recipes/" + recipe["file"].get<std::string>(), "duplicate user recipe id: " + item.first);
|
|
continue;
|
|
}
|
|
const auto existing = recipes.find(item.first);
|
|
if (existing != recipes.end()) {
|
|
if (existing->second["source"] == "user") {
|
|
LOG_WARN("recipe '%s' from %s is shadowed by an earlier model directory", item.first.c_str(), source.label.c_str());
|
|
continue;
|
|
}
|
|
LOG_INFO("user recipe '%s' from %s overrides the builtin", item.first.c_str(), source.label.c_str());
|
|
}
|
|
recipes[item.first] = std::move(item.second.front());
|
|
}
|
|
|
|
std::set<std::string> categories;
|
|
for (const auto& slot : k_model_slots)
|
|
categories.insert(slot.categories.begin(), slot.categories.end());
|
|
for (const auto& category : categories) {
|
|
try {
|
|
for (const auto& directory : model_category_directories(source, category)) {
|
|
if (!model_path_inside(root, fs::canonical(directory))) {
|
|
report(display_path(directory), "category directory escapes model directory");
|
|
continue;
|
|
}
|
|
for (fs::recursive_directory_iterator it(directory), end; it != end; ++it) {
|
|
std::error_code path_error;
|
|
const fs::path canonical = fs::canonical(it->path(), path_error);
|
|
if (path_error || !model_path_inside(root, canonical)) {
|
|
it.disable_recursion_pending();
|
|
report(display_path(it->path()), path_error ? path_error.message() : "path escapes model directory");
|
|
continue;
|
|
}
|
|
if (!it->is_regular_file() || !supported_model_file(it->path(), category) ||
|
|
!seen_files.insert(model_path_key(canonical)).second)
|
|
continue;
|
|
const auto relative = it->path().lexically_relative(root).generic_u8string();
|
|
const auto id = source.primary ? relative : source.id + "/" + relative;
|
|
files[id] = {category, it->path().lexically_relative(directory).generic_u8string(), relative, source};
|
|
}
|
|
}
|
|
} catch (const std::exception& e) {
|
|
report(source.id + "/" + category, e.what());
|
|
}
|
|
}
|
|
}
|
|
std::map<std::pair<std::string, std::string>, size_t> names;
|
|
for (const auto& entry : files)
|
|
++names[{entry.second.category, entry.second.name}];
|
|
for (auto& entry : files) {
|
|
auto& file = entry.second;
|
|
if (names[{file.category, file.name}] > 1)
|
|
file.name = (file.root.path / fs::u8path(file.relative_path)).generic_u8string();
|
|
}
|
|
models.recipes = std::move(recipes);
|
|
models.files = std::move(files);
|
|
models.catalog_errors = std::move(errors);
|
|
}
|
|
|
|
json model_status(ServerRuntime& runtime) {
|
|
std::lock_guard<std::mutex> lock(runtime.models->status_mutex);
|
|
auto& models = *runtime.models;
|
|
json status = models.status;
|
|
if (status["state"] == "loading" && !status["load_elapsed_seconds"].is_null()) {
|
|
status["load_elapsed_seconds"] = std::chrono::duration<double>(std::chrono::steady_clock::now() - models.load_started_at).count();
|
|
}
|
|
return status;
|
|
}
|
|
|
|
void initialize_server_models(ServerRuntime& runtime) {
|
|
auto& models = *runtime.models;
|
|
models.base_context = *runtime.ctx_params;
|
|
clear_model_paths(models.base_context);
|
|
models.base_defaults = *runtime.default_gen_params;
|
|
if (runtime.sd_ctx) {
|
|
models.status["state"] = "ready";
|
|
models.status["origin"] = "command_line";
|
|
json features = json::array();
|
|
for (const auto& feature : k_recipe_features)
|
|
features.push_back(feature.first);
|
|
models.active_recipe = {
|
|
{"id", "@command-line"},
|
|
{"name", "Command-line model (all features)"},
|
|
{"source", "command_line"},
|
|
{"description", "Uses the loaded command-line model and its generation defaults."},
|
|
{"slots", json::array()},
|
|
{"defaults", json::object()},
|
|
{"features", features},
|
|
};
|
|
models.active_recipe["effective_features"] = resolve_recipe_features(models.active_recipe);
|
|
models.status["runtime_recipe"] = models.active_recipe;
|
|
}
|
|
refresh_model_catalog(runtime);
|
|
}
|
|
|
|
void prepare_model_load(ServerRuntime& runtime, const json& body, SDContextParams& context, SDGenerationParams& defaults, json* loaded_recipe) {
|
|
check_keys(body, {"recipe", "models"});
|
|
const auto recipe_id = body.at("recipe").get<std::string>();
|
|
auto selection = body.at("models");
|
|
if (!selection.is_object())
|
|
throw std::runtime_error("models must be an object");
|
|
auto& models = *runtime.models;
|
|
std::lock_guard<std::mutex> lock(models.catalog_mutex);
|
|
const auto found = models.recipes.find(recipe_id);
|
|
if (found == models.recipes.end())
|
|
throw std::runtime_error("unknown recipe");
|
|
const auto& recipe = found->second;
|
|
if (loaded_recipe)
|
|
*loaded_recipe = recipe;
|
|
context = models.base_context;
|
|
defaults = models.base_defaults;
|
|
apply_recipe_context(recipe.value("context", json::object()), context);
|
|
apply_recipe_defaults(recipe.value("defaults", json::object()), defaults);
|
|
std::set<std::string> allowed;
|
|
for (const auto& slot : recipe.at("slots")) {
|
|
const auto id = slot.at("id").get<std::string>();
|
|
allowed.insert(id);
|
|
if (id == "tae" || id == "taesd") {
|
|
const std::string alias = id == "tae" ? "taesd" : "tae";
|
|
allowed.insert(alias);
|
|
if (selection.contains(alias)) {
|
|
if (selection.contains(id))
|
|
throw std::runtime_error("specify only one of tae and taesd");
|
|
selection[id] = selection.at(alias);
|
|
}
|
|
}
|
|
const std::string file_id = selection.value(id, std::string());
|
|
if (file_id.empty()) {
|
|
if (slot.at("required").get<bool>())
|
|
throw std::runtime_error("missing required model: " + id);
|
|
continue;
|
|
}
|
|
const auto file = models.files.find(file_id);
|
|
if (file == models.files.end() || !matches_model_category(*find_model_slot(id), file->second.category)) {
|
|
throw std::runtime_error("unknown model or incorrect category: " + id);
|
|
}
|
|
const auto& root = file->second.root.path;
|
|
const fs::path path = fs::canonical(root / fs::u8path(file->second.relative_path));
|
|
if (!model_path_inside(root, path) || !fs::is_regular_file(path)) {
|
|
throw std::runtime_error("model file must remain inside the model directory: " + id);
|
|
}
|
|
context.*find_model_slot(id)->field = path.u8string();
|
|
}
|
|
for (auto it = selection.begin(); it != selection.end(); ++it) {
|
|
if (!allowed.count(it.key()))
|
|
throw std::runtime_error("unknown model slot: " + it.key());
|
|
}
|
|
for (const auto& group : recipe.value("required_any", json::array())) {
|
|
if (std::none_of(group.begin(), group.end(), [&](const json& value) {
|
|
return !selection.value(value.get<std::string>(), std::string()).empty();
|
|
})) {
|
|
throw std::runtime_error("select at least one model from: " + group.dump());
|
|
}
|
|
}
|
|
if (!context.resolve_and_validate(IMG_GEN))
|
|
throw std::runtime_error("invalid model context parameters");
|
|
}
|