Skip to content

Commit 6d484cd

Browse files
stduhpfstduhpf
committed
inpaint: masked diffusion for inpainting models with inflated mask
1 parent d78db9b commit 6d484cd

2 files changed

Lines changed: 97 additions & 18 deletions

File tree

‎src/stable-diffusion.cpp‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2812,7 +2812,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
28122812
request->height / request->vae_scale_factor,
28132813
1,
28142814
1},
2815-
sd::ops::InterpolateMode::MaxPool);
2815+
sd::ops::InterpolateMode::NearestMax);
28162816

28172817
sd::Tensor<float> init_latent;
28182818
sd::Tensor<float> control_latent;
@@ -2956,10 +2956,14 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
29562956
latents.ref_images = std::move(ref_images);
29572957
latents.ref_latents = std::move(ref_latents);
29582958

2959-
if (!sd_version_is_inpaint(sd_ctx->sd->version)) {
2960-
latents.denoise_mask = std::move(latent_mask);
2961-
}
2962-
2959+
if (sd_version_is_inpaint(sd_ctx->sd->version)) {
2960+
latent_mask = sd::ops::maxPool2D(latent_mask,
2961+
{3, 3},
2962+
{1, 1},
2963+
{1, 1});
2964+
}
2965+
latents.denoise_mask = std::move(latent_mask);
2966+
29632967
return latents;
29642968
}
29652969

‎src/tensor.hpp‎

Lines changed: 88 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -819,9 +819,9 @@ namespace sd {
819819
namespace ops {
820820
enum class InterpolateMode {
821821
Nearest,
822-
MaxPool,
823-
MinPool,
824-
AvgPool,
822+
NearestMax,
823+
NearestMin,
824+
NearestAvg,
825825
};
826826

827827
inline int64_t normalize_slice_bound(int64_t index, int64_t dim_size) {
@@ -1020,9 +1020,9 @@ namespace sd {
10201020
InterpolateMode mode = InterpolateMode::Nearest,
10211021
bool align_corners = false) {
10221022
bool is_nearest_like_mode = (mode == InterpolateMode::Nearest ||
1023-
mode == InterpolateMode::MaxPool ||
1024-
mode == InterpolateMode::MinPool ||
1025-
mode == InterpolateMode::AvgPool);
1023+
mode == InterpolateMode::NearestMax ||
1024+
mode == InterpolateMode::NearestMin ||
1025+
mode == InterpolateMode::NearestAvg);
10261026
if (!is_nearest_like_mode) {
10271027
tensor_throw_invalid_argument("Only nearest-like interpolate modes are implemented, got mode=" +
10281028
std::to_string(static_cast<int>(mode)));
@@ -1083,23 +1083,23 @@ namespace sd {
10831083
}
10841084

10851085
T val;
1086-
if (mode == InterpolateMode::MaxPool) {
1086+
if (mode == InterpolateMode::NearestMax) {
10871087
val = std::numeric_limits<T>::lowest();
1088-
} else if(mode == InterpolateMode::MinPool) {
1088+
} else if(mode == InterpolateMode::NearestMin) {
10891089
val = std::numeric_limits<T>::max();
1090-
} else if(mode == InterpolateMode::AvgPool) {
1090+
} else if(mode == InterpolateMode::NearestAvg) {
10911091
val = T(0);
10921092
}
10931093

10941094
bool done_window = false;
10951095
std::vector<int64_t> current_in_coord = input_start;
10961096

10971097
while (!done_window) {
1098-
if (mode == InterpolateMode::MaxPool) {
1098+
if (mode == InterpolateMode::NearestMax) {
10991099
val = std::max(val, input.index(current_in_coord));
1100-
} else if(mode == InterpolateMode::MinPool) {
1100+
} else if(mode == InterpolateMode::NearestMin) {
11011101
val = std::min(val, input.index(current_in_coord));
1102-
} else if(mode == InterpolateMode::AvgPool) {
1102+
} else if(mode == InterpolateMode::NearestAvg) {
11031103
val += input.index(current_in_coord);
11041104
}
11051105

@@ -1113,7 +1113,7 @@ namespace sd {
11131113
}
11141114
}
11151115
}
1116-
if (mode == InterpolateMode::AvgPool) {
1116+
if (mode == InterpolateMode::NearestAvg) {
11171117
int64_t window_size = 1;
11181118
for (size_t i = 0; i < static_cast<size_t>(output.dim()); ++i) {
11191119
window_size *= (input_end[i] - input_start[i]);
@@ -1207,6 +1207,81 @@ namespace sd {
12071207
align_corners);
12081208
}
12091209

1210+
template <typename T>
1211+
inline Tensor<T> maxPool2D(const Tensor<T>& input,
1212+
std::vector<int64_t> kernel_size,
1213+
std::vector<int64_t> stride,
1214+
std::vector<int64_t> padding) {
1215+
if (input.dim() != 4) {
1216+
tensor_throw_invalid_argument("Tensor maxPool2D requires 4D input: input_dim=" +
1217+
std::to_string(input.dim()) + ", input_shape=" +
1218+
tensor_shape_to_string(input.shape()));
1219+
}
1220+
if (kernel_size.size() != 2 || stride.size() != 2 || padding.size() != 2) {
1221+
tensor_throw_invalid_argument("Tensor maxPool2D requires kernel_size, stride, and padding to have length 2");
1222+
}
1223+
for (size_t i = 0; i < 2; ++i) {
1224+
if (kernel_size[i] <= 0) {
1225+
tensor_throw_invalid_argument("Tensor maxPool2D kernel_size must be positive: kernel_size=" +
1226+
tensor_shape_to_string(kernel_size));
1227+
}
1228+
if (stride[i] <= 0) {
1229+
tensor_throw_invalid_argument("Tensor maxPool2D stride must be positive: stride=" +
1230+
tensor_shape_to_string(stride));
1231+
}
1232+
if (padding[i] < 0) {
1233+
tensor_throw_invalid_argument("Tensor maxPool2D padding must be non-negative: padding=" +
1234+
tensor_shape_to_string(padding));
1235+
}
1236+
}
1237+
1238+
const int64_t in_height = input.shape()[0];
1239+
const int64_t in_width = input.shape()[1];
1240+
const int64_t in_channels = input.shape()[2];
1241+
const int64_t batch_size = input.shape()[3];
1242+
1243+
const int64_t out_height = (in_height + 2 * padding[0] - kernel_size[0]) / stride[0] + 1;
1244+
const int64_t out_width = (in_width + 2 * padding[1] - kernel_size[1]) / stride[1] + 1;
1245+
1246+
if (out_height <= 0 || out_width <= 0) {
1247+
tensor_throw_invalid_argument("maxPool2D results in invalid output dimensions: " +
1248+
std::to_string(out_height) + "x" + std::to_string(out_width));
1249+
}
1250+
1251+
Tensor<T> output({out_height, out_width, in_channels, batch_size});
1252+
1253+
for (int64_t oh = 0; oh < out_height; ++oh) {
1254+
for (int64_t ow = 0; ow < out_width; ++ow) {
1255+
for (int64_t c = 0; c < in_channels; ++c) {
1256+
for (int64_t b = 0; b < batch_size; ++b) {
1257+
T max_val = std::numeric_limits<T>::lowest();
1258+
bool has_valid_input = false;
1259+
1260+
for (int64_t kh = 0; kh < kernel_size[0]; ++kh) {
1261+
for (int64_t kw = 0; kw < kernel_size[1]; ++kw) {
1262+
int64_t ih = oh * stride[0] + kh - padding[0];
1263+
int64_t iw = ow * stride[1] + kw - padding[1];
1264+
1265+
if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) {
1266+
T val = input.index(ih, iw, c, b);
1267+
max_val = std::max(max_val, val);
1268+
has_valid_input = true;
1269+
}
1270+
}
1271+
}
1272+
1273+
if (has_valid_input) {
1274+
output.index(oh, ow, c, b) = max_val;
1275+
} else {
1276+
output.index(oh, ow, c, b) = T(0);
1277+
}
1278+
}
1279+
}
1280+
}
1281+
}
1282+
return output;
1283+
}
1284+
12101285
template <typename T>
12111286
inline Tensor<T> concat(const Tensor<T>& lhs, const Tensor<T>& rhs, size_t dim) {
12121287
if (lhs.dim() != rhs.dim()) {

0 commit comments

Comments
 (0)