From b447a2a7194e11de2845d2dad42033125bc9bfad Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Mon, 5 Jun 2023 17:21:38 +0200
Subject: [PATCH] feat: support upscaled image generation with esrgan (#509)

---
 Makefile                        |  2 +-
 go.mod                          |  2 +-
 pkg/stablediffusion/generate.go | 12 ++++++++++++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/Makefile b/Makefile
index 4df0190..51238fd 100644
--- a/Makefile
+++ b/Makefile
@@ -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=?=
diff --git a/go.mod b/go.mod
index edaa622..d9666c3 100644
--- 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
diff --git a/pkg/stablediffusion/generate.go b/pkg/stablediffusion/generate.go
index 97989e9..cef96e8 100644
--- 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,