diff --git a/README.md b/README.md index bf722882c..a3359fcb1 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ API and command-line option may change frequently.*** - [LingBot-Video](./docs/lingbot_video.md) - [PhotoMaker](./docs/photo_maker.md) support. - Control Net support with SD 1.5 + - [ADetailer](./docs/adetailer.md) - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora) - Latent Consistency Models support (LCM/LCM-LoRA) - Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md) diff --git a/docs/adetailer.md b/docs/adetailer.md new file mode 100644 index 000000000..2bc64df2c --- /dev/null +++ b/docs/adetailer.md @@ -0,0 +1,110 @@ +# ADetailer + +`sd-cli` can run a YOLOv8 object detector on an existing or newly generated +image and perform a cropped inpaint pass for every detected object. The first +implementation supports YOLOv8 detection checkpoints. YOLOv8 segmentation and +MediaPipe models are not supported yet. + +## Convert a detector + +Ultralytics checkpoints must be converted before use. The converter fuses +BatchNorm into convolution layers and writes a safetensors file with the weight +names expected by the native GGML implementation. + +```bash +python scripts/convert_yolov8_to_safetensors.py face_yolov8n.pt face_yolov8n.safetensors +``` + +The converter requires Python packages `ultralytics`, `torch`, and +`safetensors`. +Only YOLOv8 detection checkpoints are accepted. +PyTorch checkpoints use pickle internally, so only convert `.pt` files from a +trusted source. + +## Repair an existing image + +Use the dedicated `adetailer` mode to detect and repair objects in an existing +image: + +```bash +./bin/sd-cli \ + -M adetailer \ + -m model.safetensors \ + -i input.png \ + -o repaired.png \ + -p "detailed portrait photo" \ + --negative-prompt "deformed face" \ + --steps 24 \ + --cfg-scale 6 \ + --strength 0.4 \ + --sampling-method dpm++2m \ + --scheduler karras \ + --ad-model face_yolov8n.safetensors \ + --extra-ad-args "confidence=0.3,inpaint_padding=32,mask_blur=4" +``` + +This mode reuses the normal image-generation options for the detail pass: + +- `--init-img`, `--output`, `--prompt`, and `--negative-prompt` +- `--steps`, `--cfg-scale`, `--sampling-method`, and `--scheduler` +- `--strength`, `--seed`, LoRA settings, VAE tiling, and backend assignments +- `--width` and `--height`, which also resize the input when specified + +`--ad-prompt` and `--ad-negative-prompt` optionally override the normal prompts. +Values provided in `--extra-ad-args`, such as `steps`, `cfg_scale`, +`denoising_strength`, or `inpaint_width`, take precedence over inherited values. + +## Repair generated images + +ADetailer can also run automatically after normal image generation: + +```bash +./bin/sd-cli \ + -m model.safetensors \ + -p "portrait photo" \ + --ad-model face_yolov8n.safetensors \ + --ad-prompt "[PROMPT], detailed face" \ + --ad-negative-prompt "" \ + --extra-ad-args "confidence=0.3,denoising_strength=0.4,inpaint_width=512,inpaint_height=512" +``` + +An empty ADetailer prompt inherits the main prompt. `[PROMPT]` inserts the main +prompt, `[SEP]` assigns different prompts to consecutive masks, and `[SKIP]` +skips the corresponding mask. + +All settings other than the detector path and prompts are passed through +`--extra-ad-args` as a comma-separated `key=value` list: + +| Key | Default | Description | +| --- | ---: | --- | +| `input_size` | `640` | Square YOLO input size; must be a multiple of 32 | +| `confidence` | `0.3` | Detection confidence threshold | +| `nms` | `0.45` | NMS IoU threshold | +| `max_detections` | `100` | Maximum detections retained after NMS | +| `mask_k_largest` | `0` | Keep only the largest K masks; zero keeps all | +| `mask_min_ratio` | `0` | Minimum bbox area relative to the image | +| `mask_max_ratio` | `1` | Maximum bbox area relative to the image | +| `dilate_erode` | `4` | Positive values dilate; negative values erode | +| `x_offset`, `y_offset` | `0` | Mask offset in pixels; positive Y moves upward | +| `mask_mode` | `none` | `none`, `merge`, or `merge_invert` | +| `merge_masks`, `invert_mask` | `false` | Boolean alternatives to `mask_mode` | +| `mask_blur` | `4` | Final composite feather radius | +| `inpaint_padding` | `32` | Padding around the detected region | +| `inpaint_width`, `inpaint_height` | mode-specific | `512x512` after generation; input/output size in `adetailer` mode | +| `denoising_strength` | mode-specific | `0.4` after generation; inherits `--strength` in `adetailer` mode | +| `steps` | `0` | Detail steps; zero inherits the main generation | +| `cfg_scale` | `-1` | Detail CFG; a negative value inherits the main generation | +| `sample_method` | inherited | Detail sampler name | +| `scheduler` | inherited | Detail scheduler name | +| `sort_by` | `none` | `none`, `left_to_right`, `center_to_edge`, or `area` | + +Multiple masks are processed serially. Each completed inpaint becomes the input +for the next mask, and the seed is incremented by the mask index. Use +`mask_mode=merge` to process all detections in one inpaint pass. + +The detector uses the `detector` backend module. For example, keep detection on +the CPU while diffusion runs on CUDA: + +```bash +--backend "diffusion=cuda0,detector=cpu" +``` diff --git a/docs/backend.md b/docs/backend.md index 353255371..c936b2232 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -153,6 +153,7 @@ still runs out of memory, tiling is enabled and the decode retried once. | `controlnet` | ControlNet | `controlnet`, `control` | | `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` | | `upscaler` | ESRGAN upscaler | `upscaler`, `esrgan`, `hires` | +| `detector` | ADetailer YOLOv8 detector | `detector`, `adetailer`, `yolo` | `te` is the preferred module name for text encoders. `clip` is kept as an accepted alias because many existing commands and model names use CLIP terminology. diff --git a/examples/cli/README.md b/examples/cli/README.md index e8a14098c..a7de00420 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -6,6 +6,9 @@ For detailed command-line arguments, run: ./bin/sd-cli -h ``` +For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see +[ADetailer](../../docs/adetailer.md). + Metadata mode inspects PNG/JPEG container metadata without loading any model: ```bash diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 29cb391b6..68a3148cc 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -199,7 +199,7 @@ struct SDCliParams { options.manual_options = { {"-M", "--mode", - "run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen", + "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen", on_mode_arg}, {"", "--preview", @@ -566,6 +566,65 @@ bool save_results(const SDCliParams& cli_params, return sucessful_reults != 0; } +static bool apply_adetailer(sd_ctx_t* sd_ctx, + const sd_ctx_params_t& sd_ctx_params, + const SDContextParams& ctx_params, + const SDGenerationParams& gen_params, + const sd_img_gen_params_t& img_gen_params, + SDMode mode, + SDImageVec& results, + int num_results) { + if (gen_params.ad_model_path.empty()) { + return true; + } + + sd_adetailer_params_t ad_params{}; + ad_params.prompt = gen_params.ad_prompt.empty() ? nullptr : gen_params.ad_prompt.c_str(); + ad_params.negative_prompt = gen_params.ad_negative_prompt.empty() ? nullptr : gen_params.ad_negative_prompt.c_str(); + ad_params.extra_ad_args = gen_params.extra_ad_args.c_str(); + + ADetailerCtxPtr ad_ctx(new_adetailer_ctx(gen_params.ad_model_path.c_str(), + ctx_params.n_threads, + sd_ctx_params.backend, + sd_ctx_params.params_backend)); + if (ad_ctx == nullptr) { + LOG_ERROR("new_adetailer_ctx failed"); + return false; + } + + for (int i = 0; i < num_results; ++i) { + if (results[i].data == nullptr) { + continue; + } + sd_img_gen_params_t ad_generation_params = img_gen_params; + ad_generation_params.seed = img_gen_params.seed + i; + if (mode == IMG_GEN) { + ad_generation_params.width = 512; + ad_generation_params.height = 512; + ad_generation_params.strength = 0.4f; + } + sd_image_t* detailed_images = nullptr; + int detailed_count = 0; + if (!adetail_image(ad_ctx.get(), + sd_ctx, + results[i], + &ad_params, + &ad_generation_params, + &detailed_images, + &detailed_count) || + detailed_count <= 0 || detailed_images == nullptr || detailed_images[0].data == nullptr) { + free_sd_images(detailed_images, detailed_count); + LOG_ERROR("ADetailer failed for image %d", i + 1); + return false; + } + free(results[i].data); + results[i] = detailed_images[0]; + detailed_images[0] = {0, 0, 0, nullptr}; + free_sd_images(detailed_images, detailed_count); + } + return true; +} + int main(int argc, const char* argv[]) { if (argc > 1 && std::string(argv[1]) == "--version") { std::cout << version_string() << "\n"; @@ -598,6 +657,11 @@ int main(int argc, const char* argv[]) { return 0; } + if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) { + LOG_ERROR("--ad-model is only supported in image generation and adetailer modes"); + return 1; + } + if (gen_params.video_frames > 4) { size_t last_dot_pos = cli_params.preview_path.find_last_of("."); std::string base_path = cli_params.preview_path; @@ -806,15 +870,22 @@ int main(int argc, const char* argv[]) { gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method); } - if (cli_params.mode == IMG_GEN) { - sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t(); + sd_img_gen_params_t img_gen_params{}; + const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER; + if (use_img_gen_params) { + img_gen_params = gen_params.to_sd_img_gen_params_t(); + } + if (cli_params.mode == IMG_GEN) { sd_image_t* generated_images = nullptr; if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) { generated_images = nullptr; num_results = 0; } results.adopt(generated_images, num_results); + } else if (cli_params.mode == ADETAILER) { + num_results = 1; + results.push_back(gen_params.init_image.release()); } else if (cli_params.mode == VID_GEN) { sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); sd_image_t* generated_video = nullptr; @@ -828,6 +899,18 @@ int main(int argc, const char* argv[]) { LOG_ERROR("generate failed"); return 1; } + + if (use_img_gen_params && + !apply_adetailer(sd_ctx.get(), + sd_ctx_params, + ctx_params, + gen_params, + img_gen_params, + cli_params.mode, + results, + num_results)) { + return 1; + } } int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 1dfb6aa70..4f9e9c4c3 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -30,6 +30,7 @@ namespace fs = std::filesystem; const char* const modes_str[] = { "img_gen", + "adetailer", "vid_gen", "convert", "upscale", @@ -916,6 +917,26 @@ ArgOptions SDGenerationParams::get_options() { "the negative prompt (default: \"\")", 0, &negative_prompt}, + {"", + "--ad-model", + "path to a converted YOLOv8 detection model for ADetailer", + 0, + &ad_model_path}, + {"", + "--ad-prompt", + "ADetailer prompt; empty inherits the main prompt, supports [PROMPT], [SEP], and [SKIP]", + 0, + &ad_prompt}, + {"", + "--ad-negative-prompt", + "ADetailer negative prompt; empty inherits the main negative prompt, supports [PROMPT] and [SEP]", + 0, + &ad_negative_prompt}, + {"", + "--extra-ad-args", + "extra ADetailer args, key=value list. Supports input_size, confidence, nms, max_detections, mask_k_largest, mask_min_ratio, mask_max_ratio, dilate_erode, x_offset, y_offset, mask_mode, merge_masks, invert_mask, mask_blur, inpaint_padding, inpaint_width, inpaint_height, denoising_strength, steps, cfg_scale, sample_method, scheduler, sort_by", + (int)',', + &extra_ad_args}, {"-i", "--init-img", "path to the init image", @@ -1831,6 +1852,10 @@ bool SDGenerationParams::from_json_str( load_if_exists("prompt", prompt); load_if_exists("negative_prompt", negative_prompt); + load_if_exists("ad_model", ad_model_path); + load_if_exists("ad_prompt", ad_prompt); + load_if_exists("ad_negative_prompt", ad_negative_prompt); + load_if_exists("extra_ad_args", extra_ad_args); load_if_exists("cache_mode", cache_mode); load_if_exists("cache_option", cache_option); load_if_exists("scm_mask", scm_mask); @@ -2347,13 +2372,19 @@ bool SDGenerationParams::validate(SDMode mode) { } } - if (mode == UPSCALE) { + if (mode == UPSCALE || mode == ADETAILER) { if (init_image_path.length() == 0) { - LOG_ERROR("error: upscale mode needs an init image (--init-img)\n"); + LOG_ERROR("error: %s mode needs an init image (--init-img)\n", + mode == UPSCALE ? "upscale" : "adetailer"); return false; } } + if (mode == ADETAILER && ad_model_path.empty()) { + LOG_ERROR("error: adetailer mode needs a detector model (--ad-model)\n"); + return false; + } + return true; } @@ -2561,6 +2592,10 @@ std::string SDGenerationParams::to_string() const { << " high_noise_loras: \"" << high_noise_loras_str << "\",\n" << " prompt: \"" << prompt << "\",\n" << " negative_prompt: \"" << negative_prompt << "\",\n" + << " ad_model_path: \"" << ad_model_path << "\",\n" + << " ad_prompt: \"" << ad_prompt << "\",\n" + << " ad_negative_prompt: \"" << ad_negative_prompt << "\",\n" + << " extra_ad_args: \"" << extra_ad_args << "\",\n" << " clip_skip: " << clip_skip << ",\n" << " width: " << width << ",\n" << " height: " << height << ",\n" @@ -2675,8 +2710,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, int64_t seed, SDMode mode) { json root; - root["schema"] = "sdcpp.image.params/v1"; - root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen"; + root["schema"] = "sdcpp.image.params/v1"; + root["mode"] = "img_gen"; + if (mode == VID_GEN) { + root["mode"] = "vid_gen"; + } else if (mode == ADETAILER) { + root["mode"] = "adetailer"; + } root["generator"] = { {"name", "stable-diffusion.cpp"}, {"version", safe_json_string(sd_version())}, @@ -2690,6 +2730,14 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, {"positive", gen_params.prompt}, {"negative", gen_params.negative_prompt}, }; + if (!gen_params.ad_model_path.empty()) { + root["adetailer"] = { + {"model", sd_basename(gen_params.ad_model_path)}, + {"prompt", gen_params.ad_prompt}, + {"negative_prompt", gen_params.ad_negative_prompt}, + {"extra_args", gen_params.extra_ad_args}, + }; + } root["sampling"] = build_sampling_metadata_json(gen_params.sample_params, gen_params.skip_layers, &gen_params.custom_sigmas); @@ -2844,6 +2892,18 @@ std::string get_image_params(const SDContextParams& ctx_params, if (!gen_params.extra_sample_args.empty()) { parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", "; } + if (!gen_params.ad_model_path.empty()) { + parameter_string += "ADetailer model: " + sd_basename(gen_params.ad_model_path) + ", "; + if (!gen_params.ad_prompt.empty()) { + parameter_string += "ADetailer prompt: " + gen_params.ad_prompt + ", "; + } + if (!gen_params.ad_negative_prompt.empty()) { + parameter_string += "ADetailer negative prompt: " + gen_params.ad_negative_prompt + ", "; + } + if (!gen_params.extra_ad_args.empty()) { + parameter_string += "ADetailer args: " + gen_params.extra_ad_args + ", "; + } + } parameter_string += "Seed: " + std::to_string(seed) + ", "; parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", "; parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", "; diff --git a/examples/common/common.h b/examples/common/common.h index 3a5b107bc..befab2591 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -16,10 +16,11 @@ #define BOOL_STR(b) ((b) ? "true" : "false") extern const char* const modes_str[]; -#define SD_ALL_MODES_STR "img_gen, vid_gen, convert, upscale, metadata" +#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata" enum SDMode { IMG_GEN, + ADETAILER, VID_GEN, CONVERT, UPSCALE, @@ -186,6 +187,10 @@ struct SDGenerationParams { // User-facing input fields. std::string prompt; std::string negative_prompt; + std::string ad_model_path; + std::string ad_prompt; + std::string ad_negative_prompt; + std::string extra_ad_args; int clip_skip = -1; // <= 0 represents unspecified int width = -1; int height = -1; diff --git a/examples/common/resource_owners.hpp b/examples/common/resource_owners.hpp index d47134abe..d7525a5fa 100644 --- a/examples/common/resource_owners.hpp +++ b/examples/common/resource_owners.hpp @@ -40,12 +40,21 @@ struct UpscalerCtxDeleter { } }; +struct ADetailerCtxDeleter { + void operator()(adetailer_ctx_t* ctx) const { + if (ctx != nullptr) { + free_adetailer_ctx(ctx); + } + } +}; + template using FreeUniquePtr = std::unique_ptr; -using FilePtr = std::unique_ptr; -using SDCtxPtr = std::unique_ptr; -using UpscalerCtxPtr = std::unique_ptr; +using FilePtr = std::unique_ptr; +using SDCtxPtr = std::unique_ptr; +using UpscalerCtxPtr = std::unique_ptr; +using ADetailerCtxPtr = std::unique_ptr; class SDImageOwner { private: diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index eeb59f873..8a20b1325 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -509,6 +509,27 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx, SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); +typedef struct adetailer_ctx_t adetailer_ctx_t; + +typedef struct { + const char* prompt; + const char* negative_prompt; + const char* extra_ad_args; +} sd_adetailer_params_t; + +SD_API adetailer_ctx_t* new_adetailer_ctx(const char* detector_path, + int n_threads, + const char* backend, + const char* params_backend); +SD_API void free_adetailer_ctx(adetailer_ctx_t* adetailer_ctx); +SD_API bool adetail_image(adetailer_ctx_t* adetailer_ctx, + sd_ctx_t* sd_ctx, + sd_image_t input_image, + const sd_adetailer_params_t* adetailer_params, + const sd_img_gen_params_t* inpaint_params, + sd_image_t** images_out, + int* num_images_out); + SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, diff --git a/scripts/convert_yolov8_to_safetensors.py b/scripts/convert_yolov8_to_safetensors.py new file mode 100644 index 000000000..c05f8f20c --- /dev/null +++ b/scripts/convert_yolov8_to_safetensors.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Convert an Ultralytics YOLOv8 detection checkpoint for sd.cpp ADetailer.""" + +import argparse +import json +from pathlib import Path + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Convert an Ultralytics YOLOv8 detection .pt checkpoint to safetensors." + ) + parser.add_argument("input", type=Path, help="input YOLOv8 detection checkpoint") + parser.add_argument("output", type=Path, help="output safetensors path") + parser.add_argument( + "--input-size", type=int, default=640, help="detector input size metadata (default: 640)" + ) + return parser.parse_args() + + +def main(): + args = parse_args() + if args.input_size < 32 or args.input_size % 32 != 0: + raise ValueError("--input-size must be a positive multiple of 32") + if args.output.suffix.lower() != ".safetensors": + raise ValueError("output path must use the .safetensors extension") + + try: + import torch + from safetensors.torch import save_file + from ultralytics import YOLO + from ultralytics.nn.modules.head import Detect + except ImportError as exc: + raise SystemExit("conversion requires ultralytics, torch, and safetensors") from exc + + torch_load = torch.load + + def load_trusted_checkpoint(*load_args, **load_kwargs): + load_kwargs.setdefault("weights_only", False) + return torch_load(*load_args, **load_kwargs) + + torch.load = load_trusted_checkpoint + try: + yolo = YOLO(str(args.input)) + finally: + torch.load = torch_load + network = yolo.model + if not isinstance(network.model[-1], Detect) or network.model[-1].__class__.__name__ != "Detect": + raise ValueError("only YOLOv8 detection checkpoints are supported; segmentation is not yet supported") + + network.eval() + network.fuse() + state_dict = network.state_dict() + required = { + "model.0.conv.weight", + "model.22.cv2.0.2.weight", + "model.22.cv3.0.2.weight", + } + missing = sorted(required.difference(state_dict)) + if missing: + raise ValueError(f"checkpoint does not match the supported YOLOv8 layout; missing {missing}") + + tensors = {} + for name, tensor in state_dict.items(): + if not name.startswith("model.") or ".bn." in name or name.endswith("dfl.conv.weight"): + continue + if not (name.endswith(".weight") or name.endswith(".bias")): + continue + dtype = torch.float16 if name.endswith(".weight") else torch.float32 + tensors[name] = tensor.detach().to(device="cpu", dtype=dtype).contiguous() + + metadata = { + "format": "pt", + "yolov8.variant": "detect", + "yolov8.input_size": str(args.input_size), + "yolov8.num_classes": str(int(network.model[-1].nc)), + "yolov8.reg_max": str(int(network.model[-1].reg_max)), + "yolov8.names": json.dumps(yolo.names, ensure_ascii=False), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(args.output), metadata=metadata) + print(f"wrote {args.output}: {len(tensors)} tensors") + + +if __name__ == "__main__": + main() diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index c66bb63f0..a83166438 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -83,6 +83,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m *module = SDBackendModule::UPSCALER; return true; } + if (name == "detector" || name == "adetailer" || name == "yolo") { + *module = SDBackendModule::DETECTOR; + return true; + } return false; } @@ -956,6 +960,8 @@ const char* sd_backend_module_name(SDBackendModule module) { return "photomaker"; case SDBackendModule::UPSCALER: return "upscaler"; + case SDBackendModule::DETECTOR: + return "detector"; } return "unknown"; } diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index 1f3bc8b3a..d5498e8f4 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -20,6 +20,7 @@ enum class SDBackendModule { CONTROL_NET, PHOTOMAKER, UPSCALER, + DETECTOR, }; struct SDBackendAssignment { diff --git a/src/detailer.cpp b/src/detailer.cpp new file mode 100644 index 000000000..fa9884be3 --- /dev/null +++ b/src/detailer.cpp @@ -0,0 +1,1020 @@ +#include "detailer.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/util.h" +#include "json.hpp" +#include "model_loader.h" + +struct adetailer_ctx_t { + ADetailerGGML* detailer = nullptr; +}; + +struct LetterboxInput { + sd::Tensor tensor; + float scale = 1.f; + int pad_x = 0; + int pad_y = 0; +}; + +struct Mask { + int width = 0; + int height = 0; + std::vector data; +}; + +struct CropRegion { + int x1 = 0; + int y1 = 0; + int x2 = 0; + int y2 = 0; +}; + +static std::string trim_copy(const std::string& value) { + size_t first = value.find_first_not_of(" \t\r\n"); + if (first == std::string::npos) { + return ""; + } + size_t last = value.find_last_not_of(" \t\r\n"); + return value.substr(first, last - first + 1); +} + +static std::string lower_copy(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return value; +} + +static bool parse_bool(const std::string& value, bool* result) { + const std::string lower = lower_copy(trim_copy(value)); + if (lower == "1" || lower == "true" || lower == "yes" || lower == "on") { + *result = true; + return true; + } + if (lower == "0" || lower == "false" || lower == "no" || lower == "off") { + *result = false; + return true; + } + return false; +} + +static bool parse_int(const std::string& value, int* result) { + try { + size_t used = 0; + int parsed = std::stoi(trim_copy(value), &used); + if (used != trim_copy(value).size()) { + return false; + } + *result = parsed; + return true; + } catch (...) { + return false; + } +} + +static bool parse_float(const std::string& value, float* result) { + try { + size_t used = 0; + float parsed = std::stof(trim_copy(value), &used); + if (used != trim_copy(value).size() || !std::isfinite(parsed)) { + return false; + } + *result = parsed; + return true; + } catch (...) { + return false; + } +} + +static float image_channel(const sd_image_t& image, int x, int y, int channel) { + if (image.data == nullptr || image.width == 0 || image.height == 0 || image.channel == 0) { + return 0.f; + } + x = std::clamp(x, 0, static_cast(image.width) - 1); + y = std::clamp(y, 0, static_cast(image.height) - 1); + int source_channel = image.channel == 1 ? 0 : std::min(channel, static_cast(image.channel) - 1); + size_t index = (static_cast(y) * image.width + x) * image.channel + source_channel; + return image.data[index] / 255.f; +} + +static float bilinear_channel(const sd_image_t& image, float x, float y, int channel) { + int x0 = static_cast(std::floor(x)); + int y0 = static_cast(std::floor(y)); + int x1 = x0 + 1; + int y1 = y0 + 1; + float wx = x - x0; + float wy = y - y0; + float a = image_channel(image, x0, y0, channel) * (1.f - wx) + image_channel(image, x1, y0, channel) * wx; + float b = image_channel(image, x0, y1, channel) * (1.f - wx) + image_channel(image, x1, y1, channel) * wx; + return a * (1.f - wy) + b * wy; +} + +static LetterboxInput make_letterbox_input(sd_image_t image, int input_size) { + LetterboxInput result; + result.tensor = sd::full({input_size, input_size, 3, 1}, 114.f / 255.f); + result.scale = std::min(static_cast(input_size) / image.width, + static_cast(input_size) / image.height); + int resized_width = std::max(1, static_cast(std::round(image.width * result.scale))); + int resized_height = std::max(1, static_cast(std::round(image.height * result.scale))); + result.pad_x = (input_size - resized_width) / 2; + result.pad_y = (input_size - resized_height) / 2; + + for (int y = 0; y < resized_height; ++y) { + float source_y = (y + 0.5f) / result.scale - 0.5f; + for (int x = 0; x < resized_width; ++x) { + float source_x = (x + 0.5f) / result.scale - 0.5f; + for (int c = 0; c < 3; ++c) { + result.tensor.index(x + result.pad_x, y + result.pad_y, c, 0) = bilinear_channel(image, source_x, source_y, c); + } + } + } + return result; +} + +static float logistic(float x) { + if (x >= 0.f) { + float z = std::exp(-x); + return 1.f / (1.f + z); + } + float z = std::exp(x); + return z / (1.f + z); +} + +static float dfl_expectation(const float* values, int64_t anchor_count, int channel_offset, int reg_max, int64_t anchor) { + float maximum = -std::numeric_limits::infinity(); + for (int i = 0; i < reg_max; ++i) { + maximum = std::max(maximum, values[anchor + anchor_count * (channel_offset + i)]); + } + float denominator = 0.f; + float numerator = 0.f; + for (int i = 0; i < reg_max; ++i) { + float probability = std::exp(values[anchor + anchor_count * (channel_offset + i)] - maximum); + denominator += probability; + numerator += probability * i; + } + return denominator > 0.f ? numerator / denominator : 0.f; +} + +static float box_iou(const ADetailerDetection& a, const ADetailerDetection& b) { + float x1 = std::max(a.x1, b.x1); + float y1 = std::max(a.y1, b.y1); + float x2 = std::min(a.x2, b.x2); + float y2 = std::min(a.y2, b.y2); + float intersection = std::max(0.f, x2 - x1) * std::max(0.f, y2 - y1); + float area_a = std::max(0.f, a.x2 - a.x1) * std::max(0.f, a.y2 - a.y1); + float area_b = std::max(0.f, b.x2 - b.x1) * std::max(0.f, b.y2 - b.y1); + float union_area = area_a + area_b - intersection; + return union_area > 0.f ? intersection / union_area : 0.f; +} + +static std::vector decode_detections(const sd::Tensor& raw, + const YOLOv8Config& config, + const LetterboxInput& letterbox, + sd_image_t image, + const ADetailerParams& params) { + std::vector candidates; + if (raw.empty() || raw.dim() < 2) { + return candidates; + } + int64_t anchor_count = raw.shape()[0]; + int64_t channels = raw.shape()[1]; + if (channels != config.reg_max * 4 + config.num_classes) { + LOG_ERROR("unexpected YOLOv8 output channels: %lld", static_cast(channels)); + return candidates; + } + + const std::array strides = {8, 16, 32}; + std::array offsets = {0, 0, 0, 0}; + for (int i = 0; i < 3; ++i) { + int grid = params.input_size / strides[i]; + offsets[i + 1] = offsets[i] + static_cast(grid) * grid; + } + if (offsets[3] != anchor_count) { + LOG_ERROR("unexpected YOLOv8 anchor count: %lld (expected %lld)", + static_cast(anchor_count), + static_cast(offsets[3])); + return candidates; + } + + const float* values = raw.data(); + for (int scale_index = 0; scale_index < 3; ++scale_index) { + int stride = strides[scale_index]; + int grid = params.input_size / stride; + for (int64_t anchor = offsets[scale_index]; anchor < offsets[scale_index + 1]; ++anchor) { + int64_t local = anchor - offsets[scale_index]; + int grid_x = static_cast(local % grid); + int grid_y = static_cast(local / grid); + + float confidence = 0.f; + int class_id = 0; + for (int c = 0; c < config.num_classes; ++c) { + float score = logistic(values[anchor + anchor_count * (config.reg_max * 4 + c)]); + if (score > confidence) { + confidence = score; + class_id = c; + } + } + if (confidence < params.confidence) { + continue; + } + + float left = dfl_expectation(values, anchor_count, 0 * config.reg_max, config.reg_max, anchor); + float top = dfl_expectation(values, anchor_count, 1 * config.reg_max, config.reg_max, anchor); + float right = dfl_expectation(values, anchor_count, 2 * config.reg_max, config.reg_max, anchor); + float bottom = dfl_expectation(values, anchor_count, 3 * config.reg_max, config.reg_max, anchor); + float center_x = (grid_x + 0.5f) * stride; + float center_y = (grid_y + 0.5f) * stride; + + ADetailerDetection detection; + detection.x1 = (center_x - left * stride - letterbox.pad_x) / letterbox.scale; + detection.y1 = (center_y - top * stride - letterbox.pad_y) / letterbox.scale; + detection.x2 = (center_x + right * stride - letterbox.pad_x) / letterbox.scale; + detection.y2 = (center_y + bottom * stride - letterbox.pad_y) / letterbox.scale; + detection.x1 = std::clamp(detection.x1, 0.f, static_cast(image.width)); + detection.y1 = std::clamp(detection.y1, 0.f, static_cast(image.height)); + detection.x2 = std::clamp(detection.x2, 0.f, static_cast(image.width)); + detection.y2 = std::clamp(detection.y2, 0.f, static_cast(image.height)); + detection.confidence = confidence; + detection.class_id = class_id; + if (detection.x2 > detection.x1 && detection.y2 > detection.y1) { + candidates.push_back(detection); + } + } + } + + std::sort(candidates.begin(), candidates.end(), [](const auto& a, const auto& b) { + return a.confidence > b.confidence; + }); + std::vector selected; + for (const auto& candidate : candidates) { + bool suppressed = false; + for (const auto& kept : selected) { + if (candidate.class_id == kept.class_id && box_iou(candidate, kept) > params.nms_threshold) { + suppressed = true; + break; + } + } + if (!suppressed) { + selected.push_back(candidate); + if (static_cast(selected.size()) >= params.max_detections) { + break; + } + } + } + return selected; +} + +static float detection_area(const ADetailerDetection& detection) { + return std::max(0.f, detection.x2 - detection.x1) * std::max(0.f, detection.y2 - detection.y1); +} + +static void filter_and_sort_detections(std::vector* detections, + int width, + int height, + const ADetailerParams& params) { + const float image_area = static_cast(width) * height; + detections->erase(std::remove_if(detections->begin(), detections->end(), [&](const auto& detection) { + float ratio = detection_area(detection) / image_area; + return ratio < params.mask_min_ratio || ratio > params.mask_max_ratio; + }), + detections->end()); + + if (params.mask_k_largest > 0 && static_cast(detections->size()) > params.mask_k_largest) { + std::partial_sort(detections->begin(), + detections->begin() + params.mask_k_largest, + detections->end(), + [](const auto& a, const auto& b) { return detection_area(a) > detection_area(b); }); + detections->resize(params.mask_k_largest); + } + + if (params.sort_by == ADETAILER_SORT_LEFT_TO_RIGHT) { + std::sort(detections->begin(), detections->end(), [](const auto& a, const auto& b) { return a.x1 < b.x1; }); + } else if (params.sort_by == ADETAILER_SORT_CENTER_TO_EDGE) { + const float cx = width * 0.5f; + const float cy = height * 0.5f; + std::sort(detections->begin(), detections->end(), [&](const auto& a, const auto& b) { + float ax = (a.x1 + a.x2) * 0.5f - cx; + float ay = (a.y1 + a.y2) * 0.5f - cy; + float bx = (b.x1 + b.x2) * 0.5f - cx; + float by = (b.y1 + b.y2) * 0.5f - cy; + return ax * ax + ay * ay < bx * bx + by * by; + }); + } else if (params.sort_by == ADETAILER_SORT_AREA) { + std::sort(detections->begin(), detections->end(), [](const auto& a, const auto& b) { + return detection_area(a) > detection_area(b); + }); + } +} + +static Mask bbox_mask(const ADetailerDetection& detection, int width, int height) { + Mask mask; + mask.width = width; + mask.height = height; + mask.data.assign(static_cast(width) * height, 0); + int x1 = std::clamp(static_cast(std::floor(detection.x1)), 0, width); + int y1 = std::clamp(static_cast(std::floor(detection.y1)), 0, height); + int x2 = std::clamp(static_cast(std::ceil(detection.x2)), 0, width); + int y2 = std::clamp(static_cast(std::ceil(detection.y2)), 0, height); + for (int y = y1; y < y2; ++y) { + std::fill(mask.data.begin() + static_cast(y) * width + x1, + mask.data.begin() + static_cast(y) * width + x2, + 255); + } + return mask; +} + +static Mask offset_mask(const Mask& source, int x_offset, int y_offset) { + Mask output{source.width, source.height, std::vector(source.data.size(), 0)}; + for (int y = 0; y < source.height; ++y) { + int target_y = y - y_offset; + if (target_y < 0 || target_y >= source.height) { + continue; + } + for (int x = 0; x < source.width; ++x) { + int target_x = x + x_offset; + if (target_x >= 0 && target_x < source.width) { + output.data[static_cast(target_y) * source.width + target_x] = source.data[static_cast(y) * source.width + x]; + } + } + } + return output; +} + +static Mask morphology_mask(const Mask& source, int amount) { + if (amount == 0) { + return source; + } + int kernel = std::abs(amount); + int before = kernel / 2; + int after = kernel - before - 1; + Mask output{source.width, source.height, std::vector(source.data.size(), amount > 0 ? 0 : 255)}; + for (int y = 0; y < source.height; ++y) { + for (int x = 0; x < source.width; ++x) { + uint8_t value = amount > 0 ? 0 : 255; + for (int ky = -before; ky <= after; ++ky) { + for (int kx = -before; kx <= after; ++kx) { + int sx = x + kx; + int sy = y + ky; + uint8_t sample = (sx >= 0 && sx < source.width && sy >= 0 && sy < source.height) + ? source.data[static_cast(sy) * source.width + sx] + : 0; + value = amount > 0 ? std::max(value, sample) : std::min(value, sample); + } + } + output.data[static_cast(y) * source.width + x] = value; + } + } + return output; +} + +static Mask gaussian_blur_mask(const Mask& source, int radius) { + if (radius <= 0) { + return source; + } + float sigma = static_cast(radius); + int kernel_radius = std::max(1, static_cast(std::ceil(2.5f * sigma))); + std::vector kernel(kernel_radius * 2 + 1); + float sum = 0.f; + for (int i = -kernel_radius; i <= kernel_radius; ++i) { + float value = std::exp(-(i * i) / (2.f * sigma * sigma)); + kernel[i + kernel_radius] = value; + sum += value; + } + for (float& value : kernel) { + value /= sum; + } + + std::vector horizontal(source.data.size()); + for (int y = 0; y < source.height; ++y) { + for (int x = 0; x < source.width; ++x) { + float value = 0.f; + for (int k = -kernel_radius; k <= kernel_radius; ++k) { + int sx = std::clamp(x + k, 0, source.width - 1); + value += source.data[static_cast(y) * source.width + sx] * kernel[k + kernel_radius]; + } + horizontal[static_cast(y) * source.width + x] = value; + } + } + + Mask output{source.width, source.height, std::vector(source.data.size())}; + for (int y = 0; y < source.height; ++y) { + for (int x = 0; x < source.width; ++x) { + float value = 0.f; + for (int k = -kernel_radius; k <= kernel_radius; ++k) { + int sy = std::clamp(y + k, 0, source.height - 1); + value += horizontal[static_cast(sy) * source.width + x] * kernel[k + kernel_radius]; + } + output.data[static_cast(y) * source.width + x] = static_cast(std::clamp(value, 0.f, 255.f) + 0.5f); + } + } + return output; +} + +static std::vector make_masks(const std::vector& detections, + int width, + int height, + const ADetailerParams& params) { + std::vector masks; + for (const auto& detection : detections) { + Mask mask = bbox_mask(detection, width, height); + if (params.x_offset != 0 || params.y_offset != 0) { + mask = offset_mask(mask, params.x_offset, params.y_offset); + } + mask = morphology_mask(mask, params.dilate_erode); + if (std::any_of(mask.data.begin(), mask.data.end(), [](uint8_t value) { return value != 0; })) { + masks.push_back(std::move(mask)); + } + } + if (params.merge_masks && !masks.empty()) { + Mask merged{width, height, std::vector(static_cast(width) * height, 0)}; + for (const Mask& mask : masks) { + for (size_t i = 0; i < merged.data.size(); ++i) { + merged.data[i] = std::max(merged.data[i], mask.data[i]); + } + } + masks = {std::move(merged)}; + } + if (params.invert_mask) { + for (Mask& mask : masks) { + for (uint8_t& value : mask.data) { + value = 255 - value; + } + } + } + return masks; +} + +static bool mask_bbox(const Mask& mask, CropRegion* region) { + int x1 = mask.width; + int y1 = mask.height; + int x2 = 0; + int y2 = 0; + for (int y = 0; y < mask.height; ++y) { + for (int x = 0; x < mask.width; ++x) { + if (mask.data[static_cast(y) * mask.width + x] != 0) { + x1 = std::min(x1, x); + y1 = std::min(y1, y); + x2 = std::max(x2, x + 1); + y2 = std::max(y2, y + 1); + } + } + } + if (x2 <= x1 || y2 <= y1) { + return false; + } + *region = {x1, y1, x2, y2}; + return true; +} + +static CropRegion expand_crop(CropRegion crop, + int image_width, + int image_height, + int padding, + int target_width, + int target_height) { + crop.x1 = std::max(0, crop.x1 - padding); + crop.y1 = std::max(0, crop.y1 - padding); + crop.x2 = std::min(image_width, crop.x2 + padding); + crop.y2 = std::min(image_height, crop.y2 + padding); + + float target_aspect = static_cast(target_width) / target_height; + int width = crop.x2 - crop.x1; + int height = crop.y2 - crop.y1; + int desired_width = width; + int desired_height = height; + if (static_cast(width) / height < target_aspect) { + desired_width = static_cast(std::ceil(height * target_aspect)); + } else { + desired_height = static_cast(std::ceil(width / target_aspect)); + } + desired_width = std::min(desired_width, image_width); + desired_height = std::min(desired_height, image_height); + int center_x = (crop.x1 + crop.x2) / 2; + int center_y = (crop.y1 + crop.y2) / 2; + crop.x1 = std::clamp(center_x - desired_width / 2, 0, image_width - desired_width); + crop.y1 = std::clamp(center_y - desired_height / 2, 0, image_height - desired_height); + crop.x2 = crop.x1 + desired_width; + crop.y2 = crop.y1 + desired_height; + return crop; +} + +static sd_image_t resize_crop_image(const sd_image_t& source, const CropRegion& crop, int width, int height) { + sd_image_t output = {static_cast(width), static_cast(height), 3, nullptr}; + output.data = static_cast(malloc(static_cast(width) * height * 3)); + if (output.data == nullptr) { + return {0, 0, 0, nullptr}; + } + float scale_x = static_cast(crop.x2 - crop.x1) / width; + float scale_y = static_cast(crop.y2 - crop.y1) / height; + for (int y = 0; y < height; ++y) { + float source_y = crop.y1 + (y + 0.5f) * scale_y - 0.5f; + for (int x = 0; x < width; ++x) { + float source_x = crop.x1 + (x + 0.5f) * scale_x - 0.5f; + for (int c = 0; c < 3; ++c) { + float value = bilinear_channel(source, source_x, source_y, c); + output.data[(static_cast(y) * width + x) * 3 + c] = static_cast(std::clamp(value * 255.f, 0.f, 255.f) + 0.5f); + } + } + } + return output; +} + +static sd_image_t resize_crop_mask(const Mask& source, const CropRegion& crop, int width, int height) { + sd_image_t output = {static_cast(width), static_cast(height), 1, nullptr}; + output.data = static_cast(malloc(static_cast(width) * height)); + if (output.data == nullptr) { + return {0, 0, 0, nullptr}; + } + float scale_x = static_cast(crop.x2 - crop.x1) / width; + float scale_y = static_cast(crop.y2 - crop.y1) / height; + for (int y = 0; y < height; ++y) { + int source_y = std::clamp(crop.y1 + static_cast((y + 0.5f) * scale_y), 0, source.height - 1); + for (int x = 0; x < width; ++x) { + int source_x = std::clamp(crop.x1 + static_cast((x + 0.5f) * scale_x), 0, source.width - 1); + output.data[static_cast(y) * width + x] = source.data[static_cast(source_y) * source.width + source_x]; + } + } + return output; +} + +static sd_image_t copy_as_rgb(const sd_image_t& source) { + CropRegion full{0, 0, static_cast(source.width), static_cast(source.height)}; + return resize_crop_image(source, full, source.width, source.height); +} + +static void composite_crop(sd_image_t* destination, + const sd_image_t& generated, + const Mask& feather_mask, + const CropRegion& crop) { + int crop_width = crop.x2 - crop.x1; + int crop_height = crop.y2 - crop.y1; + float scale_x = static_cast(generated.width) / crop_width; + float scale_y = static_cast(generated.height) / crop_height; + for (int y = crop.y1; y < crop.y2; ++y) { + float source_y = (y - crop.y1 + 0.5f) * scale_y - 0.5f; + for (int x = crop.x1; x < crop.x2; ++x) { + float alpha = feather_mask.data[static_cast(y) * feather_mask.width + x] / 255.f; + if (alpha <= 0.f) { + continue; + } + float source_x = (x - crop.x1 + 0.5f) * scale_x - 0.5f; + size_t destination_index = (static_cast(y) * destination->width + x) * 3; + for (int c = 0; c < 3; ++c) { + float generated_value = bilinear_channel(generated, source_x, source_y, c) * 255.f; + float original_value = destination->data[destination_index + c]; + float blended = original_value * (1.f - alpha) + generated_value * alpha; + destination->data[destination_index + c] = static_cast(std::clamp(blended, 0.f, 255.f) + 0.5f); + } + } + } +} + +static std::vector split_prompts(const std::string& prompt) { + std::vector result; + size_t start = 0; + while (true) { + size_t separator = prompt.find("[SEP]", start); + result.push_back(trim_copy(prompt.substr(start, separator == std::string::npos ? std::string::npos : separator - start))); + if (separator == std::string::npos) { + break; + } + start = separator + 5; + } + return result; +} + +static std::string resolve_prompt(const char* prompt_template, const char* base_prompt, int index) { + std::string base = base_prompt == nullptr ? "" : base_prompt; + std::string templ = prompt_template == nullptr ? "" : prompt_template; + auto prompts = split_prompts(templ); + std::string selected = prompts[std::min(index, static_cast(prompts.size()) - 1)]; + if (selected.empty()) { + return base; + } + size_t position = 0; + while ((position = selected.find("[PROMPT]", position)) != std::string::npos) { + selected.replace(position, 8, base); + position += base.size(); + } + return selected; +} + +static bool prompt_is_skip(const std::string& prompt) { + return trim_copy(prompt) == "[SKIP]"; +} + +static bool validate_params(const ADetailerParams& params) { + return params.input_size >= 32 && params.input_size % 32 == 0 && + params.confidence >= 0.f && params.confidence <= 1.f && + params.nms_threshold >= 0.f && params.nms_threshold <= 1.f && + params.max_detections > 0 && params.mask_k_largest >= 0 && + params.mask_min_ratio >= 0.f && params.mask_max_ratio <= 1.f && + params.mask_min_ratio <= params.mask_max_ratio && params.mask_blur >= 0 && + params.inpaint_padding >= 0 && params.inpaint_width > 0 && params.inpaint_height > 0 && + params.denoising_strength >= 0.f && params.denoising_strength <= 1.f && params.steps >= 0 && + (params.cfg_scale < 0.f || std::isfinite(params.cfg_scale)); +} + +static std::vector parse_class_names(const std::string& value, int class_count) { + std::vector names(static_cast(class_count)); + try { + nlohmann::json parsed = nlohmann::json::parse(value); + if (parsed.is_array()) { + for (size_t i = 0; i < names.size() && i < parsed.size(); ++i) { + if (parsed[i].is_string()) { + names[i] = parsed[i].get(); + } + } + } else if (parsed.is_object()) { + for (const auto& item : parsed.items()) { + int class_id = 0; + if (parse_int(item.key(), &class_id) && class_id >= 0 && class_id < class_count && + item.value().is_string()) { + names[static_cast(class_id)] = item.value().get(); + } + } + } + } catch (const std::exception&) { + } + return names; +} + +ADetailerGGML::ADetailerGGML(int n_threads, + std::string backend_spec, + std::string params_backend_spec) + : n_threads(n_threads > 0 ? n_threads : 1), + backend_spec(std::move(backend_spec)), + params_backend_spec(std::move(params_backend_spec)) { +} + +ADetailerGGML::~ADetailerGGML() { + model_manager.reset(); + detector.reset(); +} + +bool ADetailerGGML::load_from_file(const std::string& detector_path) { + std::string error; + if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), nullptr, &error)) { + LOG_ERROR("ADetailer backend config failed: %s", error.c_str()); + return false; + } + ggml_backend_t backend = backend_manager.runtime_backend(SDBackendModule::DETECTOR); + ggml_backend_t params_backend = backend_manager.params_backend(SDBackendModule::DETECTOR); + if (backend == nullptr || params_backend == nullptr) { + LOG_ERROR("failed to initialize detector backend"); + return false; + } + + model_manager = std::make_shared(); + model_manager->set_n_threads(n_threads); + model_manager->set_enable_mmap(false); + ModelLoader& loader = model_manager->loader(); + if (!loader.init_from_file(detector_path)) { + LOG_ERROR("failed to load ADetailer detector: '%s'", detector_path.c_str()); + return false; + } + + detector = std::make_shared(backend, loader.get_tensor_storage_map(), model_manager); + if (!detector || !detector->model || !detector->config.valid) { + LOG_ERROR("unsupported YOLOv8 detection weights: '%s'", detector_path.c_str()); + return false; + } + + class_names.assign(static_cast(detector->config.num_classes), ""); + auto names_item = loader.get_metadata().find("yolov8.names"); + if (names_item != loader.get_metadata().end()) { + class_names = parse_class_names(names_item->second, detector->config.num_classes); + } + + std::map tensors; + detector->get_param_tensors(tensors); + if (!model_manager->register_param_tensors("YOLOv8", + std::move(tensors), + backend_manager.params_backend_is_disk(SDBackendModule::DETECTOR) + ? ModelManager::ResidencyMode::Disk + : ModelManager::ResidencyMode::ParamBackend, + backend, + params_backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("failed to register YOLOv8 detector tensors"); + return false; + } + return true; +} + +std::vector ADetailerGGML::predict(sd_image_t image, + const ADetailerParams& params) { + LetterboxInput input = make_letterbox_input(image, params.input_size); + int64_t start = ggml_time_ms(); + sd::Tensor raw = detector->compute(n_threads, input.tensor); + detector->free_compute_buffer(); + if (raw.empty()) { + LOG_ERROR("YOLOv8 detector inference failed"); + return {}; + } + auto detections = decode_detections(raw, detector->config, input, image, params); + LOG_INFO("ADetailer detected %zu object(s), taking %.2fs", + detections.size(), + (ggml_time_ms() - start) / 1000.f); + for (size_t i = 0; i < detections.size(); ++i) { + const auto& detection = detections[i]; + std::string object = "class_" + std::to_string(detection.class_id); + if (detection.class_id >= 0 && static_cast(detection.class_id) < class_names.size() && + !class_names[static_cast(detection.class_id)].empty()) { + object = class_names[static_cast(detection.class_id)]; + } + LOG_INFO("ADetailer detection %zu: object=%s, class_id=%d, confidence=%.3f, bbox=[x1=%.1f, y1=%.1f, x2=%.1f, y2=%.1f]", + i + 1, + object.c_str(), + detection.class_id, + detection.confidence, + detection.x1, + detection.y1, + detection.x2, + detection.y2); + } + return detections; +} + +static bool parse_adetailer_params(const sd_adetailer_params_t& public_params, + const sd_img_gen_params_t& inpaint_params, + ADetailerParams* params) { + if (params == nullptr) { + return false; + } + *params = {}; + params->prompt = public_params.prompt; + params->negative_prompt = public_params.negative_prompt; + params->inpaint_width = inpaint_params.width; + params->inpaint_height = inpaint_params.height; + params->denoising_strength = inpaint_params.strength; + const char* extra_ad_args = public_params.extra_ad_args; + if (extra_ad_args == nullptr || extra_ad_args[0] == '\0') { + return validate_params(*params); + } + + std::stringstream stream(extra_ad_args); + std::string item; + while (std::getline(stream, item, ',')) { + item = trim_copy(item); + if (item.empty()) { + continue; + } + size_t equals = item.find('='); + if (equals == std::string::npos) { + LOG_ERROR("invalid --extra-ad-args item '%s'; expected key=value", item.c_str()); + return false; + } + std::string key = lower_copy(trim_copy(item.substr(0, equals))); + std::string value = trim_copy(item.substr(equals + 1)); + bool ok = true; + if (key == "input_size") { + ok = parse_int(value, ¶ms->input_size); + } else if (key == "confidence") { + ok = parse_float(value, ¶ms->confidence); + } else if (key == "nms" || key == "nms_threshold") { + ok = parse_float(value, ¶ms->nms_threshold); + } else if (key == "max_detections") { + ok = parse_int(value, ¶ms->max_detections); + } else if (key == "mask_k_largest") { + ok = parse_int(value, ¶ms->mask_k_largest); + } else if (key == "mask_min_ratio") { + ok = parse_float(value, ¶ms->mask_min_ratio); + } else if (key == "mask_max_ratio") { + ok = parse_float(value, ¶ms->mask_max_ratio); + } else if (key == "dilate_erode") { + ok = parse_int(value, ¶ms->dilate_erode); + } else if (key == "x_offset") { + ok = parse_int(value, ¶ms->x_offset); + } else if (key == "y_offset") { + ok = parse_int(value, ¶ms->y_offset); + } else if (key == "merge_masks") { + ok = parse_bool(value, ¶ms->merge_masks); + } else if (key == "invert_mask") { + ok = parse_bool(value, ¶ms->invert_mask); + } else if (key == "mask_blur") { + ok = parse_int(value, ¶ms->mask_blur); + } else if (key == "inpaint_padding") { + ok = parse_int(value, ¶ms->inpaint_padding); + } else if (key == "inpaint_width") { + ok = parse_int(value, ¶ms->inpaint_width); + } else if (key == "inpaint_height") { + ok = parse_int(value, ¶ms->inpaint_height); + } else if (key == "denoising_strength") { + ok = parse_float(value, ¶ms->denoising_strength); + } else if (key == "steps") { + ok = parse_int(value, ¶ms->steps); + } else if (key == "cfg_scale") { + ok = parse_float(value, ¶ms->cfg_scale); + } else if (key == "sample_method") { + params->sample_method = str_to_sample_method(value.c_str()); + ok = params->sample_method != SAMPLE_METHOD_COUNT; + } else if (key == "scheduler") { + params->scheduler = str_to_scheduler(value.c_str()); + ok = params->scheduler != SCHEDULER_COUNT; + } else if (key == "sort_by") { + std::string sort = lower_copy(value); + if (sort == "none") { + params->sort_by = ADETAILER_SORT_NONE; + } else if (sort == "left_to_right" || sort == "left-to-right") { + params->sort_by = ADETAILER_SORT_LEFT_TO_RIGHT; + } else if (sort == "center_to_edge" || sort == "center-to-edge") { + params->sort_by = ADETAILER_SORT_CENTER_TO_EDGE; + } else if (sort == "area") { + params->sort_by = ADETAILER_SORT_AREA; + } else { + ok = false; + } + } else if (key == "mask_mode") { + std::string mode = lower_copy(value); + if (mode == "none") { + params->merge_masks = false; + params->invert_mask = false; + } else if (mode == "merge") { + params->merge_masks = true; + params->invert_mask = false; + } else if (mode == "merge_invert" || mode == "merge-invert") { + params->merge_masks = true; + params->invert_mask = true; + } else { + ok = false; + } + } else { + LOG_ERROR("unknown --extra-ad-args key '%s'", key.c_str()); + return false; + } + if (!ok) { + LOG_ERROR("invalid --extra-ad-args value '%s=%s'", key.c_str(), value.c_str()); + return false; + } + } + if (!validate_params(*params)) { + LOG_ERROR("invalid --extra-ad-args parameter range"); + return false; + } + return true; +} + +adetailer_ctx_t* new_adetailer_ctx(const char* detector_path, + int n_threads, + const char* backend, + const char* params_backend) { + if (detector_path == nullptr || detector_path[0] == '\0') { + return nullptr; + } + auto* context = static_cast(calloc(1, sizeof(adetailer_ctx_t))); + if (context == nullptr) { + return nullptr; + } + context->detailer = new ADetailerGGML(n_threads, SAFE_STR(backend), SAFE_STR(params_backend)); + if (context->detailer == nullptr || !context->detailer->load_from_file(detector_path)) { + delete context->detailer; + free(context); + return nullptr; + } + return context; +} + +void free_adetailer_ctx(adetailer_ctx_t* context) { + if (context == nullptr) { + return; + } + delete context->detailer; + context->detailer = nullptr; + free(context); +} + +bool adetail_image(adetailer_ctx_t* context, + sd_ctx_t* sd_ctx, + sd_image_t input_image, + const sd_adetailer_params_t* public_params, + const sd_img_gen_params_t* inpaint_params, + sd_image_t** images_out, + int* num_images_out) { + if (images_out != nullptr) { + *images_out = nullptr; + } + if (num_images_out != nullptr) { + *num_images_out = 0; + } + if (context == nullptr || context->detailer == nullptr || sd_ctx == nullptr || public_params == nullptr || + inpaint_params == nullptr || input_image.data == nullptr) { + return false; + } + + ADetailerParams params; + if (!parse_adetailer_params(*public_params, *inpaint_params, ¶ms)) { + return false; + } + + auto detections = context->detailer->predict(input_image, params); + filter_and_sort_detections(&detections, input_image.width, input_image.height, params); + auto masks = make_masks(detections, input_image.width, input_image.height, params); + + sd_image_t current = copy_as_rgb(input_image); + if (current.data == nullptr) { + return false; + } + + int applied = 0; + for (size_t i = 0; i < masks.size(); ++i) { + std::string prompt = resolve_prompt(params.prompt, inpaint_params->prompt, static_cast(i)); + std::string negative_prompt = resolve_prompt(params.negative_prompt, + inpaint_params->negative_prompt, + static_cast(i)); + if (prompt_is_skip(prompt)) { + continue; + } + + CropRegion crop; + if (!mask_bbox(masks[i], &crop)) { + continue; + } + crop = expand_crop(crop, + current.width, + current.height, + params.inpaint_padding, + params.inpaint_width, + params.inpaint_height); + sd_image_t local_image = resize_crop_image(current, crop, params.inpaint_width, params.inpaint_height); + sd_image_t local_mask = resize_crop_mask(masks[i], crop, params.inpaint_width, params.inpaint_height); + if (local_image.data == nullptr || local_mask.data == nullptr) { + free(local_image.data); + free(local_mask.data); + free(current.data); + return false; + } + + sd_img_gen_params_t generation = *inpaint_params; + generation.prompt = prompt.c_str(); + generation.negative_prompt = negative_prompt.c_str(); + generation.init_image = local_image; + generation.mask_image = local_mask; + generation.width = params.inpaint_width; + generation.height = params.inpaint_height; + generation.strength = params.denoising_strength; + generation.seed = inpaint_params->seed + static_cast(i); + generation.batch_count = 1; + generation.ref_images = nullptr; + generation.ref_images_count = 0; + generation.control_image = {}; + generation.pm_params = {}; + generation.pulid_params = {}; + generation.hires.enabled = false; + if (params.steps > 0) { + generation.sample_params.sample_steps = params.steps; + generation.sample_params.custom_sigmas = nullptr; + generation.sample_params.custom_sigmas_count = 0; + } + if (params.cfg_scale >= 0.f) { + generation.sample_params.guidance.txt_cfg = params.cfg_scale; + } + if (params.sample_method != SAMPLE_METHOD_COUNT) { + generation.sample_params.sample_method = params.sample_method; + } + if (params.scheduler != SCHEDULER_COUNT) { + generation.sample_params.scheduler = params.scheduler; + } + + sd_image_t* generated = nullptr; + int generated_count = 0; + bool success = generate_image(sd_ctx, &generation, &generated, &generated_count); + free(local_image.data); + free(local_mask.data); + if (!success || generated_count <= 0 || generated == nullptr || generated[0].data == nullptr) { + free_sd_images(generated, generated_count); + free(current.data); + return false; + } + + Mask feather = gaussian_blur_mask(masks[i], params.mask_blur); + composite_crop(¤t, generated[0], feather, crop); + free_sd_images(generated, generated_count); + ++applied; + } + + auto* outputs = static_cast(calloc(1, sizeof(sd_image_t))); + if (outputs == nullptr) { + free(current.data); + return false; + } + outputs[0] = current; + if (images_out != nullptr) { + *images_out = outputs; + } else { + free_sd_images(outputs, 1); + } + if (num_images_out != nullptr) { + *num_images_out = 1; + } + LOG_INFO("ADetailer applied %d mask(s)", applied); + return true; +} diff --git a/src/detailer.h b/src/detailer.h new file mode 100644 index 000000000..17a81a1dd --- /dev/null +++ b/src/detailer.h @@ -0,0 +1,75 @@ +#ifndef __SD_DETAILER_H__ +#define __SD_DETAILER_H__ + +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "model/detector/yolov8.h" +#include "model_manager.h" +#include "stable-diffusion.h" + +struct ADetailerDetection { + float x1 = 0.f; + float y1 = 0.f; + float x2 = 0.f; + float y2 = 0.f; + float confidence = 0.f; + int class_id = 0; +}; + +enum ADetailerSort { + ADETAILER_SORT_NONE, + ADETAILER_SORT_LEFT_TO_RIGHT, + ADETAILER_SORT_CENTER_TO_EDGE, + ADETAILER_SORT_AREA, +}; + +struct ADetailerParams { + const char* prompt = nullptr; + const char* negative_prompt = nullptr; + int input_size = 640; + float confidence = 0.3f; + float nms_threshold = 0.45f; + int max_detections = 100; + int mask_k_largest = 0; + float mask_min_ratio = 0.f; + float mask_max_ratio = 1.f; + int dilate_erode = 4; + int x_offset = 0; + int y_offset = 0; + bool merge_masks = false; + bool invert_mask = false; + int mask_blur = 4; + int inpaint_padding = 32; + int inpaint_width = 512; + int inpaint_height = 512; + float denoising_strength = 0.4f; + int steps = 0; + float cfg_scale = -1.f; + sample_method_t sample_method = SAMPLE_METHOD_COUNT; + scheduler_t scheduler = SCHEDULER_COUNT; + ADetailerSort sort_by = ADETAILER_SORT_NONE; +}; + +struct ADetailerGGML { + SDBackendManager backend_manager; + std::shared_ptr model_manager; + std::shared_ptr detector; + std::vector class_names; + int n_threads = 1; + std::string backend_spec; + std::string params_backend_spec; + + ADetailerGGML(int n_threads, + std::string backend_spec, + std::string params_backend_spec); + ~ADetailerGGML(); + + bool load_from_file(const std::string& detector_path); + std::vector predict(sd_image_t image, + const ADetailerParams& params); +}; + +#endif // __SD_DETAILER_H__ diff --git a/src/model/detector/yolov8.h b/src/model/detector/yolov8.h new file mode 100644 index 000000000..a90fdf9a6 --- /dev/null +++ b/src/model/detector/yolov8.h @@ -0,0 +1,362 @@ +#ifndef __SD_MODEL_DETECTOR_YOLOV8_H__ +#define __SD_MODEL_DETECTOR_YOLOV8_H__ + +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "core/util.h" + +struct YOLOv8Config { + std::array out_channels{}; + std::map hidden_channels; + std::map repeats; + int detect_box_channels = 0; + int detect_cls_channels = 0; + int reg_max = 0; + int num_classes = 0; + bool valid = false; + + static YOLOv8Config detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "") { + YOLOv8Config config; + auto full_name = [&](const std::string& name) { + return prefix.empty() ? name : prefix + "." + name; + }; + auto find_weight = [&](const std::string& name) -> const TensorStorage* { + auto iter = tensor_storage_map.find(full_name(name)); + return iter == tensor_storage_map.end() ? nullptr : &iter->second; + }; + auto conv_out = [&](const std::string& name) -> int { + const TensorStorage* weight = find_weight(name); + return weight != nullptr && weight->n_dims == 4 ? static_cast(weight->ne[3]) : 0; + }; + + for (int layer : {0, 1, 3, 5, 7, 16, 19}) { + config.out_channels[layer] = conv_out("model." + std::to_string(layer) + ".conv.weight"); + } + for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) { + const std::string base = "model." + std::to_string(layer); + config.out_channels[layer] = conv_out(base + ".cv2.conv.weight"); + config.hidden_channels[layer] = conv_out(base + ".cv1.conv.weight") / 2; + + int repeat_count = 0; + while (find_weight(base + ".m." + std::to_string(repeat_count) + ".cv1.conv.weight") != nullptr) { + ++repeat_count; + } + config.repeats[layer] = repeat_count; + } + config.out_channels[9] = conv_out("model.9.cv2.conv.weight"); + + config.detect_box_channels = conv_out("model.22.cv2.0.0.conv.weight"); + config.detect_cls_channels = conv_out("model.22.cv3.0.0.conv.weight"); + const int box_outputs = conv_out("model.22.cv2.0.2.weight"); + config.num_classes = conv_out("model.22.cv3.0.2.weight"); + config.reg_max = box_outputs / 4; + + config.valid = config.out_channels[0] > 0 && config.out_channels[9] > 0 && + config.out_channels[15] > 0 && config.out_channels[18] > 0 && + config.out_channels[21] > 0 && config.detect_box_channels > 0 && + config.detect_cls_channels > 0 && box_outputs > 0 && box_outputs % 4 == 0 && + config.num_classes > 0; + for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) { + config.valid = config.valid && config.hidden_channels[layer] > 0 && config.repeats[layer] > 0; + } + + if (config.valid) { + LOG_DEBUG("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d", + config.num_classes, + config.reg_max, + config.out_channels[15], + config.out_channels[18], + config.out_channels[21]); + } + return config; + } +}; + +class YOLOConv : public UnaryBlock { + int out_channels_ = 0; + +public: + YOLOConv(int in_channels, int out_channels, int kernel, int stride = 1) + : out_channels_(out_channels) { + blocks["conv"] = std::shared_ptr(new Conv2d(in_channels, + out_channels, + {kernel, kernel}, + {stride, stride}, + {kernel / 2, kernel / 2}, + {1, 1}, + true)); + } + + int out_channels() const { + return out_channels_; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + return ggml_silu_inplace(ctx->ggml_ctx, conv->forward(ctx, x)); + } +}; + +class YOLOBottleneck : public UnaryBlock { + bool shortcut_ = false; + +public: + YOLOBottleneck(int channels, bool shortcut) + : shortcut_(shortcut) { + blocks["cv1"] = std::shared_ptr(new YOLOConv(channels, channels, 3)); + blocks["cv2"] = std::shared_ptr(new YOLOConv(channels, channels, 3)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto cv1 = std::dynamic_pointer_cast(blocks["cv1"]); + auto cv2 = std::dynamic_pointer_cast(blocks["cv2"]); + auto out = cv2->forward(ctx, cv1->forward(ctx, x)); + return shortcut_ ? ggml_add(ctx->ggml_ctx, x, out) : out; + } +}; + +class YOLOC2f : public UnaryBlock { + int hidden_channels_ = 0; + int repeats_ = 0; + +public: + YOLOC2f(int in_channels, + int out_channels, + int hidden_channels, + int repeats, + bool shortcut) + : hidden_channels_(hidden_channels), repeats_(repeats) { + blocks["cv1"] = std::shared_ptr(new YOLOConv(in_channels, hidden_channels * 2, 1)); + blocks["cv2"] = std::shared_ptr(new YOLOConv(hidden_channels * (2 + repeats), out_channels, 1)); + for (int i = 0; i < repeats; ++i) { + blocks["m." + std::to_string(i)] = std::shared_ptr(new YOLOBottleneck(hidden_channels, shortcut)); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto cv1 = std::dynamic_pointer_cast(blocks["cv1"]); + auto cv2 = std::dynamic_pointer_cast(blocks["cv2"]); + auto split = cv1->forward(ctx, x); + + // split: [N, 2*C, H, W], ggml layout [W, H, 2*C, N]. + auto y0 = ggml_view_4d(ctx->ggml_ctx, + split, + split->ne[0], + split->ne[1], + hidden_channels_, + split->ne[3], + split->nb[1], + split->nb[2], + split->nb[3], + 0); + auto y1 = ggml_view_4d(ctx->ggml_ctx, + split, + split->ne[0], + split->ne[1], + hidden_channels_, + split->ne[3], + split->nb[1], + split->nb[2], + split->nb[3], + static_cast(hidden_channels_) * split->nb[2]); + auto joined = ggml_concat(ctx->ggml_ctx, y0, y1, 2); + auto last = y1; + for (int i = 0; i < repeats_; ++i) { + auto block = std::dynamic_pointer_cast(blocks["m." + std::to_string(i)]); + last = block->forward(ctx, last); + joined = ggml_concat(ctx->ggml_ctx, joined, last, 2); + } + return cv2->forward(ctx, joined); + } +}; + +class YOLOSPPF : public UnaryBlock { +public: + YOLOSPPF(int in_channels, int out_channels) { + blocks["cv1"] = std::shared_ptr(new YOLOConv(in_channels, in_channels / 2, 1)); + blocks["cv2"] = std::shared_ptr(new YOLOConv(in_channels * 2, out_channels, 1)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto cv1 = std::dynamic_pointer_cast(blocks["cv1"]); + auto cv2 = std::dynamic_pointer_cast(blocks["cv2"]); + x = cv1->forward(ctx, x); + auto y1 = ggml_pool_2d(ctx->ggml_ctx, x, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2); + auto y2 = ggml_pool_2d(ctx->ggml_ctx, y1, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2); + auto y3 = ggml_pool_2d(ctx->ggml_ctx, y2, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2); + auto out = ggml_concat(ctx->ggml_ctx, x, y1, 2); + out = ggml_concat(ctx->ggml_ctx, out, y2, 2); + out = ggml_concat(ctx->ggml_ctx, out, y3, 2); + return cv2->forward(ctx, out); + } +}; + +class YOLODetect : public GGMLBlock { + int num_classes_ = 0; + int reg_max_ = 0; + +public: + YOLODetect(const std::array& in_channels, + int box_channels, + int cls_channels, + int reg_max, + int num_classes) + : num_classes_(num_classes), reg_max_(reg_max) { + for (int i = 0; i < 3; ++i) { + const std::string box = "cv2." + std::to_string(i); + blocks[box + ".0"] = std::shared_ptr(new YOLOConv(in_channels[i], box_channels, 3)); + blocks[box + ".1"] = std::shared_ptr(new YOLOConv(box_channels, box_channels, 3)); + blocks[box + ".2"] = std::shared_ptr(new Conv2d(box_channels, reg_max * 4, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true)); + + const std::string cls = "cv3." + std::to_string(i); + blocks[cls + ".0"] = std::shared_ptr(new YOLOConv(in_channels[i], cls_channels, 3)); + blocks[cls + ".1"] = std::shared_ptr(new YOLOConv(cls_channels, cls_channels, 3)); + blocks[cls + ".2"] = std::shared_ptr(new Conv2d(cls_channels, num_classes, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true)); + } + } + + ggml_tensor* forward_scale(GGMLRunnerContext* ctx, ggml_tensor* x, int index) { + const std::string box = "cv2." + std::to_string(index); + auto box0 = std::dynamic_pointer_cast(blocks[box + ".0"]); + auto box1 = std::dynamic_pointer_cast(blocks[box + ".1"]); + auto box2 = std::dynamic_pointer_cast(blocks[box + ".2"]); + + const std::string cls = "cv3." + std::to_string(index); + auto cls0 = std::dynamic_pointer_cast(blocks[cls + ".0"]); + auto cls1 = std::dynamic_pointer_cast(blocks[cls + ".1"]); + auto cls2 = std::dynamic_pointer_cast(blocks[cls + ".2"]); + + auto boxes = box2->forward(ctx, box1->forward(ctx, box0->forward(ctx, x))); + auto classes = cls2->forward(ctx, cls1->forward(ctx, cls0->forward(ctx, x))); + return ggml_concat(ctx->ggml_ctx, boxes, classes, 2); + } + + int output_channels() const { + return reg_max_ * 4 + num_classes_; + } +}; + +class YOLOv8Model : public GGMLBlock { + YOLOv8Config config_; + + std::shared_ptr make_c2f(int layer, int in_channels, bool shortcut) { + return std::make_shared(in_channels, + config_.out_channels[layer], + config_.hidden_channels.at(layer), + config_.repeats.at(layer), + shortcut); + } + +public: + explicit YOLOv8Model(YOLOv8Config config) + : config_(std::move(config)) { + blocks["model.0"] = std::make_shared(3, config_.out_channels[0], 3, 2); + blocks["model.1"] = std::make_shared(config_.out_channels[0], config_.out_channels[1], 3, 2); + blocks["model.2"] = make_c2f(2, config_.out_channels[1], true); + blocks["model.3"] = std::make_shared(config_.out_channels[2], config_.out_channels[3], 3, 2); + blocks["model.4"] = make_c2f(4, config_.out_channels[3], true); + blocks["model.5"] = std::make_shared(config_.out_channels[4], config_.out_channels[5], 3, 2); + blocks["model.6"] = make_c2f(6, config_.out_channels[5], true); + blocks["model.7"] = std::make_shared(config_.out_channels[6], config_.out_channels[7], 3, 2); + blocks["model.8"] = make_c2f(8, config_.out_channels[7], true); + blocks["model.9"] = std::make_shared(config_.out_channels[8], config_.out_channels[9]); + + blocks["model.12"] = make_c2f(12, config_.out_channels[9] + config_.out_channels[6], false); + blocks["model.15"] = make_c2f(15, config_.out_channels[12] + config_.out_channels[4], false); + blocks["model.16"] = std::make_shared(config_.out_channels[15], config_.out_channels[16], 3, 2); + blocks["model.18"] = make_c2f(18, config_.out_channels[16] + config_.out_channels[12], false); + blocks["model.19"] = std::make_shared(config_.out_channels[18], config_.out_channels[19], 3, 2); + blocks["model.21"] = make_c2f(21, config_.out_channels[19] + config_.out_channels[9], false); + blocks["model.22"] = std::make_shared( + std::array{config_.out_channels[15], config_.out_channels[18], config_.out_channels[21]}, + config_.detect_box_channels, + config_.detect_cls_channels, + config_.reg_max, + config_.num_classes); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto run = [&](int layer, ggml_tensor* input) { + return std::dynamic_pointer_cast(blocks["model." + std::to_string(layer)])->forward(ctx, input); + }; + + auto x0 = run(0, x); + auto x1 = run(1, x0); + auto x2 = run(2, x1); + auto x3 = run(3, x2); + auto x4 = run(4, x3); + auto x5 = run(5, x4); + auto x6 = run(6, x5); + auto x7 = run(7, x6); + auto x8 = run(8, x7); + auto x9 = run(9, x8); + + auto x12 = run(12, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x9, 2, GGML_SCALE_MODE_NEAREST), x6, 2)); + auto x15 = run(15, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x12, 2, GGML_SCALE_MODE_NEAREST), x4, 2)); + auto x16 = run(16, x15); + auto x18 = run(18, ggml_concat(ctx->ggml_ctx, x16, x12, 2)); + auto x19 = run(19, x18); + auto x21 = run(21, ggml_concat(ctx->ggml_ctx, x19, x9, 2)); + + auto detect = std::dynamic_pointer_cast(blocks["model.22"]); + auto p3 = detect->forward_scale(ctx, x15, 0); + auto p4 = detect->forward_scale(ctx, x18, 1); + auto p5 = detect->forward_scale(ctx, x21, 2); + p3 = ggml_reshape_2d(ctx->ggml_ctx, p3, p3->ne[0] * p3->ne[1], detect->output_channels()); + p4 = ggml_reshape_2d(ctx->ggml_ctx, p4, p4->ne[0] * p4->ne[1], detect->output_channels()); + p5 = ggml_reshape_2d(ctx->ggml_ctx, p5, p5->ne[0] * p5->ne[1], detect->output_channels()); + return ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, p3, p4, 0), p5, 0); + } +}; + +struct YOLOv8Runner : public GGMLRunner { + YOLOv8Config config; + std::unique_ptr model; + + YOLOv8Runner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(YOLOv8Config::detect_from_weights(tensor_storage_map)) { + if (config.valid) { + model = std::make_unique(config); + model->init(params_ctx, tensor_storage_map, ""); + } + } + + std::string get_desc() override { + return "yolov8"; + } + + void get_param_tensors(std::map& tensors) { + if (model) { + model->get_param_tensors(tensors); + } + } + + ggml_cgraph* build_graph(const sd::Tensor& input) { + if (!model) { + return nullptr; + } + ggml_cgraph* graph = new_graph_custom(1 << 16); + auto x = make_input(input); + auto runner_ctx = get_context(); + auto output = model->forward(&runner_ctx, x); + ggml_build_forward_expand(graph, output); + return graph; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& input) { + auto get_graph = [&]() { return build_graph(input); }; + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false, false, false)); + } +}; + +#endif // __SD_MODEL_DETECTOR_YOLOV8_H__ diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 223cc9670..df71eab11 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -100,7 +100,8 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) { // https://huggingface.co/docs/safetensors/index bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, - std::string* error) { + std::string* error, + std::map* metadata) { std::ifstream file(file_path, std::ios::binary); if (!file.is_open()) { set_error(error, "failed to open '" + file_path + "'"); @@ -150,6 +151,18 @@ bool read_safetensors_file(const std::string& file_path, return false; } + if (metadata != nullptr) { + metadata->clear(); + auto metadata_item = header_.find("__metadata__"); + if (metadata_item != header_.end() && metadata_item->is_object()) { + for (const auto& item : metadata_item->items()) { + if (item.value().is_string()) { + metadata->emplace(item.key(), item.value().get()); + } + } + } + } + tensor_storages.clear(); for (auto& item : header_.items()) { std::string name = item.key(); diff --git a/src/model_io/safetensors_io.h b/src/model_io/safetensors_io.h index b18f0ae75..4291b54f6 100644 --- a/src/model_io/safetensors_io.h +++ b/src/model_io/safetensors_io.h @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__ #define __SD_MODEL_IO_SAFETENSORS_IO_H__ +#include #include #include @@ -10,7 +11,8 @@ bool is_safetensors_file(const std::string& file_path); bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, - std::string* error = nullptr); + std::string* error = nullptr, + std::map* metadata = nullptr); bool read_safetensors_index_file(const std::string& file_path, std::vector& shard_paths, std::string* error = nullptr); diff --git a/src/model_loader.cpp b/src/model_loader.cpp index d22a6c5d7..853ed1440 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -317,7 +317,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const std::vector tensor_storages; std::string error; - if (!read_safetensors_file(file_path, tensor_storages, &error)) { + if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) { LOG_ERROR("%s", error.c_str()); return false; } diff --git a/src/model_loader.h b/src/model_loader.h index d311b7030..f7ebcf3ef 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -36,6 +36,7 @@ class ModelLoader { std::vector file_data; bool model_files_processed = false; String2TensorStorage tensor_storage_map; + std::map metadata_; int n_threads_; size_t add_file_path(const std::string& file_path); @@ -63,6 +64,7 @@ class ModelLoader { std::map get_vae_wtype_stat(); String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; } const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; } + const std::map& get_metadata() const { return metadata_; } void set_n_threads(int n_threads); void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = ""); void process_model_files(bool enable_mmap = false, bool writable_mmap = true);