fix: flushes after printf (#38)

.dockerignore

@@ -0,0 +1,6 @@
+build*/
+test/
+
+.cache/
+*.swp
+models/

Dockerfile

@@ -0,0 +1,17 @@
+ARG UBUNTU_VERSION=22.04
+
+FROM ubuntu:$UBUNTU_VERSION as build
+
+RUN apt-get update && apt-get install -y build-essential git cmake
+
+WORKDIR /sd.cpp
+
+COPY . .
+
+RUN mkdir build && cd build && cmake .. && cmake --build . --config Release
+
+FROM ubuntu:$UBUNTU_VERSION as runtime
+
+COPY --from=build /sd.cpp/build/bin/sd /sd
+
+ENTRYPOINT [ "/sd" ]

README.md

@@ -23,6 +23,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
     - Linux
     - Mac OS
     - Windows
+    - Android (via Termux)
 
 ### TODO
 
@@ -34,6 +35,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - [ ] LoRA support
 - [ ] k-quants support
 - [ ] Cross-platform reproducibility (perhaps ensuring consistency with the original SD)
+- [ ] Adapting to more weight formats
 
 ## Usage
 
@@ -49,6 +51,7 @@ cd stable-diffusion.cpp
 ```
 cd stable-diffusion.cpp
 git pull origin master
+git submodule init
 git submodule update
 ```
 
@@ -84,6 +87,8 @@ You can specify the output model format using the --out_type parameter
 
 ### Build
 
+#### Build from scratch
+
 ```shell
 mkdir build
 cd build
@@ -91,7 +96,7 @@ cmake ..
 cmake --build . --config Release
 ```
 
-#### Using OpenBLAS
+##### Using OpenBLAS
 
 ```
 cmake .. -DGGML_OPENBLAS=ON
@@ -149,6 +154,22 @@ Using formats of different precisions will yield results of varying quality.
   <img src="./assets/img2img_output.png" width="256x">
 </p>
 
+### Docker
+
+#### Building using Docker
+
+```shell
+docker build -t sd .
+```
+
+#### Run
+
+```shell
+docker run -v /path/to/models:/models -v /path/to/output/:/output sd [args...]
+# For example
+# docker run -v ./models:/models -v ./build:/output sd -m /models/sd-v1-4-ggml-model-f16.bin -p "a lovely cat" -v -o /output/output.png
+```
+
 ## Memory/Disk Requirements
 
 | precision | f32  | f16  |q8_0  |q5_0  |q5_1  |q4_0  |q4_1  |

examples/main.cpp

@@ -1,4 +1,5 @@
 #include <stdio.h>
+#include <ctime>
 #include <fstream>
 #include <iostream>
 #include <random>
@@ -28,18 +29,17 @@
 #define TXT2IMG "txt2img"
 #define IMG2IMG "img2img"
 
-// get_num_physical_cores is copy from 
+// get_num_physical_cores is copy from
 // https://github.com/ggerganov/llama.cpp/blob/master/examples/common.cpp
 // LICENSE: https://github.com/ggerganov/llama.cpp/blob/master/LICENSE
 int32_t get_num_physical_cores() {
 #ifdef __linux__
     // enumerate the set of thread siblings, num entries is num cores
     std::unordered_set<std::string> siblings;
-    for (uint32_t cpu=0; cpu < UINT32_MAX; ++cpu) {
-        std::ifstream thread_siblings("/sys/devices/system/cpu"
-            + std::to_string(cpu) + "/topology/thread_siblings");
+    for (uint32_t cpu = 0; cpu < UINT32_MAX; ++cpu) {
+        std::ifstream thread_siblings("/sys/devices/system/cpu" + std::to_string(cpu) + "/topology/thread_siblings");
         if (!thread_siblings.is_open()) {
-            break; // no more cpus
+            break;  // no more cpus
         }
         std::string line;
         if (std::getline(thread_siblings, line)) {
@@ -61,7 +61,7 @@ int32_t get_num_physical_cores() {
         return num_physical_cores;
     }
 #elif defined(_WIN32)
-    //TODO: Implement
+    // TODO: Implement
 #endif
     unsigned int n_threads = std::thread::hardware_concurrency();
     return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
@@ -282,6 +282,11 @@ void parse_args(int argc, const char* argv[], Option* opt) {
         fprintf(stderr, "error: can only work with strength in [0.0, 1.0]\n");
         exit(1);
     }
+
+    if (opt->seed < 0) {
+        srand((int)time(NULL));
+        opt->seed = rand();
+    }
 }
 
 int main(int argc, const char* argv[]) {

stable-diffusion.cpp

@@ -26,12 +26,16 @@ static SDLogLevel log_level = SDLogLevel::INFO;
         }                                                                                             \
         if (level == SDLogLevel::DEBUG) {                                                             \
             printf("[DEBUG] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__);          \
+            fflush(stdout);                                                                           \
         } else if (level == SDLogLevel::INFO) {                                                       \
             printf("[INFO]  %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__);          \
+            fflush(stdout);                                                                           \
         } else if (level == SDLogLevel::WARN) {                                                       \
             fprintf(stderr, "[WARN]  %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \
+            fflush(stdout);                                                                           \
         } else if (level == SDLogLevel::ERROR) {                                                      \
             fprintf(stderr, "[ERROR] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \
+            fflush(stdout);                                                                           \
         }                                                                                             \
     } while (0)
 
@@ -135,6 +139,7 @@ float ggml_tensor_get_f32(const ggml_tensor* tensor, int l, int k = 0, int j = 0
 
 void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) {
     printf("shape(%zu, %zu, %zu, %zu)\n", tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+    fflush(stdout);
     if (shape_only) {
         return;
     }
@@ -156,6 +161,7 @@ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) {
                         continue;
                     }
                     printf("  [%d, %d, %d, %d] = %f\n", i, j, k, l, ggml_tensor_get_f32(tensor, l, k, j, i));
+                    fflush(stdout);
                 }
             }
         }
@@ -251,6 +257,11 @@ void image_vec_to_ggml(const std::vector<uint8_t>& vec,
     }
 }
 
+struct ggml_tensor * ggml_group_norm_32(struct ggml_context * ctx,
+                                        struct ggml_tensor * a) {
+    return ggml_group_norm(ctx, a, 32);
+}
+
 /*================================================== CLIPTokenizer ===================================================*/
 
 const std::string UNK_TOKEN = "<|endoftext|>";
@@ -899,7 +910,7 @@ struct ResBlock {
 
         // in_layers
         // group norm 32
-        auto h = ggml_group_norm(ctx, x);
+        auto h = ggml_group_norm_32(ctx, x);
         h = ggml_add(ctx,
                      ggml_mul(ctx,
                               ggml_repeat(ctx,
@@ -929,7 +940,7 @@ struct ResBlock {
         // out_layers
         h = ggml_add(ctx, h, emb_out);
         // group norm 32
-        h = ggml_group_norm_inplace(ctx, h);
+        h = ggml_group_norm_inplace(ctx, h, 32);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, out_layer_0_w, 1, 1, out_layer_0_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, out_layer_0_b, 1, 1, out_layer_0_b->ne[0], 1), h));
@@ -1122,7 +1133,7 @@ struct SpatialTransformer {
 
         auto x_in = x;
         // group norm 32
-        x = ggml_group_norm(ctx, x);
+        x = ggml_group_norm_32(ctx, x);
         x = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_w, 1, 1, norm_w->ne[0], 1), x), x),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_b, 1, 1, norm_b->ne[0], 1), x));
@@ -1424,7 +1435,7 @@ struct UpSample {
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
         // x: [N, channels, h, w]
-        x = ggml_upscale(ctx, x);  // [N, channels, h*2, w*2]
+        x = ggml_upscale(ctx, x, 2);  // [N, channels, h*2, w*2]
         x = ggml_conv_2d(ctx, conv_w, x, 1, 1, 1, 1, 1, 1);
 
         x = ggml_add(ctx,
@@ -1815,7 +1826,7 @@ struct UNetModel {
 
         // out
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx,
                               ggml_repeat(ctx,
@@ -1919,7 +1930,7 @@ struct ResnetBlock {
         // z: [N, in_channels, h, w]
 
         // group norm 32
-        auto h = ggml_group_norm(ctx, z);
+        auto h = ggml_group_norm_32(ctx, z);
         h = ggml_mul(ctx,
                      ggml_repeat(ctx,
                                  ggml_reshape_4d(ctx, norm1_w, 1, 1, norm1_w->ne[0], 1),
@@ -1941,7 +1952,7 @@ struct ResnetBlock {
                                  h));  // [N, out_channels, h, w]
 
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm2_w, 1, 1, norm2_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm2_b, 1, 1, norm2_b->ne[0], 1), h));
@@ -2028,7 +2039,7 @@ struct AttnBlock {
         // x: [N, in_channels, h, w]
 
         // group norm 32
-        auto h_ = ggml_group_norm(ctx, x);
+        auto h_ = ggml_group_norm_32(ctx, x);
         h_ = ggml_add(ctx,
                       ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_w, 1, 1, norm_w->ne[0], 1), h_), h_),
                       ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_b, 1, 1, norm_b->ne[0], 1), h_));
@@ -2253,7 +2264,7 @@ struct Encoder {
         h = mid.block_2.forward(ctx, h);  // [N, block_in, h, w]
 
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_w, 1, 1, norm_out_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_b, 1, 1, norm_out_b->ne[0], 1), h));
@@ -2435,7 +2446,7 @@ struct Decoder {
         }
 
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_w, 1, 1, norm_out_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_b, 1, 1, norm_out_b->ne[0], 1), h));
@@ -3256,6 +3267,10 @@ class StableDiffusionGGML {
                         }
                     }
                 }
+
+#ifdef GGML_PERF
+                ggml_graph_print(&diffusion_graph);
+#endif
                 int64_t t1 = ggml_time_ms();
                 LOG_INFO("step %d sampling completed, taking %.2fs", i + 1, (t1 - t0) * 1.0f / 1000);
                 LOG_DEBUG("diffusion graph use %.2fMB runtime memory: static %.2fMB, dynamic %.2fMB",
@@ -3345,6 +3360,10 @@ class StableDiffusionGGML {
             int64_t t0 = ggml_time_ms();
             ggml_graph_compute_with_ctx(ctx, &vae_graph, n_threads);
             int64_t t1 = ggml_time_ms();
+
+#ifdef GGML_PERF
+            ggml_graph_print(&vae_graph);
+#endif
             LOG_DEBUG("computing vae graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
 
             result = ggml_dup_tensor(res_ctx, moments);
@@ -3470,6 +3489,10 @@ class StableDiffusionGGML {
             int64_t t0 = ggml_time_ms();
             ggml_graph_compute_with_ctx(ctx, &vae_graph, n_threads);
             int64_t t1 = ggml_time_ms();
+
+#ifdef GGML_PERF
+            ggml_graph_print(&vae_graph);
+#endif
             LOG_DEBUG("computing vae graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
 
             result_img = ggml_dup_tensor(res_ctx, img);