leejet · leejet · Dec 28, 2023 · Dec 5, 2023 · Dec 8, 2023 · Dec 8, 2023
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -25,6 +25,7 @@ endif()
 #option(SD_BUILD_TESTS                "sd: build tests"    ${SD_STANDALONE})
 option(SD_BUILD_EXAMPLES             "sd: build examples" ${SD_STANDALONE})
 option(SD_CUBLAS                     "sd: cuda backend" OFF)
+option(SD_METAL                      "sd: metal backend" OFF)
 option(SD_FLASH_ATTN                 "sd: use flash attention for x4 less memory usage" OFF)
 option(BUILD_SHARED_LIBS             "sd: build shared libs" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
@@ -35,6 +36,12 @@ if(SD_CUBLAS)
     add_definitions(-DSD_USE_CUBLAS)
 endif()
 
+if(SD_METAL)
+	message("Use Metal as backend stable-diffusion")
+    set(GGML_METAL ON)
+    add_definitions(-DSD_USE_METAL)
+endif()
+
 if(SD_FLASH_ATTN)
     message("Use Flash Attention for memory optimization")
     add_definitions(-DSD_USE_FLASH_ATTENTION)

diff --git a/README.md b/README.md
@@ -27,6 +27,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
 - Latent Consistency Models support (LCM/LCM-LoRA)
 - Faster and memory efficient latent decoding with [TAESD](https://github.com/madebyollin/taesd)
+- Upscale images generated with [ESRGAN](https://github.com/xinntao/Real-ESRGAN)
 - Sampling method
     - `Euler A`
     - `Euler`
@@ -52,7 +53,8 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
     - Implement Winograd Convolution 2D for 3x3 kernel filtering
 - [ ] Continuing to reduce memory usage (quantizing the weights of ggml_conv_2d)
 - [ ] Implement BPE Tokenizer
-- [ ] Implement [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN/tree/master) upscaler
+- [ ] Implement Textual Inversion (embeddings)
+- [ ] Implement Inpainting support
 - [ ] k-quants support
 
 ## Usage
@@ -135,6 +137,7 @@ arguments:
   -m, --model [MODEL]                path to model
   --vae [VAE]                        path to vae
   --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
+  -um, --upscale-model [ESRGAN_PATH] path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.
   --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
                                      If not specified, the default is the type of the weight file.
   --lora-model-dir [DIR]             lora model directory

diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -59,6 +59,7 @@ struct SDParams {
     std::string model_path;
     std::string vae_path;
     std::string taesd_path;
+    std::string esrgan_path;
     ggml_type wtype = GGML_TYPE_COUNT;
     std::string lora_model_dir;
     std::string output_path = "output.png";
@@ -67,6 +68,7 @@ struct SDParams {
     std::string prompt;
     std::string negative_prompt;
     float cfg_scale = 7.0f;
+    int clip_skip_layers = 0;
     int width       = 512;
     int height      = 512;
     int batch_count = 1;
@@ -78,6 +80,7 @@ struct SDParams {
     RNGType rng_type           = CUDA_RNG;
     int64_t seed               = 42;
     bool verbose               = false;
+    bool vae_tiling            = false;
 };
 
 void print_params(SDParams params) {
@@ -115,6 +118,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  -m, --model [MODEL]                path to model\n");
     printf("  --vae [VAE]                        path to vae\n");
     printf("  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)\n");
+    printf("  -um, --upscale-model [ESRGAN_PATH] path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.\n");
     printf("  --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)\n");
     printf("                                     If not specified, the default is the type of the weight file.\n");
     printf("  --lora-model-dir [DIR]             lora model directory\n");
@@ -134,6 +138,8 @@ void print_usage(int argc, const char* argv[]) {
     printf("  -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)\n");
     printf("  -b, --batch-count COUNT            number of images to generate.\n");
     printf("  --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)\n");
+    printf("  -cs, --clip-skip N                 number of layers to skip of clip model (default: 0)\n");
+    printf("  -vt, --vae-tiling                  process vae in tiles to reduce memory usage\n");
     printf("  -v, --verbose                      print extra info\n");
 }
 
@@ -185,6 +191,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.taesd_path = argv[i];
+        } else if (arg == "--upscale-model" || arg == "-um") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.esrgan_path = argv[i];
         } else if (arg == "--type") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -270,6 +282,14 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.sample_steps = std::stoi(argv[i]);
+        } else if (arg == "-cs" || arg == "--clip-skip") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.clip_skip_layers = std::stoi(argv[i]);
+        } else if (arg == "-vt" || arg == "--vae-tiling") {
+            params.vae_tiling = true;
         } else if (arg == "-b" || arg == "--batch-count") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -458,9 +478,9 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    StableDiffusion sd(params.n_threads, vae_decode_only, params.taesd_path, true, params.lora_model_dir, params.rng_type);
+    StableDiffusion sd(params.n_threads, vae_decode_only, params.taesd_path, params.esrgan_path, true, params.vae_tiling, params.lora_model_dir, params.rng_type);
 
-    if (!sd.load_from_file(params.model_path, params.vae_path, params.wtype, params.schedule)) {
+    if (!sd.load_from_file(params.model_path, params.vae_path, params.wtype, params.schedule, params.clip_skip_layers)) {
         return 1;
     }
 
@@ -488,6 +508,19 @@ int main(int argc, const char* argv[]) {
                              params.seed);
     }
 
+    if(params.esrgan_path.size() > 0) {
+        // TODO: support more ESRGAN models, making it easier to set up ESRGAN models.
+        /*  hardcoded scale factor because just RealESRGAN_x4plus_anime_6B is compatible
+            See also: https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py
+
+            To avoid this, the upscaler needs to be separated from the stable diffusion pipeline.
+            However, a considerable amount of work would be required for this. It might be better
+            to opt for a complete project refactoring that facilitates the easier assignment of parameters.
+        */
+        params.width *= 4;
+        params.height *= 4;
+    }
+
     if (results.size() == 0 || results.size() != params.batch_count) {
         LOG_ERROR("generate failed");
         return 1;

diff --git a/ggml b/ggml
diff --git a/model.cpp b/model.cpp
@@ -14,6 +14,10 @@
 #include "ggml/ggml-backend.h"
 #include "ggml/ggml.h"
 
+#ifdef SD_USE_METAL
+#include "ggml-metal.h"
+#endif
+
 #define ST_HEADER_SIZE_LEN 8
 
 uint64_t read_u64(uint8_t* buffer) {
@@ -1208,7 +1212,7 @@ bool ModelLoader::load_vocab(on_new_token_cb_t on_new_token_cb) {
     return true;
 }
 
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb) {
+bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend_t backend) {
     bool success = true;
     for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
         std::string file_path = file_paths_[file_index];
@@ -1296,11 +1300,13 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb) {
                 continue;
             }
 
-            ggml_backend_t backend = ggml_get_backend(dst_tensor);
-
             size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 
-            if (backend == NULL || ggml_backend_is_cpu(backend)) {
+            if (dst_tensor->buffer == NULL || ggml_backend_is_cpu(backend)
+#ifdef SD_USE_METAL
+            || ggml_backend_is_metal(model.backend)
+#endif
+            ) {
                 // for the CPU and Metal backend, we can copy directly into the tensor
                 if (tensor_storage.type == dst_tensor->type) {
                     GGML_ASSERT(ggml_nbytes(dst_tensor) == tensor_storage.nbytes());

diff --git a/model.h b/model.h
@@ -116,7 +116,7 @@ class ModelLoader {
     SDVersion get_sd_version();
     ggml_type get_sd_wtype();
     bool load_vocab(on_new_token_cb_t on_new_token_cb);
-    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb);
+    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend_t backend);
     int64_t cal_mem_size(ggml_backend_t backend);
     ~ModelLoader() = default;
 };
+4 −2		.github/workflows/ci.yml
+1 −1		.gitignore
+9 −0		CMakeLists.txt
+1 −0		README.md
+13 −0		ci/run.sh
+56 −57		examples/gpt-2/main-backend.cpp
+5 −3		examples/gpt-2/main-batched.cpp
+5 −3		examples/gpt-2/main.cpp
+1 −1		examples/starcoder/starcoder-mmap.cpp
+2 −2		examples/whisper/main.cpp
+148 −81		examples/whisper/whisper.cpp
+6 −1		examples/whisper/whisper.h
+1 −1		examples/yolo/yolov3-tiny.cpp
+7 −0		include/ggml/ggml-alloc.h
+62 −17		include/ggml/ggml-backend.h
+69 −21		include/ggml/ggml.h
+1 −0		scripts/sync-llama.sh
+59 −3		src/CMakeLists.txt
+44 −38		src/ggml-alloc.c
+46 −21		src/ggml-backend-impl.h
+563 −156		src/ggml-backend.c
+1,448 −580		src/ggml-cuda.cu
+9 −1		src/ggml-cuda.h
+1 −1		src/ggml-impl.h
+6 −0		src/ggml-metal.h
+778 −178		src/ggml-metal.m
+1,118 −262		src/ggml-metal.metal
+5 −7		src/ggml-opencl.cpp
+1 −1		src/ggml-quants.c
+722 −574		src/ggml.c
+13 −2		tests/CMakeLists.txt
+0 −254		tests/test-add-mul.cpp
+84 −0		tests/test-backend-buffer.cpp
+1,474 −0		tests/test-backend-ops.cpp
+9 −3		tests/test-conv1d.cpp
+9 −4		tests/test-conv2d.cpp
+8 −1		tests/test-mul-mat.cpp