feat: implement ESRGAN upscaler + Metal Backend (#104)

* add esrgan upscaler

* add sd_tiling

* support metal backend

* add clip_skip

---------

Co-authored-by: leejet <leejet714@gmail.com>

.clang-format

@@ -7,7 +7,7 @@ IndentCaseLabels: false
 ColumnLimit: 0
 AccessModifierOffset: -4
 NamespaceIndentation: All
-FixNamespaceComments: false 
+FixNamespaceComments: false
 AlignAfterOpenBracket: true
 AlignConsecutiveAssignments: true
 IndentCaseLabels: true

.gitignore

@@ -8,6 +8,7 @@ test/
 *.bin
 *.exe
 *.gguf
-*.log
-output.png
-models/
+output*.png
+models*
+!taesd-model.gguf
+*.log

CMakeLists.txt

@@ -25,7 +25,9 @@ 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(SD_FAST_SOFTMAX               "sd: x1.5 faster softmax, indeterministic (sometimes, same seed don't generate same image), cuda only" OFF)
 option(BUILD_SHARED_LIBS             "sd: build shared libs" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
@@ -33,6 +35,15 @@ if(SD_CUBLAS)
 	message("Use CUBLAS as backend stable-diffusion")
     set(GGML_CUBLAS ON)
     add_definitions(-DSD_USE_CUBLAS)
+    if(SD_FAST_SOFTMAX)
+        set(GGML_CUDA_FAST_SOFTMAX ON)
+    endif()
+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)

README.md

@@ -9,22 +9,26 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 ## Features
 
 - Plain C/C++ implementation based on [ggml](https://github.com/ggerganov/ggml), working in the same way as [llama.cpp](https://github.com/ggerganov/llama.cpp)
-- Super lightweight and without external dependencies.
+- Super lightweight and without external dependencies
 - SD1.x and SD2.x support
+- [SD-Turbo](https://huggingface.co/stabilityai/sd-turbo) support
 - 16-bit, 32-bit float support
 - 4-bit, 5-bit and 8-bit integer quantization support
 - Accelerated memory-efficient CPU inference
     - Only requires ~2.3GB when using txt2img with fp16 precision to generate a 512x512 image, enabling Flash Attention just requires ~1.8GB.
 - AVX, AVX2 and AVX512 support for x86 architectures
-- Full CUDA backend for GPU acceleration, for now just for float16 and float32 models. There are some issues with quantized models and CUDA; it will be fixed in the future.
-- Can load ckpt, safetensors and diffusers models/checkpoints. Standalone VAEs models.
+- Full CUDA and Metal backend for GPU acceleration.
+- Can load ckpt, safetensors and diffusers models/checkpoints. Standalone VAEs models
     - No need to convert to `.ggml` or `.gguf` anymore!
-- Flash Attention for memory usage optimization (only cpu for now).
+- Flash Attention for memory usage optimization (only cpu for now)
 - Original `txt2img` and `img2img` mode
 - Negative prompt
 - [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) style tokenizer (not all the features, only token weighting for now)
 - 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)
+- VAE tiling processing for reduce memory usage
 - Sampling method
     - `Euler A`
     - `Euler`
@@ -47,9 +51,10 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - [ ] More sampling methods
 - [ ] Make inference faster
     - The current implementation of ggml_conv_2d is slow and has high memory usage
+    - Implement Winograd Convolution 2D for 3x3 kernel filtering
 - [ ] Continuing to reduce memory usage (quantizing the weights of ggml_conv_2d)
-- [ ] Implement BPE Tokenizer
-- [ ] Add [TAESD](https://github.com/madebyollin/taesd) for faster VAE decoding
+- [ ] Implement Textual Inversion (embeddings)
+- [ ] Implement Inpainting support
 - [ ] k-quants support
 
 ## Usage
@@ -80,7 +85,7 @@ git submodule update
     ```shell
     curl -L -O https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt
     # curl -L -O https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
-    # curl -L -O https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-nonema-pruned.safetensors
+    # curl -L -O https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-nonema-pruned.safetensors
     ```
 
 ### Build
@@ -110,6 +115,15 @@ cmake .. -DSD_CUBLAS=ON
 cmake --build . --config Release
 ```
 
+##### Using Metal
+
+Using Metal makes the computation run on the GPU. Currently, there are some issues with Metal when performing operations on very large matrices, making it highly inefficient at the moment. Performance improvements are expected in the near future.
+
+```
+cmake .. -DSD_METAL=ON
+cmake --build . --config Release
+```
+
 ### Using Flash Attention
 
 Enabling flash attention reduces memory usage by at least 400 MB. At the moment, it is not supported when CUBLAS is enabled because the kernel implementation is missing.
@@ -131,8 +145,11 @@ arguments:
                                      If threads <= 0, then threads will be set to the number of CPU physical cores
   -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)
+  --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  
+                                     If not specified, the default is the type of the weight file.
+  --lora-model-dir [DIR]             lora model directory
   -i, --init-img [IMAGE]             path to the input image, required by img2img
   -o, --output OUTPUT                path to write result image to (default: ./output.png)
   -p, --prompt [PROMPT]              the prompt to render
@@ -149,6 +166,8 @@ arguments:
   -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)
   -b, --batch-count COUNT            number of images to generate.
   --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)
+  --clip-skip N                      number of layers to skip of clip model (default: 0)
+  --vae-tiling                       process vae in tiles to reduce memory usage
   -v, --verbose                      print extra info
 ```
 
@@ -218,6 +237,33 @@ Here's a simple example:
 | ----  |----    |
 | ![](./assets/without_lcm.png) |![](./assets/with_lcm.png)  |
 
+## Using TAESD to faster decoding
+
+You can use TAESD to accelerate the decoding of latent images by following these steps:
+
+- Download the model [weights](https://huggingface.co/madebyollin/taesd/blob/main/diffusion_pytorch_model.safetensors).
+
+Or curl
+
+```bash
+curl -L -O https://huggingface.co/madebyollin/taesd/blob/main/diffusion_pytorch_model.safetensors
+```
+
+- Specify the model path using the `--taesd PATH` parameter. example:
+
+```bash
+sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --taesd ../models/diffusion_pytorch_model.safetensors
+```
+
+## Using ESRGAN to upscale results
+
+You can use ESRGAN to upscale the generated images. At the moment, only the [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth) model is supported. Support for more models of this architecture will be added soon.
+
+- Specify the model path using the `--upscale-model PATH` parameter. example:
+
+```bash
+sd -m ../models/v1-5-pruned-emaonly.safetensors -p "a lovely cat" --upscale-model ../models/RealESRGAN_x4plus_anime_6B.pth
+```
 
 ### Docker
 

common/zip.h

@@ -1,509 +0,0 @@
-/*
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
- * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
- * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
- * IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- */
-
-#pragma once
-#ifndef ZIP_H
-#define ZIP_H
-
-#include <stdint.h>
-#include <string.h>
-#include <sys/types.h>
-
-#ifndef ZIP_SHARED
-#define ZIP_EXPORT
-#else
-#ifdef _WIN32
-#ifdef ZIP_BUILD_SHARED
-#define ZIP_EXPORT __declspec(dllexport)
-#else
-#define ZIP_EXPORT __declspec(dllimport)
-#endif
-#else
-#define ZIP_EXPORT __attribute__((visibility("default")))
-#endif
-#endif
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-#if !defined(_POSIX_C_SOURCE) && defined(_MSC_VER)
-// 64-bit Windows is the only mainstream platform
-// where sizeof(long) != sizeof(void*)
-#ifdef _WIN64
-typedef long long ssize_t; /* byte count or error */
-#else
-typedef long ssize_t; /* byte count or error */
-#endif
-#endif
-
-/**
- * @mainpage
- *
- * Documentation for @ref zip.
- */
-
-/**
- * @addtogroup zip
- * @{
- */
-
-/**
- * Default zip compression level.
- */
-#define ZIP_DEFAULT_COMPRESSION_LEVEL 6
-
-/**
- * Error codes
- */
-#define ZIP_ENOINIT -1      // not initialized
-#define ZIP_EINVENTNAME -2  // invalid entry name
-#define ZIP_ENOENT -3       // entry not found
-#define ZIP_EINVMODE -4     // invalid zip mode
-#define ZIP_EINVLVL -5      // invalid compression level
-#define ZIP_ENOSUP64 -6     // no zip 64 support
-#define ZIP_EMEMSET -7      // memset error
-#define ZIP_EWRTENT -8      // cannot write data to entry
-#define ZIP_ETDEFLINIT -9   // cannot initialize tdefl compressor
-#define ZIP_EINVIDX -10     // invalid index
-#define ZIP_ENOHDR -11      // header not found
-#define ZIP_ETDEFLBUF -12   // cannot flush tdefl buffer
-#define ZIP_ECRTHDR -13     // cannot create entry header
-#define ZIP_EWRTHDR -14     // cannot write entry header
-#define ZIP_EWRTDIR -15     // cannot write to central dir
-#define ZIP_EOPNFILE -16    // cannot open file
-#define ZIP_EINVENTTYPE -17 // invalid entry type
-#define ZIP_EMEMNOALLOC -18 // extracting data using no memory allocation
-#define ZIP_ENOFILE -19     // file not found
-#define ZIP_ENOPERM -20     // no permission
-#define ZIP_EOOMEM -21      // out of memory
-#define ZIP_EINVZIPNAME -22 // invalid zip archive name
-#define ZIP_EMKDIR -23      // make dir error
-#define ZIP_ESYMLINK -24    // symlink error
-#define ZIP_ECLSZIP -25     // close archive error
-#define ZIP_ECAPSIZE -26    // capacity size too small
-#define ZIP_EFSEEK -27      // fseek error
-#define ZIP_EFREAD -28      // fread error
-#define ZIP_EFWRITE -29     // fwrite error
-#define ZIP_ERINIT -30      // cannot initialize reader
-#define ZIP_EWINIT -31      // cannot initialize writer
-#define ZIP_EWRINIT -32     // cannot initialize writer from reader
-
-/**
- * Looks up the error message string corresponding to an error number.
- * @param errnum error number
- * @return error message string corresponding to errnum or NULL if error is not
- * found.
- */
-extern ZIP_EXPORT const char *zip_strerror(int errnum);
-
-/**
- * @struct zip_t
- *
- * This data structure is used throughout the library to represent zip archive -
- * forward declaration.
- */
-struct zip_t;
-
-/**
- * Opens zip archive with compression level using the given mode.
- *
- * @param zipname zip archive file name.
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *zip_open(const char *zipname, int level,
-                                         char mode);
-
-/**
- * Opens zip archive with compression level using the given mode.
- * The function additionally returns @param errnum -
- *
- * @param zipname zip archive file name.
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- * @param errnum 0 on success, negative number (< 0) on error.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *
-zip_openwitherror(const char *zipname, int level, char mode, int *errnum);
-
-/**
- * Closes the zip archive, releases resources - always finalize.
- *
- * @param zip zip archive handler.
- */
-extern ZIP_EXPORT void zip_close(struct zip_t *zip);
-
-/**
- * Determines if the archive has a zip64 end of central directory headers.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - 1 (true), 0 (false), negative number (< 0) on
- *         error.
- */
-extern ZIP_EXPORT int zip_is64(struct zip_t *zip);
-
-/**
- * Opens an entry by name in the zip archive.
- *
- * For zip archive opened in 'w' or 'a' mode the function will append
- * a new entry. In readonly mode the function tries to locate the entry
- * in global dictionary.
- *
- * @param zip zip archive handler.
- * @param entryname an entry name in local dictionary.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_open(struct zip_t *zip, const char *entryname);
-
-/**
- * Opens an entry by name in the zip archive.
- *
- * For zip archive opened in 'w' or 'a' mode the function will append
- * a new entry. In readonly mode the function tries to locate the entry
- * in global dictionary (case sensitive).
- *
- * @param zip zip archive handler.
- * @param entryname an entry name in local dictionary (case sensitive).
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_opencasesensitive(struct zip_t *zip,
-                                                  const char *entryname);
-
-/**
- * Opens a new entry by index in the zip archive.
- *
- * This function is only valid if zip archive was opened in 'r' (readonly) mode.
- *
- * @param zip zip archive handler.
- * @param index index in local dictionary.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_openbyindex(struct zip_t *zip, size_t index);
-
-/**
- * Closes a zip entry, flushes buffer and releases resources.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_close(struct zip_t *zip);
-
-/**
- * Returns a local name of the current zip entry.
- *
- * The main difference between user's entry name and local entry name
- * is optional relative path.
- * Following .ZIP File Format Specification - the path stored MUST not contain
- * a drive or device letter, or a leading slash.
- * All slashes MUST be forward slashes '/' as opposed to backwards slashes '\'
- * for compatibility with Amiga and UNIX file systems etc.
- *
- * @param zip: zip archive handler.
- *
- * @return the pointer to the current zip entry name, or NULL on error.
- */
-extern ZIP_EXPORT const char *zip_entry_name(struct zip_t *zip);
-
-/**
- * Returns an index of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the index on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entry_index(struct zip_t *zip);
-
-/**
- * Determines if the current zip entry is a directory entry.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - 1 (true), 0 (false), negative number (< 0) on
- *         error.
- */
-extern ZIP_EXPORT int zip_entry_isdir(struct zip_t *zip);
-
-/**
- * Returns the uncompressed size of the current zip entry.
- * Alias for zip_entry_uncomp_size (for backward compatibility).
- *
- * @param zip zip archive handler.
- *
- * @return the uncompressed size in bytes.
- */
-extern ZIP_EXPORT unsigned long long zip_entry_size(struct zip_t *zip);
-
-/**
- * Returns the uncompressed size of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the uncompressed size in bytes.
- */
-extern ZIP_EXPORT unsigned long long zip_entry_uncomp_size(struct zip_t *zip);
-
-/**
- * Returns the compressed size of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the compressed size in bytes.
- */
-extern ZIP_EXPORT unsigned long long zip_entry_comp_size(struct zip_t *zip);
-
-/**
- * Returns CRC-32 checksum of the current zip entry.
- *
- * @param zip zip archive handler.
- *
- * @return the CRC-32 checksum.
- */
-extern ZIP_EXPORT unsigned int zip_entry_crc32(struct zip_t *zip);
-
-/**
- * Compresses an input buffer for the current zip entry.
- *
- * @param zip zip archive handler.
- * @param buf input buffer.
- * @param bufsize input buffer size (in bytes).
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_write(struct zip_t *zip, const void *buf,
-                                      size_t bufsize);
-
-/**
- * Compresses a file for the current zip entry.
- *
- * @param zip zip archive handler.
- * @param filename input file.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_fwrite(struct zip_t *zip, const char *filename);
-
-/**
- * Extracts the current zip entry into output buffer.
- *
- * The function allocates sufficient memory for a output buffer.
- *
- * @param zip zip archive handler.
- * @param buf output buffer.
- * @param bufsize output buffer size (in bytes).
- *
- * @note remember to release memory allocated for a output buffer.
- *       for large entries, please take a look at zip_entry_extract function.
- *
- * @return the return code - the number of bytes actually read on success.
- *         Otherwise a negative number (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entry_read(struct zip_t *zip, void **buf,
-                                         size_t *bufsize);
-
-/**
- * Extracts the current zip entry into a memory buffer using no memory
- * allocation.
- *
- * @param zip zip archive handler.
- * @param buf preallocated output buffer.
- * @param bufsize output buffer size (in bytes).
- *
- * @note ensure supplied output buffer is large enough.
- *       zip_entry_size function (returns uncompressed size for the current
- *       entry) can be handy to estimate how big buffer is needed.
- *       For large entries, please take a look at zip_entry_extract function.
- *
- * @return the return code - the number of bytes actually read on success.
- *         Otherwise a negative number (< 0) on error (e.g. bufsize is not large
- * enough).
- */
-extern ZIP_EXPORT ssize_t zip_entry_noallocread(struct zip_t *zip, void *buf,
-                                                size_t bufsize);
-
-/**
- * Extracts the current zip entry into output file.
- *
- * @param zip zip archive handler.
- * @param filename output file.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_entry_fread(struct zip_t *zip, const char *filename);
-
-/**
- * Extracts the current zip entry using a callback function (on_extract).
- *
- * @param zip zip archive handler.
- * @param on_extract callback function.
- * @param arg opaque pointer (optional argument, which you can pass to the
- *        on_extract callback)
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int
-zip_entry_extract(struct zip_t *zip,
-                  size_t (*on_extract)(void *arg, uint64_t offset,
-                                       const void *data, size_t size),
-                  void *arg);
-
-/**
- * Returns the number of all entries (files and directories) in the zip archive.
- *
- * @param zip zip archive handler.
- *
- * @return the return code - the number of entries on success, negative number
- *         (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entries_total(struct zip_t *zip);
-
-/**
- * Deletes zip archive entries.
- *
- * @param zip zip archive handler.
- * @param entries array of zip archive entries to be deleted.
- * @param len the number of entries to be deleted.
- * @return the number of deleted entries, or negative number (< 0) on error.
- */
-extern ZIP_EXPORT ssize_t zip_entries_delete(struct zip_t *zip,
-                                             char *const entries[], size_t len);
-
-/**
- * Extracts a zip archive stream into directory.
- *
- * If on_extract is not NULL, the callback will be called after
- * successfully extracted each zip entry.
- * Returning a negative value from the callback will cause abort and return an
- * error. The last argument (void *arg) is optional, which you can use to pass
- * data to the on_extract callback.
- *
- * @param stream zip archive stream.
- * @param size stream size.
- * @param dir output directory.
- * @param on_extract on extract callback.
- * @param arg opaque pointer.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int
-zip_stream_extract(const char *stream, size_t size, const char *dir,
-                   int (*on_extract)(const char *filename, void *arg),
-                   void *arg);
-
-/**
- * Opens zip archive stream into memory.
- *
- * @param stream zip archive stream.
- * @param size stream size.
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *zip_stream_open(const char *stream, size_t size,
-                                                int level, char mode);
-
-/**
- * Opens zip archive stream into memory.
- * The function additionally returns @param errnum -
- *
- * @param stream zip archive stream.
- * @param size stream size.*
- * @param level compression level (0-9 are the standard zlib-style levels).
- * @param mode file access mode.
- *        - 'r': opens a file for reading/extracting (the file must exists).
- *        - 'w': creates an empty file for writing.
- *        - 'a': appends to an existing archive.
- * @param errnum 0 on success, negative number (< 0) on error.
- *
- * @return the zip archive handler or NULL on error
- */
-extern ZIP_EXPORT struct zip_t *zip_stream_openwitherror(const char *stream,
-                                                         size_t size, int level,
-                                                         char mode,
-                                                         int *errnum);
-
-/**
- * Copy zip archive stream output buffer.
- *
- * @param zip zip archive handler.
- * @param buf output buffer. User should free buf.
- * @param bufsize output buffer size (in bytes).
- *
- * @return copy size
- */
-extern ZIP_EXPORT ssize_t zip_stream_copy(struct zip_t *zip, void **buf,
-                                          size_t *bufsize);
-
-/**
- * Close zip archive releases resources.
- *
- * @param zip zip archive handler.
- *
- * @return
- */
-extern ZIP_EXPORT void zip_stream_close(struct zip_t *zip);
-
-/**
- * Creates a new archive and puts files into a single zip archive.
- *
- * @param zipname zip archive file.
- * @param filenames input files.
- * @param len: number of input files.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_create(const char *zipname, const char *filenames[],
-                                 size_t len);
-
-/**
- * Extracts a zip archive file into directory.
- *
- * If on_extract_entry is not NULL, the callback will be called after
- * successfully extracted each zip entry.
- * Returning a negative value from the callback will cause abort and return an
- * error. The last argument (void *arg) is optional, which you can use to pass
- * data to the on_extract_entry callback.
- *
- * @param zipname zip archive file.
- * @param dir output directory.
- * @param on_extract_entry on extract callback.
- * @param arg opaque pointer.
- *
- * @return the return code - 0 on success, negative number (< 0) on error.
- */
-extern ZIP_EXPORT int zip_extract(const char *zipname, const char *dir,
-                                  int (*on_extract_entry)(const char *filename,
-                                                          void *arg),
-                                  void *arg);
-/** @} */
-#ifdef __cplusplus
-}
-#endif
-
-#endif

examples/cli/main.cpp

@@ -58,6 +58,8 @@ 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";
@@ -66,6 +68,7 @@ struct SDParams {
     std::string prompt;
     std::string negative_prompt;
     float cfg_scale = 7.0f;
+    int clip_skip   = -1;  // <= 0 represents unspecified
     int width       = 512;
     int height      = 512;
     int batch_count = 1;
@@ -77,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) {
@@ -86,11 +90,14 @@ void print_params(SDParams params) {
     printf("    model_path:        %s\n", params.model_path.c_str());
     printf("    wtype:             %s\n", params.wtype < GGML_TYPE_COUNT ? ggml_type_name(params.wtype) : "unspecified");
     printf("    vae_path:          %s\n", params.vae_path.c_str());
+    printf("    taesd_path:        %s\n", params.taesd_path.c_str());
+    printf("    esrgan_path:       %s\n", params.esrgan_path.c_str());
     printf("    output_path:       %s\n", params.output_path.c_str());
     printf("    init_img:          %s\n", params.input_path.c_str());
     printf("    prompt:            %s\n", params.prompt.c_str());
     printf("    negative_prompt:   %s\n", params.negative_prompt.c_str());
     printf("    cfg_scale:         %.2f\n", params.cfg_scale);
+    printf("    clip_skip:         %d\n", params.clip_skip);
     printf("    width:             %d\n", params.width);
     printf("    height:            %d\n", params.height);
     printf("    sample_method:     %s\n", sample_method_str[params.sample_method]);
@@ -100,6 +107,7 @@ void print_params(SDParams params) {
     printf("    rng:               %s\n", rng_type_to_str[params.rng_type]);
     printf("    seed:              %ld\n", params.seed);
     printf("    batch_count:       %d\n", params.batch_count);
+    printf("    vae_tiling:        %s\n", params.vae_tiling ? "true" : "false");
 }
 
 void print_usage(int argc, const char* argv[]) {
@@ -112,8 +120,10 @@ void print_usage(int argc, const char* argv[]) {
     printf("                                     If threads <= 0, then threads will be set to the number of CPU physical cores\n");
     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("  --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.");
+    printf("                                     If not specified, the default is the type of the weight file.\n");
     printf("  --lora-model-dir [DIR]             lora model directory\n");
     printf("  -i, --init-img [IMAGE]             path to the input image, required by img2img\n");
     printf("  -o, --output OUTPUT                path to write result image to (default: ./output.png)\n");
@@ -131,6 +141,9 @@ 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("  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n");
+    printf("                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x\n");
+    printf("  --vae-tiling                       process vae in tiles to reduce memory usage\n");
     printf("  -v, --verbose                      print extra info\n");
 }
 
@@ -176,6 +189,18 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.vae_path = argv[i];
+        } else if (arg == "--taesd") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.taesd_path = argv[i];
+        } else if (arg == "--upscale-model") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.esrgan_path = argv[i];
         } else if (arg == "--type") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -261,6 +286,14 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.sample_steps = std::stoi(argv[i]);
+        } else if (arg == "--clip-skip") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.clip_skip = std::stoi(argv[i]);
+        } else if (arg == "--vae-tiling") {
+            params.vae_tiling = true;
         } else if (arg == "-b" || arg == "--batch-count") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -449,8 +482,9 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    StableDiffusion sd(params.n_threads, vae_decode_only, true, params.lora_model_dir, params.rng_type);
-    if (!sd.load_from_file(params.model_path, params.vae_path, params.wtype, params.schedule)) {
+    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, params.clip_skip)) {
         return 1;
     }
 
