feat: add ADetailer support (#1785)

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)

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"
+```

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.
 

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

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

examples/common/common.cpp

@@ -30,6 +30,7 @@ namespace fs = std::filesystem;
 
 const char* const modes_str[] = {
     "img_gen",
+    "adetailer",
     "vid_gen",
     "convert",
     "upscale",
@@ -922,6 +923,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",
@@ -1842,6 +1863,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);
@@ -2358,13 +2383,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;
 }
 
@@ -2587,6 +2618,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"
@@ -2701,8 +2736,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())},
@@ -2716,6 +2756,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);
@@ -2870,6 +2918,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) + ", ";

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,
@@ -187,6 +188,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;

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 <typename T>
 using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>;
 
-using FilePtr        = std::unique_ptr<FILE, FileCloser>;
-using SDCtxPtr       = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
-using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
+using FilePtr         = std::unique_ptr<FILE, FileCloser>;
+using SDCtxPtr        = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
+using UpscalerCtxPtr  = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
+using ADetailerCtxPtr = std::unique_ptr<adetailer_ctx_t, ADetailerCtxDeleter>;
 
 class SDImageOwner {
 private:

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,

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()

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";
 }

src/core/ggml_extend_backend.h

@@ -20,6 +20,7 @@ enum class SDBackendModule {
     CONTROL_NET,
     PHOTOMAKER,
     UPSCALER,
+    DETECTOR,
 };
 
 struct SDBackendAssignment {

src/detailer.h

@@ -0,0 +1,75 @@
+#ifndef __SD_DETAILER_H__
+#define __SD_DETAILER_H__
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#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<ModelManager> model_manager;
+    std::shared_ptr<YOLOv8Runner> detector;
+    std::vector<std::string> 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<ADetailerDetection> predict(sd_image_t image,
+                                            const ADetailerParams& params);
+};
+
+#endif  // __SD_DETAILER_H__

src/model/detector/yolov8.h

@@ -0,0 +1,362 @@
+#ifndef __SD_MODEL_DETECTOR_YOLOV8_H__
+#define __SD_MODEL_DETECTOR_YOLOV8_H__
+
+#include <algorithm>
+#include <array>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "core/ggml_extend.hpp"
+#include "core/util.h"
+
+struct YOLOv8Config {
+    std::array<int, 23> out_channels{};
+    std::map<int, int> hidden_channels;
+    std::map<int, int> 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<int>(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<GGMLBlock>(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<Conv2d>(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<GGMLBlock>(new YOLOConv(channels, channels, 3));
+        blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(channels, channels, 3));
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
+        auto cv2 = std::dynamic_pointer_cast<YOLOConv>(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<GGMLBlock>(new YOLOConv(in_channels, hidden_channels * 2, 1));
+        blocks["cv2"] = std::shared_ptr<GGMLBlock>(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<GGMLBlock>(new YOLOBottleneck(hidden_channels, shortcut));
+        }
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        auto cv1   = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
+        auto cv2   = std::dynamic_pointer_cast<YOLOConv>(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<size_t>(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<YOLOBottleneck>(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<GGMLBlock>(new YOLOConv(in_channels, in_channels / 2, 1));
+        blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels * 2, out_channels, 1));
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
+        auto cv2 = std::dynamic_pointer_cast<YOLOConv>(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<int, 3>& 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<GGMLBlock>(new YOLOConv(in_channels[i], box_channels, 3));
+            blocks[box + ".1"]    = std::shared_ptr<GGMLBlock>(new YOLOConv(box_channels, box_channels, 3));
+            blocks[box + ".2"]    = std::shared_ptr<GGMLBlock>(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<GGMLBlock>(new YOLOConv(in_channels[i], cls_channels, 3));
+            blocks[cls + ".1"]    = std::shared_ptr<GGMLBlock>(new YOLOConv(cls_channels, cls_channels, 3));
+            blocks[cls + ".2"]    = std::shared_ptr<GGMLBlock>(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<YOLOConv>(blocks[box + ".0"]);
+        auto box1             = std::dynamic_pointer_cast<YOLOConv>(blocks[box + ".1"]);
+        auto box2             = std::dynamic_pointer_cast<Conv2d>(blocks[box + ".2"]);
+
+        const std::string cls = "cv3." + std::to_string(index);
+        auto cls0             = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".0"]);
+        auto cls1             = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".1"]);
+        auto cls2             = std::dynamic_pointer_cast<Conv2d>(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<YOLOC2f> make_c2f(int layer, int in_channels, bool shortcut) {
+        return std::make_shared<YOLOC2f>(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<YOLOConv>(3, config_.out_channels[0], 3, 2);
+        blocks["model.1"] = std::make_shared<YOLOConv>(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<YOLOConv>(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<YOLOConv>(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<YOLOConv>(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<YOLOSPPF>(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<YOLOConv>(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<YOLOConv>(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<YOLODetect>(
+            std::array<int, 3>{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<UnaryBlock>(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<YOLODetect>(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<YOLOv8Model> model;
+
+    YOLOv8Runner(ggml_backend_t backend,
+                 const String2TensorStorage& tensor_storage_map,
+                 std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager),
+          config(YOLOv8Config::detect_from_weights(tensor_storage_map)) {
+        if (config.valid) {
+            model = std::make_unique<YOLOv8Model>(config);
+            model->init(params_ctx, tensor_storage_map, "");
+        }
+    }
+
+    std::string get_desc() override {
+        return "yolov8";
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+        if (model) {
+            model->get_param_tensors(tensors);
+        }
+    }
+
+    ggml_cgraph* build_graph(const sd::Tensor<float>& 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<float> compute(int n_threads, const sd::Tensor<float>& input) {
+        auto get_graph = [&]() { return build_graph(input); };
+        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false));
+    }
+};
+
+#endif  // __SD_MODEL_DETECTOR_YOLOV8_H__

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<TensorStorage>& tensor_storages,
-                           std::string* error) {
+                           std::string* error,
+                           std::map<std::string, std::string>* 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<std::string>());
+                }
+            }
+        }
+    }
+
     tensor_storages.clear();
     for (auto& item : header_.items()) {
         std::string name           = item.key();

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 <map>
 #include <string>
 #include <vector>
 
@@ -10,7 +11,8 @@
 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);
+                           std::string* error                           = nullptr,
+                           std::map<std::string, std::string>* metadata = nullptr);
 bool read_safetensors_index_file(const std::string& file_path,
                                  std::vector<std::string>& shard_paths,
                                  std::string* error = nullptr);

src/model_loader.cpp

@@ -317,7 +317,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
 
     std::vector<TensorStorage> 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;
     }

src/model_loader.h

@@ -36,6 +36,7 @@ protected:
     std::vector<ModelFileData> file_data;
     bool model_files_processed = false;
     String2TensorStorage tensor_storage_map;
+    std::map<std::string, std::string> metadata_;
     int n_threads_;
 
     size_t add_file_path(const std::string& file_path);
@@ -63,6 +64,7 @@ public:
     std::map<ggml_type, uint32_t> 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<std::string, std::string>& 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);