feat: add ADetailer support (#1785)

Author: leejet

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) + ", ";