upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-08-02 00:00:41 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add ADetailer support (#1785)

Browse Source
This commit is contained in:
leejet
2026-07-14 23:59:11 +08:00
committed by GitHub
parent c00a9e956d
commit a8a91b24cd
19 changed files with 1874 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

362
src/model/detector/yolov8.h Normal file
View File

@@ -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__