xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-21 15:36:49 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Added validation for number of generated images

Browse Source 
and defined image generation ratios
This commit is contained in:
ckt1031
2023-11-15 20:09:48 +08:00
parent 7b5efd9f3e
commit 6f5655ec7a
3 changed files with 22 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
controller/relay-image.go
View File

@@ -14,6 +14,17 @@ import (
"github.com/gin-gonic/gin"
)
func isWithinRange(element string, value int) bool {
if _, ok := common.DalleGenerationImageAmountRatios[element]; !ok {
return false
}
min := common.DalleGenerationImageAmountRatios[element][0]
max := common.DalleGenerationImageAmountRatios[element][1]
return value >= min && value <= max
}
func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
imageModel := "dall-e-2"
requestSize := "1024x1024"
@@ -49,6 +60,11 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
return errorWrapper(errors.New("prompt is required"), "required_field_missing", http.StatusBadRequest)
}
// Number of generated images validation
if isWithinRange(imageModel, imageRequest.N) == false {
return errorWrapper(errors.New("invalud value of n"), "number_of_generated_images_not_within_range", http.StatusBadRequest)
}
// map model name
modelMapping := c.GetString("model_mapping")
isModelMapped := false