fixed img2img for different sizes

LostRuins · Apr 8, 2024 · 5c323a0 · 5c323a0
1 parent 2bc11e0
commit 5c323a0
Show file tree

Hide file tree

Showing 2 changed files with 10,386 additions and 5 deletions.
diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp
@@ -30,6 +30,9 @@
 #define STB_IMAGE_WRITE_STATIC
 #include "stb_image_write.h"
 
+#define STB_IMAGE_RESIZE_IMPLEMENTATION
+#include "stb_image_resize2.h"
+
 const char* rng_type_to_str[] = {
     "std_default",
     "cuda",
@@ -282,6 +285,10 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
     sd_image_t input_image = {0,0,0,nullptr};
     std::vector<uint8_t> image_buffer;
     int nx, ny, nc;
+    int img2imgW = inputs.width; //for img2img input
+    int img2imgH = inputs.height;
+    int img2imgC = 3; // Assuming RGB image
+    std::vector<uint8_t> resized_image_buf(img2imgW * img2imgH * img2imgC);
 
     if(!is_quiet)
     {
@@ -372,7 +379,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
 
         input_image_buffer = stbi_load_from_memory(image_buffer.data(), image_buffer.size(), &nx, &ny, &nc, 3);
 
-        if (nx < 64 || ny < 64 || nc!= 3) {
+        if (nx < 64 || ny < 64 || nx > 1024 || ny > 1024 || nc!= 3) {
             printf("\nKCPP SD: bad input image dimensions %d x %d!\n",nx,ny);
             output.data = "";
             output.status = 0;
@@ -385,10 +392,19 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
             return output;
         }
 
-        input_image.width = nx;
-        input_image.height = ny;
-        input_image.channel = nc;
-        input_image.data = input_image_buffer;
+        // Resize the image
+        unsigned char * resok = stbir_resize_uint8_linear(input_image_buffer, nx, ny, 0, resized_image_buf.data(), img2imgW, img2imgH, 0, (stbir_pixel_layout)img2imgC);
+        if (!resok) {
+            printf("\nKCPP SD: resize image failed!\n");
+            output.data = "";
+            output.status = 0;
+            return output;
+        }
+
+        input_image.width = img2imgW;
+        input_image.height = img2imgH;
+        input_image.channel = img2imgC;
+        input_image.data = resized_image_buf.data();
 
         if(!is_quiet && sddebugmode==1)
         {