diff --git a/common/model-ratio.go b/common/model-ratio.go index f4cc33ec..457e0013 100644 --- a/common/model-ratio.go +++ b/common/model-ratio.go @@ -18,6 +18,12 @@ var DalleSizeRatios = map[string]map[string]float64{ }, } +// Array in values, of maxmium and minimum +var DalleGenerationImageAmountRatios = map[string][2]int{ + "dall-e-2": {1, 10}, + "dall-e-3": {1, 1}, // OpenAI allows n=1 currently. +} + // ModelRatio // https://platform.openai.com/docs/models/model-endpoint-compatibility // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf diff --git a/controller/relay-image.go b/controller/relay-image.go index 3ff7dc18..d97e5e38 100644 --- a/controller/relay-image.go +++ b/controller/relay-image.go @@ -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 diff --git a/one-api b/one-api new file mode 100755 index 00000000..231e75df Binary files /dev/null and b/one-api differ