Merge pull request #264 from wenerme/wener

feat: default to dall-e-3, support more image res, add image style

Author: Leizhenpeng

code/handlers/card_pic_action.go

@@ -17,6 +17,10 @@ func NewPicResolutionHandler(cardMsg CardMsg, m MessageHandler) CardHandlerFunc
 			CommonProcessPicResolution(cardMsg, cardAction, m.sessionCache)
 			return nil, nil
 		}
+		if cardMsg.Kind == PicStyleKind {
+			CommonProcessPicStyle(cardMsg, cardAction, m.sessionCache)
+			return nil, nil
+		}
 		return nil, ErrNextHandler
 	}
 }
@@ -57,13 +61,25 @@ func CommonProcessPicResolution(msg CardMsg,
 		&msg.MsgId)
 }
 
+func CommonProcessPicStyle(msg CardMsg,
+	cardAction *larkcard.CardAction,
+	cache services.SessionServiceCacheInterface) {
+	option := cardAction.Action.Option
+	fmt.Println(larkcore.Prettify(msg))
+	cache.SetPicStyle(msg.SessionId, services.PicStyle(option))
+	//send text
+	replyMsg(context.Background(), "已更新图片风格为"+option,
+		&msg.MsgId)
+}
+
 func (m MessageHandler) CommonProcessPicMore(msg CardMsg) {
 	resolution := m.sessionCache.GetPicResolution(msg.SessionId)
+	style := m.sessionCache.GetPicStyle(msg.SessionId)
 
 	logger.Debugf("resolution: %v", resolution)
 	logger.Debug("msg: %v", msg)
 	question := msg.Value.(string)
-	bs64, _ := m.gpt.GenerateOneImage(question, resolution)
+	bs64, _ := m.gpt.GenerateOneImage(question, resolution, style)
 	replayImageCardByBase64(context.Background(), bs64, &msg.MsgId,
 		&msg.SessionId, question)
 }

code/handlers/event_pic_action.go

@@ -29,7 +29,7 @@ func (*PicAction) Execute(a *ActionInfo) bool {
 		a.handler.sessionCache.SetMode(*a.info.sessionId,
 			services.ModePicCreate)
 		a.handler.sessionCache.SetPicResolution(*a.info.sessionId,
-			services.Resolution256)
+			services.Resolution1024)
 		sendPicCreateInstructionCard(*a.ctx, a.info.sessionId,
 			a.info.msgId)
 		return false
