From 68eb6f5ce1e9d25f6557b92512cf70aed13dfca0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?c=E8=8F=8C?= <littlecjun@gmail.com>
Date: Thu, 23 Nov 2023 18:30:24 +0800
Subject: [PATCH] fix bug

---
 pkg/openai/dall-e.go       | 10 +++++-----
 pkg/tokenizer/tokenizer.go |  8 ++++----
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/pkg/openai/dall-e.go b/pkg/openai/dall-e.go
index 280ab87..3f4e9bd 100644
--- a/pkg/openai/dall-e.go
+++ b/pkg/openai/dall-e.go
@@ -51,7 +51,6 @@ func DalleHandler(c *gin.Context) {
 	model := dalleRequest.Model
 
 	var chatlog store.Tokens
-	chatlog.Model = model
 	chatlog.CompletionCount = dalleRequest.N
 
 	if model == "dall-e" {
@@ -59,7 +58,7 @@ func DalleHandler(c *gin.Context) {
 	}
 	model = model + "." + dalleRequest.Size
 
-	if model == "dall-e-2" {
+	if dalleRequest.Model == "dall-e-2" || dalleRequest.Model == "dall-e" {
 		if !slice.Contain([]string{"256x256", "512x512", "1024x1024"}, dalleRequest.Size) {
 			c.JSON(http.StatusBadRequest, gin.H{
 				"error": gin.H{
@@ -68,7 +67,7 @@ func DalleHandler(c *gin.Context) {
 			})
 			return
 		}
-	} else if model == "dall-e-3" {
+	} else if dalleRequest.Model == "dall-e-3" {
 		if !slice.Contain([]string{"256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"}, dalleRequest.Size) {
 			c.JSON(http.StatusBadRequest, gin.H{
 				"error": gin.H{
@@ -77,8 +76,8 @@ func DalleHandler(c *gin.Context) {
 			})
 			return
 		}
-		if dalleRequest.Quality == "HD" {
-			model = model + ".HD"
+		if dalleRequest.Quality == "hd" {
+			model = model + ".hd"
 		}
 	} else {
 		c.JSON(http.StatusBadRequest, gin.H{
@@ -88,6 +87,7 @@ func DalleHandler(c *gin.Context) {
 		})
 		return
 	}
+	chatlog.Model = model
 
 	token, _ := c.Get("localuser")
 
diff --git a/pkg/tokenizer/tokenizer.go b/pkg/tokenizer/tokenizer.go
index c7240e0..ac67bae 100644
--- a/pkg/tokenizer/tokenizer.go
+++ b/pkg/tokenizer/tokenizer.go
@@ -116,13 +116,13 @@ func Cost(model string, promptCount, completionCount int) float64 {
 		cost = float64(0.04 * completion)
 	case "dall-e-3.1024x1792", "dall-e-3.1792x1024":
 		cost = float64(0.08 * completion)
-	case "dall-e-3.256x256.HD":
+	case "dall-e-3.256x256.hd":
 		cost = float64(0.08 * completion)
-	case "dall-e-3.512x512.HD":
+	case "dall-e-3.512x512.hd":
 		cost = float64(0.08 * completion)
-	case "dall-e-3.1024x1024.HD":
+	case "dall-e-3.1024x1024.hd":
 		cost = float64(0.08 * completion)
-	case "dall-e-3.1024x1792.HD", "dall-e-3.1792x1024.HD":
+	case "dall-e-3.1024x1792.hd", "dall-e-3.1792x1024.hd":
 		cost = float64(0.12 * completion)
 
 	// claude /million tokens