Skip to content

Commit

Permalink
feat: add tiny dream stable diffusion support
Browse files Browse the repository at this point in the history
Signed-off-by: Gianluca Boiano <morf3089@gmail.com>
  • Loading branch information
M0Rf30 committed Nov 13, 2023
1 parent 5546118 commit c7b5907
Show file tree
Hide file tree
Showing 14 changed files with 209 additions and 6 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ go-llama
go-llama-stable
/gpt4all
go-stable-diffusion
go-tiny-dream
go-piper
/go-bert
go-ggllm
Expand Down
2 changes: 1 addition & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 29 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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?=
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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/
Expand Down
37 changes: 37 additions & 0 deletions api/api_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
15 changes: 13 additions & 2 deletions api/openai/image.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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")
Expand Down
23 changes: 23 additions & 0 deletions cmd/grpc/tinydream/main.go
Original file line number Diff line number Diff line change
@@ -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)
}
}
3 changes: 2 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -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=
Expand Down
32 changes: 32 additions & 0 deletions pkg/backend/image/tinydream.go
Original file line number Diff line number Diff line change
@@ -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)
}
2 changes: 2 additions & 0 deletions pkg/model/initializers.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ const (
RwkvBackend = "rwkv"
WhisperBackend = "whisper"
StableDiffusionBackend = "stablediffusion"
TinyDreamBackend = "tinydream"
PiperBackend = "piper"
LCHuggingFaceBackend = "langchain-huggingface"
)
Expand All @@ -56,6 +57,7 @@ var AutoLoadBackends []string = []string{
RwkvBackend,
WhisperBackend,
StableDiffusionBackend,
TinyDreamBackend,
PiperBackend,
}

Expand Down
36 changes: 36 additions & 0 deletions pkg/tinydream/generate.go
Original file line number Diff line number Diff line change
@@ -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,
)
}
10 changes: 10 additions & 0 deletions pkg/tinydream/generate_unsupported.go
Original file line number Diff line number Diff line change
@@ -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")
}
20 changes: 20 additions & 0 deletions pkg/tinydream/tinydream.go
Original file line number Diff line number Diff line change
@@ -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)
}

0 comments on commit c7b5907

Please sign in to comment.