From 2a8e935f64c1bf2046f54e4b281d8a988749c77c Mon Sep 17 00:00:00 2001
From: cktsun1031 <65409152+cktsun1031@users.noreply.github.com>
Date: Tue, 7 Nov 2023 14:49:56 +0800
Subject: [PATCH] Updated ImageRequest struct, added new models to
 OpenAIModels, and updated DalleSizeRatios and ModelRatio maps

---
 common/model-ratio.go     | 16 ++++++++++++++-
 controller/model.go       | 13 +++++++++++--
 controller/relay-image.go | 41 ++++++++++++++++++++++-----------------
 controller/relay.go       | 10 +++++++---
 4 files changed, 56 insertions(+), 24 deletions(-)

diff --git a/common/model-ratio.go b/common/model-ratio.go
index 8f4be8c3..eab6ddfd 100644
--- a/common/model-ratio.go
+++ b/common/model-ratio.go
@@ -5,6 +5,19 @@ import (
 	"strings"
 )
 
+var DalleSizeRatios = map[string]map[string]float64{
+	"dall-e-2": {
+		"256x256":   1,
+		"512x512":   1.125,
+		"1024x1024": 1.25,
+	},
+	"dall-e-3": {
+		"1024x1024": 1,
+		"1024×1792": 2,
+		"1792×1024": 2,
+	},
+}
+
 // ModelRatio
 // https://platform.openai.com/docs/models/model-endpoint-compatibility
 // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf
@@ -41,7 +54,8 @@ var ModelRatio = map[string]float64{
 	"text-search-ada-doc-001":   10,
 	"text-moderation-stable":    0.1,
 	"text-moderation-latest":    0.1,
-	"dall-e":                    8,
+	"dall-e-2":                  8,      // $0.016 - $0.020 / image
+	"dall-e-3":                  20,     // ￥0.040 - ￥0.120 / image
 	"claude-instant-1":          0.815,  // $1.63 / 1M tokens
 	"claude-2":                  5.51,   // $11.02 / 1M tokens
 	"ERNIE-Bot":                 0.8572, // ￥0.012 / 1k tokens
diff --git a/controller/model.go b/controller/model.go
index 2a7dc538..42541e35 100644
--- a/controller/model.go
+++ b/controller/model.go
@@ -55,12 +55,21 @@ func init() {
 	// https://platform.openai.com/docs/models/model-endpoint-compatibility
 	openAIModels = []OpenAIModels{
 		{
-			Id:         "dall-e",
+			Id:         "dall-e-2",
 			Object:     "model",
 			Created:    1677649963,
 			OwnedBy:    "openai",
 			Permission: permission,
-			Root:       "dall-e",
+			Root:       "dall-e-2",
+			Parent:     nil,
+		},
+		{
+			Id:         "dall-e-3",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "dall-e-3",
 			Parent:     nil,
 		},
 		{
diff --git a/controller/relay-image.go b/controller/relay-image.go
index 193e9d30..8cfc63da 100644
--- a/controller/relay-image.go
+++ b/controller/relay-image.go
@@ -15,7 +15,8 @@ import (
 )
 
 func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
-	imageModel := "dall-e"
+	imageModel := "dall-e-2"
+	requestSize := "1024x1024"
 
 	tokenId := c.GetInt("token_id")
 	channelType := c.GetInt("channel")
@@ -25,6 +26,7 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	group := c.GetString("group")
 
 	var imageRequest ImageRequest
+
 	if consumeQuota {
 		err := common.UnmarshalBodyReusable(c, &imageRequest)
 		if err != nil {
@@ -32,21 +34,21 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		}
 	}
 
+	// Model validation
+	if imageRequest.Model != "" || imageRequest.Model != "dall-e-3" {
+		imageModel = "dall-e-2"
+	}
+
+	// Size validation
+	if imageRequest.Size != "" {
+		requestSize = imageRequest.Size
+	}
+
 	// Prompt validation
 	if imageRequest.Prompt == "" {
 		return errorWrapper(errors.New("prompt is required"), "required_field_missing", http.StatusBadRequest)
 	}
 
-	// Not "256x256", "512x512", or "1024x1024"
-	if imageRequest.Size != "" && imageRequest.Size != "256x256" && imageRequest.Size != "512x512" && imageRequest.Size != "1024x1024" {
-		return errorWrapper(errors.New("size must be one of 256x256, 512x512, or 1024x1024"), "invalid_field_value", http.StatusBadRequest)
-	}
-
-	// N should between 1 and 10
-	if imageRequest.N != 0 && (imageRequest.N < 1 || imageRequest.N > 10) {
-		return errorWrapper(errors.New("n must be between 1 and 10"), "invalid_field_value", http.StatusBadRequest)
-	}
-
 	// map model name
 	modelMapping := c.GetString("model_mapping")
 	isModelMapped := false
@@ -88,14 +90,17 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	}
 
 	sizeRatio := 1.0
-	// Size
-	if imageRequest.Size == "256x256" {
-		sizeRatio = 1
-	} else if imageRequest.Size == "512x512" {
-		sizeRatio = 1.125
-	} else if imageRequest.Size == "1024x1024" {
-		sizeRatio = 1.25
+
+	if ratios, ok := common.DalleSizeRatios[imageModel]; ok {
+		if ratio, ok := ratios[requestSize]; ok {
+			sizeRatio = ratio
+
+			if imageRequest.Quality == "hd" {
+				sizeRatio = ratio * 2
+			}
+		}
 	}
+
 	quota := int(ratio*sizeRatio*1000) * imageRequest.N
 
 	if consumeQuota && userQuota-quota < 0 {
diff --git a/controller/relay.go b/controller/relay.go
index 54557a01..fdac0d88 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -79,9 +79,13 @@ type TextRequest struct {
 }
 
 type ImageRequest struct {
-	Prompt string `json:"prompt"`
-	N      int    `json:"n"`
-	Size   string `json:"size"`
+	Model          string `json:"model"`
+	Prompt         string `json:"prompt" binding:"required"`
+	N              int    `json:"n"`
+	Size           string `json:"size"`
+	Quality        string `json:"quality"`
+	ResponseFormat string `json:"response_format"`
+	Style          string `json:"style"`
 }
 
 type AudioResponse struct {