sync stable-diffusion.cpp version,fix bugs, update go mod version

README.md

@@ -21,11 +21,11 @@ go get github.com/seasonjs/stable-diffusion
 
 See `deps` folder for dylib compatibility, push request is welcome.
 
-| platform | x32         | x64                 | arm         |
-|----------|-------------|---------------------|-------------|
-| windows  | not support | support             | not support |
-| linux    | not support | not support         | not support |
-| darwin   | not support | support  (not test) | support     |
+| platform | x32         | x64                     | arm         |
+|----------|-------------|-------------------------|-------------|
+| windows  | not support | support avx/avx2/avx512 | not support |
+| linux    | not support | support                 | not support |
+| darwin   | not support | support  (no test)      | support     |
 
 ## Usage
 
@@ -101,7 +101,7 @@ func main()  {
 		panic(err)
 	}
 	defer outfile.Close()
-	err = model.ImagePredict(inFile, "pink cat", outfile)
+	err = model.ImagePredict(inFile, "the cat that wears shoes", outfile)
 }
 ```
 If `NewStableDiffusionAutoModel` can't automatic loading of dynamic library, please use `NewStableDiffusionModel` method load manually.
@@ -159,6 +159,12 @@ See detail at [stable-diffusion-doc](https://pkg.go.dev/github.com/seasonjs/stab
 * [ggml.cpp](https://github.com/leejet/ggml.cpp)
 * [purego](https://github.com/ebitengine/purego)
 
+## Successful Examples
+
+<p align="center">
+  <img src="./assets/love_cat2.png" width="512x">
+</p>
+
 ## License
 
 Copyright (c) seasonjs. All rights reserved.

binding.go

@@ -4,7 +4,10 @@
 package sd
 
 import (
+	"encoding/base64"
+	"errors"
 	"github.com/ebitengine/purego"
+	"runtime"
 	"unsafe"
 )
 
@@ -36,6 +39,7 @@ const (
 	DPMPP2S_A                     = "DPMPP2S_A"
 	DPMPP2M                       = "DPMPP2M"
 	DPMPP2Mv2                     = "DPMPP2Mv2"
+	LCM                           = "LCM"
 	N_SAMPLE_METHODS              = "N_SAMPLE_METHODS"
 )
 
@@ -75,6 +79,7 @@ const (
 
 	cSetStableDiffusionLogLevel   = "set_stable_diffusion_log_level"
 	cGetStableDiffusionSystemInfo = "get_stable_diffusion_system_info"
+	cFreeBuffer                   = "free_buffer"
 )
 
 type CStableDiffusion struct {
@@ -90,7 +95,7 @@ type CStableDiffusion struct {
 	cSetTxt2imgSeed           func(options uintptr, seed int64)
 
 	cNewSdImg2imgOptions      func() uintptr
-	cSetImg2imgInitImg        func(options uintptr, initImg uintptr, size int64)
+	cSetImg2imgInitImg        func(options uintptr, base64Str string)
 	cSetImg2imgPrompt         func(options uintptr, prompt string)
 	cSetImg2imgNegativePrompt func(options uintptr, negativePrompt string)
 	cSetImg2imgCfgScale       func(options uintptr, cfgScale float32)
@@ -100,14 +105,16 @@ type CStableDiffusion struct {
 	cSetImg2imgStrength       func(options uintptr, strength float32)
 	cSetImg2imgSeed           func(options uintptr, seed int64)
 
-	cCreateStableDiffusion  func(nThreads int, vaeDecodeOnly bool, freeParamsImmediately bool, rngType string) uintptr
+	cCreateStableDiffusion  func(nThreads int, vaeDecodeOnly bool, freeParamsImmediately bool, loraModelDir string, rngType string) uintptr
 	cDestroyStableDiffusion func(sd uintptr)
 	cLoadFromFile           func(sd uintptr, path string, schedule string)
-	cTxt2img                func(sd uintptr, options uintptr) *byte
-	cImg2img                func(sd uintptr, options uintptr) *byte
+	cTxt2img                func(sd uintptr, options uintptr, byteSize *int64) *byte
+	cImg2img                func(sd uintptr, options uintptr, byte2 *int64) *byte
 
 	cSetStableDiffusionLogLevel   func(level string)
 	cGetStableDiffusionSystemInfo func() string
+
+	cFreeBuffer func(buffer uintptr)
 }
 
 func NewCStableDiffusion(libraryPath string) (*CStableDiffusion, error) {
@@ -126,7 +133,7 @@ func NewCStableDiffusion(libraryPath string) (*CStableDiffusion, error) {
 		setTxt2imgSeed           func(options uintptr, seed int64)
 
 		newSdImg2imgOptions      func() uintptr
-		setImg2imgInitImg        func(options uintptr, initImg uintptr, size int64)
+		setImg2imgInitImg        func(options uintptr, base64Str string)
 		setImg2imgPrompt         func(options uintptr, prompt string)
 		setImg2imgNegativePrompt func(options uintptr, negativePrompt string)
 		setImg2imgCfgScale       func(options uintptr, cfgScale float32)
@@ -136,14 +143,16 @@ func NewCStableDiffusion(libraryPath string) (*CStableDiffusion, error) {
 		setImg2imgStrength       func(options uintptr, strength float32)
 		setImg2imgSeed           func(options uintptr, seed int64)
 
-		createStableDiffusion  func(nThreads int, vaeDecodeOnly bool, freeParamsImmediately bool, rngType string) uintptr
+		createStableDiffusion  func(nThreads int, vaeDecodeOnly bool, freeParamsImmediately bool, loraModelDir string, rngType string) uintptr
 		destroyStableDiffusion func(sd uintptr)
 		loadFromFile           func(sd uintptr, path string, schedule string)
-		txt2img                func(sd uintptr, options uintptr) *byte
-		img2img                func(sd uintptr, options uintptr) *byte
+		txt2img                func(sd uintptr, options uintptr, byteSize *int64) *byte
+		img2img                func(sd uintptr, options uintptr, byte2 *int64) *byte
 
 		setStableDiffusionLogLevel   func(level string)
 		getStableDiffusionSystemInfo func() string
+
+		freeBuffer func(buffer uintptr)
 	)
 	purego.RegisterLibFunc(&newSdTxt2imgOptions, libSd, cNewSdTxt2imgOptions)
 	purego.RegisterLibFunc(&setTxt2imgPrompt, libSd, cSetTxt2imgPrompt)
@@ -174,6 +183,8 @@ func NewCStableDiffusion(libraryPath string) (*CStableDiffusion, error) {
 	purego.RegisterLibFunc(&setStableDiffusionLogLevel, libSd, cSetStableDiffusionLogLevel)
 	purego.RegisterLibFunc(&getStableDiffusionSystemInfo, libSd, cGetStableDiffusionSystemInfo)
 
+	purego.RegisterLibFunc(&freeBuffer, libSd, cFreeBuffer)
+
 	return &CStableDiffusion{
 		libSd,
 
@@ -205,11 +216,12 @@ func NewCStableDiffusion(libraryPath string) (*CStableDiffusion, error) {
 
 		setStableDiffusionLogLevel,
 		getStableDiffusionSystemInfo,
+		freeBuffer,
 	}, nil
 }
 
-func (cSD *CStableDiffusion) NewStableDiffusionCtx(nThreads int, vaeDecodeOnly bool, freeParamsImmediately bool, rngType RNGType) *CSDCtx {
-	ctx := cSD.cCreateStableDiffusion(nThreads, vaeDecodeOnly, freeParamsImmediately, string(rngType))
+func (cSD *CStableDiffusion) NewStableDiffusionCtx(nThreads int, vaeDecodeOnly bool, freeParamsImmediately bool, loraModelDir string, rngType RNGType) *CSDCtx {
+	ctx := cSD.cCreateStableDiffusion(nThreads, vaeDecodeOnly, freeParamsImmediately, loraModelDir, string(rngType))
 	return &CSDCtx{ctx: ctx, csd: cSD}
 }
 
@@ -230,7 +242,13 @@ func (c *CSDCtx) StableDiffusionLoadFromFile(path string, schedule Schedule) {
 	c.csd.cLoadFromFile(c.ctx, path, string(schedule))
 }
 
-func (c *CSDCtx) StableDiffusionTextToImage(prompt string, negativePrompt string, cfgScale float32, width int, height int, sampleMethod SampleMethod, sampleSteps int, seed int64) []byte {
+func (c *CSDCtx) StableDiffusionTextToImage(prompt string, negativePrompt string, cfgScale float32, width int, height int, sampleMethod SampleMethod, sampleSteps int, seed int64) ([]byte, error) {
+	if width <= 0 {
+		return nil, errors.New("width must be greater than 0")
+	}
+	if height <= 0 {
+		return nil, errors.New("height must be greater than 0")
+	}
 	options := c.csd.cNewSdImg2imgOptions()
 	c.csd.cSetTxt2imgPrompt(options, prompt)
 	c.csd.cSetTxt2imgNegativePrompt(options, negativePrompt)
@@ -239,16 +257,28 @@ func (c *CSDCtx) StableDiffusionTextToImage(prompt string, negativePrompt string
 	c.csd.cSetTxt2imgSampleMethod(options, string(sampleMethod))
 	c.csd.cSetTxt2imgSampleSteps(options, sampleSteps)
 	c.csd.cSetTxt2imgSeed(options, seed)
-	output := c.csd.cTxt2img(c.ctx, options)
-	//runtime.KeepAlive(output)
-	size := width * height * 3 //width*height*rgb
+	minSize := int64(width * height * 3)
+	var size int64
+	runtime.KeepAlive(size)
+	output := c.csd.cTxt2img(c.ctx, options, &size)
+	if size < minSize {
+		size = minSize
+	}
 	data := unsafe.Slice(output, size)
-	return data
+	size = 0
+	c.csd.cFreeBuffer(uintptr(unsafe.Pointer(output)))
+	return data, nil
 }
 
-func (c *CSDCtx) StableDiffusionImageToImage(initImg []byte, prompt string, negativePrompt string, cfgScale float32, width int, height int, sampleMethod SampleMethod, sampleSteps int, strength float32, seed int64) []byte {
+func (c *CSDCtx) StableDiffusionImageToImage(initImg []byte, prompt string, negativePrompt string, cfgScale float32, width int, height int, sampleMethod SampleMethod, sampleSteps int, strength float32, seed int64) ([]byte, error) {
+	if width <= 0 || width%64 != 0 {
+		return nil, errors.New("width must be greater than 0 and must be a multiple of 64")
+	}
+	if height <= 0 || height%64 != 0 {
+		return nil, errors.New("height must be greater than 0 and must be a multiple of 64")
+	}
 	options := c.csd.cNewSdImg2imgOptions()
-	c.csd.cSetImg2imgInitImg(options, uintptr(unsafe.Pointer(&initImg[0])), int64(len(initImg)))
+	c.csd.cSetImg2imgInitImg(options, base64.StdEncoding.EncodeToString(initImg))
 	c.csd.cSetImg2imgPrompt(options, prompt)
 	c.csd.cSetImg2imgNegativePrompt(options, negativePrompt)
 	c.csd.cSetImg2imgCfgScale(options, cfgScale)
@@ -257,11 +287,17 @@ func (c *CSDCtx) StableDiffusionImageToImage(initImg []byte, prompt string, nega
 	c.csd.cSetImg2imgSampleSteps(options, sampleSteps)
 	c.csd.cSetImg2imgStrength(options, strength)
 	c.csd.cSetImg2imgSeed(options, seed)
-	output := c.csd.cImg2img(c.ctx, options)
-	//runtime.KeepAlive(output)
-	size := width * height * 3 //width*height*rgb
+	minSize := int64(width * height * 3)
+	var size int64
+	runtime.KeepAlive(size)
+	output := c.csd.cImg2img(c.ctx, options, &size)
+	if size < minSize {
+		size = minSize
+	}
 	data := unsafe.Slice(output, size)
-	return data
+	size = 0
+	c.csd.cFreeBuffer(uintptr(unsafe.Pointer(output)))
+	return data, nil
 }
 
 func (c *CSDCtx) Close() {

binding_test.go

@@ -4,6 +4,7 @@
 package sd
 
 import (
+	"encoding/base64"
 	"fmt"
 	"image"
 	"image/color"
@@ -29,7 +30,7 @@ func getLibrary() string {
 // int n_threads = -1;
 // std::string mode = TXT2IMG;
 // std::string model_path;
-// std::string output_path = "output.png";
+// std::string output_path = "love_cat2.png";
 // std::string init_img;
 // std::string prompt;
 // std::string negative_prompt;
@@ -105,22 +106,28 @@ func TestStableDiffusionTextToImage(t *testing.T) {
 	if err != nil {
 		t.Log(err)
 	}
-	ctx := sd.NewStableDiffusionCtx(8, true, true, CUDA_RNG)
+	ctx := sd.NewStableDiffusionCtx(8, true, true, "", CUDA_RNG)
 	defer ctx.Close()
 	ctx.StableDiffusionLoadFromFile("./models/miniSD-ggml-model-q5_0.bin", DEFAULT)
-	data := ctx.StableDiffusionTextToImage("A lovely cat, high quality", "", 7.0, 256, 256, EULER_A, 20, 42)
-	writeToFile(t, data, 256, 256, "./data/output.png")
+	data, _ := ctx.StableDiffusionTextToImage("A lovely cat, high quality", "", 7.0, 256, 256, EULER_A, 20, 42)
+	writeToFile(t, data, 256, 256, "./data/love_cat2.png")
 }
 
 func TestStableDiffusionImgToImage(t *testing.T) {
 	sd, err := NewCStableDiffusion(getLibrary())
 	if err != nil {
 		t.Log(err)
 	}
-	ctx := sd.NewStableDiffusionCtx(8, false, true, CUDA_RNG)
+	ctx := sd.NewStableDiffusionCtx(8, false, true, "", CUDA_RNG)
 	defer ctx.Close()
 	ctx.StableDiffusionLoadFromFile("./models/miniSD-ggml-model-q5_0.bin", DEFAULT)
-	img := readFromFile(t, "./data/output.png")
-	data := ctx.StableDiffusionImageToImage(img, "A lovely cat that theme pink", "", 7.0, 256, 256, EULER_A, 20, 0.4, 42)
+	img := readFromFile(t, "./data/love_cat2.png")
+	data, _ := ctx.StableDiffusionImageToImage(img, "A lovely cat that theme pink", "", 7.0, 256, 256, EULER_A, 20, 0.4, 42)
 	writeToFile(t, data, 256, 256, "./data/output1.png")
 }
+
+func TestBase64(t *testing.T) {
+	img := readFromFile(t, "./assets/love_cat2.png")
+	imgBase64 := base64.StdEncoding.EncodeToString(img)
+	t.Log(imgBase64)
+}

deps/stable-diffusion-abi.h

@@ -0,0 +1,96 @@
+#ifndef STABLE_DIFFUSION_CPP_STABLE_DIFFUSION_ABI_H
+#define STABLE_DIFFUSION_CPP_STABLE_DIFFUSION_ABI_H
+
+
+#include <cstdint>
+
+#ifdef STABLE_DIFFUSION_SHARED
+#if defined(_WIN32) && !defined(__MINGW32__)
+#ifdef STABLE_DIFFUSION_BUILD
+#define STABLE_DIFFUSION_API __declspec(dllexport)
+#else
+#define STABLE_DIFFUSION_API __declspec(dllimport)
+#endif
+#else
+#define STABLE_DIFFUSION_API __attribute__((visibility("default")))
+#endif
+#else
+#define STABLE_DIFFUSION_API
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+struct sd_txt2img_options;
+
+struct sd_img2img_options;
+
+// These methods are used in binding in other languages,golang, python,etc.
+
+//==============================sd_txt2img_options===============================
+STABLE_DIFFUSION_API sd_txt2img_options* new_sd_txt2img_options();
+
+STABLE_DIFFUSION_API void set_txt2img_prompt(sd_txt2img_options* opt, const char* prompt);
+
+STABLE_DIFFUSION_API void set_txt2img_negative_prompt(sd_txt2img_options* opt, const char* negative_prompt);
+
+STABLE_DIFFUSION_API void set_txt2img_cfg_scale(sd_txt2img_options* opt, float cfg_scale);
+
+STABLE_DIFFUSION_API void set_txt2img_size(sd_txt2img_options* opt, int width, int height);
+
+STABLE_DIFFUSION_API void set_txt2img_sample_method(sd_txt2img_options* opt, const char* sample_method);
+
+STABLE_DIFFUSION_API void set_txt2img_sample_steps(sd_txt2img_options* opt, int sample_steps);
+
+STABLE_DIFFUSION_API void set_txt2img_seed(sd_txt2img_options* opt, int64_t seed);
+//================================================================================
+
+//==============================sd_img2img_options===============================
+STABLE_DIFFUSION_API sd_img2img_options* new_sd_img2img_options();
+
+STABLE_DIFFUSION_API void set_img2img_init_img(sd_img2img_options* opt, const char* init_img);
+
+STABLE_DIFFUSION_API void set_img2img_prompt(sd_img2img_options* opt, const char* prompt);
+
+STABLE_DIFFUSION_API void set_img2img_negative_prompt(sd_img2img_options* opt, const char* negative_prompt);
+
+STABLE_DIFFUSION_API void set_img2img_cfg_scale(sd_img2img_options* opt, float cfg_scale);
+
+STABLE_DIFFUSION_API void set_img2img_size(sd_img2img_options* opt, int width, int height);
+
+STABLE_DIFFUSION_API void set_img2img_sample_method(sd_img2img_options* opt, const char* sample_method);
+
+STABLE_DIFFUSION_API void set_img2img_sample_steps(sd_img2img_options* opt, int sample_steps);
+
+STABLE_DIFFUSION_API void set_img2img_strength(sd_img2img_options* opt, float strength);
+
+STABLE_DIFFUSION_API void set_img2img_seed(sd_img2img_options* opt, int64_t seed);
+//================================================================================
+
+STABLE_DIFFUSION_API void* create_stable_diffusion(
+    int n_threads,
+    bool vae_decode_only,
+    bool free_params_immediately,
+    const char* lora_model_dir,
+    const char* rng_type);
+
+STABLE_DIFFUSION_API void destroy_stable_diffusion(void* sd);
+
+STABLE_DIFFUSION_API bool load_from_file(void* sd, const char* file_path, const char* schedule);
+
+STABLE_DIFFUSION_API uint8_t* txt2img(void* sd, const sd_txt2img_options* opt, int64_t* output_size);
+
+STABLE_DIFFUSION_API uint8_t* img2img(void* sd, const sd_img2img_options* opt, int64_t* output_size);
+
+STABLE_DIFFUSION_API void set_stable_diffusion_log_level(const char* level);
+
+STABLE_DIFFUSION_API const char* get_stable_diffusion_system_info();
+
+STABLE_DIFFUSION_API void free_buffer(const uint8_t* buffer);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif //STABLE_DIFFUSION_CPP_STABLE_DIFFUSION_C_H