feat: add tiny dream stable diffusion support (#1283)

Signed-off-by: Gianluca Boiano <morf3089@gmail.com>

pkg/model/initializers.go

@@ -40,6 +40,7 @@ const (
 	RwkvBackend            = "rwkv"
 	WhisperBackend         = "whisper"
 	StableDiffusionBackend = "stablediffusion"
+	TinyDreamBackend       = "tinydream"
 	PiperBackend           = "piper"
 	LCHuggingFaceBackend   = "langchain-huggingface"
 
@@ -64,6 +65,7 @@ var AutoLoadBackends []string = []string{
 	RwkvBackend,
 	WhisperBackend,
 	StableDiffusionBackend,
+	TinyDreamBackend,
 	PiperBackend,
 }
 

pkg/tinydream/generate.go

@@ -0,0 +1,36 @@
+//go:build tinydream
+// +build tinydream
+
+package tinydream
+
+import (
+	"fmt"
+	"path/filepath"
+
+	tinyDream "github.com/M0Rf30/go-tiny-dream"
+)
+
+func GenerateImage(height, width, step, seed int, positive_prompt, negative_prompt, dst, asset_dir string) error {
+	fmt.Println(dst)
+	if height > 512 || width > 512 {
+		return tinyDream.GenerateImage(
+			1,
+			step,
+			seed,
+			positive_prompt,
+			negative_prompt,
+			filepath.Dir(dst),
+			asset_dir,
+		)
+	}
+
+	return tinyDream.GenerateImage(
+		0,
+		step,
+		seed,
+		positive_prompt,
+		negative_prompt,
+		filepath.Dir(dst),
+		asset_dir,
+	)
+}

pkg/tinydream/generate_unsupported.go

@@ -0,0 +1,10 @@
+//go:build !tinydream
+// +build !tinydream
+
+package tinydream
+
+import "fmt"
+
+func GenerateImage(height, width, step, seed int, positive_prompt, negative_prompt, dst, asset_dir string) error {
+	return fmt.Errorf("This version of LocalAI was built without the tinytts tag")
+}

pkg/tinydream/tinydream.go

@@ -0,0 +1,20 @@
+package tinydream
+
+import "os"
+
+type TinyDream struct {
+	assetDir string
+}
+
+func New(assetDir string) (*TinyDream, error) {
+	if _, err := os.Stat(assetDir); err != nil {
+		return nil, err
+	}
+	return &TinyDream{
+		assetDir: assetDir,
+	}, nil
+}
+
+func (td *TinyDream) GenerateImage(height, width, step, seed int, positive_prompt, negative_prompt, dst string) error {
+	return GenerateImage(height, width, step, seed, positive_prompt, negative_prompt, dst, td.assetDir)
+}