feat: support upscaled image generation with esrgan (#509)

2 years ago · b447a2a719
parent ec4fd1d219
commit b447a2a719
3 changed files with 14 additions and 2 deletions
--- a/2
+++ b/2
@ -15,7 +15,7 @@ BLOOMZ_VERSION?=1834e77b83faafe912ad4092ccf7f77937349e2f
 export BUILD_TYPE?=
 CGO_LDFLAGS?=
 CUDA_LIBPATH?=/usr/local/cuda/lib64/
-STABLEDIFFUSION_VERSION?=c0748eca3642d58bcf9521108bcee46959c647dc
+STABLEDIFFUSION_VERSION?=d89260f598afb809279bc72aa0107b4292587632
 GO_TAGS?=
 BUILD_ID?=git
 LD_FLAGS=?=
--- a/go.mod
+++ b/go.mod
@ -14,7 +14,7 @@ require (
 	github.com/google/uuid v1.3.0
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/imdario/mergo v0.3.16
-	github.com/mudler/go-stable-diffusion v0.0.0-20230516152536-c0748eca3642
+	github.com/mudler/go-stable-diffusion v0.0.0-20230605114250-a6706a426a90
 	github.com/nomic-ai/gpt4all/gpt4all-bindings/golang v0.0.0-20230604125924-bbe195ee0207
 	github.com/onsi/ginkgo/v2 v2.9.7
 	github.com/onsi/gomega v1.27.7
--- a/pkg/stablediffusion/generate.go
+++ b/pkg/stablediffusion/generate.go
@ -8,6 +8,18 @@ import (
 )

 func GenerateImage(height, width, mode, step, seed int, positive_prompt, negative_prompt, dst, asset_dir string) error {
+	if height > 512 || width > 512 {
+		return stableDiffusion.GenerateImageUpscaled(
+			height,
+			width,
+			step,
+			seed,
+			positive_prompt,
+			negative_prompt,
+			dst,
+			asset_dir,
+		)
+	}
 	return stableDiffusion.GenerateImage(
 		height,
 		width,