mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-05-17 12:39:09 +00:00
sd: sync to master-591-331cfa5 (#2155)
* sd: sync to master-585-44cca3d * sd: sync to master-587-b8bdffc * sd: sync to master-591-331cfa5
This commit is contained in:
Wagner Bruna
GitHub
parent
61478cbf4a
commit
e2bdd6d7aa
33 changed files with 3703 additions and 1135 deletions
4
Makefile
4
Makefile
|
|
@ -679,7 +679,7 @@ llama-impl.o: src/llama-impl.cpp src/llama-impl.h
|
|||
budget.o: common/reasoning-budget.cpp common/reasoning-budget.h
|
||||
$(CXX) $(CXXFLAGS) -c $< -o $@
|
||||
|
||||
SDCPP_COMMON_BASENAMES := stable-diffusion.h stable-diffusion.cpp sample-cache.h sample-cache.cpp util.cpp upscaler.cpp model.cpp name_conversion.cpp tokenizers/bpe_tokenizer.cpp tokenizers/bpe_tokenizer.h tokenizers/clip_tokenizer.cpp tokenizers/clip_tokenizer.h tokenizers/mistral_tokenizer.cpp tokenizers/mistral_tokenizer.h tokenizers/qwen2_tokenizer.cpp tokenizers/qwen2_tokenizer.h tokenizers/t5_unigram_tokenizer.cpp tokenizers/t5_unigram_tokenizer.h tokenizers/tokenizer.cpp tokenizers/tokenizer.h tokenizers/tokenize_util.cpp tokenizers/tokenize_util.h thirdparty/zip.c
|
||||
SDCPP_COMMON_BASENAMES := stable-diffusion.h stable-diffusion.cpp sample-cache.h sample-cache.cpp util.cpp upscaler.h upscaler.cpp model.cpp name_conversion.cpp model_io/gguf_io.cpp model_io/gguf_io.h model_io/gguf_reader_ext.h model_io/pickle_io.cpp model_io/safetensors_io.cpp model_io/safetensors_io.h model_io/tensor_storage.h model_io/torch_legacy_io.cpp model_io/torch_zip_io.cpp tokenizers/bpe_tokenizer.cpp tokenizers/bpe_tokenizer.h tokenizers/clip_tokenizer.cpp tokenizers/clip_tokenizer.h tokenizers/mistral_tokenizer.cpp tokenizers/mistral_tokenizer.h tokenizers/qwen2_tokenizer.cpp tokenizers/qwen2_tokenizer.h tokenizers/t5_unigram_tokenizer.cpp tokenizers/t5_unigram_tokenizer.h tokenizers/tokenizer.cpp tokenizers/tokenizer.h tokenizers/tokenize_util.cpp tokenizers/tokenize_util.h thirdparty/zip.c
|
||||
SDCPP_COMMON_SOURCES := $(foreach f,$(SDCPP_COMMON_BASENAMES),otherarch/sdcpp/$(f))
|
||||
SDCPP_FLAGS := -I./vendor/nlohmann
|
||||
|
||||
|
|
@ -736,7 +736,7 @@ mainvk: tools/completion/completion.cpp common/arg.cpp common/speculative.cpp co
|
|||
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
fitparams: tools/fit-params/fit-params.cpp common/arg.cpp common/speculative.cpp common/ngram-cache.cpp common/ngram-map.cpp common/ngram-mod.cpp common/chat.cpp common/preset.cpp common/download.cpp build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_vulkan.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
|
||||
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
sdmain: $(SDCPP_COMMON_SOURCES) otherarch/sdcpp/main.cpp otherarch/sdcpp/image_metadata.cpp otherarch/sdcpp/common/log.cpp otherarch/sdcpp/common/media_io.cpp otherarch/sdcpp/common/common.cpp otherarch/sdcpp/version.cpp otherarch/sdcpp/tokenizers/vocab/vocab.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
sdmain: $(SDCPP_COMMON_SOURCES) otherarch/sdcpp/main.cpp otherarch/sdcpp/image_metadata.cpp otherarch/sdcpp/convert.cpp otherarch/sdcpp/common/log.cpp otherarch/sdcpp/common/media_io.cpp otherarch/sdcpp/common/common.cpp otherarch/sdcpp/version.cpp otherarch/sdcpp/tokenizers/vocab/vocab.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
$(CXX) $(CXXFLAGS) $(SDCPP_FLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
whispermain: otherarch/whispercpp/main.cpp otherarch/whispercpp/whisper.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
|
|
|
|||
|
|
@ -107,47 +107,60 @@ static bool is_absolute_path(const std::string& p) {
|
|||
|
||||
std::string ArgOptions::wrap_text(const std::string& text, size_t width, size_t indent) {
|
||||
std::ostringstream oss;
|
||||
size_t line_len = 0;
|
||||
size_t pos = 0;
|
||||
size_t line_len = 0;
|
||||
|
||||
while (pos < text.size()) {
|
||||
// Preserve manual newlines
|
||||
if (text[pos] == '\n') {
|
||||
oss << '\n'
|
||||
<< std::string(indent, ' ');
|
||||
line_len = indent;
|
||||
line_len = 0;
|
||||
++pos;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Add the character
|
||||
oss << text[pos];
|
||||
++line_len;
|
||||
++pos;
|
||||
if (std::isspace(static_cast<unsigned char>(text[pos]))) {
|
||||
++pos;
|
||||
continue;
|
||||
}
|
||||
|
||||
// If the current line exceeds width, try to break at the last space
|
||||
if (line_len >= width) {
|
||||
std::string current = oss.str();
|
||||
size_t back = current.size();
|
||||
size_t word_start = pos;
|
||||
while (pos < text.size() &&
|
||||
text[pos] != '\n' &&
|
||||
!std::isspace(static_cast<unsigned char>(text[pos]))) {
|
||||
++pos;
|
||||
}
|
||||
|
||||
// Find the last space (for a clean break)
|
||||
while (back > 0 && current[back - 1] != ' ' && current[back - 1] != '\n')
|
||||
--back;
|
||||
|
||||
// If found a space to break on
|
||||
if (back > 0 && current[back - 1] != '\n') {
|
||||
std::string before = current.substr(0, back - 1);
|
||||
std::string after = current.substr(back);
|
||||
oss.str("");
|
||||
oss.clear();
|
||||
oss << before << "\n"
|
||||
<< std::string(indent, ' ') << after;
|
||||
} else {
|
||||
// If no space found, just break at width
|
||||
oss << "\n"
|
||||
<< std::string(indent, ' ');
|
||||
std::string word = text.substr(word_start, pos - word_start);
|
||||
while (!word.empty()) {
|
||||
size_t separator_len = line_len == 0 ? 0 : 1;
|
||||
if (line_len + separator_len + word.size() <= width) {
|
||||
if (separator_len > 0) {
|
||||
oss << ' ';
|
||||
++line_len;
|
||||
}
|
||||
oss << word;
|
||||
line_len += word.size();
|
||||
word.clear();
|
||||
continue;
|
||||
}
|
||||
|
||||
if (line_len > 0) {
|
||||
oss << '\n'
|
||||
<< std::string(indent, ' ');
|
||||
line_len = 0;
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t chunk_len = std::min(width, word.size());
|
||||
oss << word.substr(0, chunk_len);
|
||||
line_len = chunk_len;
|
||||
word.erase(0, chunk_len);
|
||||
if (!word.empty()) {
|
||||
oss << '\n'
|
||||
<< std::string(indent, ' ');
|
||||
line_len = 0;
|
||||
}
|
||||
line_len = indent;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -351,7 +364,10 @@ ArgOptions SDContextParams::get_options() {
|
|||
"--lora-model-dir",
|
||||
"lora model directory",
|
||||
&lora_model_dir},
|
||||
|
||||
{"",
|
||||
"--hires-upscalers-dir",
|
||||
"highres fix upscaler model directory",
|
||||
&hires_upscalers_dir},
|
||||
{"",
|
||||
"--tensor-type-rules",
|
||||
"weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
|
||||
|
|
@ -649,6 +665,7 @@ std::string SDContextParams::to_string() const {
|
|||
<< " wtype: " << sd_type_name(wtype) << ",\n"
|
||||
<< " tensor_type_rules: \"" << tensor_type_rules << "\",\n"
|
||||
<< " lora_model_dir: \"" << lora_model_dir << "\",\n"
|
||||
<< " hires_upscalers_dir: \"" << hires_upscalers_dir << "\",\n"
|
||||
<< " photo_maker_path: \"" << photo_maker_path << "\",\n"
|
||||
<< " rng_type: " << sd_rng_type_name(rng_type) << ",\n"
|
||||
<< " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
|
||||
|
|
@ -777,6 +794,12 @@ ArgOptions SDGenerationParams::get_options() {
|
|||
"--pm-id-embed-path",
|
||||
"path to PHOTOMAKER v2 id embed",
|
||||
&pm_id_embed_path},
|
||||
{"",
|
||||
"--hires-upscaler",
|
||||
"highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent (nearest-exact), "
|
||||
"Latent (antialiased), Latent (bicubic), Latent (bicubic antialiased), or a model name "
|
||||
"under --hires-upscalers-dir (default: Latent)",
|
||||
&hires_upscaler},
|
||||
};
|
||||
|
||||
options.int_options = {
|
||||
|
|
@ -826,6 +849,22 @@ ArgOptions SDGenerationParams::get_options() {
|
|||
"--upscale-tile-size",
|
||||
"tile size for ESRGAN upscaling (default: 128)",
|
||||
&upscale_tile_size},
|
||||
{"",
|
||||
"--hires-width",
|
||||
"highres fix target width, 0 to use --hires-scale (default: 0)",
|
||||
&hires_width},
|
||||
{"",
|
||||
"--hires-height",
|
||||
"highres fix target height, 0 to use --hires-scale (default: 0)",
|
||||
&hires_height},
|
||||
{"",
|
||||
"--hires-steps",
|
||||
"highres fix second pass sample steps, 0 to reuse --steps (default: 0)",
|
||||
&hires_steps},
|
||||
{"",
|
||||
"--hires-upscale-tile-size",
|
||||
"highres fix upscaler tile size, reserved for model-backed upscalers (default: 128)",
|
||||
&hires_upscale_tile_size},
|
||||
};
|
||||
|
||||
options.float_options = {
|
||||
|
|
@ -913,6 +952,14 @@ ArgOptions SDGenerationParams::get_options() {
|
|||
"--vae-tile-overlap",
|
||||
"tile overlap for vae tiling, in fraction of tile size (default: 0.5)",
|
||||
&vae_tiling_params.target_overlap},
|
||||
{"",
|
||||
"--hires-scale",
|
||||
"highres fix scale when target size is not set (default: 2.0)",
|
||||
&hires_scale},
|
||||
{"",
|
||||
"--hires-denoising-strength",
|
||||
"highres fix second pass denoising strength (default: 0.7)",
|
||||
&hires_denoising_strength},
|
||||
};
|
||||
|
||||
options.bool_options = {
|
||||
|
|
@ -936,6 +983,11 @@ ArgOptions SDGenerationParams::get_options() {
|
|||
"process vae in tiles to reduce memory usage",
|
||||
true,
|
||||
&vae_tiling_params.enabled},
|
||||
{"",
|
||||
"--hires",
|
||||
"enable highres fix",
|
||||
true,
|
||||
&hires_enabled},
|
||||
};
|
||||
|
||||
auto on_seed_arg = [&](int argc, const char** argv, int index) {
|
||||
|
|
@ -1424,6 +1476,37 @@ static bool parse_lora_json_field(const json& parent,
|
|||
return true;
|
||||
}
|
||||
|
||||
static bool resolve_model_file_from_dir(const std::string& model_name,
|
||||
const std::string& model_dir,
|
||||
const std::vector<std::string>& valid_ext,
|
||||
const char* label,
|
||||
std::string& resolved_path) {
|
||||
if (model_dir.empty()) {
|
||||
LOG_ERROR("%s directory is empty", label);
|
||||
return false;
|
||||
}
|
||||
if (model_name.empty() ||
|
||||
model_name.find('/') != std::string::npos ||
|
||||
model_name.find('\\') != std::string::npos ||
|
||||
fs::path(model_name).has_root_path() ||
|
||||
fs::path(model_name).has_extension()) {
|
||||
LOG_ERROR("%s must be a model name without path or extension: %s", label, model_name.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
fs::path model_dir_path = model_dir;
|
||||
for (const auto& ext : valid_ext) {
|
||||
fs::path try_path = model_dir_path / (model_name + ext);
|
||||
if (fs::exists(try_path) && fs::is_regular_file(try_path)) {
|
||||
resolved_path = try_path.lexically_normal().string();
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
LOG_ERROR("can not find %s %s in %s", label, model_name.c_str(), model_dir_path.lexically_normal().string().c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::from_json_str(
|
||||
const std::string& json_str,
|
||||
const std::function<std::string(const std::string&)>& lora_path_resolver) {
|
||||
|
|
@ -1487,6 +1570,34 @@ bool SDGenerationParams::from_json_str(
|
|||
load_if_exists("increase_ref_index", increase_ref_index);
|
||||
load_if_exists("embed_image_metadata", embed_image_metadata);
|
||||
|
||||
if (j.contains("hires") && j["hires"].is_object()) {
|
||||
const json& hires_json = j["hires"];
|
||||
if (hires_json.contains("enabled") && hires_json["enabled"].is_boolean()) {
|
||||
hires_enabled = hires_json["enabled"];
|
||||
}
|
||||
if (hires_json.contains("upscaler") && hires_json["upscaler"].is_string()) {
|
||||
hires_upscaler = hires_json["upscaler"];
|
||||
}
|
||||
if (hires_json.contains("scale") && hires_json["scale"].is_number()) {
|
||||
hires_scale = hires_json["scale"];
|
||||
}
|
||||
if (hires_json.contains("target_width") && hires_json["target_width"].is_number_integer()) {
|
||||
hires_width = hires_json["target_width"];
|
||||
}
|
||||
if (hires_json.contains("target_height") && hires_json["target_height"].is_number_integer()) {
|
||||
hires_height = hires_json["target_height"];
|
||||
}
|
||||
if (hires_json.contains("steps") && hires_json["steps"].is_number_integer()) {
|
||||
hires_steps = hires_json["steps"];
|
||||
}
|
||||
if (hires_json.contains("denoising_strength") && hires_json["denoising_strength"].is_number()) {
|
||||
hires_denoising_strength = hires_json["denoising_strength"];
|
||||
}
|
||||
if (hires_json.contains("upscale_tile_size") && hires_json["upscale_tile_size"].is_number_integer()) {
|
||||
hires_upscale_tile_size = hires_json["upscale_tile_size"];
|
||||
}
|
||||
}
|
||||
|
||||
auto parse_sample_params_json = [&](const json& sample_json,
|
||||
sd_sample_params_t& target_params,
|
||||
std::vector<int>& target_skip_layers,
|
||||
|
|
@ -1800,7 +1911,7 @@ bool SDGenerationParams::initialize_cache_params() {
|
|||
return true;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::resolve(const std::string& lora_model_dir, bool strict) {
|
||||
bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) {
|
||||
if (high_noise_sample_params.sample_steps <= 0) {
|
||||
high_noise_sample_params.sample_steps = -1;
|
||||
}
|
||||
|
|
@ -1819,6 +1930,27 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, bool strict)
|
|||
sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100);
|
||||
}
|
||||
|
||||
hires_upscaler_model_path.clear();
|
||||
if (hires_enabled) {
|
||||
if (hires_upscaler.empty()) {
|
||||
hires_upscaler = "Latent";
|
||||
}
|
||||
resolved_hires_upscaler = str_to_sd_hires_upscaler(hires_upscaler.c_str());
|
||||
if (resolved_hires_upscaler == SD_HIRES_UPSCALER_NONE) {
|
||||
hires_enabled = false;
|
||||
} else if (resolved_hires_upscaler == SD_HIRES_UPSCALER_COUNT) {
|
||||
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".pth"};
|
||||
if (!resolve_model_file_from_dir(hires_upscaler,
|
||||
hires_upscalers_dir,
|
||||
valid_ext,
|
||||
"hires upscaler",
|
||||
hires_upscaler_model_path)) {
|
||||
return false;
|
||||
}
|
||||
resolved_hires_upscaler = SD_HIRES_UPSCALER_MODEL;
|
||||
}
|
||||
}
|
||||
|
||||
prompt_with_lora = prompt;
|
||||
if (!lora_model_dir.empty()) {
|
||||
extract_and_remove_lora(lora_model_dir);
|
||||
|
|
@ -1883,6 +2015,29 @@ bool SDGenerationParams::validate(SDMode mode) {
|
|||
return false;
|
||||
}
|
||||
|
||||
if (hires_enabled) {
|
||||
if (hires_width < 0 || hires_height < 0) {
|
||||
LOG_ERROR("error: hires target width and height must be >= 0");
|
||||
return false;
|
||||
}
|
||||
if (hires_scale <= 0.f && hires_width <= 0 && hires_height <= 0) {
|
||||
LOG_ERROR("error: hires scale must be positive when target size is not set");
|
||||
return false;
|
||||
}
|
||||
if (hires_steps < 0) {
|
||||
LOG_ERROR("error: hires steps must be >= 0");
|
||||
return false;
|
||||
}
|
||||
if (hires_denoising_strength <= 0.f || hires_denoising_strength > 1.f) {
|
||||
LOG_ERROR("error: hires denoising strength must be in (0.0, 1.0]");
|
||||
return false;
|
||||
}
|
||||
if (hires_upscale_tile_size < 1) {
|
||||
LOG_ERROR("error: hires upscale tile size must be positive");
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if (mode == UPSCALE) {
|
||||
if (init_image_path.length() == 0) {
|
||||
LOG_ERROR("error: upscale mode needs an init image (--init-img)\n");
|
||||
|
|
@ -1893,8 +2048,11 @@ bool SDGenerationParams::validate(SDMode mode) {
|
|||
return true;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict) {
|
||||
if (!resolve(lora_model_dir, strict)) {
|
||||
bool SDGenerationParams::resolve_and_validate(SDMode mode,
|
||||
const std::string& lora_model_dir,
|
||||
const std::string& hires_upscalers_dir,
|
||||
bool strict) {
|
||||
if (!resolve(lora_model_dir, hires_upscalers_dir, strict)) {
|
||||
return false;
|
||||
}
|
||||
if (!validate(mode)) {
|
||||
|
|
@ -1965,6 +2123,16 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
|||
params.pm_params = pm_params;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
|
||||
params.hires.enabled = hires_enabled;
|
||||
params.hires.upscaler = resolved_hires_upscaler;
|
||||
params.hires.model_path = hires_upscaler_model_path.empty() ? nullptr : hires_upscaler_model_path.c_str();
|
||||
params.hires.scale = hires_scale;
|
||||
params.hires.target_width = hires_width;
|
||||
params.hires.target_height = hires_height;
|
||||
params.hires.steps = hires_steps;
|
||||
params.hires.denoising_strength = hires_denoising_strength;
|
||||
params.hires.upscale_tile_size = hires_upscale_tile_size;
|
||||
return params;
|
||||
}
|
||||
|
||||
|
|
@ -2089,6 +2257,15 @@ std::string SDGenerationParams::to_string() const {
|
|||
<< " seed: " << seed << ",\n"
|
||||
<< " upscale_repeats: " << upscale_repeats << ",\n"
|
||||
<< " upscale_tile_size: " << upscale_tile_size << ",\n"
|
||||
<< " hires: { enabled: " << (hires_enabled ? "true" : "false")
|
||||
<< ", upscaler: \"" << hires_upscaler << "\""
|
||||
<< ", model_path: \"" << hires_upscaler_model_path << "\""
|
||||
<< ", scale: " << hires_scale
|
||||
<< ", target_width: " << hires_width
|
||||
<< ", target_height: " << hires_height
|
||||
<< ", steps: " << hires_steps
|
||||
<< ", denoising_strength: " << hires_denoising_strength
|
||||
<< ", upscale_tile_size: " << hires_upscale_tile_size << " },\n"
|
||||
<< " vae_tiling_params: { "
|
||||
<< vae_tiling_params.enabled << ", "
|
||||
<< vae_tiling_params.tile_size_x << ", "
|
||||
|
|
@ -2104,7 +2281,192 @@ std::string version_string() {
|
|||
return std::string("stable-diffusion.cpp version ") + sd_version() + ", commit " + sd_commit();
|
||||
}
|
||||
|
||||
std::string get_image_params(const SDContextParams& ctx_params, const SDGenerationParams& gen_params, int64_t seed) {
|
||||
static std::string safe_json_string(const char* value) {
|
||||
return value ? value : "";
|
||||
}
|
||||
|
||||
static void set_json_basename_if_not_empty(json& target, const char* key, const std::string& path) {
|
||||
if (!path.empty()) {
|
||||
target[key] = sd_basename(path);
|
||||
}
|
||||
}
|
||||
|
||||
static json build_sampling_metadata_json(const sd_sample_params_t& sample_params,
|
||||
const std::vector<int>& skip_layers,
|
||||
const std::vector<float>* custom_sigmas = nullptr) {
|
||||
json sampling = {
|
||||
{"steps", sample_params.sample_steps},
|
||||
{"eta", sample_params.eta},
|
||||
{"shifted_timestep", sample_params.shifted_timestep},
|
||||
{"flow_shift", sample_params.flow_shift},
|
||||
{"guidance",
|
||||
{
|
||||
{"txt_cfg", sample_params.guidance.txt_cfg},
|
||||
{"img_cfg", sample_params.guidance.img_cfg},
|
||||
{"distilled_guidance", sample_params.guidance.distilled_guidance},
|
||||
{"slg",
|
||||
{
|
||||
{"scale", sample_params.guidance.slg.scale},
|
||||
{"layers", skip_layers},
|
||||
{"start", sample_params.guidance.slg.layer_start},
|
||||
{"end", sample_params.guidance.slg.layer_end},
|
||||
}},
|
||||
}},
|
||||
};
|
||||
if (sample_params.sample_method != SAMPLE_METHOD_COUNT) {
|
||||
sampling["method"] = safe_json_string(sd_sample_method_name(sample_params.sample_method));
|
||||
}
|
||||
if (sample_params.scheduler != SCHEDULER_COUNT) {
|
||||
sampling["scheduler"] = safe_json_string(sd_scheduler_name(sample_params.scheduler));
|
||||
}
|
||||
if (custom_sigmas != nullptr) {
|
||||
sampling["custom_sigmas"] = *custom_sigmas;
|
||||
}
|
||||
return sampling;
|
||||
}
|
||||
|
||||
std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
const SDGenerationParams& gen_params,
|
||||
int64_t seed,
|
||||
SDMode mode) {
|
||||
json root;
|
||||
root["schema"] = "sdcpp.image.params/v1";
|
||||
root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen";
|
||||
root["generator"] = {
|
||||
{"name", "stable-diffusion.cpp"},
|
||||
{"version", safe_json_string(sd_version())},
|
||||
{"commit", safe_json_string(sd_commit())},
|
||||
};
|
||||
root["seed"] = seed;
|
||||
root["width"] = gen_params.get_resolved_width();
|
||||
root["height"] = gen_params.get_resolved_height();
|
||||
|
||||
root["prompt"] = {
|
||||
{"positive", gen_params.prompt},
|
||||
{"negative", gen_params.negative_prompt},
|
||||
};
|
||||
root["sampling"] = build_sampling_metadata_json(gen_params.sample_params,
|
||||
gen_params.skip_layers,
|
||||
&gen_params.custom_sigmas);
|
||||
|
||||
json models;
|
||||
set_json_basename_if_not_empty(models, "model", ctx_params.model_path);
|
||||
set_json_basename_if_not_empty(models, "clip_l", ctx_params.clip_l_path);
|
||||
set_json_basename_if_not_empty(models, "clip_g", ctx_params.clip_g_path);
|
||||
set_json_basename_if_not_empty(models, "clip_vision", ctx_params.clip_vision_path);
|
||||
set_json_basename_if_not_empty(models, "t5xxl", ctx_params.t5xxl_path);
|
||||
set_json_basename_if_not_empty(models, "llm", ctx_params.llm_path);
|
||||
set_json_basename_if_not_empty(models, "llm_vision", ctx_params.llm_vision_path);
|
||||
set_json_basename_if_not_empty(models, "diffusion_model", ctx_params.diffusion_model_path);
|
||||
set_json_basename_if_not_empty(models, "high_noise_diffusion_model", ctx_params.high_noise_diffusion_model_path);
|
||||
set_json_basename_if_not_empty(models, "vae", ctx_params.vae_path);
|
||||
set_json_basename_if_not_empty(models, "taesd", ctx_params.taesd_path);
|
||||
set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path);
|
||||
root["models"] = std::move(models);
|
||||
|
||||
root["clip_skip"] = gen_params.clip_skip;
|
||||
root["strength"] = gen_params.strength;
|
||||
root["control_strength"] = gen_params.control_strength;
|
||||
root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
|
||||
root["increase_ref_index"] = gen_params.increase_ref_index;
|
||||
if (mode == VID_GEN) {
|
||||
root["video"] = {
|
||||
{"frame_count", gen_params.video_frames},
|
||||
{"fps", gen_params.fps},
|
||||
};
|
||||
root["moe_boundary"] = gen_params.moe_boundary;
|
||||
root["vace_strength"] = gen_params.vace_strength;
|
||||
root["high_noise_sampling"] = build_sampling_metadata_json(gen_params.high_noise_sample_params,
|
||||
gen_params.high_noise_skip_layers);
|
||||
}
|
||||
|
||||
root["rng"] = safe_json_string(sd_rng_type_name(ctx_params.rng_type));
|
||||
if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) {
|
||||
root["sampler_rng"] = safe_json_string(sd_rng_type_name(ctx_params.sampler_rng_type));
|
||||
}
|
||||
|
||||
json loras = json::array();
|
||||
for (const auto& entry : gen_params.lora_map) {
|
||||
loras.push_back({
|
||||
{"name", sd_basename(entry.first)},
|
||||
{"multiplier", entry.second},
|
||||
{"is_high_noise", false},
|
||||
});
|
||||
}
|
||||
for (const auto& entry : gen_params.high_noise_lora_map) {
|
||||
loras.push_back({
|
||||
{"name", sd_basename(entry.first)},
|
||||
{"multiplier", entry.second},
|
||||
{"is_high_noise", true},
|
||||
});
|
||||
}
|
||||
if (!loras.empty()) {
|
||||
root["loras"] = std::move(loras);
|
||||
}
|
||||
|
||||
if (gen_params.hires_enabled) {
|
||||
root["hires"] = {
|
||||
{"enabled", gen_params.hires_enabled},
|
||||
{"upscaler", gen_params.hires_upscaler},
|
||||
{"model", gen_params.hires_upscaler_model_path.empty() ? "" : sd_basename(gen_params.hires_upscaler_model_path)},
|
||||
{"scale", gen_params.hires_scale},
|
||||
{"target_width", gen_params.hires_width},
|
||||
{"target_height", gen_params.hires_height},
|
||||
{"steps", gen_params.hires_steps},
|
||||
{"denoising_strength", gen_params.hires_denoising_strength},
|
||||
{"upscale_tile_size", gen_params.hires_upscale_tile_size},
|
||||
};
|
||||
}
|
||||
|
||||
if (gen_params.cache_params.mode != SD_CACHE_DISABLED) {
|
||||
root["cache"] = {
|
||||
{"requested_mode", gen_params.cache_mode},
|
||||
{"requested_option", gen_params.cache_option},
|
||||
{"mode", gen_params.cache_params.mode},
|
||||
{"scm_mask", gen_params.scm_mask},
|
||||
{"scm_policy_dynamic", gen_params.scm_policy_dynamic},
|
||||
{"reuse_threshold", gen_params.cache_params.reuse_threshold},
|
||||
{"start_percent", gen_params.cache_params.start_percent},
|
||||
{"end_percent", gen_params.cache_params.end_percent},
|
||||
{"error_decay_rate", gen_params.cache_params.error_decay_rate},
|
||||
{"use_relative_threshold", gen_params.cache_params.use_relative_threshold},
|
||||
{"reset_error_on_compute", gen_params.cache_params.reset_error_on_compute},
|
||||
{"Fn_compute_blocks", gen_params.cache_params.Fn_compute_blocks},
|
||||
{"Bn_compute_blocks", gen_params.cache_params.Bn_compute_blocks},
|
||||
{"residual_diff_threshold", gen_params.cache_params.residual_diff_threshold},
|
||||
{"max_warmup_steps", gen_params.cache_params.max_warmup_steps},
|
||||
{"max_cached_steps", gen_params.cache_params.max_cached_steps},
|
||||
{"max_continuous_cached_steps", gen_params.cache_params.max_continuous_cached_steps},
|
||||
{"taylorseer_n_derivatives", gen_params.cache_params.taylorseer_n_derivatives},
|
||||
{"taylorseer_skip_interval", gen_params.cache_params.taylorseer_skip_interval},
|
||||
{"spectrum_w", gen_params.cache_params.spectrum_w},
|
||||
{"spectrum_m", gen_params.cache_params.spectrum_m},
|
||||
{"spectrum_lam", gen_params.cache_params.spectrum_lam},
|
||||
{"spectrum_window_size", gen_params.cache_params.spectrum_window_size},
|
||||
{"spectrum_flex_window", gen_params.cache_params.spectrum_flex_window},
|
||||
{"spectrum_warmup_steps", gen_params.cache_params.spectrum_warmup_steps},
|
||||
{"spectrum_stop_percent", gen_params.cache_params.spectrum_stop_percent},
|
||||
};
|
||||
}
|
||||
|
||||
if (gen_params.vae_tiling_params.enabled) {
|
||||
root["vae_tiling"] = {
|
||||
{"enabled", gen_params.vae_tiling_params.enabled},
|
||||
{"tile_size_x", gen_params.vae_tiling_params.tile_size_x},
|
||||
{"tile_size_y", gen_params.vae_tiling_params.tile_size_y},
|
||||
{"target_overlap", gen_params.vae_tiling_params.target_overlap},
|
||||
{"rel_size_x", gen_params.vae_tiling_params.rel_size_x},
|
||||
{"rel_size_y", gen_params.vae_tiling_params.rel_size_y},
|
||||
};
|
||||
}
|
||||
|
||||
return root.dump();
|
||||
}
|
||||
|
||||
std::string get_image_params(const SDContextParams& ctx_params,
|
||||
const SDGenerationParams& gen_params,
|
||||
int64_t seed,
|
||||
SDMode mode) {
|
||||
std::string parameter_string;
|
||||
if (gen_params.prompt_with_lora.size() != 0) {
|
||||
parameter_string += gen_params.prompt_with_lora + "\n";
|
||||
|
|
@ -2117,7 +2479,7 @@ std::string get_image_params(const SDContextParams& ctx_params, const SDGenerati
|
|||
parameter_string += "Steps: " + std::to_string(gen_params.sample_params.sample_steps) + ", ";
|
||||
parameter_string += "CFG scale: " + std::to_string(gen_params.sample_params.guidance.txt_cfg) + ", ";
|
||||
if (gen_params.sample_params.guidance.slg.scale != 0 && gen_params.skip_layers.size() != 0) {
|
||||
parameter_string += "SLG scale: " + std::to_string(gen_params.sample_params.guidance.txt_cfg) + ", ";
|
||||
parameter_string += "SLG scale: " + std::to_string(gen_params.sample_params.guidance.slg.scale) + ", ";
|
||||
parameter_string += "Skip layers: [";
|
||||
for (const auto& layer : gen_params.skip_layers) {
|
||||
parameter_string += std::to_string(layer) + ", ";
|
||||
|
|
@ -2162,6 +2524,14 @@ std::string get_image_params(const SDContextParams& ctx_params, const SDGenerati
|
|||
if (gen_params.clip_skip != -1) {
|
||||
parameter_string += "Clip skip: " + std::to_string(gen_params.clip_skip) + ", ";
|
||||
}
|
||||
if (gen_params.hires_enabled) {
|
||||
parameter_string += "Hires upscale: " + gen_params.hires_upscaler + ", ";
|
||||
parameter_string += "Hires scale: " + std::to_string(gen_params.hires_scale) + ", ";
|
||||
parameter_string += "Hires resize: " + std::to_string(gen_params.hires_width) + "x" + std::to_string(gen_params.hires_height) + ", ";
|
||||
parameter_string += "Hires steps: " + std::to_string(gen_params.hires_steps) + ", ";
|
||||
parameter_string += "Denoising strength: " + std::to_string(gen_params.hires_denoising_strength) + ", ";
|
||||
}
|
||||
parameter_string += "Version: stable-diffusion.cpp";
|
||||
parameter_string += ", SDCPP: " + build_sdcpp_image_metadata_json(ctx_params, gen_params, seed, mode);
|
||||
return parameter_string;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -101,6 +101,7 @@ struct SDContextParams {
|
|||
sd_type_t wtype = SD_TYPE_COUNT;
|
||||
std::string tensor_type_rules;
|
||||
std::string lora_model_dir = ".";
|
||||
std::string hires_upscalers_dir;
|
||||
|
||||
std::map<std::string, std::string> embedding_map;
|
||||
std::vector<sd_embedding_t> embedding_vec;
|
||||
|
|
@ -190,12 +191,23 @@ struct SDGenerationParams {
|
|||
int upscale_repeats = 1;
|
||||
int upscale_tile_size = 128;
|
||||
|
||||
bool hires_enabled = false;
|
||||
std::string hires_upscaler = "Latent";
|
||||
std::string hires_upscaler_model_path;
|
||||
float hires_scale = 2.f;
|
||||
int hires_width = 0;
|
||||
int hires_height = 0;
|
||||
int hires_steps = 0;
|
||||
float hires_denoising_strength = 0.7f;
|
||||
int hires_upscale_tile_size = 128;
|
||||
|
||||
std::map<std::string, float> lora_map;
|
||||
std::map<std::string, float> high_noise_lora_map;
|
||||
|
||||
// Derived and normalized fields.
|
||||
std::string prompt_with_lora; // for metadata record only
|
||||
std::vector<sd_lora_t> lora_vec;
|
||||
sd_hires_upscaler_t resolved_hires_upscaler;
|
||||
|
||||
// Owned execution payload.
|
||||
SDImageOwner init_image;
|
||||
|
|
@ -225,15 +237,25 @@ struct SDGenerationParams {
|
|||
void set_width_and_height_if_unset(int w, int h);
|
||||
int get_resolved_width() const;
|
||||
int get_resolved_height() const;
|
||||
bool resolve(const std::string& lora_model_dir, bool strict = false);
|
||||
bool resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict = false);
|
||||
bool validate(SDMode mode);
|
||||
bool resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict = false);
|
||||
bool resolve_and_validate(SDMode mode,
|
||||
const std::string& lora_model_dir,
|
||||
const std::string& hires_upscalers_dir,
|
||||
bool strict = false);
|
||||
sd_img_gen_params_t to_sd_img_gen_params_t();
|
||||
sd_vid_gen_params_t to_sd_vid_gen_params_t();
|
||||
std::string to_string() const;
|
||||
};
|
||||
|
||||
std::string version_string();
|
||||
std::string get_image_params(const SDContextParams& ctx_params, const SDGenerationParams& gen_params, int64_t seed);
|
||||
std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
const SDGenerationParams& gen_params,
|
||||
int64_t seed,
|
||||
SDMode mode = IMG_GEN);
|
||||
std::string get_image_params(const SDContextParams& ctx_params,
|
||||
const SDGenerationParams& gen_params,
|
||||
int64_t seed,
|
||||
SDMode mode = IMG_GEN);
|
||||
|
||||
#endif // __EXAMPLES_COMMON_COMMON_H__
|
||||
|
|
|
|||
138
otherarch/sdcpp/convert.cpp
Normal file
138
otherarch/sdcpp/convert.cpp
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
#include <cstring>
|
||||
#include <mutex>
|
||||
#include <regex>
|
||||
#include <vector>
|
||||
|
||||
#include "model.h"
|
||||
#include "model_io/gguf_io.h"
|
||||
#include "model_io/safetensors_io.h"
|
||||
#include "util.h"
|
||||
|
||||
#include "ggml-cpu.h"
|
||||
|
||||
static ggml_type get_export_tensor_type(ModelLoader& model_loader,
|
||||
const TensorStorage& tensor_storage,
|
||||
ggml_type type,
|
||||
const TensorTypeRules& tensor_type_rules) {
|
||||
const std::string& name = tensor_storage.name;
|
||||
ggml_type tensor_type = tensor_storage.type;
|
||||
ggml_type dst_type = type;
|
||||
|
||||
for (const auto& tensor_type_rule : tensor_type_rules) {
|
||||
std::regex pattern(tensor_type_rule.first);
|
||||
if (std::regex_search(name, pattern)) {
|
||||
dst_type = tensor_type_rule.second;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (model_loader.tensor_should_be_converted(tensor_storage, dst_type)) {
|
||||
tensor_type = dst_type;
|
||||
}
|
||||
|
||||
return tensor_type;
|
||||
}
|
||||
|
||||
static bool load_tensors_for_export(ModelLoader& model_loader,
|
||||
ggml_context* ggml_ctx,
|
||||
ggml_type type,
|
||||
const TensorTypeRules& tensor_type_rules,
|
||||
std::vector<TensorWriteInfo>& tensors) {
|
||||
std::mutex tensor_mutex;
|
||||
auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
|
||||
const std::string& name = tensor_storage.name;
|
||||
ggml_type tensor_type = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
|
||||
|
||||
std::lock_guard<std::mutex> lock(tensor_mutex);
|
||||
ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
|
||||
if (tensor == nullptr) {
|
||||
LOG_ERROR("ggml_new_tensor failed");
|
||||
return false;
|
||||
}
|
||||
ggml_set_name(tensor, name.c_str());
|
||||
|
||||
if (!tensor->data) {
|
||||
GGML_ASSERT(ggml_nelements(tensor) == 0);
|
||||
// Avoid crashing writers by setting a dummy pointer for zero-sized tensors.
|
||||
LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
|
||||
tensor->data = ggml_get_mem_buffer(ggml_ctx);
|
||||
}
|
||||
|
||||
TensorWriteInfo write_info;
|
||||
write_info.tensor = tensor;
|
||||
write_info.n_dims = tensor_storage.n_dims;
|
||||
for (int i = 0; i < tensor_storage.n_dims; ++i) {
|
||||
write_info.ne[i] = tensor_storage.ne[i];
|
||||
}
|
||||
|
||||
*dst_tensor = tensor;
|
||||
tensors.push_back(std::move(write_info));
|
||||
|
||||
return true;
|
||||
};
|
||||
|
||||
bool success = model_loader.load_tensors(on_new_tensor_cb);
|
||||
LOG_INFO("load tensors done");
|
||||
return success;
|
||||
}
|
||||
|
||||
bool convert(const char* input_path,
|
||||
const char* vae_path,
|
||||
const char* output_path,
|
||||
sd_type_t output_type,
|
||||
const char* tensor_type_rules,
|
||||
bool convert_name) {
|
||||
ModelLoader model_loader;
|
||||
|
||||
if (!model_loader.init_from_file(input_path)) {
|
||||
LOG_ERROR("init model loader from file failed: '%s'", input_path);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (vae_path != nullptr && strlen(vae_path) > 0) {
|
||||
if (!model_loader.init_from_file(vae_path, "vae.")) {
|
||||
LOG_ERROR("init model loader from file failed: '%s'", vae_path);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if (convert_name) {
|
||||
model_loader.convert_tensors_name();
|
||||
}
|
||||
|
||||
ggml_type type = (ggml_type)output_type;
|
||||
bool output_is_safetensors = ends_with(output_path, ".safetensors");
|
||||
TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
|
||||
|
||||
auto backend = ggml_backend_cpu_init();
|
||||
size_t mem_size = 1 * 1024 * 1024; // for padding
|
||||
mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
|
||||
mem_size += model_loader.get_params_mem_size(backend, type);
|
||||
LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
|
||||
ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
|
||||
|
||||
if (ggml_ctx == nullptr) {
|
||||
LOG_ERROR("ggml_init failed for converter");
|
||||
ggml_backend_free(backend);
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<TensorWriteInfo> tensors;
|
||||
bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors);
|
||||
ggml_backend_free(backend);
|
||||
|
||||
std::string error;
|
||||
if (success) {
|
||||
if (output_is_safetensors) {
|
||||
success = write_safetensors_file(output_path, tensors, &error);
|
||||
} else {
|
||||
success = write_gguf_file(output_path, tensors, &error);
|
||||
}
|
||||
}
|
||||
|
||||
if (!success && !error.empty()) {
|
||||
LOG_ERROR("%s", error.c_str());
|
||||
}
|
||||
|
||||
ggml_free(ggml_ctx);
|
||||
return success;
|
||||
}
|
||||
|
|
@ -1523,12 +1523,10 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
|
|||
const std::vector<float>& sigmas,
|
||||
std::shared_ptr<RNG> rng,
|
||||
float eta) {
|
||||
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
|
||||
float sigma = sigmas[i];
|
||||
float sigma_to = sigmas[i + 1];
|
||||
float sigma = sigmas[i];
|
||||
float sigma_to = sigmas[i + 1];
|
||||
|
||||
auto model_output_opt = model(x, sigma, i + 1);
|
||||
if (model_output_opt.empty()) {
|
||||
|
|
@ -1551,12 +1549,11 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
|
|||
float std_dev_t = eta * std::sqrt(variance);
|
||||
|
||||
x = pred_original_sample +
|
||||
std::sqrt((1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2))/ alpha_prod_t_prev) * model_output;
|
||||
std::sqrt((1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2)) / alpha_prod_t_prev) * model_output;
|
||||
|
||||
if (eta > 0) {
|
||||
x+= std_dev_t / std::sqrt(alpha_prod_t_prev) * sd::Tensor<float>::randn_like(x, rng);
|
||||
x += std_dev_t / std::sqrt(alpha_prod_t_prev) * sd::Tensor<float>::randn_like(x, rng);
|
||||
}
|
||||
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
|
@ -1584,8 +1581,10 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
|
|||
|
||||
auto get_timestep_from_sigma = [&](float s) -> int {
|
||||
auto it = std::lower_bound(compvis_sigmas.begin(), compvis_sigmas.end(), s);
|
||||
if (it == compvis_sigmas.begin()) return 0;
|
||||
if (it == compvis_sigmas.end()) return TIMESTEPS - 1;
|
||||
if (it == compvis_sigmas.begin())
|
||||
return 0;
|
||||
if (it == compvis_sigmas.end())
|
||||
return TIMESTEPS - 1;
|
||||
int idx_high = static_cast<int>(std::distance(compvis_sigmas.begin(), it));
|
||||
int idx_low = idx_high - 1;
|
||||
if (std::abs(compvis_sigmas[idx_high] - s) < std::abs(compvis_sigmas[idx_low] - s)) {
|
||||
|
|
@ -1596,7 +1595,6 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
|
|||
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
|
||||
float sigma_to = sigmas[i + 1];
|
||||
int prev_timestep = get_timestep_from_sigma(sigma_to);
|
||||
int timestep_s = (int)floor((1 - eta) * prev_timestep);
|
||||
|
|
@ -1626,7 +1624,6 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
|
|||
x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
|
||||
std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
|
||||
}
|
||||
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2758,16 +2758,16 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
|
|||
bool is_conv,
|
||||
WeightAdapter::ForwardParams::conv2d_params_t conv_params,
|
||||
float scale) {
|
||||
GGML_ASSERT((w1 != NULL || (w1a != NULL && w1b != NULL)));
|
||||
GGML_ASSERT((w2 != NULL || (w2a != NULL && w2b != NULL)));
|
||||
GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr)));
|
||||
GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr)));
|
||||
|
||||
int uq = (w1 != NULL) ? (int)w1->ne[0] : (int)w1a->ne[0];
|
||||
int up = (w1 != NULL) ? (int)w1->ne[1] : (int)w1b->ne[1];
|
||||
int uq = (w1 != nullptr) ? (int)w1->ne[0] : (int)w1a->ne[0];
|
||||
int up = (w1 != nullptr) ? (int)w1->ne[1] : (int)w1b->ne[1];
|
||||
|
||||
int q_actual = is_conv ? (int)h->ne[2] : (int)h->ne[0];
|
||||
int vq = q_actual / uq;
|
||||
|
||||
int vp = (w2 != NULL) ? (is_conv ? (int)w2->ne[3] : (int)w2->ne[1])
|
||||
int vp = (w2 != nullptr) ? (is_conv ? (int)w2->ne[3] : (int)w2->ne[1])
|
||||
: (int)w2a->ne[1];
|
||||
GGML_ASSERT(q_actual == (uq * vq) && "Input dimension mismatch for LoKR split");
|
||||
|
||||
|
|
@ -2803,7 +2803,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
|
|||
#endif
|
||||
|
||||
ggml_tensor* h_split = ggml_reshape_3d(ctx, h, vq, uq * merge_batch_uq, batch / merge_batch_uq);
|
||||
if (w2 != NULL) {
|
||||
if (w2 != nullptr) {
|
||||
hb = ggml_mul_mat(ctx, w2, h_split);
|
||||
} else {
|
||||
hb = ggml_mul_mat(ctx, w2b, ggml_mul_mat(ctx, w2a, h_split));
|
||||
|
|
@ -2816,7 +2816,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
|
|||
hb_t = ggml_reshape_3d(ctx, hb_t, uq, vp * merge_batch_vp, batch / merge_batch_vp);
|
||||
|
||||
ggml_tensor* hc_t;
|
||||
if (w1 != NULL) {
|
||||
if (w1 != nullptr) {
|
||||
hc_t = ggml_mul_mat(ctx, w1, hb_t);
|
||||
} else {
|
||||
hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_t));
|
||||
|
|
@ -2834,7 +2834,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
|
|||
// 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch]
|
||||
ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch);
|
||||
|
||||
if (w2 != NULL) {
|
||||
if (w2 != nullptr) {
|
||||
hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr,
|
||||
conv_params.s0,
|
||||
conv_params.s1,
|
||||
|
|
@ -2902,7 +2902,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
|
|||
ggml_tensor* hb_merged = ggml_reshape_2d(ctx, hb, w_out * h_out * vp, uq * batch);
|
||||
ggml_tensor* hc_t;
|
||||
ggml_tensor* hb_merged_t = ggml_cont(ctx, ggml_transpose(ctx, hb_merged));
|
||||
if (w1 != NULL) {
|
||||
if (w1 != nullptr) {
|
||||
// Would be great to be able to transpose w1 instead to avoid transposing both hb and hc
|
||||
hc_t = ggml_mul_mat(ctx, w1, hb_merged_t);
|
||||
} else {
|
||||
|
|
|
|||
|
|
@ -278,7 +278,9 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
|
|||
bool valid = cli_params.resolve_and_validate();
|
||||
if (valid && cli_params.mode != METADATA) {
|
||||
valid = ctx_params.resolve_and_validate(cli_params.mode) &&
|
||||
gen_params.resolve_and_validate(cli_params.mode, ctx_params.lora_model_dir);
|
||||
gen_params.resolve_and_validate(cli_params.mode,
|
||||
ctx_params.lora_model_dir,
|
||||
ctx_params.hires_upscalers_dir);
|
||||
}
|
||||
|
||||
if (!valid) {
|
||||
|
|
@ -431,10 +433,11 @@ bool save_results(const SDCliParams& cli_params,
|
|||
if (!img.data)
|
||||
return false;
|
||||
|
||||
std::string params = gen_params.embed_image_metadata
|
||||
? get_image_params(ctx_params, gen_params, gen_params.seed + idx)
|
||||
: "";
|
||||
const bool ok = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, 90);
|
||||
const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx;
|
||||
std::string params = gen_params.embed_image_metadata
|
||||
? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
|
||||
: "";
|
||||
const bool ok = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, 90);
|
||||
LOG_INFO("save result image %d to '%s' (%s)", idx, path.string().c_str(), ok ? "success" : "failure");
|
||||
return ok;
|
||||
};
|
||||
|
|
@ -688,6 +691,13 @@ int main(int argc, const char* argv[]) {
|
|||
vae_decode_only = false;
|
||||
}
|
||||
|
||||
if (gen_params.hires_enabled &&
|
||||
(gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_MODEL ||
|
||||
gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_LANCZOS ||
|
||||
gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_NEAREST)) {
|
||||
vae_decode_only = false;
|
||||
}
|
||||
|
||||
sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview);
|
||||
|
||||
SDImageVec results;
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -5,20 +5,13 @@
|
|||
#include <map>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
#include "gguf.h"
|
||||
#include "json.hpp"
|
||||
#include "model_io/tensor_storage.h"
|
||||
#include "ordered_map.hpp"
|
||||
#include "zip.h"
|
||||
|
||||
#define SD_MAX_DIMS 5
|
||||
|
||||
enum SDVersion {
|
||||
VERSION_SD1,
|
||||
|
|
@ -195,116 +188,10 @@ enum PMVersion {
|
|||
PM_VERSION_2,
|
||||
};
|
||||
|
||||
struct TensorStorage {
|
||||
std::string name;
|
||||
ggml_type type = GGML_TYPE_F32;
|
||||
ggml_type expected_type = GGML_TYPE_COUNT;
|
||||
bool is_f8_e4m3 = false;
|
||||
bool is_f8_e5m2 = false;
|
||||
bool is_f64 = false;
|
||||
bool is_i64 = false;
|
||||
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
||||
int n_dims = 0;
|
||||
|
||||
size_t file_index = 0;
|
||||
int index_in_zip = -1; // >= means stored in a zip file
|
||||
uint64_t offset = 0; // offset in file
|
||||
|
||||
TensorStorage() = default;
|
||||
|
||||
TensorStorage(std::string name, ggml_type type, const int64_t* ne, int n_dims, size_t file_index, size_t offset = 0)
|
||||
: name(std::move(name)), type(type), n_dims(n_dims), file_index(file_index), offset(offset) {
|
||||
for (int i = 0; i < n_dims; i++) {
|
||||
this->ne[i] = ne[i];
|
||||
}
|
||||
}
|
||||
|
||||
int64_t nelements() const {
|
||||
int64_t n = 1;
|
||||
for (int i = 0; i < SD_MAX_DIMS; i++) {
|
||||
n *= ne[i];
|
||||
}
|
||||
return n;
|
||||
}
|
||||
|
||||
int64_t nbytes() const {
|
||||
return nelements() * ggml_type_size(type) / ggml_blck_size(type);
|
||||
}
|
||||
|
||||
int64_t nbytes_to_read() const {
|
||||
if (is_f8_e4m3 || is_f8_e5m2) {
|
||||
return nbytes() / 2;
|
||||
} else if (is_f64 || is_i64) {
|
||||
return nbytes() * 2;
|
||||
} else {
|
||||
return nbytes();
|
||||
}
|
||||
}
|
||||
|
||||
void unsqueeze() {
|
||||
if (n_dims == 2) {
|
||||
n_dims = 4;
|
||||
ne[3] = ne[1];
|
||||
ne[2] = ne[0];
|
||||
ne[1] = 1;
|
||||
ne[0] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<TensorStorage> chunk(size_t n) {
|
||||
std::vector<TensorStorage> chunks;
|
||||
uint64_t chunk_size = nbytes_to_read() / n;
|
||||
// printf("%d/%d\n", chunk_size, nbytes_to_read());
|
||||
reverse_ne();
|
||||
for (size_t i = 0; i < n; i++) {
|
||||
TensorStorage chunk_i = *this;
|
||||
chunk_i.ne[0] = ne[0] / n;
|
||||
chunk_i.offset = offset + i * chunk_size;
|
||||
chunk_i.reverse_ne();
|
||||
chunks.push_back(chunk_i);
|
||||
}
|
||||
reverse_ne();
|
||||
return chunks;
|
||||
}
|
||||
|
||||
void reverse_ne() {
|
||||
int64_t new_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
||||
for (int i = 0; i < n_dims; i++) {
|
||||
new_ne[i] = ne[n_dims - 1 - i];
|
||||
}
|
||||
for (int i = 0; i < n_dims; i++) {
|
||||
ne[i] = new_ne[i];
|
||||
}
|
||||
}
|
||||
|
||||
std::string to_string() const {
|
||||
std::stringstream ss;
|
||||
const char* type_name = ggml_type_name(type);
|
||||
if (is_f8_e4m3) {
|
||||
type_name = "f8_e4m3";
|
||||
} else if (is_f8_e5m2) {
|
||||
type_name = "f8_e5m2";
|
||||
} else if (is_f64) {
|
||||
type_name = "f64";
|
||||
} else if (is_i64) {
|
||||
type_name = "i64";
|
||||
}
|
||||
ss << name << " | " << type_name << " | ";
|
||||
ss << n_dims << " [";
|
||||
for (int i = 0; i < SD_MAX_DIMS; i++) {
|
||||
ss << ne[i];
|
||||
if (i != SD_MAX_DIMS - 1) {
|
||||
ss << ", ";
|
||||
}
|
||||
}
|
||||
ss << "]";
|
||||
return ss.str();
|
||||
}
|
||||
};
|
||||
|
||||
typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
|
||||
|
||||
typedef OrderedMap<std::string, TensorStorage> String2TensorStorage;
|
||||
using TensorTypeRules = std::vector<std::pair<std::string, ggml_type>>;
|
||||
|
||||
TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
|
||||
|
||||
class ModelLoader {
|
||||
protected:
|
||||
|
|
@ -314,16 +201,10 @@ protected:
|
|||
|
||||
void add_tensor_storage(const TensorStorage& tensor_storage);
|
||||
|
||||
bool parse_data_pkl(uint8_t* buffer,
|
||||
size_t buffer_size,
|
||||
zip_t* zip,
|
||||
std::string dir,
|
||||
size_t file_index,
|
||||
const std::string prefix);
|
||||
|
||||
bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_ckpt_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
|
||||
|
||||
public:
|
||||
|
|
@ -354,7 +235,6 @@ public:
|
|||
return names;
|
||||
}
|
||||
|
||||
bool save_to_gguf_file(const std::string& file_path, ggml_type type, const std::string& tensor_type_rules);
|
||||
bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
|
||||
int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
|
||||
~ModelLoader() = default;
|
||||
|
|
|
|||
57
otherarch/sdcpp/model_io/binary_io.h
Normal file
57
otherarch/sdcpp/model_io/binary_io.h
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
#ifndef __SD_MODEL_IO_BINARY_IO_H__
|
||||
#define __SD_MODEL_IO_BINARY_IO_H__
|
||||
|
||||
#include <cstdint>
|
||||
#include <ostream>
|
||||
|
||||
namespace model_io {
|
||||
|
||||
inline int32_t read_int(const uint8_t* buffer) {
|
||||
uint32_t value = 0;
|
||||
value |= static_cast<uint32_t>(buffer[3]) << 24;
|
||||
value |= static_cast<uint32_t>(buffer[2]) << 16;
|
||||
value |= static_cast<uint32_t>(buffer[1]) << 8;
|
||||
value |= static_cast<uint32_t>(buffer[0]);
|
||||
return static_cast<int32_t>(value);
|
||||
}
|
||||
|
||||
inline uint16_t read_short(const uint8_t* buffer) {
|
||||
uint16_t value = 0;
|
||||
value |= static_cast<uint16_t>(buffer[1]) << 8;
|
||||
value |= static_cast<uint16_t>(buffer[0]);
|
||||
return value;
|
||||
}
|
||||
|
||||
inline uint64_t read_u64(const uint8_t* buffer) {
|
||||
uint64_t value = 0;
|
||||
value |= static_cast<uint64_t>(buffer[7]) << 56;
|
||||
value |= static_cast<uint64_t>(buffer[6]) << 48;
|
||||
value |= static_cast<uint64_t>(buffer[5]) << 40;
|
||||
value |= static_cast<uint64_t>(buffer[4]) << 32;
|
||||
value |= static_cast<uint64_t>(buffer[3]) << 24;
|
||||
value |= static_cast<uint64_t>(buffer[2]) << 16;
|
||||
value |= static_cast<uint64_t>(buffer[1]) << 8;
|
||||
value |= static_cast<uint64_t>(buffer[0]);
|
||||
return value;
|
||||
}
|
||||
|
||||
inline void write_u64(std::ostream& stream, uint64_t value) {
|
||||
uint8_t buffer[8];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
buffer[i] = static_cast<uint8_t>((value >> (8 * i)) & 0xFF);
|
||||
}
|
||||
stream.write((const char*)buffer, sizeof(buffer));
|
||||
}
|
||||
|
||||
inline int find_char(const uint8_t* buffer, int len, char c) {
|
||||
for (int pos = 0; pos < len; pos++) {
|
||||
if (buffer[pos] == (uint8_t)c) {
|
||||
return pos;
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
} // namespace model_io
|
||||
|
||||
#endif // __SD_MODEL_IO_BINARY_IO_H__
|
||||
123
otherarch/sdcpp/model_io/gguf_io.cpp
Normal file
123
otherarch/sdcpp/model_io/gguf_io.cpp
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
#include "gguf_io.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <fstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "gguf.h"
|
||||
#include "gguf_reader_ext.h"
|
||||
#include "util.h"
|
||||
|
||||
static void set_error(std::string* error, const std::string& message) {
|
||||
if (error != nullptr) {
|
||||
*error = message;
|
||||
}
|
||||
}
|
||||
|
||||
bool is_gguf_file(const std::string& file_path) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
char magic[4];
|
||||
|
||||
file.read(magic, sizeof(magic));
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
for (uint32_t i = 0; i < sizeof(magic); i++) {
|
||||
if (magic[i] != GGUF_MAGIC[i]) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool read_gguf_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error) {
|
||||
tensor_storages.clear();
|
||||
|
||||
gguf_context* ctx_gguf_ = nullptr;
|
||||
ggml_context* ctx_meta_ = nullptr;
|
||||
|
||||
ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_});
|
||||
if (!ctx_gguf_) {
|
||||
GGUFReader gguf_reader;
|
||||
if (!gguf_reader.load(file_path)) {
|
||||
set_error(error, "failed to open '" + file_path + "' with GGUFReader");
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t data_offset = gguf_reader.data_offset();
|
||||
for (const auto& gguf_tensor_info : gguf_reader.tensors()) {
|
||||
TensorStorage tensor_storage(
|
||||
gguf_tensor_info.name,
|
||||
gguf_tensor_info.type,
|
||||
gguf_tensor_info.shape.data(),
|
||||
static_cast<int>(gguf_tensor_info.shape.size()),
|
||||
0,
|
||||
data_offset + gguf_tensor_info.offset);
|
||||
|
||||
tensor_storages.push_back(tensor_storage);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
int n_tensors = static_cast<int>(gguf_get_n_tensors(ctx_gguf_));
|
||||
|
||||
size_t data_offset = gguf_get_data_offset(ctx_gguf_);
|
||||
for (int i = 0; i < n_tensors; i++) {
|
||||
std::string name = gguf_get_tensor_name(ctx_gguf_, i);
|
||||
ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
|
||||
size_t offset = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
|
||||
|
||||
TensorStorage tensor_storage(name, dummy->type, dummy->ne, ggml_n_dims(dummy), 0, offset);
|
||||
|
||||
if (ggml_nbytes(dummy) != tensor_storage.nbytes()) {
|
||||
gguf_free(ctx_gguf_);
|
||||
ggml_free(ctx_meta_);
|
||||
set_error(error, "size mismatch for tensor '" + name + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
tensor_storages.push_back(tensor_storage);
|
||||
}
|
||||
|
||||
gguf_free(ctx_gguf_);
|
||||
ggml_free(ctx_meta_);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool write_gguf_file(const std::string& file_path,
|
||||
const std::vector<TensorWriteInfo>& tensors,
|
||||
std::string* error) {
|
||||
gguf_context* gguf_ctx = gguf_init_empty();
|
||||
if (gguf_ctx == nullptr) {
|
||||
set_error(error, "gguf_init_empty failed");
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const TensorWriteInfo& write_tensor : tensors) {
|
||||
ggml_tensor* tensor = write_tensor.tensor;
|
||||
if (tensor == nullptr) {
|
||||
set_error(error, "null tensor cannot be written to GGUF");
|
||||
gguf_free(gguf_ctx);
|
||||
return false;
|
||||
}
|
||||
gguf_add_tensor(gguf_ctx, tensor);
|
||||
}
|
||||
|
||||
LOG_INFO("trying to save tensors to %s", file_path.c_str());
|
||||
bool success = gguf_write_to_file(gguf_ctx, file_path.c_str(), false);
|
||||
if (!success) {
|
||||
set_error(error, "failed to write GGUF file '" + file_path + "'");
|
||||
}
|
||||
gguf_free(gguf_ctx);
|
||||
return success;
|
||||
}
|
||||
17
otherarch/sdcpp/model_io/gguf_io.h
Normal file
17
otherarch/sdcpp/model_io/gguf_io.h
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
#ifndef __SD_MODEL_IO_GGUF_IO_H__
|
||||
#define __SD_MODEL_IO_GGUF_IO_H__
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "tensor_storage.h"
|
||||
|
||||
bool is_gguf_file(const std::string& file_path);
|
||||
bool read_gguf_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error = nullptr);
|
||||
bool write_gguf_file(const std::string& file_path,
|
||||
const std::vector<TensorWriteInfo>& tensors,
|
||||
std::string* error = nullptr);
|
||||
|
||||
#endif // __SD_MODEL_IO_GGUF_IO_H__
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
#ifndef __GGUF_READER_HPP__
|
||||
#define __GGUF_READER_HPP__
|
||||
#ifndef __SD_MODEL_IO_GGUF_READER_EXT_H__
|
||||
#define __SD_MODEL_IO_GGUF_READER_EXT_H__
|
||||
|
||||
#include <cstdint>
|
||||
#include <fstream>
|
||||
|
|
@ -231,4 +231,4 @@ public:
|
|||
size_t data_offset() const { return data_offset_; }
|
||||
};
|
||||
|
||||
#endif // __GGUF_READER_HPP__
|
||||
#endif // __SD_MODEL_IO_GGUF_READER_EXT_H__
|
||||
1064
otherarch/sdcpp/model_io/pickle_io.cpp
Normal file
1064
otherarch/sdcpp/model_io/pickle_io.cpp
Normal file
File diff suppressed because it is too large
Load diff
21
otherarch/sdcpp/model_io/pickle_io.h
Normal file
21
otherarch/sdcpp/model_io/pickle_io.h
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
#ifndef __SD_MODEL_IO_PICKLE_IO_H__
|
||||
#define __SD_MODEL_IO_PICKLE_IO_H__
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "tensor_storage.h"
|
||||
|
||||
bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* object_size);
|
||||
bool pickle_object_is_torch_magic_number(const uint8_t* buffer, size_t buffer_size);
|
||||
bool parse_pickle_uint32_object(const uint8_t* buffer, size_t buffer_size, uint32_t* value);
|
||||
bool parse_torch_state_dict_pickle(const uint8_t* buffer,
|
||||
size_t buffer_size,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::unordered_map<std::string, uint64_t>& storage_nbytes,
|
||||
std::string* error = nullptr);
|
||||
|
||||
#endif // __SD_MODEL_IO_PICKLE_IO_H__
|
||||
316
otherarch/sdcpp/model_io/safetensors_io.cpp
Normal file
316
otherarch/sdcpp/model_io/safetensors_io.cpp
Normal file
|
|
@ -0,0 +1,316 @@
|
|||
#include "safetensors_io.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <exception>
|
||||
#include <fstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "binary_io.h"
|
||||
#include "json.hpp"
|
||||
#include "util.h"
|
||||
|
||||
static constexpr size_t ST_HEADER_SIZE_LEN = 8;
|
||||
|
||||
static void set_error(std::string* error, const std::string& message) {
|
||||
if (error != nullptr) {
|
||||
*error = message;
|
||||
}
|
||||
}
|
||||
|
||||
bool is_safetensors_file(const std::string& file_path) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// get file size
|
||||
file.seekg(0, file.end);
|
||||
size_t file_size_ = file.tellg();
|
||||
file.seekg(0, file.beg);
|
||||
|
||||
// read header size
|
||||
if (file_size_ <= ST_HEADER_SIZE_LEN) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
|
||||
file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t header_size_ = model_io::read_u64(header_size_buf);
|
||||
if (header_size_ >= file_size_ || header_size_ <= 2) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// read header
|
||||
std::vector<char> header_buf;
|
||||
header_buf.resize(header_size_ + 1);
|
||||
header_buf[header_size_] = '\0';
|
||||
file.read(header_buf.data(), header_size_);
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
try {
|
||||
nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
|
||||
ggml_type ttype = GGML_TYPE_COUNT;
|
||||
if (dtype == "F16") {
|
||||
ttype = GGML_TYPE_F16;
|
||||
} else if (dtype == "BF16") {
|
||||
ttype = GGML_TYPE_BF16;
|
||||
} else if (dtype == "F32") {
|
||||
ttype = GGML_TYPE_F32;
|
||||
} else if (dtype == "F64") {
|
||||
ttype = GGML_TYPE_F32;
|
||||
} else if (dtype == "F8_E4M3") {
|
||||
ttype = GGML_TYPE_F16;
|
||||
} else if (dtype == "F8_E5M2") {
|
||||
ttype = GGML_TYPE_F16;
|
||||
} else if (dtype == "I32") {
|
||||
ttype = GGML_TYPE_I32;
|
||||
} else if (dtype == "I64") {
|
||||
ttype = GGML_TYPE_I32;
|
||||
}
|
||||
return ttype;
|
||||
}
|
||||
|
||||
// https://huggingface.co/docs/safetensors/index
|
||||
bool read_safetensors_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
set_error(error, "failed to open '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
// get file size
|
||||
file.seekg(0, file.end);
|
||||
size_t file_size_ = file.tellg();
|
||||
file.seekg(0, file.beg);
|
||||
|
||||
// read header size
|
||||
if (file_size_ <= ST_HEADER_SIZE_LEN) {
|
||||
set_error(error, "invalid safetensor file '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
|
||||
file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
|
||||
if (!file) {
|
||||
set_error(error, "read safetensors header size failed: '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t header_size_ = model_io::read_u64(header_size_buf);
|
||||
if (header_size_ >= file_size_) {
|
||||
set_error(error, "invalid safetensor file '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
// read header
|
||||
std::vector<char> header_buf;
|
||||
header_buf.resize(header_size_ + 1);
|
||||
header_buf[header_size_] = '\0';
|
||||
file.read(header_buf.data(), header_size_);
|
||||
if (!file) {
|
||||
set_error(error, "read safetensors header failed: '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
nlohmann::json header_;
|
||||
try {
|
||||
header_ = nlohmann::json::parse(header_buf.data());
|
||||
} catch (const std::exception&) {
|
||||
set_error(error, "parsing safetensors header failed: '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
tensor_storages.clear();
|
||||
for (auto& item : header_.items()) {
|
||||
std::string name = item.key();
|
||||
nlohmann::json tensor_info = item.value();
|
||||
// LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str());
|
||||
|
||||
if (name == "__metadata__") {
|
||||
continue;
|
||||
}
|
||||
|
||||
std::string dtype = tensor_info["dtype"];
|
||||
nlohmann::json shape = tensor_info["shape"];
|
||||
|
||||
if (dtype == "U8") {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t begin = tensor_info["data_offsets"][0].get<size_t>();
|
||||
size_t end = tensor_info["data_offsets"][1].get<size_t>();
|
||||
|
||||
ggml_type type = safetensors_dtype_to_ggml_type(dtype);
|
||||
if (type == GGML_TYPE_COUNT) {
|
||||
set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (shape.size() > SD_MAX_DIMS) {
|
||||
set_error(error, "invalid tensor '" + name + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
int n_dims = (int)shape.size();
|
||||
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
||||
for (int i = 0; i < n_dims; i++) {
|
||||
ne[i] = shape[i].get<int64_t>();
|
||||
}
|
||||
|
||||
if (n_dims == 5) {
|
||||
n_dims = 4;
|
||||
ne[0] = ne[0] * ne[1];
|
||||
ne[1] = ne[2];
|
||||
ne[2] = ne[3];
|
||||
ne[3] = ne[4];
|
||||
}
|
||||
|
||||
// ggml_n_dims returns 1 for scalars
|
||||
if (n_dims == 0) {
|
||||
n_dims = 1;
|
||||
}
|
||||
|
||||
TensorStorage tensor_storage(name, type, ne, n_dims, 0, ST_HEADER_SIZE_LEN + header_size_ + begin);
|
||||
tensor_storage.reverse_ne();
|
||||
|
||||
size_t tensor_data_size = end - begin;
|
||||
|
||||
bool tensor_size_ok;
|
||||
if (dtype == "F8_E4M3") {
|
||||
tensor_storage.is_f8_e4m3 = true;
|
||||
// f8 -> f16
|
||||
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
|
||||
} else if (dtype == "F8_E5M2") {
|
||||
tensor_storage.is_f8_e5m2 = true;
|
||||
// f8 -> f16
|
||||
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
|
||||
} else if (dtype == "F64") {
|
||||
tensor_storage.is_f64 = true;
|
||||
// f64 -> f32
|
||||
tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
|
||||
} else if (dtype == "I64") {
|
||||
tensor_storage.is_i64 = true;
|
||||
// i64 -> i32
|
||||
tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
|
||||
} else {
|
||||
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
|
||||
}
|
||||
if (!tensor_size_ok) {
|
||||
set_error(error, "size mismatch for tensor '" + name + "' (" + dtype + ")");
|
||||
return false;
|
||||
}
|
||||
|
||||
tensor_storages.push_back(tensor_storage);
|
||||
|
||||
// LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str());
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_F16:
|
||||
*dtype = "F16";
|
||||
return true;
|
||||
case GGML_TYPE_BF16:
|
||||
*dtype = "BF16";
|
||||
return true;
|
||||
case GGML_TYPE_F32:
|
||||
*dtype = "F32";
|
||||
return true;
|
||||
case GGML_TYPE_I32:
|
||||
*dtype = "I32";
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool write_safetensors_file(const std::string& file_path,
|
||||
const std::vector<TensorWriteInfo>& tensors,
|
||||
std::string* error) {
|
||||
nlohmann::ordered_json header = nlohmann::ordered_json::object();
|
||||
|
||||
uint64_t data_offset = 0;
|
||||
for (const TensorWriteInfo& write_tensor : tensors) {
|
||||
ggml_tensor* tensor = write_tensor.tensor;
|
||||
if (tensor == nullptr) {
|
||||
set_error(error, "null tensor cannot be written to safetensors");
|
||||
return false;
|
||||
}
|
||||
|
||||
const std::string name = ggml_get_name(tensor);
|
||||
std::string dtype;
|
||||
if (!ggml_type_to_safetensors_dtype(tensor->type, &dtype)) {
|
||||
set_error(error,
|
||||
"unsupported safetensors dtype '" + std::string(ggml_type_name(tensor->type)) +
|
||||
"' for tensor '" + name + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
const uint64_t tensor_nbytes = ggml_nbytes(tensor);
|
||||
|
||||
nlohmann::ordered_json json_tensor_info = nlohmann::ordered_json::object();
|
||||
json_tensor_info["dtype"] = dtype;
|
||||
|
||||
nlohmann::ordered_json shape = nlohmann::ordered_json::array();
|
||||
for (int i = 0; i < write_tensor.n_dims; ++i) {
|
||||
shape.push_back(write_tensor.ne[write_tensor.n_dims - 1 - i]);
|
||||
}
|
||||
json_tensor_info["shape"] = shape;
|
||||
|
||||
nlohmann::ordered_json data_offsets = nlohmann::ordered_json::array();
|
||||
data_offsets.push_back(data_offset);
|
||||
data_offsets.push_back(data_offset + tensor_nbytes);
|
||||
json_tensor_info["data_offsets"] = data_offsets;
|
||||
|
||||
header[name] = json_tensor_info;
|
||||
data_offset += tensor_nbytes;
|
||||
}
|
||||
|
||||
const std::string header_str = header.dump();
|
||||
|
||||
std::ofstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
set_error(error, "failed to open '" + file_path + "' for writing");
|
||||
return false;
|
||||
}
|
||||
|
||||
LOG_INFO("trying to save tensors to %s", file_path.c_str());
|
||||
model_io::write_u64(file, header_str.size());
|
||||
file.write(header_str.data(), header_str.size());
|
||||
if (!file) {
|
||||
set_error(error, "failed to write safetensors header to '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const TensorWriteInfo& write_tensor : tensors) {
|
||||
ggml_tensor* tensor = write_tensor.tensor;
|
||||
const std::string name = ggml_get_name(tensor);
|
||||
const size_t tensor_nbytes = ggml_nbytes(tensor);
|
||||
file.write((const char*)tensor->data, tensor_nbytes);
|
||||
if (!file) {
|
||||
set_error(error,
|
||||
"failed to write tensor '" + name + "' to '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
17
otherarch/sdcpp/model_io/safetensors_io.h
Normal file
17
otherarch/sdcpp/model_io/safetensors_io.h
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
|
||||
#define __SD_MODEL_IO_SAFETENSORS_IO_H__
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "tensor_storage.h"
|
||||
|
||||
bool is_safetensors_file(const std::string& file_path);
|
||||
bool read_safetensors_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error = nullptr);
|
||||
bool write_safetensors_file(const std::string& file_path,
|
||||
const std::vector<TensorWriteInfo>& tensors,
|
||||
std::string* error = nullptr);
|
||||
|
||||
#endif // __SD_MODEL_IO_SAFETENSORS_IO_H__
|
||||
132
otherarch/sdcpp/model_io/tensor_storage.h
Normal file
132
otherarch/sdcpp/model_io/tensor_storage.h
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
#ifndef __SD_TENSOR_STORAGE_H__
|
||||
#define __SD_TENSOR_STORAGE_H__
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <functional>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml.h"
|
||||
|
||||
#define SD_MAX_DIMS 5
|
||||
|
||||
struct TensorStorage {
|
||||
std::string name;
|
||||
ggml_type type = GGML_TYPE_F32;
|
||||
ggml_type expected_type = GGML_TYPE_COUNT;
|
||||
bool is_f8_e4m3 = false;
|
||||
bool is_f8_e5m2 = false;
|
||||
bool is_f64 = false;
|
||||
bool is_i64 = false;
|
||||
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
||||
int n_dims = 0;
|
||||
|
||||
std::string storage_key;
|
||||
size_t file_index = 0;
|
||||
int index_in_zip = -1; // >= means stored in a zip file
|
||||
uint64_t offset = 0; // offset in file
|
||||
|
||||
TensorStorage() = default;
|
||||
|
||||
TensorStorage(std::string name, ggml_type type, const int64_t* ne, int n_dims, size_t file_index, size_t offset = 0)
|
||||
: name(std::move(name)), type(type), n_dims(n_dims), file_index(file_index), offset(offset) {
|
||||
for (int i = 0; i < n_dims; i++) {
|
||||
this->ne[i] = ne[i];
|
||||
}
|
||||
}
|
||||
|
||||
int64_t nelements() const {
|
||||
int64_t n = 1;
|
||||
for (int i = 0; i < SD_MAX_DIMS; i++) {
|
||||
n *= ne[i];
|
||||
}
|
||||
return n;
|
||||
}
|
||||
|
||||
int64_t nbytes() const {
|
||||
return nelements() * ggml_type_size(type) / ggml_blck_size(type);
|
||||
}
|
||||
|
||||
int64_t nbytes_to_read() const {
|
||||
if (is_f8_e4m3 || is_f8_e5m2) {
|
||||
return nbytes() / 2;
|
||||
} else if (is_f64 || is_i64) {
|
||||
return nbytes() * 2;
|
||||
} else {
|
||||
return nbytes();
|
||||
}
|
||||
}
|
||||
|
||||
void unsqueeze() {
|
||||
if (n_dims == 2) {
|
||||
n_dims = 4;
|
||||
ne[3] = ne[1];
|
||||
ne[2] = ne[0];
|
||||
ne[1] = 1;
|
||||
ne[0] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<TensorStorage> chunk(size_t n) {
|
||||
std::vector<TensorStorage> chunks;
|
||||
uint64_t chunk_size = nbytes_to_read() / n;
|
||||
// printf("%d/%d\n", chunk_size, nbytes_to_read());
|
||||
reverse_ne();
|
||||
for (size_t i = 0; i < n; i++) {
|
||||
TensorStorage chunk_i = *this;
|
||||
chunk_i.ne[0] = ne[0] / n;
|
||||
chunk_i.offset = offset + i * chunk_size;
|
||||
chunk_i.reverse_ne();
|
||||
chunks.push_back(chunk_i);
|
||||
}
|
||||
reverse_ne();
|
||||
return chunks;
|
||||
}
|
||||
|
||||
void reverse_ne() {
|
||||
int64_t new_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
||||
for (int i = 0; i < n_dims; i++) {
|
||||
new_ne[i] = ne[n_dims - 1 - i];
|
||||
}
|
||||
for (int i = 0; i < n_dims; i++) {
|
||||
ne[i] = new_ne[i];
|
||||
}
|
||||
}
|
||||
|
||||
std::string to_string() const {
|
||||
std::stringstream ss;
|
||||
const char* type_name = ggml_type_name(type);
|
||||
if (is_f8_e4m3) {
|
||||
type_name = "f8_e4m3";
|
||||
} else if (is_f8_e5m2) {
|
||||
type_name = "f8_e5m2";
|
||||
} else if (is_f64) {
|
||||
type_name = "f64";
|
||||
} else if (is_i64) {
|
||||
type_name = "i64";
|
||||
}
|
||||
ss << name << " | " << type_name << " | ";
|
||||
ss << n_dims << " [";
|
||||
for (int i = 0; i < SD_MAX_DIMS; i++) {
|
||||
ss << ne[i];
|
||||
if (i != SD_MAX_DIMS - 1) {
|
||||
ss << ", ";
|
||||
}
|
||||
}
|
||||
ss << "]";
|
||||
return ss.str();
|
||||
}
|
||||
};
|
||||
|
||||
struct TensorWriteInfo {
|
||||
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
||||
int n_dims = 0;
|
||||
ggml_tensor* tensor = nullptr;
|
||||
};
|
||||
|
||||
typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
|
||||
|
||||
#endif // __SD_TENSOR_STORAGE_H__
|
||||
252
otherarch/sdcpp/model_io/torch_legacy_io.cpp
Normal file
252
otherarch/sdcpp/model_io/torch_legacy_io.cpp
Normal file
|
|
@ -0,0 +1,252 @@
|
|||
#include "torch_legacy_io.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <fstream>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "pickle_io.h"
|
||||
#include "util.h"
|
||||
|
||||
// torch.save format background:
|
||||
//
|
||||
// - Before PyTorch 1.6.0, torch.save used this legacy non-zip format by
|
||||
// default.
|
||||
// - Since PyTorch 1.6.0, torch.save defaults to an uncompressed ZIP64 archive
|
||||
// containing data.pkl, data/, version, and, since PyTorch 2.1.0, byteorder.
|
||||
// - The old format can still be produced explicitly with:
|
||||
// torch.save(obj, path, _use_new_zipfile_serialization=False)
|
||||
//
|
||||
// Whether obj is a state_dict or a whole nn.Module does not change the outer
|
||||
// container format selected by torch.save. It changes the pickled object inside:
|
||||
//
|
||||
// - state_dict: usually an OrderedDict[str, Tensor]. pickle_io.cpp supports a
|
||||
// restricted subset of this layout because tensor metadata and raw storages
|
||||
// can be recovered without executing pickle callables.
|
||||
// - whole module/checkpoint object: arbitrary Python object graph. This may
|
||||
// require importing user classes and executing pickle GLOBAL/REDUCE rebuild
|
||||
// logic, so it is intentionally not supported here.
|
||||
//
|
||||
// Legacy non-zip PyTorch files are not a single pickle object:
|
||||
//
|
||||
// 1. pickle object: PyTorch legacy magic number
|
||||
// 2. pickle object: legacy protocol version, expected to be 1001
|
||||
// 3. pickle object: sys_info metadata, ignored by this reader
|
||||
// 4. pickle object: state_dict metadata, parsed by pickle_io.cpp
|
||||
// 5. pickle object: serialized storage key list, skipped here
|
||||
// 6. raw storage data payloads
|
||||
// - PyTorch writes storages after the pickles, ordered by storage key
|
||||
// - each storage has an 8-byte legacy storage header followed by raw bytes
|
||||
static constexpr size_t LEGACY_STORAGE_HEADER_SIZE = 8;
|
||||
|
||||
static void set_error(std::string* error, const std::string& message) {
|
||||
if (error != nullptr) {
|
||||
*error = message;
|
||||
}
|
||||
}
|
||||
|
||||
static std::string bytes_to_hex(const std::vector<uint8_t>& bytes) {
|
||||
static const char* hex = "0123456789ABCDEF";
|
||||
std::string result;
|
||||
result.reserve(bytes.size() * 3);
|
||||
for (size_t i = 0; i < bytes.size(); ++i) {
|
||||
if (i > 0) {
|
||||
result.push_back('-');
|
||||
}
|
||||
result.push_back(hex[(bytes[i] >> 4) & 0x0F]);
|
||||
result.push_back(hex[bytes[i] & 0x0F]);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
static bool is_probably_tar_file(const std::vector<uint8_t>& header) {
|
||||
return header.size() >= 262 &&
|
||||
header[257] == 'u' &&
|
||||
header[258] == 's' &&
|
||||
header[259] == 't' &&
|
||||
header[260] == 'a' &&
|
||||
header[261] == 'r';
|
||||
}
|
||||
|
||||
static std::string torch_legacy_diagnostics(const std::string& file_path, const std::vector<uint8_t>& buffer) {
|
||||
if (!ends_with(file_path, ".pt") && !ends_with(file_path, ".pth")) {
|
||||
return "";
|
||||
}
|
||||
if (buffer.empty()) {
|
||||
return "unsupported PyTorch file '" + file_path + "': empty file";
|
||||
}
|
||||
|
||||
size_t short_len = std::min<size_t>(buffer.size(), 32);
|
||||
std::vector<uint8_t> short_header(buffer.begin(), buffer.begin() + short_len);
|
||||
const bool raw_pickle = buffer[0] == 0x80;
|
||||
const bool tar_file = is_probably_tar_file(buffer);
|
||||
|
||||
std::string message = "unsupported PyTorch file '" + file_path + "': first bytes " +
|
||||
bytes_to_hex(short_header) +
|
||||
", raw_pickle=" + (raw_pickle ? "true" : "false") +
|
||||
", tar=" + (tar_file ? "true" : "false");
|
||||
if (raw_pickle) {
|
||||
message += "; raw pickle did not match the restricted state_dict layouts currently supported";
|
||||
} else if (tar_file) {
|
||||
message += "; legacy tar PyTorch checkpoints are not supported yet";
|
||||
}
|
||||
return message;
|
||||
}
|
||||
|
||||
bool read_torch_legacy_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
set_error(error, "failed to open '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
file.seekg(0, file.end);
|
||||
size_t file_size = (size_t)file.tellg();
|
||||
file.seekg(0, file.beg);
|
||||
if (file_size == 0) {
|
||||
set_error(error, "empty file '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> buffer(file_size);
|
||||
file.read((char*)buffer.data(), file_size);
|
||||
if (!file) {
|
||||
set_error(error, "failed to read '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
auto finalize_tensor_offsets = [&](size_t storage_data_offset,
|
||||
const std::unordered_map<std::string, uint64_t>& legacy_storage_map) -> bool {
|
||||
if (storage_data_offset > file_size) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<std::string> storage_keys;
|
||||
storage_keys.reserve(legacy_storage_map.size());
|
||||
for (const auto& [storage_key, _] : legacy_storage_map) {
|
||||
storage_keys.push_back(storage_key);
|
||||
}
|
||||
std::sort(storage_keys.begin(), storage_keys.end());
|
||||
|
||||
std::unordered_map<std::string, uint64_t> storage_offsets;
|
||||
uint64_t current_offset = storage_data_offset;
|
||||
for (const auto& storage_key : storage_keys) {
|
||||
auto it = legacy_storage_map.find(storage_key);
|
||||
if (it == legacy_storage_map.end()) {
|
||||
return false;
|
||||
}
|
||||
if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
|
||||
return false;
|
||||
}
|
||||
storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
|
||||
current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
|
||||
}
|
||||
|
||||
for (auto& tensor_storage : tensor_storages) {
|
||||
if (tensor_storage.storage_key.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
auto it_offset = storage_offsets.find(tensor_storage.storage_key);
|
||||
auto it_size = legacy_storage_map.find(tensor_storage.storage_key);
|
||||
if (it_offset == storage_offsets.end() || it_size == legacy_storage_map.end()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t base_offset = it_offset->second;
|
||||
uint64_t storage_nbytes = it_size->second;
|
||||
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||
if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
|
||||
return false;
|
||||
}
|
||||
|
||||
tensor_storage.offset = base_offset + tensor_storage.offset;
|
||||
tensor_storage.storage_key.clear();
|
||||
}
|
||||
|
||||
return true;
|
||||
};
|
||||
|
||||
auto parse_state_dict_at = [&](size_t state_dict_offset, size_t state_dict_size, size_t* storage_data_offset) -> bool {
|
||||
tensor_storages.clear();
|
||||
std::unordered_map<std::string, uint64_t> legacy_storage_map;
|
||||
if (!parse_torch_state_dict_pickle(buffer.data() + state_dict_offset,
|
||||
state_dict_size,
|
||||
tensor_storages,
|
||||
legacy_storage_map,
|
||||
error)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t offset_after_state_dict = state_dict_offset + state_dict_size;
|
||||
size_t storage_keys_size = 0;
|
||||
if (!skip_pickle_object(buffer.data() + offset_after_state_dict,
|
||||
buffer.size() - offset_after_state_dict,
|
||||
&storage_keys_size)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
*storage_data_offset = offset_after_state_dict + storage_keys_size;
|
||||
return finalize_tensor_offsets(*storage_data_offset, legacy_storage_map);
|
||||
};
|
||||
|
||||
size_t object_size_1 = 0;
|
||||
size_t offset = 0;
|
||||
|
||||
if (skip_pickle_object(buffer.data(), buffer.size(), &object_size_1) &&
|
||||
pickle_object_is_torch_magic_number(buffer.data(), object_size_1)) {
|
||||
offset += object_size_1;
|
||||
|
||||
size_t object_size_2 = 0;
|
||||
if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &object_size_2)) {
|
||||
set_error(error, torch_legacy_diagnostics(file_path, buffer));
|
||||
return false;
|
||||
}
|
||||
uint32_t protocol_version = 0;
|
||||
if (!parse_pickle_uint32_object(buffer.data() + offset, object_size_2, &protocol_version) || protocol_version != 1001) {
|
||||
set_error(error, torch_legacy_diagnostics(file_path, buffer));
|
||||
return false;
|
||||
}
|
||||
offset += object_size_2;
|
||||
|
||||
size_t object_size_3 = 0;
|
||||
if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &object_size_3)) {
|
||||
set_error(error, torch_legacy_diagnostics(file_path, buffer));
|
||||
return false;
|
||||
}
|
||||
offset += object_size_3;
|
||||
|
||||
size_t state_dict_size = 0;
|
||||
if (!skip_pickle_object(buffer.data() + offset, buffer.size() - offset, &state_dict_size)) {
|
||||
set_error(error, torch_legacy_diagnostics(file_path, buffer));
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t storage_data_offset = 0;
|
||||
if (parse_state_dict_at(offset, state_dict_size, &storage_data_offset)) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (error != nullptr && error->empty()) {
|
||||
set_error(error, torch_legacy_diagnostics(file_path, buffer));
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t state_dict_size = 0;
|
||||
if (skip_pickle_object(buffer.data(), buffer.size(), &state_dict_size)) {
|
||||
size_t storage_data_offset = 0;
|
||||
if (parse_state_dict_at(0, state_dict_size, &storage_data_offset)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
if (error != nullptr && error->empty()) {
|
||||
set_error(error, torch_legacy_diagnostics(file_path, buffer));
|
||||
}
|
||||
return false;
|
||||
}
|
||||
13
otherarch/sdcpp/model_io/torch_legacy_io.h
Normal file
13
otherarch/sdcpp/model_io/torch_legacy_io.h
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
#ifndef __SD_MODEL_IO_TORCH_LEGACY_IO_H__
|
||||
#define __SD_MODEL_IO_TORCH_LEGACY_IO_H__
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "tensor_storage.h"
|
||||
|
||||
bool read_torch_legacy_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error = nullptr);
|
||||
|
||||
#endif // __SD_MODEL_IO_TORCH_LEGACY_IO_H__
|
||||
140
otherarch/sdcpp/model_io/torch_zip_io.cpp
Normal file
140
otherarch/sdcpp/model_io/torch_zip_io.cpp
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
#include "torch_zip_io.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "pickle_io.h"
|
||||
|
||||
#include "zip.h"
|
||||
|
||||
static void set_error(std::string* error, const std::string& message) {
|
||||
if (error != nullptr) {
|
||||
*error = message;
|
||||
}
|
||||
}
|
||||
|
||||
bool is_torch_zip_file(const std::string& file_path) {
|
||||
zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
|
||||
if (zip == nullptr) {
|
||||
return false;
|
||||
}
|
||||
zip_close(zip);
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool find_zip_entry(zip_t* zip, const std::string& entry_name, int* index, uint64_t* size) {
|
||||
size_t n = zip_entries_total(zip);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
zip_entry_openbyindex(zip, i);
|
||||
std::string name = zip_entry_name(zip);
|
||||
if (name == entry_name) {
|
||||
*index = (int)i;
|
||||
*size = zip_entry_size(zip);
|
||||
zip_entry_close(zip);
|
||||
return true;
|
||||
}
|
||||
zip_entry_close(zip);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static bool parse_zip_data_pkl(const uint8_t* buffer,
|
||||
size_t buffer_size,
|
||||
zip_t* zip,
|
||||
const std::string& dir,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error) {
|
||||
std::vector<TensorStorage> parsed_tensors;
|
||||
std::unordered_map<std::string, uint64_t> storage_nbytes;
|
||||
if (!parse_torch_state_dict_pickle(buffer, buffer_size, parsed_tensors, storage_nbytes, error)) {
|
||||
if (error != nullptr && error->empty()) {
|
||||
*error = "failed to parse torch zip pickle metadata";
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
for (auto& tensor_storage : parsed_tensors) {
|
||||
if (tensor_storage.storage_key.empty()) {
|
||||
set_error(error, "tensor '" + tensor_storage.name + "' has no storage key");
|
||||
return false;
|
||||
}
|
||||
|
||||
const std::string entry_name = dir + "data/" + tensor_storage.storage_key;
|
||||
int zip_index = -1;
|
||||
uint64_t entry_size = 0;
|
||||
if (!find_zip_entry(zip, entry_name, &zip_index, &entry_size)) {
|
||||
set_error(error, "storage entry '" + entry_name + "' was not found");
|
||||
return false;
|
||||
}
|
||||
|
||||
auto it_storage_size = storage_nbytes.find(tensor_storage.storage_key);
|
||||
if (it_storage_size != storage_nbytes.end() && entry_size < it_storage_size->second) {
|
||||
set_error(error, "storage entry '" + entry_name + "' is smaller than pickle metadata");
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||
if (tensor_storage.offset + tensor_nbytes > entry_size) {
|
||||
set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
tensor_storage.index_in_zip = zip_index;
|
||||
tensor_storage.storage_key.clear();
|
||||
tensor_storages.push_back(tensor_storage);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool read_torch_zip_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error) {
|
||||
zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
|
||||
if (zip == nullptr) {
|
||||
set_error(error, "failed to open '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
tensor_storages.clear();
|
||||
bool success = true;
|
||||
bool found_data_pkl = false;
|
||||
int n = (int)zip_entries_total(zip);
|
||||
for (int i = 0; i < n; ++i) {
|
||||
zip_entry_openbyindex(zip, i);
|
||||
std::string name = zip_entry_name(zip);
|
||||
size_t pos = name.find("data.pkl");
|
||||
if (pos != std::string::npos) {
|
||||
found_data_pkl = true;
|
||||
std::string dir = name.substr(0, pos);
|
||||
void* pkl_data = nullptr;
|
||||
size_t pkl_size = 0;
|
||||
zip_entry_read(zip, &pkl_data, &pkl_size);
|
||||
|
||||
if (pkl_data == nullptr || pkl_size == 0) {
|
||||
set_error(error, "failed to read '" + name + "' from '" + file_path + "'");
|
||||
success = false;
|
||||
} else if (!parse_zip_data_pkl((const uint8_t*)pkl_data, pkl_size, zip, dir, tensor_storages, error)) {
|
||||
success = false;
|
||||
}
|
||||
|
||||
free(pkl_data);
|
||||
}
|
||||
zip_entry_close(zip);
|
||||
|
||||
if (!success) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (success && !found_data_pkl) {
|
||||
set_error(error, "data.pkl was not found in '" + file_path + "'");
|
||||
success = false;
|
||||
}
|
||||
|
||||
zip_close(zip);
|
||||
return success;
|
||||
}
|
||||
14
otherarch/sdcpp/model_io/torch_zip_io.h
Normal file
14
otherarch/sdcpp/model_io/torch_zip_io.h
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
#ifndef __SD_MODEL_IO_TORCH_ZIP_IO_H__
|
||||
#define __SD_MODEL_IO_TORCH_ZIP_IO_H__
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "tensor_storage.h"
|
||||
|
||||
bool is_torch_zip_file(const std::string& file_path);
|
||||
bool read_torch_zip_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error = nullptr);
|
||||
|
||||
#endif // __SD_MODEL_IO_TORCH_ZIP_IO_H__
|
||||
|
|
@ -21,7 +21,31 @@
|
|||
#include "util.cpp"
|
||||
#include "name_conversion.cpp"
|
||||
#include "upscaler.cpp"
|
||||
|
||||
#include "zip.c"
|
||||
#include "model_io/binary_io.h"
|
||||
namespace pickle {
|
||||
#include "model_io/pickle_io.cpp"
|
||||
}
|
||||
namespace gguf {
|
||||
#include "model_io/gguf_io.cpp"
|
||||
}
|
||||
namespace safetensors {
|
||||
#include "model_io/safetensors_io.cpp"
|
||||
}
|
||||
using namespace pickle;
|
||||
namespace torch_legacy {
|
||||
#include "model_io/torch_legacy_io.cpp"
|
||||
}
|
||||
namespace torch_zip {
|
||||
#include "model_io/torch_zip_io.cpp"
|
||||
}
|
||||
using namespace gguf;
|
||||
using namespace safetensors;
|
||||
using namespace torch_legacy;
|
||||
using namespace torch_zip;
|
||||
#include "model.cpp"
|
||||
|
||||
#include "tokenizers/bpe_tokenizer.cpp"
|
||||
#include "tokenizers/clip_tokenizer.cpp"
|
||||
#include "tokenizers/mistral_tokenizer.cpp"
|
||||
|
|
@ -29,7 +53,6 @@
|
|||
#include "tokenizers/t5_unigram_tokenizer.cpp"
|
||||
#include "tokenizers/tokenizer.cpp"
|
||||
#include "tokenizers/tokenize_util.cpp"
|
||||
#include "zip.c"
|
||||
|
||||
#include "otherarch/utils.h"
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@
|
|||
#include "pmid.hpp"
|
||||
#include "sample-cache.h"
|
||||
#include "tae.hpp"
|
||||
#include "upscaler.h"
|
||||
#include "vae.hpp"
|
||||
|
||||
#include "latent-preview.h"
|
||||
|
|
@ -234,11 +235,11 @@ public:
|
|||
device = 0;
|
||||
}
|
||||
if (device >= device_count) {
|
||||
LOG_WARN("Cannot find targeted vulkan device (%llu). Falling back to device 0.", device);
|
||||
LOG_WARN("Cannot find targeted vulkan device (%zu). Falling back to device 0.", device);
|
||||
device = 0;
|
||||
}
|
||||
}
|
||||
LOG_INFO("Vulkan: Using device %llu", device);
|
||||
LOG_INFO("Vulkan: Using device %zu", device);
|
||||
backend = ggml_backend_vk_init(device);
|
||||
}
|
||||
if (!backend) {
|
||||
|
|
@ -2380,6 +2381,35 @@ enum lora_apply_mode_t str_to_lora_apply_mode(const char* str) {
|
|||
return LORA_APPLY_MODE_COUNT;
|
||||
}
|
||||
|
||||
const char* hires_upscaler_to_str[] = {
|
||||
"None",
|
||||
"Latent",
|
||||
"Latent (nearest)",
|
||||
"Latent (nearest-exact)",
|
||||
"Latent (antialiased)",
|
||||
"Latent (bicubic)",
|
||||
"Latent (bicubic antialiased)",
|
||||
"Lanczos",
|
||||
"Nearest",
|
||||
"Model",
|
||||
};
|
||||
|
||||
const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler) {
|
||||
if (upscaler >= SD_HIRES_UPSCALER_NONE && upscaler < SD_HIRES_UPSCALER_COUNT) {
|
||||
return hires_upscaler_to_str[upscaler];
|
||||
}
|
||||
return NONE_STR;
|
||||
}
|
||||
|
||||
enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) {
|
||||
for (int i = 0; i < SD_HIRES_UPSCALER_COUNT; i++) {
|
||||
if (!strcmp(str, hires_upscaler_to_str[i])) {
|
||||
return (enum sd_hires_upscaler_t)i;
|
||||
}
|
||||
}
|
||||
return SD_HIRES_UPSCALER_COUNT;
|
||||
}
|
||||
|
||||
void sd_cache_params_init(sd_cache_params_t* cache_params) {
|
||||
*cache_params = {};
|
||||
cache_params->mode = SD_CACHE_DISABLED;
|
||||
|
|
@ -2408,6 +2438,19 @@ void sd_cache_params_init(sd_cache_params_t* cache_params) {
|
|||
cache_params->spectrum_stop_percent = 0.9f;
|
||||
}
|
||||
|
||||
void sd_hires_params_init(sd_hires_params_t* hires_params) {
|
||||
*hires_params = {};
|
||||
hires_params->enabled = false;
|
||||
hires_params->upscaler = SD_HIRES_UPSCALER_LATENT;
|
||||
hires_params->model_path = nullptr;
|
||||
hires_params->scale = 2.0f;
|
||||
hires_params->target_width = 0;
|
||||
hires_params->target_height = 0;
|
||||
hires_params->steps = 0;
|
||||
hires_params->denoising_strength = 0.7f;
|
||||
hires_params->upscale_tile_size = 128;
|
||||
}
|
||||
|
||||
void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
|
||||
*sd_ctx_params = {};
|
||||
sd_ctx_params->vae_decode_only = true;
|
||||
|
|
@ -2577,6 +2620,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
|
|||
sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f};
|
||||
sd_img_gen_params->vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
|
||||
sd_cache_params_init(&sd_img_gen_params->cache);
|
||||
sd_hires_params_init(&sd_img_gen_params->hires);
|
||||
}
|
||||
|
||||
char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
|
||||
|
|
@ -2603,7 +2647,8 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
|
|||
"increase_ref_index: %s\n"
|
||||
"control_strength: %.2f\n"
|
||||
"photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
|
||||
"VAE tiling: %s\n",
|
||||
"VAE tiling: %s\n"
|
||||
"hires: {enabled=%s, upscaler=%s, model_path=%s, scale=%.2f, target=%dx%d, steps=%d, denoising_strength=%.2f}\n",
|
||||
SAFE_STR(sd_img_gen_params->prompt),
|
||||
SAFE_STR(sd_img_gen_params->negative_prompt),
|
||||
sd_img_gen_params->clip_skip,
|
||||
|
|
@ -2620,7 +2665,15 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
|
|||
sd_img_gen_params->pm_params.style_strength,
|
||||
sd_img_gen_params->pm_params.id_images_count,
|
||||
SAFE_STR(sd_img_gen_params->pm_params.id_embed_path),
|
||||
BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled));
|
||||
BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled),
|
||||
BOOL_STR(sd_img_gen_params->hires.enabled),
|
||||
sd_hires_upscaler_name(sd_img_gen_params->hires.upscaler),
|
||||
SAFE_STR(sd_img_gen_params->hires.model_path),
|
||||
sd_img_gen_params->hires.scale,
|
||||
sd_img_gen_params->hires.target_width,
|
||||
sd_img_gen_params->hires.target_height,
|
||||
sd_img_gen_params->hires.steps,
|
||||
sd_img_gen_params->hires.denoising_strength);
|
||||
const char* cache_mode_str = "disabled";
|
||||
if (sd_img_gen_params->cache.mode == SD_CACHE_EASYCACHE) {
|
||||
cache_mode_str = "easycache";
|
||||
|
|
@ -2724,8 +2777,10 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
|
|||
return EXPONENTIAL_SCHEDULER;
|
||||
}
|
||||
}
|
||||
if (sample_method == LCM_SAMPLE_METHOD) {
|
||||
if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) {
|
||||
return LCM_SCHEDULER;
|
||||
} else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
|
||||
return SIMPLE_SCHEDULER;
|
||||
}
|
||||
return DISCRETE_SCHEDULER;
|
||||
}
|
||||
|
|
@ -2799,6 +2854,7 @@ struct GenerationRequest {
|
|||
sd_guidance_params_t guidance = {};
|
||||
sd_guidance_params_t high_noise_guidance = {};
|
||||
sd_pm_params_t pm_params = {};
|
||||
sd_hires_params_t hires = {};
|
||||
int frames = -1;
|
||||
float vace_strength = 1.f;
|
||||
|
||||
|
|
@ -2820,6 +2876,7 @@ struct GenerationRequest {
|
|||
auto_resize_ref_image = sd_img_gen_params->auto_resize_ref_image;
|
||||
guidance = sd_img_gen_params->sample_params.guidance;
|
||||
pm_params = sd_img_gen_params->pm_params;
|
||||
hires = sd_img_gen_params->hires;
|
||||
cache_params = &sd_img_gen_params->cache;
|
||||
resolve(sd_ctx);
|
||||
}
|
||||
|
|
@ -2842,26 +2899,76 @@ struct GenerationRequest {
|
|||
}
|
||||
|
||||
void align_generation_request_size() {
|
||||
align_image_size(&width, &height, "generation request");
|
||||
}
|
||||
|
||||
void align_image_size(int* target_width, int* target_height, const char* label) {
|
||||
int spatial_multiple = vae_scale_factor * diffusion_model_down_factor;
|
||||
int width_offset = align_up_offset(width, spatial_multiple);
|
||||
int height_offset = align_up_offset(height, spatial_multiple);
|
||||
int width_offset = align_up_offset(*target_width, spatial_multiple);
|
||||
int height_offset = align_up_offset(*target_height, spatial_multiple);
|
||||
if (width_offset <= 0 && height_offset <= 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
int original_width = width;
|
||||
int original_height = height;
|
||||
int original_width = *target_width;
|
||||
int original_height = *target_height;
|
||||
|
||||
width += width_offset;
|
||||
height += height_offset;
|
||||
LOG_WARN("align up %dx%d to %dx%d (multiple=%d)",
|
||||
*target_width += width_offset;
|
||||
*target_height += height_offset;
|
||||
LOG_WARN("align %s up %dx%d to %dx%d (multiple=%d)",
|
||||
label,
|
||||
original_width,
|
||||
original_height,
|
||||
width,
|
||||
height,
|
||||
*target_width,
|
||||
*target_height,
|
||||
spatial_multiple);
|
||||
}
|
||||
|
||||
void resolve_hires() {
|
||||
if (!hires.enabled) {
|
||||
return;
|
||||
}
|
||||
if (hires.upscaler == SD_HIRES_UPSCALER_NONE) {
|
||||
hires.enabled = false;
|
||||
return;
|
||||
}
|
||||
if (hires.upscaler < SD_HIRES_UPSCALER_NONE || hires.upscaler >= SD_HIRES_UPSCALER_COUNT) {
|
||||
LOG_WARN("hires upscaler '%d' is invalid, disabling hires", hires.upscaler);
|
||||
hires.enabled = false;
|
||||
return;
|
||||
}
|
||||
if (hires.upscaler == SD_HIRES_UPSCALER_MODEL && strlen(SAFE_STR(hires.model_path)) == 0) {
|
||||
LOG_WARN("hires model upscaler requires a model path, disabling hires");
|
||||
hires.enabled = false;
|
||||
return;
|
||||
}
|
||||
if (hires.scale <= 0.f && hires.target_width <= 0 && hires.target_height <= 0) {
|
||||
LOG_WARN("hires scale must be positive when no target size is set, disabling hires");
|
||||
hires.enabled = false;
|
||||
return;
|
||||
}
|
||||
hires.denoising_strength = std::clamp(hires.denoising_strength, 0.0001f, 1.f);
|
||||
hires.steps = std::max(0, hires.steps);
|
||||
|
||||
if (hires.target_width > 0 && hires.target_height > 0) {
|
||||
// pass
|
||||
} else if (hires.target_width > 0) {
|
||||
hires.target_height = hires.target_width;
|
||||
} else if (hires.target_height > 0) {
|
||||
hires.target_width = hires.target_height;
|
||||
} else {
|
||||
hires.target_width = static_cast<int>(std::round(width * hires.scale));
|
||||
hires.target_height = static_cast<int>(std::round(height * hires.scale));
|
||||
}
|
||||
|
||||
if (hires.target_width <= 0 || hires.target_height <= 0) {
|
||||
LOG_WARN("hires target size is not positive, disabling hires");
|
||||
hires.enabled = false;
|
||||
return;
|
||||
}
|
||||
align_image_size(&hires.target_width, &hires.target_height, "hires target");
|
||||
}
|
||||
|
||||
static void resolve_guidance(sd_ctx_t* sd_ctx,
|
||||
sd_guidance_params_t* guidance,
|
||||
bool* use_uncond,
|
||||
|
|
@ -2902,6 +3009,7 @@ struct GenerationRequest {
|
|||
|
||||
void resolve(sd_ctx_t* sd_ctx) {
|
||||
align_generation_request_size();
|
||||
resolve_hires();
|
||||
seed = resolve_seed(seed);
|
||||
|
||||
resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_cond);
|
||||
|
|
@ -3392,7 +3500,7 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
|
|||
}
|
||||
decoded_images.push_back(std::move(image));
|
||||
int64_t t2 = ggml_time_ms();
|
||||
LOG_INFO("latent %" PRId64 " decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000);
|
||||
LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000);
|
||||
}
|
||||
|
||||
int64_t t4 = ggml_time_ms();
|
||||
|
|
@ -3414,6 +3522,135 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
|
|||
return result_images;
|
||||
}
|
||||
|
||||
static sd::Tensor<float> upscale_hires_latent(sd_ctx_t* sd_ctx,
|
||||
const sd::Tensor<float>& latent,
|
||||
const GenerationRequest& request,
|
||||
UpscalerGGML* upscaler) {
|
||||
auto get_hires_latent_target_shape = [&]() {
|
||||
std::vector<int64_t> target_shape = latent.shape();
|
||||
if (target_shape.size() < 2) {
|
||||
target_shape.clear();
|
||||
return target_shape;
|
||||
}
|
||||
target_shape[0] = request.hires.target_width / request.vae_scale_factor;
|
||||
target_shape[1] = request.hires.target_height / request.vae_scale_factor;
|
||||
return target_shape;
|
||||
};
|
||||
|
||||
if (request.hires.upscaler == SD_HIRES_UPSCALER_LATENT ||
|
||||
request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST ||
|
||||
request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT ||
|
||||
request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_ANTIALIASED ||
|
||||
request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC ||
|
||||
request.hires.upscaler == SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED) {
|
||||
std::vector<int64_t> target_shape = get_hires_latent_target_shape();
|
||||
if (target_shape.empty()) {
|
||||
LOG_ERROR("latent has invalid shape for hires upscale");
|
||||
return {};
|
||||
}
|
||||
|
||||
sd::ops::InterpolateMode mode = sd::ops::InterpolateMode::Nearest;
|
||||
bool antialias = false;
|
||||
switch (request.hires.upscaler) {
|
||||
case SD_HIRES_UPSCALER_LATENT:
|
||||
mode = sd::ops::InterpolateMode::Bilinear;
|
||||
break;
|
||||
case SD_HIRES_UPSCALER_LATENT_NEAREST:
|
||||
mode = sd::ops::InterpolateMode::Nearest;
|
||||
break;
|
||||
case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT:
|
||||
mode = sd::ops::InterpolateMode::NearestExact;
|
||||
break;
|
||||
case SD_HIRES_UPSCALER_LATENT_ANTIALIASED:
|
||||
mode = sd::ops::InterpolateMode::Bilinear;
|
||||
antialias = true;
|
||||
break;
|
||||
case SD_HIRES_UPSCALER_LATENT_BICUBIC:
|
||||
mode = sd::ops::InterpolateMode::Bicubic;
|
||||
break;
|
||||
case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED:
|
||||
mode = sd::ops::InterpolateMode::Bicubic;
|
||||
antialias = true;
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
|
||||
LOG_INFO("hires %s upscale %" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64,
|
||||
sd_hires_upscaler_name(request.hires.upscaler),
|
||||
latent.shape()[0],
|
||||
latent.shape()[1],
|
||||
target_shape[0],
|
||||
target_shape[1]);
|
||||
|
||||
return sd::ops::interpolate(latent, target_shape, mode, false, antialias);
|
||||
} else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL ||
|
||||
request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS ||
|
||||
request.hires.upscaler == SD_HIRES_UPSCALER_NEAREST) {
|
||||
if (sd_ctx->sd->vae_decode_only) {
|
||||
LOG_ERROR("hires %s upscaler requires VAE encoder weights; create the context with vae_decode_only=false",
|
||||
sd_hires_upscaler_name(request.hires.upscaler));
|
||||
return {};
|
||||
}
|
||||
if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL && upscaler == nullptr) {
|
||||
LOG_ERROR("hires model upscaler context is null");
|
||||
return {};
|
||||
}
|
||||
|
||||
sd::Tensor<float> decoded = sd_ctx->sd->decode_first_stage(latent);
|
||||
if (decoded.empty()) {
|
||||
LOG_ERROR("decode_first_stage failed before hires %s upscale",
|
||||
sd_hires_upscaler_name(request.hires.upscaler));
|
||||
return {};
|
||||
}
|
||||
|
||||
sd::Tensor<float> upscaled_tensor;
|
||||
if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
|
||||
upscaled_tensor = upscaler->upscale_tensor(decoded);
|
||||
if (upscaled_tensor.empty()) {
|
||||
LOG_ERROR("hires model upscale failed");
|
||||
return {};
|
||||
}
|
||||
|
||||
if (upscaled_tensor.shape()[0] != request.hires.target_width ||
|
||||
upscaled_tensor.shape()[1] != request.hires.target_height) {
|
||||
upscaled_tensor = sd::ops::interpolate(upscaled_tensor,
|
||||
{request.hires.target_width,
|
||||
request.hires.target_height,
|
||||
upscaled_tensor.shape()[2],
|
||||
upscaled_tensor.shape()[3]});
|
||||
}
|
||||
} else {
|
||||
sd::ops::InterpolateMode mode = request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS
|
||||
? sd::ops::InterpolateMode::Lanczos
|
||||
: sd::ops::InterpolateMode::Nearest;
|
||||
LOG_INFO("hires %s image upscale %" PRId64 "x%" PRId64 " -> %dx%d",
|
||||
sd_hires_upscaler_name(request.hires.upscaler),
|
||||
decoded.shape()[0],
|
||||
decoded.shape()[1],
|
||||
request.hires.target_width,
|
||||
request.hires.target_height);
|
||||
upscaled_tensor = sd::ops::interpolate(decoded,
|
||||
{request.hires.target_width,
|
||||
request.hires.target_height,
|
||||
decoded.shape()[2],
|
||||
decoded.shape()[3]},
|
||||
mode);
|
||||
upscaled_tensor = sd::ops::clamp(upscaled_tensor, 0.0f, 1.0f);
|
||||
}
|
||||
|
||||
sd::Tensor<float> upscaled_latent = sd_ctx->sd->encode_first_stage(upscaled_tensor);
|
||||
if (upscaled_latent.empty()) {
|
||||
LOG_ERROR("encode_first_stage failed after hires %s upscale",
|
||||
sd_hires_upscaler_name(request.hires.upscaler));
|
||||
}
|
||||
return upscaled_latent;
|
||||
}
|
||||
|
||||
LOG_ERROR("unsupported hires upscaler '%s'", sd_hires_upscaler_name(request.hires.upscaler));
|
||||
return {};
|
||||
}
|
||||
|
||||
SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
|
||||
if (sd_ctx == nullptr || sd_img_gen_params == nullptr) {
|
||||
return nullptr;
|
||||
|
|
@ -3501,14 +3738,139 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
|
|||
}
|
||||
return nullptr;
|
||||
}
|
||||
if (sd_ctx->sd->free_params_immediately) {
|
||||
if (sd_ctx->sd->free_params_immediately && !request.hires.enabled) {
|
||||
sd_ctx->sd->diffusion_model->free_params_buffer();
|
||||
}
|
||||
int64_t denoise_end = ggml_time_ms();
|
||||
LOG_INFO("generating %" PRId64 " latent images completed, taking %.2fs",
|
||||
LOG_INFO("generating %zu latent images completed, taking %.2fs",
|
||||
final_latents.size(),
|
||||
(denoise_end - denoise_start) * 1.0f / 1000);
|
||||
|
||||
if (request.hires.enabled && request.hires.target_width > 0) {
|
||||
LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
|
||||
|
||||
std::unique_ptr<UpscalerGGML> hires_upscaler;
|
||||
if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
|
||||
LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
|
||||
hires_upscaler = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
|
||||
false,
|
||||
request.hires.upscale_tile_size);
|
||||
if (!hires_upscaler->load_from_file(request.hires.model_path,
|
||||
sd_ctx->sd->offload_params_to_cpu,
|
||||
sd_ctx->sd->n_threads)) {
|
||||
LOG_ERROR("load hires model upscaler failed");
|
||||
if (sd_ctx->sd->free_params_immediately) {
|
||||
sd_ctx->sd->diffusion_model->free_params_buffer();
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
int hires_steps = request.hires.steps > 0 ? request.hires.steps : plan.sample_steps;
|
||||
|
||||
// sd-webui behavior: scale up total steps so trimming by denoising_strength yields exactly hires_steps effective steps,
|
||||
// unlike img2img which trims from a fixed step count
|
||||
hires_steps = static_cast<int>(hires_steps / request.hires.denoising_strength);
|
||||
|
||||
std::vector<float> hires_sigmas = sd_ctx->sd->denoiser->get_sigmas(
|
||||
hires_steps,
|
||||
sd_ctx->sd->get_image_seq_len(request.hires.target_height, request.hires.target_width),
|
||||
sd_img_gen_params->sample_params.scheduler,
|
||||
sd_ctx->sd->version);
|
||||
|
||||
size_t t_enc = static_cast<size_t>(hires_steps * request.hires.denoising_strength);
|
||||
if (t_enc >= static_cast<size_t>(hires_steps)) {
|
||||
t_enc = static_cast<size_t>(hires_steps) - 1;
|
||||
}
|
||||
std::vector<float> hires_sigma_sched(hires_sigmas.begin() + hires_steps - static_cast<int>(t_enc) - 1,
|
||||
hires_sigmas.end());
|
||||
LOG_INFO("hires fix: %d steps, denoising_strength=%.2f, sigma_sched_size=%zu",
|
||||
hires_steps,
|
||||
request.hires.denoising_strength,
|
||||
hires_sigma_sched.size());
|
||||
|
||||
std::vector<sd::Tensor<float>> hires_final_latents;
|
||||
int64_t hires_denoise_start = ggml_time_ms();
|
||||
for (int b = 0; b < (int)final_latents.size(); b++) {
|
||||
int64_t cur_seed = request.seed + b;
|
||||
sd_ctx->sd->rng->manual_seed(cur_seed);
|
||||
sd_ctx->sd->sampler_rng->manual_seed(cur_seed);
|
||||
|
||||
sd::Tensor<float> upscaled = upscale_hires_latent(sd_ctx,
|
||||
final_latents[b],
|
||||
request,
|
||||
hires_upscaler.get());
|
||||
if (upscaled.empty()) {
|
||||
if (sd_ctx->sd->free_params_immediately) {
|
||||
sd_ctx->sd->diffusion_model->free_params_buffer();
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd_ctx->sd->rng);
|
||||
|
||||
sd::Tensor<float> hires_denoise_mask;
|
||||
if (!latents.denoise_mask.empty()) {
|
||||
std::vector<int64_t> mask_shape = latents.denoise_mask.shape();
|
||||
mask_shape[0] = upscaled.shape()[0];
|
||||
mask_shape[1] = upscaled.shape()[1];
|
||||
hires_denoise_mask = sd::ops::interpolate(latents.denoise_mask,
|
||||
mask_shape,
|
||||
sd::ops::InterpolateMode::NearestMax);
|
||||
}
|
||||
|
||||
int64_t hires_sample_start = ggml_time_ms();
|
||||
sd::Tensor<float> x_0 = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model,
|
||||
true,
|
||||
upscaled,
|
||||
std::move(noise),
|
||||
embeds.cond,
|
||||
embeds.uncond,
|
||||
embeds.img_cond,
|
||||
embeds.id_cond,
|
||||
latents.control_image,
|
||||
request.control_strength,
|
||||
request.guidance,
|
||||
plan.eta,
|
||||
request.shifted_timestep,
|
||||
plan.sample_method,
|
||||
sd_ctx->sd->is_flow_denoiser(),
|
||||
hires_sigma_sched,
|
||||
plan.start_merge_step,
|
||||
latents.ref_latents,
|
||||
request.increase_ref_index,
|
||||
hires_denoise_mask,
|
||||
sd::Tensor<float>(),
|
||||
1.f,
|
||||
request.cache_params);
|
||||
int64_t hires_sample_end = ggml_time_ms();
|
||||
if (!x_0.empty()) {
|
||||
LOG_INFO("hires sampling %d/%d completed, taking %.2fs",
|
||||
b + 1,
|
||||
(int)final_latents.size(),
|
||||
(hires_sample_end - hires_sample_start) * 1.0f / 1000);
|
||||
hires_final_latents.push_back(std::move(x_0));
|
||||
continue;
|
||||
}
|
||||
|
||||
LOG_ERROR("hires sampling for image %d/%d failed after %.2fs",
|
||||
b + 1,
|
||||
(int)final_latents.size(),
|
||||
(hires_sample_end - hires_sample_start) * 1.0f / 1000);
|
||||
if (sd_ctx->sd->free_params_immediately) {
|
||||
sd_ctx->sd->diffusion_model->free_params_buffer();
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
if (sd_ctx->sd->free_params_immediately) {
|
||||
sd_ctx->sd->diffusion_model->free_params_buffer();
|
||||
}
|
||||
int64_t hires_denoise_end = ggml_time_ms();
|
||||
LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000);
|
||||
|
||||
final_latents = std::move(hires_final_latents);
|
||||
}
|
||||
|
||||
auto result = decode_image_outputs(sd_ctx, request, final_latents);
|
||||
if (result == nullptr) {
|
||||
return nullptr;
|
||||
|
|
|
|||
|
|
@ -290,6 +290,32 @@ typedef struct {
|
|||
const char* path;
|
||||
} sd_lora_t;
|
||||
|
||||
enum sd_hires_upscaler_t {
|
||||
SD_HIRES_UPSCALER_NONE,
|
||||
SD_HIRES_UPSCALER_LATENT,
|
||||
SD_HIRES_UPSCALER_LATENT_NEAREST,
|
||||
SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT,
|
||||
SD_HIRES_UPSCALER_LATENT_ANTIALIASED,
|
||||
SD_HIRES_UPSCALER_LATENT_BICUBIC,
|
||||
SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED,
|
||||
SD_HIRES_UPSCALER_LANCZOS,
|
||||
SD_HIRES_UPSCALER_NEAREST,
|
||||
SD_HIRES_UPSCALER_MODEL,
|
||||
SD_HIRES_UPSCALER_COUNT,
|
||||
};
|
||||
|
||||
typedef struct {
|
||||
bool enabled;
|
||||
enum sd_hires_upscaler_t upscaler;
|
||||
const char* model_path;
|
||||
float scale;
|
||||
int target_width;
|
||||
int target_height;
|
||||
int steps;
|
||||
float denoising_strength;
|
||||
int upscale_tile_size;
|
||||
} sd_hires_params_t;
|
||||
|
||||
typedef struct {
|
||||
const sd_lora_t* loras;
|
||||
uint32_t lora_count;
|
||||
|
|
@ -313,6 +339,7 @@ typedef struct {
|
|||
sd_pm_params_t pm_params;
|
||||
sd_tiling_params_t vae_tiling_params;
|
||||
sd_cache_params_t cache;
|
||||
sd_hires_params_t hires;
|
||||
} sd_img_gen_params_t;
|
||||
|
||||
typedef struct {
|
||||
|
|
@ -366,8 +393,11 @@ SD_API const char* sd_preview_name(enum preview_t preview);
|
|||
SD_API enum preview_t str_to_preview(const char* str);
|
||||
SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode);
|
||||
SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str);
|
||||
SD_API const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler);
|
||||
SD_API enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str);
|
||||
|
||||
SD_API void sd_cache_params_init(sd_cache_params_t* cache_params);
|
||||
SD_API void sd_hires_params_init(sd_hires_params_t* hires_params);
|
||||
|
||||
SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params);
|
||||
SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);
|
||||
|
|
|
|||
|
|
@ -815,11 +815,202 @@ namespace sd {
|
|||
namespace ops {
|
||||
enum class InterpolateMode {
|
||||
Nearest,
|
||||
NearestExact,
|
||||
NearestMax,
|
||||
NearestMin,
|
||||
NearestAvg,
|
||||
Bilinear,
|
||||
Bicubic,
|
||||
Lanczos,
|
||||
};
|
||||
|
||||
inline bool is_nearest_like_interpolate_mode(InterpolateMode mode) {
|
||||
return mode == InterpolateMode::Nearest ||
|
||||
mode == InterpolateMode::NearestExact ||
|
||||
mode == InterpolateMode::NearestMax ||
|
||||
mode == InterpolateMode::NearestMin ||
|
||||
mode == InterpolateMode::NearestAvg;
|
||||
}
|
||||
|
||||
inline bool is_2d_filter_interpolate_mode(InterpolateMode mode) {
|
||||
return mode == InterpolateMode::Bilinear ||
|
||||
mode == InterpolateMode::Bicubic ||
|
||||
mode == InterpolateMode::Lanczos;
|
||||
}
|
||||
|
||||
inline int64_t nearest_exact_interpolate_index(int64_t output_index,
|
||||
int64_t input_size,
|
||||
int64_t output_size) {
|
||||
const double scale = static_cast<double>(input_size) / static_cast<double>(output_size);
|
||||
const double center = (static_cast<double>(output_index) + 0.5) * scale - 0.5;
|
||||
return std::min(std::max<int64_t>(static_cast<int64_t>(std::floor(center + 0.5)), 0), input_size - 1);
|
||||
}
|
||||
|
||||
inline double linear_interpolate_weight(double x) {
|
||||
x = std::abs(x);
|
||||
return x < 1.0 ? 1.0 - x : 0.0;
|
||||
}
|
||||
|
||||
inline double cubic_interpolate_weight(double x) {
|
||||
constexpr double a = -0.75; // Match PyTorch bicubic interpolation.
|
||||
x = std::abs(x);
|
||||
if (x <= 1.0) {
|
||||
return ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0;
|
||||
}
|
||||
if (x < 2.0) {
|
||||
return ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a;
|
||||
}
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
inline double sinc(double x) {
|
||||
constexpr double pi = 3.14159265358979323846;
|
||||
if (std::abs(x) < 1e-12) {
|
||||
return 1.0;
|
||||
}
|
||||
const double pix = pi * x;
|
||||
return std::sin(pix) / pix;
|
||||
}
|
||||
|
||||
inline double lanczos_interpolate_weight(double x) {
|
||||
constexpr double radius = 3.0;
|
||||
x = std::abs(x);
|
||||
if (x >= radius) {
|
||||
return 0.0;
|
||||
}
|
||||
return sinc(x) * sinc(x / radius);
|
||||
}
|
||||
|
||||
struct InterpolateContributor {
|
||||
int64_t index;
|
||||
double weight;
|
||||
};
|
||||
|
||||
inline std::vector<std::vector<InterpolateContributor>> make_interpolate_contributors(
|
||||
int64_t input_size,
|
||||
int64_t output_size,
|
||||
InterpolateMode mode,
|
||||
bool antialias) {
|
||||
std::vector<std::vector<InterpolateContributor>> contributors(static_cast<size_t>(output_size));
|
||||
const double scale = static_cast<double>(input_size) / static_cast<double>(output_size);
|
||||
const double filter_scale = antialias ? std::max(1.0, scale) : 1.0;
|
||||
|
||||
for (int64_t out = 0; out < output_size; ++out) {
|
||||
const double center = (static_cast<double>(out) + 0.5) * scale - 0.5;
|
||||
int64_t start = 0;
|
||||
int64_t end = 0;
|
||||
|
||||
if (mode == InterpolateMode::Bilinear) {
|
||||
const double support = filter_scale;
|
||||
start = static_cast<int64_t>(std::ceil(center - support));
|
||||
end = static_cast<int64_t>(std::floor(center + support));
|
||||
} else if (mode == InterpolateMode::Bicubic) {
|
||||
const double support = 2.0 * filter_scale;
|
||||
start = static_cast<int64_t>(std::ceil(center - support));
|
||||
end = static_cast<int64_t>(std::floor(center + support));
|
||||
} else if (mode == InterpolateMode::Lanczos) {
|
||||
const double support = 3.0 * filter_scale;
|
||||
start = static_cast<int64_t>(std::ceil(center - support));
|
||||
end = static_cast<int64_t>(std::floor(center + support));
|
||||
} else {
|
||||
tensor_throw_invalid_argument("Unsupported 2D filter interpolate mode: mode=" +
|
||||
std::to_string(static_cast<int>(mode)));
|
||||
}
|
||||
|
||||
double weight_sum = 0.0;
|
||||
std::vector<InterpolateContributor>& axis_contributors = contributors[static_cast<size_t>(out)];
|
||||
axis_contributors.reserve(static_cast<size_t>(end - start + 1));
|
||||
|
||||
for (int64_t in = start; in <= end; ++in) {
|
||||
double weight = 0.0;
|
||||
if (mode == InterpolateMode::Bilinear) {
|
||||
weight = linear_interpolate_weight((center - static_cast<double>(in)) / filter_scale);
|
||||
} else if (mode == InterpolateMode::Bicubic) {
|
||||
weight = cubic_interpolate_weight((center - static_cast<double>(in)) / filter_scale);
|
||||
} else {
|
||||
weight = lanczos_interpolate_weight((center - static_cast<double>(in)) / filter_scale);
|
||||
}
|
||||
|
||||
if (weight == 0.0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t clamped_index = std::min(std::max<int64_t>(in, 0), input_size - 1);
|
||||
axis_contributors.push_back({clamped_index, weight});
|
||||
weight_sum += weight;
|
||||
}
|
||||
|
||||
if ((antialias || mode == InterpolateMode::Lanczos) &&
|
||||
std::abs(weight_sum) > 1e-12) {
|
||||
for (auto& contributor : axis_contributors) {
|
||||
contributor.weight /= weight_sum;
|
||||
}
|
||||
}
|
||||
|
||||
if (axis_contributors.empty()) {
|
||||
const int64_t nearest = std::min(
|
||||
std::max<int64_t>(static_cast<int64_t>(std::floor(center + 0.5)), 0),
|
||||
input_size - 1);
|
||||
axis_contributors.push_back({nearest, 1.0});
|
||||
}
|
||||
}
|
||||
|
||||
return contributors;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline Tensor<T> interpolate_2d_filter(const Tensor<T>& input,
|
||||
const std::vector<int64_t>& output_shape,
|
||||
InterpolateMode mode,
|
||||
bool antialias) {
|
||||
if (input.dim() < 2) {
|
||||
tensor_throw_invalid_argument("2D filter interpolate requires rank >= 2: input_shape=" +
|
||||
tensor_shape_to_string(input.shape()) + ", output_shape=" +
|
||||
tensor_shape_to_string(output_shape));
|
||||
}
|
||||
for (size_t i = 2; i < output_shape.size(); ++i) {
|
||||
if (input.shape()[i] != output_shape[i]) {
|
||||
tensor_throw_invalid_argument("2D filter interpolate only supports resizing dimensions 0 and 1: input_shape=" +
|
||||
tensor_shape_to_string(input.shape()) + ", output_shape=" +
|
||||
tensor_shape_to_string(output_shape));
|
||||
}
|
||||
}
|
||||
|
||||
Tensor<T> output(output_shape);
|
||||
const int64_t input_width = input.shape()[0];
|
||||
const int64_t input_height = input.shape()[1];
|
||||
const int64_t output_width = output_shape[0];
|
||||
const int64_t output_height = output_shape[1];
|
||||
const int64_t input_plane = input_width * input_height;
|
||||
const int64_t output_plane = output_width * output_height;
|
||||
const int64_t plane_count = input.numel() / input_plane;
|
||||
|
||||
auto x_contributors = make_interpolate_contributors(input_width, output_width, mode, antialias);
|
||||
auto y_contributors = make_interpolate_contributors(input_height, output_height, mode, antialias);
|
||||
|
||||
for (int64_t plane = 0; plane < plane_count; ++plane) {
|
||||
const int64_t input_plane_offset = plane * input_plane;
|
||||
const int64_t output_plane_offset = plane * output_plane;
|
||||
for (int64_t y = 0; y < output_height; ++y) {
|
||||
const auto& y_axis = y_contributors[static_cast<size_t>(y)];
|
||||
for (int64_t x = 0; x < output_width; ++x) {
|
||||
const auto& x_axis = x_contributors[static_cast<size_t>(x)];
|
||||
double value = 0.0;
|
||||
for (const auto& yc : y_axis) {
|
||||
const int64_t input_row_offset = input_plane_offset + yc.index * input_width;
|
||||
for (const auto& xc : x_axis) {
|
||||
value += static_cast<double>(input.data()[input_row_offset + xc.index]) *
|
||||
xc.weight * yc.weight;
|
||||
}
|
||||
}
|
||||
output.data()[output_plane_offset + y * output_width + x] = static_cast<T>(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
inline int64_t normalize_slice_bound(int64_t index, int64_t dim_size) {
|
||||
if (index < 0) {
|
||||
index += dim_size;
|
||||
|
|
@ -1014,17 +1205,20 @@ namespace sd {
|
|||
inline Tensor<T> interpolate(const Tensor<T>& input,
|
||||
std::vector<int64_t> output_shape,
|
||||
InterpolateMode mode = InterpolateMode::Nearest,
|
||||
bool align_corners = false) {
|
||||
const bool is_nearest_like_mode = (mode == InterpolateMode::Nearest ||
|
||||
mode == InterpolateMode::NearestMax ||
|
||||
mode == InterpolateMode::NearestMin ||
|
||||
mode == InterpolateMode::NearestAvg);
|
||||
if (!is_nearest_like_mode) {
|
||||
tensor_throw_invalid_argument("Only nearest-like interpolate modes are implemented, got mode=" +
|
||||
bool align_corners = false,
|
||||
bool antialias = false) {
|
||||
const bool is_nearest_like_mode = is_nearest_like_interpolate_mode(mode);
|
||||
const bool is_2d_filter_mode = is_2d_filter_interpolate_mode(mode);
|
||||
if (!is_nearest_like_mode && !is_2d_filter_mode) {
|
||||
tensor_throw_invalid_argument("Unsupported interpolate mode: mode=" +
|
||||
std::to_string(static_cast<int>(mode)));
|
||||
}
|
||||
if (antialias && !is_2d_filter_mode) {
|
||||
tensor_throw_invalid_argument("Tensor interpolate antialias requires a 2D filter mode: mode=" +
|
||||
std::to_string(static_cast<int>(mode)));
|
||||
}
|
||||
if (align_corners) {
|
||||
tensor_throw_invalid_argument("align_corners is not supported for nearest-like interpolate: input_shape=" +
|
||||
tensor_throw_invalid_argument("align_corners is not supported for tensor interpolate: input_shape=" +
|
||||
tensor_shape_to_string(input.shape()) + ", output_shape=" +
|
||||
tensor_shape_to_string(output_shape));
|
||||
}
|
||||
|
|
@ -1051,6 +1245,10 @@ namespace sd {
|
|||
}
|
||||
}
|
||||
|
||||
if (is_2d_filter_mode) {
|
||||
return interpolate_2d_filter(input, output_shape, mode, antialias);
|
||||
}
|
||||
|
||||
bool has_downsampling = false;
|
||||
for (int64_t i = 0; i < input.dim(); ++i) {
|
||||
if (input.shape()[i] > output_shape[i]) {
|
||||
|
|
@ -1060,12 +1258,20 @@ namespace sd {
|
|||
}
|
||||
|
||||
Tensor<T> output(std::move(output_shape));
|
||||
if (mode == InterpolateMode::Nearest || !has_downsampling) {
|
||||
if (mode == InterpolateMode::Nearest ||
|
||||
mode == InterpolateMode::NearestExact ||
|
||||
!has_downsampling) {
|
||||
for (int64_t flat = 0; flat < output.numel(); ++flat) {
|
||||
std::vector<int64_t> output_coord = tensor_unravel_index(flat, output.shape());
|
||||
std::vector<int64_t> input_coord(static_cast<size_t>(input.dim()), 0);
|
||||
for (size_t i = 0; i < static_cast<size_t>(input.dim()); ++i) {
|
||||
input_coord[i] = output_coord[i] * input.shape()[i] / output.shape()[i];
|
||||
if (mode == InterpolateMode::NearestExact) {
|
||||
input_coord[i] = nearest_exact_interpolate_index(output_coord[i],
|
||||
input.shape()[i],
|
||||
output.shape()[i]);
|
||||
} else {
|
||||
input_coord[i] = output_coord[i] * input.shape()[i] / output.shape()[i];
|
||||
}
|
||||
}
|
||||
output[flat] = input.index(input_coord);
|
||||
}
|
||||
|
|
@ -1083,6 +1289,12 @@ namespace sd {
|
|||
return T(0);
|
||||
case InterpolateMode::Nearest:
|
||||
return T(0);
|
||||
case InterpolateMode::NearestExact:
|
||||
return T(0);
|
||||
case InterpolateMode::Bilinear:
|
||||
case InterpolateMode::Bicubic:
|
||||
case InterpolateMode::Lanczos:
|
||||
break;
|
||||
}
|
||||
|
||||
tensor_throw_invalid_argument("Unsupported interpolate mode: mode=" +
|
||||
|
|
@ -1102,6 +1314,12 @@ namespace sd {
|
|||
break;
|
||||
case InterpolateMode::Nearest:
|
||||
break;
|
||||
case InterpolateMode::NearestExact:
|
||||
break;
|
||||
case InterpolateMode::Bilinear:
|
||||
case InterpolateMode::Bicubic:
|
||||
case InterpolateMode::Lanczos:
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -1157,17 +1375,20 @@ namespace sd {
|
|||
const std::optional<std::vector<int64_t>>& size,
|
||||
const std::optional<std::vector<double>>& scale_factor,
|
||||
InterpolateMode mode = InterpolateMode::Nearest,
|
||||
bool align_corners = false) {
|
||||
const bool is_nearest_like_mode = (mode == InterpolateMode::Nearest ||
|
||||
mode == InterpolateMode::NearestMax ||
|
||||
mode == InterpolateMode::NearestMin ||
|
||||
mode == InterpolateMode::NearestAvg);
|
||||
if (!is_nearest_like_mode) {
|
||||
tensor_throw_invalid_argument("Only nearest-like interpolate modes are implemented, got mode=" +
|
||||
bool align_corners = false,
|
||||
bool antialias = false) {
|
||||
const bool is_nearest_like_mode = is_nearest_like_interpolate_mode(mode);
|
||||
const bool is_2d_filter_mode = is_2d_filter_interpolate_mode(mode);
|
||||
if (!is_nearest_like_mode && !is_2d_filter_mode) {
|
||||
tensor_throw_invalid_argument("Unsupported interpolate mode: mode=" +
|
||||
std::to_string(static_cast<int>(mode)));
|
||||
}
|
||||
if (antialias && !is_2d_filter_mode) {
|
||||
tensor_throw_invalid_argument("Tensor interpolate antialias requires a 2D filter mode: mode=" +
|
||||
std::to_string(static_cast<int>(mode)));
|
||||
}
|
||||
if (align_corners) {
|
||||
tensor_throw_invalid_argument("align_corners is not supported for nearest-like interpolate: input_shape=" +
|
||||
tensor_throw_invalid_argument("align_corners is not supported for tensor interpolate: input_shape=" +
|
||||
tensor_shape_to_string(input.shape()));
|
||||
}
|
||||
if (size.has_value() == scale_factor.has_value()) {
|
||||
|
|
@ -1211,7 +1432,7 @@ namespace sd {
|
|||
}
|
||||
}
|
||||
|
||||
return interpolate(input, std::move(output_shape), mode, align_corners);
|
||||
return interpolate(input, std::move(output_shape), mode, align_corners, antialias);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
|
|
@ -1219,12 +1440,14 @@ namespace sd {
|
|||
const std::optional<std::vector<int64_t>>& size,
|
||||
double scale_factor,
|
||||
InterpolateMode mode = InterpolateMode::Nearest,
|
||||
bool align_corners = false) {
|
||||
bool align_corners = false,
|
||||
bool antialias = false) {
|
||||
return interpolate(input,
|
||||
size,
|
||||
std::vector<double>(size.has_value() ? size->size() : input.dim(), scale_factor),
|
||||
mode,
|
||||
align_corners);
|
||||
align_corners,
|
||||
antialias);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
|
|
|
|||
|
|
@ -62,7 +62,7 @@ void CLIPTokenizer::load_from_merges(const std::string& merges_utf8_str) {
|
|||
}
|
||||
vocab.push_back(utf8_to_utf32("<|startoftext|>"));
|
||||
vocab.push_back(utf8_to_utf32("<|endoftext|>"));
|
||||
LOG_DEBUG("vocab size: %llu", vocab.size());
|
||||
LOG_DEBUG("vocab size: %zu", vocab.size());
|
||||
int i = 0;
|
||||
for (const auto& token : vocab) {
|
||||
encoder[token] = i;
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, cons
|
|||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
|
||||
LOG_DEBUG("merges size %llu", merges.size());
|
||||
LOG_DEBUG("merges size %zu", merges.size());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) {
|
|||
}
|
||||
|
||||
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
|
||||
LOG_DEBUG("merges size %llu", merges.size());
|
||||
LOG_DEBUG("merges size %zu", merges.size());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
|
|
|
|||
|
|
@ -1,125 +1,115 @@
|
|||
#include "esrgan.hpp"
|
||||
#include "upscaler.h"
|
||||
#include "ggml_extend.hpp"
|
||||
#include "model.h"
|
||||
#include "stable-diffusion.h"
|
||||
#include "util.h"
|
||||
|
||||
struct UpscalerGGML {
|
||||
ggml_backend_t backend = nullptr; // general backend
|
||||
ggml_type model_data_type = GGML_TYPE_F16;
|
||||
std::shared_ptr<ESRGAN> esrgan_upscaler;
|
||||
std::string esrgan_path;
|
||||
int n_threads;
|
||||
bool direct = false;
|
||||
int tile_size = 128;
|
||||
UpscalerGGML::UpscalerGGML(int n_threads,
|
||||
bool direct,
|
||||
int tile_size)
|
||||
: n_threads(n_threads),
|
||||
direct(direct),
|
||||
tile_size(tile_size) {
|
||||
}
|
||||
|
||||
UpscalerGGML(int n_threads,
|
||||
bool direct = false,
|
||||
int tile_size = 128)
|
||||
: n_threads(n_threads),
|
||||
direct(direct),
|
||||
tile_size(tile_size) {
|
||||
}
|
||||
|
||||
bool load_from_file(const std::string& esrgan_path,
|
||||
bool offload_params_to_cpu,
|
||||
int n_threads) {
|
||||
ggml_log_set(ggml_log_callback_default, nullptr);
|
||||
bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
|
||||
bool offload_params_to_cpu,
|
||||
int n_threads) {
|
||||
ggml_log_set(ggml_log_callback_default, nullptr);
|
||||
#ifdef SD_USE_CUDA
|
||||
LOG_DEBUG("Using CUDA backend");
|
||||
backend = ggml_backend_cuda_init(0);
|
||||
LOG_DEBUG("Using CUDA backend");
|
||||
backend = ggml_backend_cuda_init(0);
|
||||
#endif
|
||||
#ifdef SD_USE_METAL
|
||||
LOG_DEBUG("Using Metal backend");
|
||||
backend = ggml_backend_metal_init();
|
||||
LOG_DEBUG("Using Metal backend");
|
||||
backend = ggml_backend_metal_init();
|
||||
#endif
|
||||
#ifdef SD_USE_VULKAN
|
||||
LOG_DEBUG("Using Vulkan backend");
|
||||
backend = ggml_backend_vk_init(0);
|
||||
LOG_DEBUG("Using Vulkan backend");
|
||||
backend = ggml_backend_vk_init(0);
|
||||
#endif
|
||||
#ifdef SD_USE_OPENCL
|
||||
LOG_DEBUG("Using OpenCL backend");
|
||||
backend = ggml_backend_opencl_init();
|
||||
LOG_DEBUG("Using OpenCL backend");
|
||||
backend = ggml_backend_opencl_init();
|
||||
#endif
|
||||
#ifdef SD_USE_SYCL
|
||||
LOG_DEBUG("Using SYCL backend");
|
||||
backend = ggml_backend_sycl_init(0);
|
||||
LOG_DEBUG("Using SYCL backend");
|
||||
backend = ggml_backend_sycl_init(0);
|
||||
#endif
|
||||
ModelLoader model_loader;
|
||||
if (!model_loader.init_from_file_and_convert_name(esrgan_path)) {
|
||||
LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str());
|
||||
}
|
||||
model_loader.set_wtype_override(model_data_type);
|
||||
if (!backend) {
|
||||
LOG_DEBUG("Using CPU backend");
|
||||
backend = ggml_backend_cpu_init();
|
||||
}
|
||||
LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
|
||||
esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
|
||||
if (direct) {
|
||||
esrgan_upscaler->set_conv2d_direct_enabled(true);
|
||||
}
|
||||
if (!esrgan_upscaler->load_from_file(esrgan_path, n_threads)) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
ModelLoader model_loader;
|
||||
if (!model_loader.init_from_file_and_convert_name(esrgan_path)) {
|
||||
LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str());
|
||||
}
|
||||
|
||||
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor) {
|
||||
sd::Tensor<float> upscaled;
|
||||
if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
|
||||
upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
|
||||
} else {
|
||||
auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
|
||||
auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
|
||||
if (output_tile.empty()) {
|
||||
LOG_ERROR("esrgan compute failed while processing a tile");
|
||||
return {};
|
||||
}
|
||||
return output_tile;
|
||||
};
|
||||
|
||||
upscaled = process_tiles_2d(input_tensor,
|
||||
static_cast<int>(input_tensor.shape()[0] * esrgan_upscaler->scale),
|
||||
static_cast<int>(input_tensor.shape()[1] * esrgan_upscaler->scale),
|
||||
esrgan_upscaler->scale,
|
||||
tile_size,
|
||||
tile_size,
|
||||
0.25f,
|
||||
false,
|
||||
false,
|
||||
on_processing);
|
||||
}
|
||||
esrgan_upscaler->free_compute_buffer();
|
||||
if (upscaled.empty()) {
|
||||
LOG_ERROR("esrgan compute failed");
|
||||
return {};
|
||||
}
|
||||
return upscaled;
|
||||
model_loader.set_wtype_override(model_data_type);
|
||||
if (!backend) {
|
||||
LOG_DEBUG("Using CPU backend");
|
||||
backend = ggml_backend_cpu_init();
|
||||
}
|
||||
LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
|
||||
esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
|
||||
if (direct) {
|
||||
esrgan_upscaler->set_conv2d_direct_enabled(true);
|
||||
}
|
||||
if (!esrgan_upscaler->load_from_file(esrgan_path, n_threads)) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor) {
|
||||
// upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
|
||||
sd_image_t upscaled_image = {0, 0, 0, nullptr};
|
||||
int output_width = (int)input_image.width * esrgan_upscaler->scale;
|
||||
int output_height = (int)input_image.height * esrgan_upscaler->scale;
|
||||
LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
|
||||
input_image.width, input_image.height, output_width, output_height);
|
||||
sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
|
||||
sd::Tensor<float> upscaled;
|
||||
if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
|
||||
upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
|
||||
} else {
|
||||
auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
|
||||
auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
|
||||
if (output_tile.empty()) {
|
||||
LOG_ERROR("esrgan compute failed while processing a tile");
|
||||
return {};
|
||||
}
|
||||
return output_tile;
|
||||
};
|
||||
|
||||
sd::Tensor<float> input_tensor = sd_image_to_tensor(input_image);
|
||||
sd::Tensor<float> upscaled;
|
||||
int64_t t0 = ggml_time_ms();
|
||||
upscaled = upscale_tensor(input_tensor);
|
||||
if (upscaled.empty()) {
|
||||
return upscaled_image;
|
||||
}
|
||||
sd_image_t upscaled_data = tensor_to_sd_image(upscaled);
|
||||
int64_t t3 = ggml_time_ms();
|
||||
LOG_INFO("input_image_tensor upscaled, taking %.2fs", (t3 - t0) / 1000.0f);
|
||||
upscaled_image = upscaled_data;
|
||||
upscaled = process_tiles_2d(input_tensor,
|
||||
static_cast<int>(input_tensor.shape()[0] * esrgan_upscaler->scale),
|
||||
static_cast<int>(input_tensor.shape()[1] * esrgan_upscaler->scale),
|
||||
esrgan_upscaler->scale,
|
||||
tile_size,
|
||||
tile_size,
|
||||
0.25f,
|
||||
false,
|
||||
false,
|
||||
on_processing);
|
||||
}
|
||||
esrgan_upscaler->free_compute_buffer();
|
||||
if (upscaled.empty()) {
|
||||
LOG_ERROR("esrgan compute failed");
|
||||
return {};
|
||||
}
|
||||
return upscaled;
|
||||
}
|
||||
|
||||
sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) {
|
||||
// upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
|
||||
sd_image_t upscaled_image = {0, 0, 0, nullptr};
|
||||
int output_width = (int)input_image.width * esrgan_upscaler->scale;
|
||||
int output_height = (int)input_image.height * esrgan_upscaler->scale;
|
||||
LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
|
||||
input_image.width, input_image.height, output_width, output_height);
|
||||
|
||||
sd::Tensor<float> input_tensor = sd_image_to_tensor(input_image);
|
||||
sd::Tensor<float> upscaled;
|
||||
int64_t t0 = ggml_time_ms();
|
||||
upscaled = upscale_tensor(input_tensor);
|
||||
if (upscaled.empty()) {
|
||||
return upscaled_image;
|
||||
}
|
||||
};
|
||||
sd_image_t upscaled_data = tensor_to_sd_image(upscaled);
|
||||
int64_t t3 = ggml_time_ms();
|
||||
LOG_INFO("input_image_tensor upscaled, taking %.2fs", (t3 - t0) / 1000.0f);
|
||||
upscaled_image = upscaled_data;
|
||||
return upscaled_image;
|
||||
}
|
||||
|
||||
struct upscaler_ctx_t {
|
||||
UpscalerGGML* upscaler = nullptr;
|
||||
|
|
|
|||
31
otherarch/sdcpp/upscaler.h
Normal file
31
otherarch/sdcpp/upscaler.h
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
#ifndef __SD_UPSCALER_H__
|
||||
#define __SD_UPSCALER_H__
|
||||
|
||||
#include "esrgan.hpp"
|
||||
#include "stable-diffusion.h"
|
||||
#include "tensor.hpp"
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
struct UpscalerGGML {
|
||||
ggml_backend_t backend = nullptr; // general backend
|
||||
ggml_type model_data_type = GGML_TYPE_F16;
|
||||
std::shared_ptr<ESRGAN> esrgan_upscaler;
|
||||
std::string esrgan_path;
|
||||
int n_threads;
|
||||
bool direct = false;
|
||||
int tile_size = 128;
|
||||
|
||||
UpscalerGGML(int n_threads,
|
||||
bool direct = false,
|
||||
int tile_size = 128);
|
||||
|
||||
bool load_from_file(const std::string& esrgan_path,
|
||||
bool offload_params_to_cpu,
|
||||
int n_threads);
|
||||
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
|
||||
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
|
||||
};
|
||||
|
||||
#endif // __SD_UPSCALER_H__
|
||||
|
|
@ -128,10 +128,10 @@ std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
|
|||
filename.c_str(),
|
||||
GENERIC_READ,
|
||||
FILE_SHARE_READ,
|
||||
NULL,
|
||||
nullptr,
|
||||
OPEN_EXISTING,
|
||||
FILE_ATTRIBUTE_NORMAL,
|
||||
NULL);
|
||||
nullptr);
|
||||
|
||||
if (file_handle == INVALID_HANDLE_VALUE) {
|
||||
return nullptr;
|
||||
|
|
@ -145,16 +145,16 @@ std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
|
|||
|
||||
file_size = static_cast<size_t>(size.QuadPart);
|
||||
|
||||
HANDLE mapping_handle = CreateFileMapping(file_handle, NULL, PAGE_READONLY, 0, 0, NULL);
|
||||
HANDLE mapping_handle = CreateFileMapping(file_handle, nullptr, PAGE_READONLY, 0, 0, nullptr);
|
||||
|
||||
if (mapping_handle == NULL) {
|
||||
if (mapping_handle == nullptr) {
|
||||
CloseHandle(file_handle);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
mapped_data = MapViewOfFile(mapping_handle, FILE_MAP_READ, 0, 0, file_size);
|
||||
|
||||
if (mapped_data == NULL) {
|
||||
if (mapped_data == nullptr) {
|
||||
CloseHandle(mapping_handle);
|
||||
CloseHandle(file_handle);
|
||||
return nullptr;
|
||||
|
|
@ -217,7 +217,7 @@ std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
|
|||
|
||||
size_t file_size = sb.st_size;
|
||||
|
||||
void* mapped_data = mmap(NULL, file_size, PROT_READ, mmap_flags, file_descriptor, 0);
|
||||
void* mapped_data = mmap(nullptr, file_size, PROT_READ, mmap_flags, file_descriptor, 0);
|
||||
|
||||
close(file_descriptor);
|
||||
|
||||
|
|
|
|||
|
|
@ -142,9 +142,10 @@ public:
|
|||
"vae encode compute failed while processing a tile");
|
||||
} else {
|
||||
output = _compute(n_threads, input, false);
|
||||
free_compute_buffer();
|
||||
}
|
||||
|
||||
free_compute_buffer();
|
||||
|
||||
if (output.empty()) {
|
||||
LOG_ERROR("vae encode compute failed");
|
||||
return {};
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue