Fix Flex 2 inpaint

Author: stduhpf

ggml_extend.hpp

@@ -375,18 +375,31 @@ __STATIC_INLINE__ void sd_mask_to_tensor(const uint8_t* image_data,
 
 __STATIC_INLINE__ void sd_apply_mask(struct ggml_tensor* image_data,
                                      struct ggml_tensor* mask,
-                                     struct ggml_tensor* output) {
+                                     struct ggml_tensor* output,
+                                     float masked_value = 0.5f) {
     int64_t width    = output->ne[0];
     int64_t height   = output->ne[1];
     int64_t channels = output->ne[2];
+    for (int ix = 0; ix < mask->ne[0]; ix++) {
+        for (int iy = 0; iy < mask->ne[1]; iy++) {
+            float m = ggml_tensor_get_f32(mask, ix, iy);
+            m       = round(m);  // inpaint models need binary masks
+            ggml_tensor_set_f32(mask, m, ix, iy);
+        }
+    }
+    float rescale_mx = mask->ne[0]/output->ne[0];
+    float rescale_my = mask->ne[1]/output->ne[1];
     GGML_ASSERT(output->type == GGML_TYPE_F32);
     for (int ix = 0; ix < width; ix++) {
         for (int iy = 0; iy < height; iy++) {
-            float m = ggml_tensor_get_f32(mask, ix, iy);
+            int mx = (int)(ix * rescale_mx);
+            int my = (int)(iy * rescale_my);
+            float m = ggml_tensor_get_f32(mask, mx, my);
             m       = round(m);  // inpaint models need binary masks
-            ggml_tensor_set_f32(mask, m, ix, iy);
+            ggml_tensor_set_f32(mask, m, mx, my);
             for (int k = 0; k < channels; k++) {
-                float value = (1 - m) * (ggml_tensor_get_f32(image_data, ix, iy, k) - .5) + .5;
+                float value = ggml_tensor_get_f32(image_data, ix, iy, k);
+                value = (1 - m) * (value - masked_value) + masked_value;
                 ggml_tensor_set_f32(output, value, ix, iy, k);
             }
         }

stable-diffusion.cpp

@@ -1821,16 +1821,23 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
         } else if (sd_ctx->sd->version == VERSION_FLEX_2) {
             mask_channels = 1 + init_latent->ne[2];
         }
-        ggml_tensor* masked_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
-        // Restore init_img (encode_first_stage has side effects) TODO: remove the side effects?
-        sd_image_to_tensor(init_image.data, init_img);
-        sd_apply_mask(init_img, mask_img, masked_img);
         ggml_tensor* masked_latent = NULL;
-        if (!sd_ctx->sd->use_tiny_autoencoder) {
-            ggml_tensor* moments = sd_ctx->sd->encode_first_stage(work_ctx, masked_img);
-            masked_latent        = sd_ctx->sd->get_first_stage_encoding(work_ctx, moments);
+        if (sd_ctx->sd->version != VERSION_FLEX_2) {
+            // most inpaint models mask before vae
+            ggml_tensor* masked_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
+            // Restore init_img (encode_first_stage has side effects) TODO: remove the side effects?
+            sd_image_to_tensor(init_image.data, init_img);
+            sd_apply_mask(init_img, mask_img, masked_img);
+            if (!sd_ctx->sd->use_tiny_autoencoder) {
+                ggml_tensor* moments = sd_ctx->sd->encode_first_stage(work_ctx, masked_img);
+                masked_latent      = sd_ctx->sd->get_first_stage_encoding(work_ctx, moments);
+            } else {
+                masked_latent = sd_ctx->sd->encode_first_stage(work_ctx, masked_img);
+            }
         } else {
-            masked_latent = sd_ctx->sd->encode_first_stage(work_ctx, masked_img);
+            // mask after vae
+            masked_latent = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, init_latent->ne[0], init_latent->ne[1], init_latent->ne[2], 1);
+            sd_apply_mask(init_latent, mask_img, masked_latent, 0.);
         }
         concat_latent = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, masked_latent->ne[0], masked_latent->ne[1], mask_channels + masked_latent->ne[2], 1);
         for (int ix = 0; ix < masked_latent->ne[0]; ix++) {