@@ -478,6 +512,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;

model.cpp

@@ -14,6 +14,12 @@
 #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) {
     // little endian
     uint64_t value = 0;
@@ -52,7 +58,6 @@ std::string self_attn_names[] = {
     "self_attn.q_proj.weight",
     "self_attn.k_proj.weight",
     "self_attn.v_proj.weight",
-
     "self_attn.q_proj.bias",
     "self_attn.k_proj.bias",
     "self_attn.v_proj.bias",
@@ -73,6 +78,8 @@ const char* unused_tensors[] = {
     "cond_stage_model.transformer.text_model.embeddings.position_ids",
     "cond_stage_model.model.logit_scale",
     "cond_stage_model.model.text_projection",
+    "conditioner.embedders.0.model.logit_scale",
+    "conditioner.embedders.0.model.text_projection",
     "model.diffusion_model.time_embedding.cond_proj.weight",
     "unet.time_embedding.cond_proj.weight",
     "model_ema.decay",
@@ -80,6 +87,7 @@ const char* unused_tensors[] = {
     "model_ema.diffusion_model",
     "control_model",
     "embedding_manager",
+    "denoiser.sigmas",
 };
 
 bool is_unused_tensor(std::string name) {
@@ -124,16 +132,19 @@ std::unordered_map<std::string, std::string> vae_decoder_name_map = {
 };
 
 std::string convert_open_clip_to_hf_clip(const std::string& name) {
-    std::string new_name                  = name;
+    std::string new_name = name;
+    if (starts_with(new_name, "conditioner.embedders.0.")) {
+        new_name = "cond_stage_model." + new_name.substr(strlen("conditioner.embedders.0."));
+    }
     std::string open_clip_resblock_prefix = "cond_stage_model.model.transformer.resblocks.";
     std::string hf_clip_resblock_prefix   = "cond_stage_model.transformer.text_model.encoder.layers.";
 
-    if (open_clip_to_hf_clip_model.find(name) != open_clip_to_hf_clip_model.end()) {
-        new_name = open_clip_to_hf_clip_model[name];
+    if (open_clip_to_hf_clip_model.find(new_name) != open_clip_to_hf_clip_model.end()) {
+        new_name = open_clip_to_hf_clip_model[new_name];
     }
 
-    if (name.find(open_clip_resblock_prefix) == 0) {
-        std::string remain = name.substr(open_clip_resblock_prefix.length());
+    if (new_name.find(open_clip_resblock_prefix) == 0) {
+        std::string remain = new_name.substr(open_clip_resblock_prefix.length());
         std::string idx    = remain.substr(0, remain.find("."));
         std::string suffix = remain.substr(idx.length() + 1);
 
@@ -347,7 +358,7 @@ std::string convert_diffusers_name_to_compvis(const std::string& key, char seq)
 
 std::string convert_tensor_name(const std::string& name) {
     std::string new_name;
-    if (starts_with(name, "cond_stage_model.model")) {
+    if (starts_with(name, "cond_stage_model.model") || starts_with(name, "conditioner.embedders.0.model")) {
         new_name = convert_open_clip_to_hf_clip(name);
     } else if (starts_with(name, "first_stage_model.decoder")) {
         new_name = convert_vae_decoder_name(name);
@@ -533,233 +544,102 @@ std::map<char, int> unicode_to_byte() {
     return byte_decoder;
 }
 
+bool is_zip_file(const std::string& file_path) {
+    struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
+    if (zip == NULL) {
+        return false;
+    }
+    zip_close(zip);
+    return true;
+}
+
+bool is_gguf_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    char magic[4];
+
+    file.read(magic, sizeof(magic));
+    if (!file) {
+        return false;
+    }
+    for (uint32_t i = 0; i < sizeof(magic); i++) {
+        if (magic[i] != GGUF_MAGIC[i]) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool is_safetensors_file(const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    // get file size
+    file.seekg(0, file.end);
+    size_t file_size_ = file.tellg();
+    file.seekg(0, file.beg);
+
+    // read header size
+    if (file_size_ <= ST_HEADER_SIZE_LEN) {
+        return false;
+    }
+
+    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
+    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
+    if (!file) {
+        return false;
+    }
+
+    size_t header_size_ = read_u64(header_size_buf);
+    if (header_size_ >= file_size_) {
+        return false;
+    }
+
+    // read header
+    std::vector<char> header_buf;
+    header_buf.resize(header_size_ + 1);
+    header_buf[header_size_] = '\0';
+    file.read(header_buf.data(), header_size_);
+    if (!file) {
+        return false;
+    }
+    nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
+    if (header_.is_discarded()) {
+        return false;
+    }
+    return true;
+}
+
 bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    file_paths_.push_back(file_path);
-    return true;
-}
-
-bool ModelLoader::init_from_files(const std::vector<std::string>& file_paths) {
-    for (auto& file_path : file_paths) {
-        if (!init_from_file(file_path)) {
-            return false;
-        }
+    if (is_directory(file_path)) {
+        LOG_INFO("load %s using diffusers format", file_path.c_str());
+        return init_from_diffusers_file(file_path, prefix);
+    } else if (is_gguf_file(file_path)) {
+        LOG_INFO("load %s using gguf format", file_path.c_str());
+        return init_from_gguf_file(file_path, prefix);
+    } else if (is_safetensors_file(file_path)) {
+        LOG_INFO("load %s using safetensors format", file_path.c_str());
+        return init_from_safetensors_file(file_path, prefix);
+    } else if (is_zip_file(file_path)) {
+        LOG_INFO("load %s using checkpoint format", file_path.c_str());
+        return init_from_ckpt_file(file_path, prefix);
+    } else {
+        LOG_WARN("unknown format %s", file_path.c_str());
+        return false;
     }
-    return true;
-}
-
-SDVersion ModelLoader::get_sd_version() {
-    TensorStorage token_embedding_weight;
-    for (auto& tensor_storage : tensor_storages) {
-        if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
-            tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
-            tensor_storage.name == "text_model.embeddings.token_embedding.weight" ||
-            tensor_storage.name == "te.text_model.embeddings.token_embedding.weight") {
-            token_embedding_weight = tensor_storage;
-            break;
-        }
-    }
-    if (token_embedding_weight.ne[0] == 768) {
-        return VERSION_1_x;
-    } else if (token_embedding_weight.ne[0] == 1024) {
-        return VERSION_2_x;
-    }
-    return VERSION_COUNT;
-}
-
-ggml_type ModelLoader::get_sd_wtype() {
-    for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
-        if (tensor_storage.name.find(".weight") != std::string::npos &&
-            tensor_storage.name.find("time_embed") != std::string::npos) {
-            return tensor_storage.type;
-        }
-    }
-    return GGML_TYPE_COUNT;
-}
-
-bool ModelLoader::load_vocab(on_new_token_cb_t on_new_token_cb) {
-    char* vocab_buffer          = reinterpret_cast<char*>(vocab_json);
-    nlohmann::json vocab        = nlohmann::json::parse(vocab_buffer);
-    std::map<char, int> decoder = unicode_to_byte();
-    for (auto& it : vocab.items()) {
-        int token_id          = it.value();
-        std::string token_str = it.key();
-        std::string token     = "";
-        for (char c : token_str) {
-            token += decoder[c];
-        }
-        on_new_token_cb(token, token_id);
-    }
-    return true;
-}
-
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb) {
-    bool success = true;
-    for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
-        std::string file_path = file_paths_[file_index];
-        LOG_DEBUG("loading tensors from %s", file_path.c_str());
-
-        std::ifstream file(file_path, std::ios::binary);
-        if (!file.is_open()) {
-            LOG_ERROR("failed to open '%s'", file_path.c_str());
-            return false;
-        }
-
-        bool is_zip = false;
-        for (auto& tensor_storage : tensor_storages) {
-            if (tensor_storage.index_in_zip >= 0) {
-                is_zip = true;
-                break;
-            }
-        }
-
-        struct zip_t* zip = NULL;
-        if (is_zip) {
-            zip = zip_open(file_path.c_str(), 0, 'r');
-            if (zip == NULL) {
-                LOG_ERROR("failed to open zip '%s'", file_path.c_str());
-                return false;
-            }
-        }
-
-        std::vector<uint8_t> read_buffer;
-        std::vector<uint8_t> convert_buffer;
-
-        auto read_data = [&](const TensorStorage& tensor_storage, char* buf, size_t n) {
-            if (zip != NULL) {
-                zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
-                size_t entry_size = zip_entry_size(zip);
-                if (entry_size != n) {
-                    read_buffer.resize(entry_size);
-                    zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
-                    memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
-                } else {
-                    zip_entry_noallocread(zip, (void*)buf, n);
-                }
-                zip_entry_close(zip);
-            } else {
-                file.seekg(tensor_storage.offset);
-                file.read(buf, n);
-                if (!file) {
-                    LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                    return false;
-                }
-            }
-            return true;
-        };
-
-        std::vector<TensorStorage> processed_tensor_storages;
-        for (auto& tensor_storage : tensor_storages) {
-            if (tensor_storage.file_index != file_index) {
-                continue;
-            }
-
-            // LOG_DEBUG("%s", name.c_str());
-
-            if (is_unused_tensor(tensor_storage.name)) {
-                continue;
-            }
-
-            preprocess_tensor(tensor_storage, processed_tensor_storages);
-        }
-
-        for (auto& tensor_storage : processed_tensor_storages) {
-            // LOG_DEBUG("%s", name.c_str());
-
-            ggml_tensor* dst_tensor = NULL;
-
-            success = on_new_tensor_cb(tensor_storage, &dst_tensor);
-            if (!success) {
-                LOG_WARN("process tensor failed: '%s'", tensor_storage.name.c_str());
-                break;
-            }
-
-            if (dst_tensor == NULL) {
-                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)) {
-                // 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) == nbytes_to_read);
-                    read_data(tensor_storage, (char*)dst_tensor->data, nbytes_to_read);
-
-                    if (tensor_storage.is_bf16) {
-                        // inplace op
-                        bf16_to_f32_vec((uint16_t*)dst_tensor->data, (float*)dst_tensor->data, tensor_storage.nelements());
-                    }
-                } else {
-                    read_buffer.resize(tensor_storage.nbytes());
-                    read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
-
-                    if (tensor_storage.is_bf16) {
-                        // inplace op
-                        bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                    }
-
-                    convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
-                                   dst_tensor->type, (int)tensor_storage.nelements());
-                }
-            } else {
-                read_buffer.resize(tensor_storage.nbytes());
-                read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
-
-                if (tensor_storage.is_bf16) {
-                    // inplace op
-                    bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
-                }
-
-                if (tensor_storage.type == dst_tensor->type) {
-                    // copy to device memory
-                    ggml_backend_tensor_set(dst_tensor, read_buffer.data(), 0, ggml_nbytes(dst_tensor));
-                } else {
-                    // convert first, then copy to device memory
-                    convert_buffer.resize(ggml_nbytes(dst_tensor));
-                    convert_tensor((void*)read_buffer.data(), tensor_storage.type,
-                                   (void*)convert_buffer.data(), dst_tensor->type,
-                                   (int)tensor_storage.nelements());
-                    ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
-                }
-            }
-        }
-
-        if (zip != NULL) {
-            zip_close(zip);
-        }
-
-        if (!success) {
-            break;
-        }
-    }
-    return success;
-}
-
-int64_t ModelLoader::cal_mem_size() {
-    int64_t mem_size = 0;
-    for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
-        mem_size += tensor_storage.nbytes();
-        mem_size += GGML_MEM_ALIGN * 2;  // for lora alphas
-    }
-
-    return mem_size + 10 * 1024 * 1024;
 }
 
 /*================================================= GGUFModelLoader ==================================================*/
 
-bool GGUFModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    LOG_INFO("loading model from '%s'", file_path.c_str());
-    ModelLoader::init_from_file(file_path, prefix);
+bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from '%s'", file_path.c_str());
+    file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
 
     gguf_context* ctx_gguf_ = NULL;
@@ -796,8 +676,6 @@ bool GGUFModelLoader::init_from_file(const std::string& file_path, const std::st
 
 /*================================================= SafeTensorsModelLoader ==================================================*/
 
-#define ST_HEADER_SIZE_LEN 8
-
 ggml_type str_to_ggml_type(const std::string& dtype) {
     ggml_type ttype = GGML_TYPE_COUNT;
     if (dtype == "F16") {
@@ -811,8 +689,9 @@ ggml_type str_to_ggml_type(const std::string& dtype) {
 }
 
 // https://huggingface.co/docs/safetensors/index
-bool SafeTensorsModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    ModelLoader::init_from_file(file_path, prefix);
+bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from '%s'", file_path.c_str());
+    file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
@@ -913,21 +792,18 @@ bool SafeTensorsModelLoader::init_from_file(const std::string& file_path, const
 
 /*================================================= DiffusersModelLoader ==================================================*/
 
-bool DiffusersModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    if (!is_directory(file_path)) {
-        return SafeTensorsModelLoader::init_from_file(file_path, prefix);
-    }
+bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const std::string& prefix) {
     std::string unet_path = path_join(file_path, "unet/diffusion_pytorch_model.safetensors");
     std::string vae_path  = path_join(file_path, "vae/diffusion_pytorch_model.safetensors");
     std::string clip_path = path_join(file_path, "text_encoder/model.safetensors");
 
-    if (!SafeTensorsModelLoader::init_from_file(unet_path, "unet.")) {
+    if (!init_from_safetensors_file(unet_path, "unet.")) {
         return false;
     }
-    if (!SafeTensorsModelLoader::init_from_file(vae_path, "vae.")) {
+    if (!init_from_safetensors_file(vae_path, "vae.")) {
         return false;
     }
-    if (!SafeTensorsModelLoader::init_from_file(clip_path, "te.")) {
+    if (!init_from_safetensors_file(clip_path, "te.")) {
         return false;
     }
     return true;
@@ -1127,12 +1003,12 @@ int find_char(uint8_t* buffer, int len, char c) {
 
 #define MAX_STRING_BUFFER 512
 
-bool CkptModelLoader::parse_data_pkl(uint8_t* buffer,
-                                     size_t buffer_size,
-                                     zip_t* zip,
-                                     std::string dir,
-                                     size_t file_index,
-                                     const std::string& prefix) {
+bool ModelLoader::parse_data_pkl(uint8_t* buffer,
+                                 size_t buffer_size,
+                                 zip_t* zip,
+                                 std::string dir,
+                                 size_t file_index,
+                                 const std::string& prefix) {
     uint8_t* buffer_end = buffer + buffer_size;
     if (buffer[0] == 0x80) {  // proto
         if (buffer[1] != 2) {
@@ -1235,6 +1111,7 @@ bool CkptModelLoader::parse_data_pkl(uint8_t* buffer,
                         reader.tensor_storage.file_index = file_index;
                         reader.tensor_storage.name       = prefix + reader.tensor_storage.name;
                         tensor_storages.push_back(reader.tensor_storage);
+                        // LOG_DEBUG("%s", reader.tensor_storage.name.c_str());
                         // reset
                         reader = PickleTensorReader();
                     }
@@ -1250,8 +1127,9 @@ bool CkptModelLoader::parse_data_pkl(uint8_t* buffer,
     return true;
 }
 
-bool CkptModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
-    ModelLoader::init_from_file(file_path, prefix);
+bool ModelLoader::init_from_ckpt_file(const std::string& file_path, const std::string& prefix) {
+    LOG_DEBUG("init from '%s'", file_path.c_str());
+    file_paths_.push_back(file_path);
     size_t file_index = file_paths_.size() - 1;
 
     struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
@@ -1271,7 +1149,7 @@ bool CkptModelLoader::init_from_file(const std::string& file_path, const std::st
                 size_t pkl_size;
                 zip_entry_read(zip, &pkl_data, &pkl_size);
 
-                LOG_DEBUG("%lld", pkl_size);
+                // LOG_DEBUG("%lld", pkl_size);
 
                 parse_data_pkl((uint8_t*)pkl_data, pkl_size, zip, dir, file_index, prefix);
 
@@ -1284,29 +1162,212 @@ bool CkptModelLoader::init_from_file(const std::string& file_path, const std::st
     return true;
 }
 
-/*================================================= init_model_loader_from_file ==================================================*/
+SDVersion ModelLoader::get_sd_version() {
+    TensorStorage token_embedding_weight;
+    for (auto& tensor_storage : tensor_storages) {
+        if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
+            tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
+            tensor_storage.name == "text_model.embeddings.token_embedding.weight" ||
+            tensor_storage.name == "te.text_model.embeddings.token_embedding.weight" ||
+            tensor_storage.name == "conditioner.embedders.0.model.token_embedding.weight") {
+            token_embedding_weight = tensor_storage;
+            break;
+        }
+    }
+    if (token_embedding_weight.ne[0] == 768) {
+        return VERSION_1_x;
+    } else if (token_embedding_weight.ne[0] == 1024) {
+        return VERSION_2_x;
+    }
+    return VERSION_COUNT;
+}
 
-ModelLoader* init_model_loader_from_file(const std::string& file_path) {
-    ModelLoader* model_loader = NULL;
-    if (is_directory(file_path)) {
-        LOG_DEBUG("load %s using diffusers format", file_path.c_str());
-        model_loader = new DiffusersModelLoader();
-    } else if (ends_with(file_path, ".gguf")) {
-        LOG_DEBUG("load %s using gguf format", file_path.c_str());
-        model_loader = new GGUFModelLoader();
-    } else if (ends_with(file_path, ".safetensors")) {
-        LOG_DEBUG("load %s using safetensors format", file_path.c_str());
-        model_loader = new SafeTensorsModelLoader();
-    } else if (ends_with(file_path, ".ckpt")) {
-        LOG_DEBUG("load %s using checkpoint format", file_path.c_str());
-        model_loader = new CkptModelLoader();
-    } else {
-        LOG_DEBUG("unknown format %s", file_path.c_str());
-        return NULL;
+ggml_type ModelLoader::get_sd_wtype() {
+    for (auto& tensor_storage : tensor_storages) {
+        if (is_unused_tensor(tensor_storage.name)) {
+            continue;
+        }
+
+        if (tensor_storage.name.find(".weight") != std::string::npos &&
+            tensor_storage.name.find("time_embed") != std::string::npos) {
+            return tensor_storage.type;
+        }
     }
-    if (!model_loader->init_from_file(file_path)) {
-        delete model_loader;
-        model_loader = NULL;
+    return GGML_TYPE_COUNT;
+}
+
+std::string ModelLoader::load_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(merges_utf8_c_str), sizeof(merges_utf8_c_str));
+    return merges_utf8_str;
+}
+
+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];
+        LOG_DEBUG("loading tensors from %s", file_path.c_str());
+
+        std::ifstream file(file_path, std::ios::binary);
+        if (!file.is_open()) {
+            LOG_ERROR("failed to open '%s'", file_path.c_str());
+            return false;
+        }
+
+        bool is_zip = false;
+        for (auto& tensor_storage : tensor_storages) {
+            if (tensor_storage.file_index != file_index) {
+                continue;
+            }
+            if (tensor_storage.index_in_zip >= 0) {
+                is_zip = true;
+                break;
+            }
+        }
+
+        struct zip_t* zip = NULL;
+        if (is_zip) {
+            zip = zip_open(file_path.c_str(), 0, 'r');
+            if (zip == NULL) {
+                LOG_ERROR("failed to open zip '%s'", file_path.c_str());
+                return false;
+            }
+        }
+
+        std::vector<uint8_t> read_buffer;
+        std::vector<uint8_t> convert_buffer;
+
+        auto read_data = [&](const TensorStorage& tensor_storage, char* buf, size_t n) {
+            if (zip != NULL) {
+                zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
+                size_t entry_size = zip_entry_size(zip);
+                if (entry_size != n) {
+                    read_buffer.resize(entry_size);
+                    zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
+                    memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
+                } else {
+                    zip_entry_noallocread(zip, (void*)buf, n);
+                }
+                zip_entry_close(zip);
+            } else {
+                file.seekg(tensor_storage.offset);
+                file.read(buf, n);
+                if (!file) {
+                    LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
+                    return false;
+                }
+            }
+            return true;
+        };
+
+        std::vector<TensorStorage> processed_tensor_storages;
+        for (auto& tensor_storage : tensor_storages) {
+            if (tensor_storage.file_index != file_index) {
+                continue;
+            }
+
+            // LOG_DEBUG("%s", name.c_str());
+
+            if (is_unused_tensor(tensor_storage.name)) {
+                continue;
+            }
+
+            preprocess_tensor(tensor_storage, processed_tensor_storages);
+        }
+
+        for (auto& tensor_storage : processed_tensor_storages) {
+            // LOG_DEBUG("%s", name.c_str());
+
+            ggml_tensor* dst_tensor = NULL;
+
+            success = on_new_tensor_cb(tensor_storage, &dst_tensor);
+            if (!success) {
+                LOG_WARN("process tensor failed: '%s'", tensor_storage.name.c_str());
+                break;
+            }
+
+            if (dst_tensor == NULL) {
+                continue;
+            }
+
+            size_t nbytes_to_read = tensor_storage.nbytes_to_read();
+
+            if (dst_tensor->buffer == NULL || ggml_backend_is_cpu(backend)
+#ifdef SD_USE_METAL
+                || ggml_backend_is_metal(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());
+                    read_data(tensor_storage, (char*)dst_tensor->data, nbytes_to_read);
+
+                    if (tensor_storage.is_bf16) {
+                        // inplace op
+                        bf16_to_f32_vec((uint16_t*)dst_tensor->data, (float*)dst_tensor->data, tensor_storage.nelements());
+                    }
+                } else {
+                    read_buffer.resize(tensor_storage.nbytes());
+                    read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
+
+                    if (tensor_storage.is_bf16) {
+                        // inplace op
+                        bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                    }
+
+                    convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
+                                   dst_tensor->type, (int)tensor_storage.nelements());
+                }
+            } else {
+                read_buffer.resize(tensor_storage.nbytes());
+                read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
+
+                if (tensor_storage.is_bf16) {
+                    // inplace op
+                    bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
+                }
+
+                if (tensor_storage.type == dst_tensor->type) {
+                    // copy to device memory
+                    ggml_backend_tensor_set(dst_tensor, read_buffer.data(), 0, ggml_nbytes(dst_tensor));
+                } else {
+                    // convert first, then copy to device memory
+                    convert_buffer.resize(ggml_nbytes(dst_tensor));
+                    convert_tensor((void*)read_buffer.data(), tensor_storage.type,
+                                   (void*)convert_buffer.data(), dst_tensor->type,
+                                   (int)tensor_storage.nelements());
+                    ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
+                }
+            }
+        }
+
+        if (zip != NULL) {
+            zip_close(zip);
+        }
+
+        if (!success) {
+            break;
+        }
     }
-    return model_loader;
-}
+    return success;
+}
+
+int64_t ModelLoader::cal_mem_size(ggml_backend_t backend) {
+    size_t alignment = 128;
+    if (backend != NULL) {
+        alignment = ggml_backend_get_alignment(backend);
+    }
+    int64_t mem_size = 0;
+    std::vector<TensorStorage> processed_tensor_storages;
+    for (auto& tensor_storage : tensor_storages) {
+        if (is_unused_tensor(tensor_storage.name)) {
+            continue;
+        }
+        preprocess_tensor(tensor_storage, processed_tensor_storages);
+    }
+
+    for (auto& tensor_storage : processed_tensor_storages) {
+        mem_size += tensor_storage.nbytes() + alignment;
+    }
+
+    return mem_size;
+}

model.h

@@ -7,6 +7,7 @@
 #include <string>
 #include <vector>
 
+#include "ggml/ggml-backend.h"
 #include "ggml/ggml.h"
 #include "json.hpp"
 #include "zip.h"
@@ -98,29 +99,6 @@ protected:
     std::vector<std::string> file_paths_;
     std::vector<TensorStorage> tensor_storages;
 
-public:
-    virtual bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-    virtual bool init_from_files(const std::vector<std::string>& file_paths);
-    virtual SDVersion get_sd_version();
-    virtual ggml_type get_sd_wtype();
-    virtual bool load_vocab(on_new_token_cb_t on_new_token_cb);
-    virtual bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb);
-    virtual int64_t cal_mem_size();
-    virtual ~ModelLoader() = default;
-};
-
-class GGUFModelLoader : public ModelLoader {
-public:
-    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-};
-
-class SafeTensorsModelLoader : public ModelLoader {
-public:
-    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-};
-
-class CkptModelLoader : public ModelLoader {
-private:
     bool parse_data_pkl(uint8_t* buffer,
                         size_t buffer_size,
                         zip_t* zip,
@@ -128,15 +106,18 @@ private:
                         size_t file_index,
                         const std::string& prefix);
 
+    bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_ckpt_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
+
 public:
     bool init_from_file(const std::string& file_path, const std::string& prefix = "");
+    SDVersion get_sd_version();
+    ggml_type get_sd_wtype();
+    std::string load_merges();
+    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;
 };
-
-class DiffusersModelLoader : public SafeTensorsModelLoader {
-public:
-    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-};
-
-ModelLoader* init_model_loader_from_file(const std::string& file_path);
-
 #endif  // __MODEL_H__

stable-diffusion.h

@@ -1,9 +1,12 @@
 #ifndef __STABLE_DIFFUSION_H__
 #define __STABLE_DIFFUSION_H__
 
-#include <string>
 #include <memory>
+#include <string>
 #include <vector>
+#include "ggml/ggml.h"
+
+#include "ggml/ggml.h"
 
 enum RNGType {
     STD_DEFAULT_RNG,
@@ -38,13 +41,19 @@ private:
 public:
     StableDiffusion(int n_threads                = -1,
                     bool vae_decode_only         = false,
+                    std::string taesd_path       = "",
+                    std::string esrgan_path      = "",
                     bool free_params_immediately = false,
+                    bool vae_tiling              = false,
                     std::string lora_model_dir   = "",
                     RNGType rng_type             = STD_DEFAULT_RNG);
+
     bool load_from_file(const std::string& model_path,
                         const std::string& vae_path,
                         ggml_type wtype,
-                        Schedule d = DEFAULT);
+                        Schedule d = DEFAULT,
+                        int clip_skip = -1);
+
     std::vector<uint8_t*> txt2img(
         std::string prompt,
         std::string negative_prompt,

util.cpp

@@ -1,7 +1,9 @@
 #include "util.h"
 
 #include <stdarg.h>
+#include <codecvt>
 #include <fstream>
+#include <locale>
 #include <thread>
 #include <unordered_set>
 #include <vector>
@@ -119,6 +121,21 @@ int32_t get_num_physical_cores() {
     return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
 }
 
+std::u32string utf8_to_utf32(const std::string& utf8_str) {
+    std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
+    return converter.from_bytes(utf8_str);
+}
+
+std::string utf32_to_utf8(const std::u32string& utf32_str) {
+    std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
+    return converter.to_bytes(utf32_str);
+}
+
+std::u32string unicode_value_to_utf32(int unicode_value) {
+    std::u32string utf32_string = {static_cast<char32_t>(unicode_value)};
+    return utf32_string;
+}
+
 std::string basename(const std::string& path) {
     size_t pos = path.find_last_of('/');
     if (pos != std::string::npos) {

util.h

@@ -2,6 +2,7 @@
 #define __UTIL_H__
 
 #include <string>
+#include <cstdint>
 
 bool ends_with(const std::string& str, const std::string& ending);
 bool starts_with(const std::string& str, const std::string& start);
@@ -13,6 +14,10 @@ void replace_all_chars(std::string& str, char target, char replacement);
 bool file_exists(const std::string& filename);
 bool is_directory(const std::string& path);
 
+std::u32string utf8_to_utf32(const std::string& utf8_str);
+std::string utf32_to_utf8(const std::u32string& utf32_str);
+std::u32string unicode_value_to_utf32(int unicode_value);
+
 std::string basename(const std::string& path);
 
 std::string path_join(const std::string& p1, const std::string& p2);
@@ -34,4 +39,4 @@ void log_printf(SDLogLevel level, const char* file, int line, const char* format
 #define LOG_INFO(format, ...) log_printf(SDLogLevel::INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SDLogLevel::WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_ERROR(format, ...) log_printf(SDLogLevel::ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
-#endif  // __UTIL_H__
+#endif  // __UTIL_H__