From b59fbafbc33355abf4823f4bb3a2549efc23c210 Mon Sep 17 00:00:00 2001
From: ckt1031 <65409152+ckt1031@users.noreply.github.com>
Date: Wed, 15 Nov 2023 20:18:27 +0800
Subject: [PATCH] Updated model validation and removed size not supported error
 in relayImageHelper function

---
 controller/relay-image.go | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/controller/relay-image.go b/controller/relay-image.go
index 84893ebf..a7941ba5 100644
--- a/controller/relay-image.go
+++ b/controller/relay-image.go
@@ -46,8 +46,13 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	}
 
 	// Model validation
-	if imageRequest.Model != "" && imageRequest.Model != "dall-e-3" {
-		imageModel = "dall-e-2"
+	if imageRequest.Model != "" {
+		// Check if model is supported
+		if _, ok := common.DalleSizeRatios[imageRequest.Model]; !ok {
+			return errorWrapper(errors.New("model not supported"), "model_not_supported", http.StatusBadRequest)
+		}
+
+		imageModel = imageRequest.Model
 	}
 
 	// Size validation
@@ -117,8 +122,6 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 				sizeRatio *= 1.5
 			}
 		}
-	} else {
-		return errorWrapper(errors.New("size not supported"), "size_not_supported", http.StatusBadRequest)
 	}
 
 	quota := int(ratio*sizeRatio*1000) * imageRequest.N