From c7b5907adbe9e92e557c7b75f24ac4c812bc145f Mon Sep 17 00:00:00 2001 From: M0Rf30 Date: Mon, 13 Nov 2023 02:39:55 +0100 Subject: [PATCH] feat: add tiny dream stable diffusion support Signed-off-by: Gianluca Boiano --- .github/workflows/test.yml | 2 +- .gitignore | 1 + Dockerfile | 2 +- Makefile | 30 +++++++++++++++++++++- api/api_test.go | 37 +++++++++++++++++++++++++++ api/openai/image.go | 15 +++++++++-- cmd/grpc/tinydream/main.go | 23 +++++++++++++++++ go.mod | 3 ++- go.sum | 2 ++ pkg/backend/image/tinydream.go | 32 +++++++++++++++++++++++ pkg/model/initializers.go | 2 ++ pkg/tinydream/generate.go | 36 ++++++++++++++++++++++++++ pkg/tinydream/generate_unsupported.go | 10 ++++++++ pkg/tinydream/tinydream.go | 20 +++++++++++++++ 14 files changed, 209 insertions(+), 6 deletions(-) create mode 100644 cmd/grpc/tinydream/main.go create mode 100644 pkg/backend/image/tinydream.go create mode 100644 pkg/tinydream/generate.go create mode 100644 pkg/tinydream/generate_unsupported.go create mode 100644 pkg/tinydream/tinydream.go diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2ece8a57be5e..3dc4999f46c7 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -94,7 +94,7 @@ jobs: ../.. && sudo make -j12 install - name: Test run: | - GO_TAGS="stablediffusion tts" make test + GO_TAGS="stablediffusion tinydream tts" make test tests-apple: runs-on: macOS-latest diff --git a/.gitignore b/.gitignore index ed4f26439b78..31ba062eb98a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ go-llama go-llama-stable /gpt4all go-stable-diffusion +go-tiny-dream go-piper /go-bert go-ggllm diff --git a/Dockerfile b/Dockerfile index 2091e2e3ccb1..14975442eecf 100644 --- a/Dockerfile +++ b/Dockerfile @@ -14,7 +14,7 @@ ARG TARGETVARIANT ENV BUILD_TYPE=${BUILD_TYPE} ENV EXTERNAL_GRPC_BACKENDS="huggingface-embeddings:/build/extra/grpc/huggingface/run.sh,autogptq:/build/extra/grpc/autogptq/run.sh,bark:/build/extra/grpc/bark/run.sh,diffusers:/build/extra/grpc/diffusers/run.sh,exllama:/build/extra/grpc/exllama/run.sh,vall-e-x:/build/extra/grpc/vall-e-x/run.sh,vllm:/build/extra/grpc/vllm/run.sh" ENV GALLERIES='[{"name":"model-gallery", "url":"github:go-skynet/model-gallery/index.yaml"}, {"url": "github:go-skynet/model-gallery/huggingface.yaml","name":"huggingface"}]' -ARG GO_TAGS="stablediffusion tts" +ARG GO_TAGS="stablediffusion tinydream tts" RUN apt-get update && \ apt-get install -y ca-certificates curl patch pip cmake && apt-get clean diff --git a/Makefile b/Makefile index 1b03de09bee4..1a1a0ff46d76 100644 --- a/Makefile +++ b/Makefile @@ -33,6 +33,9 @@ PIPER_VERSION?=736f6fb639ab8e3397356e48eeb6bdcb9da88a78 # stablediffusion version STABLEDIFFUSION_VERSION?=d89260f598afb809279bc72aa0107b4292587632 +# tiny-dream version +TINYDREAM_VERSION?=0308386bc80d9576109835ecb363f54f18406676 + export BUILD_TYPE?= export STABLE_BUILD_TYPE?=$(BUILD_TYPE) export CMAKE_ARGS?= @@ -116,6 +119,11 @@ ifeq ($(findstring stablediffusion,$(GO_TAGS)),stablediffusion) OPTIONAL_GRPC+=backend-assets/grpc/stablediffusion endif +ifeq ($(findstring tinydream,$(GO_TAGS)),tinydream) +# OPTIONAL_TARGETS+=go-stable-diffusion/libstablediffusion.a + OPTIONAL_GRPC+=backend-assets/grpc/tinydream +endif + ifeq ($(findstring tts,$(GO_TAGS)),tts) # OPTIONAL_TARGETS+=go-piper/libpiper_binding.a # OPTIONAL_TARGETS+=backend-assets/espeak-ng-data @@ -159,6 +167,11 @@ go-stable-diffusion: go-stable-diffusion/libstablediffusion.a: $(MAKE) -C go-stable-diffusion libstablediffusion.a +## tiny-dream +go-tiny-dream: + git clone --recurse-submodules https://github.com/M0Rf30/go-tiny-dream go-tiny-dream + cd go-tiny-dream && git checkout -b build $(TINYDREAM_VERSION) && git submodule update --init --recursive --depth 1 + ## RWKV go-rwkv: git clone --recurse-submodules $(RWKV_REPO) go-rwkv @@ -216,7 +229,10 @@ go-llama-stable/libbinding.a: go-llama-stable go-piper/libpiper_binding.a: go-piper $(MAKE) -C go-piper libpiper_binding.a example/main -get-sources: go-llama go-llama-stable go-ggml-transformers gpt4all go-piper go-rwkv whisper.cpp go-bert go-stable-diffusion +go-tiny-dream/libtinydream_binding.a: go-tiny-dream + $(MAKE) -C go-tiny-dream libtinydream_binding.a example/main + +get-sources: go-llama go-llama-stable go-ggml-transformers gpt4all go-piper go-rwkv whisper.cpp go-bert go-stable-diffusion go-tiny-dream touch $@ replace: @@ -226,6 +242,7 @@ replace: $(GOCMD) mod edit -replace github.com/ggerganov/whisper.cpp=$(shell pwd)/whisper.cpp $(GOCMD) mod edit -replace github.com/go-skynet/go-bert.cpp=$(shell pwd)/go-bert $(GOCMD) mod edit -replace github.com/mudler/go-stable-diffusion=$(shell pwd)/go-stable-diffusion + $(GOCMD) mod edit -replace github.com/M0Rf30/go-tiny-dream=$(shell pwd)/go-tiny-dream $(GOCMD) mod edit -replace github.com/mudler/go-piper=$(shell pwd)/go-piper prepare-sources: get-sources replace @@ -243,6 +260,7 @@ rebuild: ## Rebuilds the project $(MAKE) -C go-stable-diffusion clean $(MAKE) -C go-bert clean $(MAKE) -C go-piper clean + $(MAKE) -C go-tiny-dream clean $(MAKE) build prepare: prepare-sources $(OPTIONAL_TARGETS) @@ -256,6 +274,7 @@ clean: ## Remove build related file rm -rf ./go-llama-stable rm -rf ./go-gpt2 rm -rf ./go-stable-diffusion + rm -rf ./go-tiny-dream rm -rf ./go-ggml-transformers rm -rf ./backend-assets rm -rf ./go-rwkv @@ -313,6 +332,7 @@ test: prepare test-models/testmodel grpcs $(MAKE) test-llama-gguf $(MAKE) test-tts $(MAKE) test-stablediffusion + $(MAKE) test-tinydream prepare-e2e: mkdir -p $(TEST_DIR) @@ -354,6 +374,10 @@ test-stablediffusion: prepare-test TEST_DIR=$(abspath ./)/test-dir/ FIXTURES=$(abspath ./)/tests/fixtures CONFIG_FILE=$(abspath ./)/test-models/config.yaml MODELS_PATH=$(abspath ./)/test-models \ $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter="stablediffusion" --flake-attempts 1 -v -r ./api ./pkg +test-tinydream: prepare-test + TEST_DIR=$(abspath ./)/test-dir/ FIXTURES=$(abspath ./)/tests/fixtures CONFIG_FILE=$(abspath ./)/test-models/config.yaml MODELS_PATH=$(abspath ./)/test-models \ + $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter="tinydream" --flake-attempts 1 -v -r ./api ./pkg + test-container: docker build --target requirements -t local-ai-test-container . docker run -ti --rm --entrypoint /bin/bash -ti -v $(abspath ./):/build local-ai-test-container @@ -499,6 +523,10 @@ backend-assets/grpc/stablediffusion: backend-assets/grpc $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o backend-assets/grpc/stablediffusion ./cmd/grpc/stablediffusion/; \ fi +backend-assets/grpc/tinydream: backend-assets/grpc go-tiny-dream/libtinydream_binding.a + CGO_LDFLAGS="$(CGO_LDFLAGS)" LIBRARY_PATH=$(shell pwd)/go-tiny-dream \ + $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o backend-assets/grpc/tinydream ./cmd/grpc/tinydream/ + backend-assets/grpc/piper: backend-assets/grpc backend-assets/espeak-ng-data go-piper/libpiper_binding.a CGO_CXXFLAGS="$(PIPER_CGO_CXXFLAGS)" CGO_LDFLAGS="$(PIPER_CGO_LDFLAGS)" LIBRARY_PATH=$(shell pwd)/go-piper \ $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o backend-assets/grpc/piper ./cmd/grpc/piper/ diff --git a/api/api_test.go b/api/api_test.go index 6623e54277e0..f8914fd18047 100644 --- a/api/api_test.go +++ b/api/api_test.go @@ -548,6 +548,43 @@ var _ = Describe("API test", func() { Expect(resp.StatusCode).To(Equal(200), fmt.Sprint(string(dat))) Expect(resp.Header.Get("Content-Type")).To(Equal("audio/x-wav")) }) + It("installs and is capable to generate images", Label("tinydream"), func() { + if runtime.GOOS != "linux" { + Skip("test supported only on linux") + } + + response := postModelApplyRequest("http://127.0.0.1:9090/models/apply", modelApplyRequest{ + ID: "model-gallery@tinydream", + Overrides: map[string]interface{}{ + "parameters": map[string]interface{}{"model": "tinydream_assets"}, + }, + }) + + Expect(response["uuid"]).ToNot(BeEmpty(), fmt.Sprint(response)) + + uuid := response["uuid"].(string) + + Eventually(func() bool { + response := getModelStatus("http://127.0.0.1:9090/models/jobs/" + uuid) + fmt.Println(response) + return response["processed"].(bool) + }, "360s", "10s").Should(Equal(true)) + + resp, err := http.Post( + "http://127.0.0.1:9090/v1/images/generations", + "application/json", + bytes.NewBuffer([]byte(`{ + "prompt": "floating hair, portrait, ((loli)), ((one girl)), cute face, hidden hands, asymmetrical bangs, beautiful detailed eyes, eye shadow, hair ornament, ribbons, bowties, buttons, pleated skirt, (((masterpiece))), ((best quality)), colorful|((part of the head)), ((((mutated hands and fingers)))), deformed, blurry, bad anatomy, disfigured, poorly drawn face, mutation, mutated, extra limb, ugly, poorly drawn hands, missing limb, blurry, floating limbs, disconnected limbs, malformed hands, blur, out of focus, long neck, long body, Octane renderer, lowres, bad anatomy, bad hands, text", + "mode": 2, "seed":9000, + "size": "256x256", "n":2}`))) + // The response should contain an URL + Expect(err).ToNot(HaveOccurred(), fmt.Sprint(resp)) + dat, err := io.ReadAll(resp.Body) + Expect(err).ToNot(HaveOccurred(), string(dat)) + Expect(string(dat)).To(ContainSubstring("http://127.0.0.1:9090/"), string(dat)) + Expect(string(dat)).To(ContainSubstring(".png"), string(dat)) + + }) It("installs and is capable to generate images", Label("stablediffusion"), func() { if runtime.GOOS != "linux" { Skip("test supported only on linux") diff --git a/api/openai/image.go b/api/openai/image.go index 8243e167b63b..e5bde8bedbda 100644 --- a/api/openai/image.go +++ b/api/openai/image.go @@ -44,9 +44,15 @@ func ImageEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fiber.Ctx return fmt.Errorf("failed reading parameters from request:%w", err) } - if m == "" { + switch m { + case "stablediffusion": m = model.StableDiffusionBackend + case "tinydream": + m = model.TinyDreamBackend + default: + m = model.TinyDreamBackend } + log.Debug().Msgf("Loading model: %+v", m) config, input, err := readConfig(m, input, cm, o.Loader, o.Debug, 0, 0, false) @@ -82,8 +88,13 @@ func ImageEndpoint(cm *config.ConfigLoader, o *options.Option) func(c *fiber.Ctx log.Debug().Msgf("Parameter Config: %+v", config) // XXX: Only stablediffusion is supported for now - if config.Backend == "" { + switch config.Backend { + case "stablediffusion": config.Backend = model.StableDiffusionBackend + case "tinydream": + config.Backend = model.TinyDreamBackend + default: + config.Backend = model.TinyDreamBackend } sizeParts := strings.Split(input.Size, "x") diff --git a/cmd/grpc/tinydream/main.go b/cmd/grpc/tinydream/main.go new file mode 100644 index 000000000000..992580f08ab9 --- /dev/null +++ b/cmd/grpc/tinydream/main.go @@ -0,0 +1,23 @@ +package main + +// Note: this is started internally by LocalAI and a server is allocated for each model + +import ( + "flag" + + image "github.com/go-skynet/LocalAI/pkg/backend/image" + + grpc "github.com/go-skynet/LocalAI/pkg/grpc" +) + +var ( + addr = flag.String("addr", "localhost:50051", "the address to connect to") +) + +func main() { + flag.Parse() + + if err := grpc.StartServer(*addr, &image.TinyDream{}); err != nil { + panic(err) + } +} diff --git a/go.mod b/go.mod index 3262a68b6187..8891225eefc3 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/go-skynet/LocalAI -go 1.21 +go 1.21.4 require ( github.com/donomii/go-rwkv.cpp v0.0.0-20230715075832-c898cd0f62df @@ -54,6 +54,7 @@ require ( ) require ( + github.com/M0Rf30/go-tiny-dream v0.0.0-20231112234303-0f4e1ec72a2a // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/dlclark/regexp2 v1.8.1 // indirect diff --git a/go.sum b/go.sum index 239bb85d4615..996538fd1ba8 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/M0Rf30/go-tiny-dream v0.0.0-20231112234303-0f4e1ec72a2a h1:j+b1s71F6/ciNqaud9g4Erycbb67jvisYiV0jFx3bpk= +github.com/M0Rf30/go-tiny-dream v0.0.0-20231112234303-0f4e1ec72a2a/go.mod h1:wIYLWg3Y1Di5AYWZIsMDnH0X2HTDKTcZvlvtAIZCNZY= github.com/andybalholm/brotli v1.0.1/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y= github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs= github.com/andybalholm/brotli v1.0.5/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= diff --git a/pkg/backend/image/tinydream.go b/pkg/backend/image/tinydream.go new file mode 100644 index 000000000000..207d71562149 --- /dev/null +++ b/pkg/backend/image/tinydream.go @@ -0,0 +1,32 @@ +package image + +// This is a wrapper to statisfy the GRPC service interface +// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc) +import ( + "github.com/go-skynet/LocalAI/pkg/grpc/base" + pb "github.com/go-skynet/LocalAI/pkg/grpc/proto" + "github.com/go-skynet/LocalAI/pkg/tinydream" +) + +type TinyDream struct { + base.SingleThread + tinydream *tinydream.TinyDream +} + +func (sd *TinyDream) Load(opts *pb.ModelOptions) error { + var err error + // Note: the Model here is a path to a directory containing the model files + sd.tinydream, err = tinydream.New(opts.ModelFile) + return err +} + +func (td *TinyDream) GenerateImage(opts *pb.GenerateImageRequest) error { + return td.tinydream.GenerateImage( + int(opts.Height), + int(opts.Width), + int(opts.Step), + int(opts.Seed), + opts.PositivePrompt, + opts.NegativePrompt, + opts.Dst) +} diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index 535e21e50909..05dc00899654 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -35,6 +35,7 @@ const ( RwkvBackend = "rwkv" WhisperBackend = "whisper" StableDiffusionBackend = "stablediffusion" + TinyDreamBackend = "tinydream" PiperBackend = "piper" LCHuggingFaceBackend = "langchain-huggingface" ) @@ -56,6 +57,7 @@ var AutoLoadBackends []string = []string{ RwkvBackend, WhisperBackend, StableDiffusionBackend, + TinyDreamBackend, PiperBackend, } diff --git a/pkg/tinydream/generate.go b/pkg/tinydream/generate.go new file mode 100644 index 000000000000..cfcd23ccae03 --- /dev/null +++ b/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, + ) +} diff --git a/pkg/tinydream/generate_unsupported.go b/pkg/tinydream/generate_unsupported.go new file mode 100644 index 000000000000..4ffd421af62f --- /dev/null +++ b/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") +} diff --git a/pkg/tinydream/tinydream.go b/pkg/tinydream/tinydream.go new file mode 100644 index 000000000000..a316e6413938 --- /dev/null +++ b/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) +}