@@ -92,8 +92,10 @@ func (*PicAction) Execute(a *ActionInfo) bool {
 	if mode == services.ModePicCreate {
 		resolution := a.handler.sessionCache.GetPicResolution(*a.
 			info.sessionId)
+		style := a.handler.sessionCache.GetPicStyle(*a.
+			info.sessionId)
 		bs64, err := a.handler.gpt.GenerateOneImage(a.info.qParsed,
-			resolution)
+			resolution, style)
 		if err != nil {
 			replyMsg(*a.ctx, fmt.Sprintf(
 				"🤖️：图片生成失败，请稍后再试～\n错误信息: %v", err), a.info.msgId)

code/handlers/msg.go

@@ -24,6 +24,7 @@ var (
 	ClearCardKind      = CardKind("clear")            // 清空上下文
 	PicModeChangeKind  = CardKind("pic_mode_change")  // 切换图片创作模式
 	PicResolutionKind  = CardKind("pic_resolution")   // 图片分辨率调整
+	PicStyleKind       = CardKind("pic_style")        // 图片风格调整
 	PicTextMoreKind    = CardKind("pic_text_more")    // 重新根据文本生成图片
 	PicVarMoreKind     = CardKind("pic_var_more")     // 变量图片
 	RoleTagsChooseKind = CardKind("role_tags_choose") // 内置角色所属标签选择
@@ -324,29 +325,56 @@ func withOneBtn(btn *larkcard.MessageCardEmbedButton) larkcard.
 
 func withPicResolutionBtn(sessionID *string) larkcard.
 	MessageCardElement {
-	cancelMenu := newMenu("默认分辨率",
+	resolutionMenu := newMenu("默认分辨率",
 		map[string]interface{}{
 			"value":     "0",
 			"kind":      PicResolutionKind,
 			"sessionId": *sessionID,
 			"msgId":     *sessionID,
 		},
+		// dall-e-2 256, 512, 1024
+		//MenuOption{
+		//	label: "256x256",
+		//	value: string(services.Resolution256),
+		//},
+		//MenuOption{
+		//	label: "512x512",
+		//	value: string(services.Resolution512),
+		//},
+		// dall-e-3
 		MenuOption{
-			label: "256x256",
-			value: string(services.Resolution256),
+			label: "1024x1024",
+			value: string(services.Resolution1024),
 		},
 		MenuOption{
-			label: "512x512",
-			value: string(services.Resolution512),
+			label: "1024x1792",
+			value: string(services.Resolution10241792),
 		},
 		MenuOption{
-			label: "1024x1024",
-			value: string(services.Resolution1024),
+			label: "1792x1024",
+			value: string(services.Resolution17921024),
+		},
+	)
+
+	styleMenu := newMenu("风格",
+		map[string]interface{}{
+			"value":     "0",
+			"kind":      PicStyleKind,
+			"sessionId": *sessionID,
+			"msgId":     *sessionID,
+		},
+		MenuOption{
+			label: "生动风格",
+			value: string(services.PicStyleVivid),
+		},
+		MenuOption{
+			label: "自然风格",
+			value: string(services.PicStyleNatural),
 		},
 	)
 
 	actions := larkcard.NewMessageCardAction().
-		Actions([]larkcard.MessageCardActionElement{cancelMenu}).
+		Actions([]larkcard.MessageCardActionElement{resolutionMenu, styleMenu}).
 		Layout(larkcard.MessageCardActionLayoutFlow.Ptr()).
 		Build()
 	return actions

code/services/openai/gpt3_test.go

@@ -28,7 +28,7 @@ func TestGenerateOneImage(t *testing.T) {
 	gpt := NewChatGPT(*config)
 	prompt := "a red apple"
 	size := "256x256"
-	imageURL, err := gpt.GenerateOneImage(prompt, size)
+	imageURL, err := gpt.GenerateOneImage(prompt, size, "")
 	if err != nil {
 		t.Errorf("TestGenerateOneImage failed with error: %v", err)
 	}

code/services/openai/picture.go

@@ -16,6 +16,8 @@ type ImageGenerationRequestBody struct {
 	N              int    `json:"n"`
 	Size           string `json:"size"`
 	ResponseFormat string `json:"response_format"`
+	Model          string `json:"model,omitempty"`
+	Style          string `json:"style,omitempty"`
 }
 
 type ImageResponseBody struct {
@@ -33,12 +35,14 @@ type ImageVariantRequestBody struct {
 }
 
 func (gpt *ChatGPT) GenerateImage(prompt string, size string,
-	n int) ([]string, error) {
+	n int, style string) ([]string, error) {
 	requestBody := ImageGenerationRequestBody{
 		Prompt:         prompt,
 		N:              n,
 		Size:           size,
 		ResponseFormat: "b64_json",
+		Model:          "dall-e-3",
+		Style:          style,
 	}
 
 	imageResponseBody := &ImageResponseBody{}
@@ -57,8 +61,8 @@ func (gpt *ChatGPT) GenerateImage(prompt string, size string,
 }
 
 func (gpt *ChatGPT) GenerateOneImage(prompt string,
-	size string) (string, error) {
-	b64s, err := gpt.GenerateImage(prompt, size, 1)
+	size string, style string) (string, error) {
+	b64s, err := gpt.GenerateImage(prompt, size, 1, style)
 	if err != nil {
 		return "", err
 	}
@@ -67,7 +71,8 @@ func (gpt *ChatGPT) GenerateOneImage(prompt string,
 
 func (gpt *ChatGPT) GenerateOneImageWithDefaultSize(
 	prompt string) (string, error) {
-	return gpt.GenerateOneImage(prompt, "512x512")
+	// works for dall-e 2&3
+	return gpt.GenerateOneImage(prompt, "1024x1024", "")
 }
 
 func (gpt *ChatGPT) GenerateImageVariation(images string,

code/services/sessionCache.go

@@ -13,8 +13,10 @@ type SessionService struct {
 }
 type PicSetting struct {
 	resolution Resolution
+	style      PicStyle
 }
 type Resolution string
+type PicStyle string
 
 type SessionMeta struct {
 	Mode       SessionMode       `json:"mode"`
@@ -24,9 +26,15 @@ type SessionMeta struct {
 }
 
 const (
-	Resolution256  Resolution = "256x256"
-	Resolution512  Resolution = "512x512"
-	Resolution1024 Resolution = "1024x1024"
+	Resolution256      Resolution = "256x256"
+	Resolution512      Resolution = "512x512"
+	Resolution1024     Resolution = "1024x1024"
+	Resolution10241792 Resolution = "1024x1792"
+	Resolution17921024 Resolution = "1792x1024"
+)
+const (
+	PicStyleVivid   PicStyle = "vivid"
+	PicStyleNatural PicStyle = "natural"
 )
 const (
 	ModePicCreate SessionMode = "pic_create"
@@ -44,7 +52,9 @@ type SessionServiceCacheInterface interface {
 	GetAIMode(sessionId string) openai.AIMode
 	SetAIMode(sessionId string, aiMode openai.AIMode)
 	SetPicResolution(sessionId string, resolution Resolution)
+	SetPicStyle(sessionId string, resolution PicStyle)
 	GetPicResolution(sessionId string) string
+	GetPicStyle(sessionId string) string
 	Clear(sessionId string)
 }
 
@@ -141,16 +151,45 @@ func (s *SessionService) SetMsg(sessionId string, msg []openai.Messages) {
 	s.cache.Set(sessionId, sessionMeta, maxCacheTime)
 }
 
+func (s *SessionService) SetPicStyle(sessionId string, style PicStyle) {
+	maxCacheTime := time.Hour * 12
+
+	switch style {
+	case PicStyleVivid, PicStyleNatural:
+	default:
+		style = PicStyleVivid
+	}
+
+	sessionContext, ok := s.cache.Get(sessionId)
+	if !ok {
+		sessionMeta := &SessionMeta{PicSetting: PicSetting{style: style}}
+		s.cache.Set(sessionId, sessionMeta, maxCacheTime)
+		return
+	}
+	sessionMeta := sessionContext.(*SessionMeta)
+	sessionMeta.PicSetting.style = style
+	s.cache.Set(sessionId, sessionMeta, maxCacheTime)
+}
+
+func (s *SessionService) GetPicStyle(sessionId string) string {
+	sessionContext, ok := s.cache.Get(sessionId)
+	if !ok {
+		return string(PicStyleVivid)
+	}
+	sessionMeta := sessionContext.(*SessionMeta)
+	return string(sessionMeta.PicSetting.style)
+}
+
 func (s *SessionService) SetPicResolution(sessionId string,
 	resolution Resolution) {
 	maxCacheTime := time.Hour * 12
 
 	//if not in [Resolution256, Resolution512, Resolution1024] then set
 	//to Resolution256
 	switch resolution {
-	case Resolution256, Resolution512, Resolution1024:
+	case Resolution256, Resolution512, Resolution1024, Resolution10241792, Resolution17921024:
 	default:
-		resolution = Resolution256
+		resolution = Resolution1024
 	}
 
 	sessionContext, ok := s.cache.Get(sessionId)