fix ImagePredict channel error

seasonjs · Jan 16, 2024 · 40a0c2b · 40a0c2b
1 parent cd5b9c0
commit 40a0c2b
Show file tree

Hide file tree

Showing 8 changed files with 86 additions and 7 deletions.
diff --git a/.idea/deployment.xml b/.idea/deployment.xml
diff --git a/.idea/dictionaries/cyberhan.xml b/.idea/dictionaries/cyberhan.xml
diff --git a/.idea/remote-targets.xml b/.idea/remote-targets.xml
diff --git a/.idea/sshConfigs.xml b/.idea/sshConfigs.xml
diff --git a/.idea/stable-diffusion.iml b/.idea/stable-diffusion.iml
diff --git a/.idea/webServers.xml b/.idea/webServers.xml
diff --git a/sd.go b/sd.go
@@ -317,9 +317,10 @@ func imageToBytes(decode image.Image) Image {
 		}
 	}
 	return Image{
-		Width:  uint32(width),
-		Height: uint32(height),
-		Data:   bytesImg,
+		Width:   uint32(width),
+		Height:  uint32(height),
+		Data:    bytesImg,
+		Channel: 3,
 	}
 }
 

diff --git a/sd_test.go b/sd_test.go
@@ -60,6 +60,9 @@ func TestNewStableDiffusionAutoModelImagePredict(t *testing.T) {
 		return
 	}
 	defer model.Close()
+	model.SetLogCallback(func(level sd.LogLevel, msg string) {
+		t.Log(msg)
+	})
 	err = model.LoadFromFile("./models/mysd.safetensors")
 	if err != nil {
 		t.Error(err)
@@ -75,7 +78,7 @@ func TestNewStableDiffusionAutoModelImagePredict(t *testing.T) {
 	var writers []io.Writer
 	filenames := []string{
 		"./assets/love_cat0_m.png",
-		"./assets/love_cat1_m.png",
+		//"./assets/love_cat1_m.png",
 		//"./assets/love_cat5.png",
 		//"./assets/love_cat6.png"
 	}
@@ -88,8 +91,12 @@ func TestNewStableDiffusionAutoModelImagePredict(t *testing.T) {
 		defer file.Close()
 		writers = append(writers, file)
 	}
-
-	err = model.ImagePredict(inFile, "the cat that wears shoe", sd.DefaultFullParams, writers)
+	params := sd.DefaultFullParams
+	params.BatchCount = 1
+	params.Width = 256
+	params.Height = 256
+	params.NegativePrompt = ""
+	err = model.ImagePredict(inFile, "dogs", params, writers)
 	if err != nil {
 		t.Error(err)
 		return