From 0504a4643c08da3ef0fda6e66a558733faaf957b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=A1=9C=E8=8F=AF?= <huan0016@gmail.com>
Date: Tue, 20 Jun 2023 19:01:18 +0800
Subject: [PATCH] =?UTF-8?q?=E8=BE=93=E5=85=A5=E6=A3=80=E6=9F=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 routers/images.go | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/routers/images.go b/routers/images.go
index a9d6c1c..8854228 100644
--- a/routers/images.go
+++ b/routers/images.go
@@ -8,6 +8,7 @@ import (
 	_ "image/gif"
 	_ "image/jpeg"
 	_ "image/png"
+	"regexp"
 
 	"io/ioutil"
 	"log"
@@ -82,9 +83,7 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 	models.AccountRead(w, r, func(account *models.Account) {
 
 		// 通过模型推理生成图像, 为图像标记任务批次
-		if r.Header.Get("Content-Type") == "application/json" || r.Header.Get("Content-Type") == "application/json; charset=utf-8" {
-
-			// 接收模板参数
+		if match, _ := regexp.MatchString("application/json", r.Header.Get("Content-Type")); match {
 			template := &struct {
 				FromImage         int     `json:"from_image"`          // 来源图片(图生图时使用)
 				Prompt            string  `json:"prompt"`              // 提示词
@@ -92,7 +91,7 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 				NumInferenceSteps int     `json:"num_inference_steps"` // 推理步数
 				GuidanceScale     float32 `json:"guidance_scale"`      // 引导比例
 				Scheduler         string  `json:"scheduler"`           // 调度器
-				Seed              int     `json:"seed"`                // 随机种子(单张图生成时使用)
+				Seed              string  `json:"seed"`                // 随机种子(单张图生成时使用)
 				Number            int     `json:"number"`              // 生成数量
 			}{}
 			body, err := ioutil.ReadAll(r.Body)
@@ -106,6 +105,20 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 				return
 			}
 
+			// 输入检查
+			if template.Number <= 0 {
+				template.Number = 1
+			}
+			if template.NumInferenceSteps <= 0 {
+				template.NumInferenceSteps = 20
+			}
+			if template.GuidanceScale <= 0 {
+				template.GuidanceScale = 1
+			}
+			if template.GuidanceScale > 20 {
+				template.GuidanceScale = 20
+			}
+
 			// TODO: 创建任务获得任务编号, 多张图时期望可以流式推理
 			task := uuid.New().String()
 
@@ -134,6 +147,7 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 
 			w.Header().Set("Content-Type", "application/json; charset=utf-8")
 			json.NewEncoder(w).Encode(image_list)
+			//w.Write(utils.ToJSON({"task": task, "list": image_list}))
 			return
 		}