From dfc3dadc53292b09b02f737bb70c2421a44e4adb Mon Sep 17 00:00:00 2001
From: Chen Tao <70054568+eeee0717@users.noreply.github.com>
Date: Sat, 26 Apr 2025 22:15:48 +0800
Subject: [PATCH] feat: Enable image generation in assistant based on model
 selection (#5364)

* feat: Enable image generation in assistant based on model selection

* chore: remove generate image button
---
 src/renderer/src/pages/home/Inputbar/Inputbar.tsx | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
index b7c1f15e05..8515cc58aa 100644
--- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
+++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
@@ -55,7 +55,6 @@ import styled from 'styled-components'
 import NarrowLayout from '../Messages/NarrowLayout'
 import AttachmentButton, { AttachmentButtonRef } from './AttachmentButton'
 import AttachmentPreview from './AttachmentPreview'
-import GenerateImageButton from './GenerateImageButton'
 import KnowledgeBaseButton, { KnowledgeBaseButtonRef } from './KnowledgeBaseButton'
 import KnowledgeBaseInput from './KnowledgeBaseInput'
 import MCPToolsButton, { MCPToolsButtonRef } from './MCPToolsButton'
@@ -796,10 +795,6 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic, topic }) =
     updateAssistant({ ...assistant, enableWebSearch: !assistant.enableWebSearch })
   }
 
-  const onEnableGenerateImage = () => {
-    updateAssistant({ ...assistant, enableGenerateImage: !assistant.enableGenerateImage })
-  }
-
   useEffect(() => {
     if (!isWebSearchModel(model) && !WebSearchService.isWebSearchEnabled() && assistant.enableWebSearch) {
       updateAssistant({ ...assistant, enableWebSearch: false })
@@ -807,6 +802,9 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic, topic }) =
     if (!isGenerateImageModel(model) && assistant.enableGenerateImage) {
       updateAssistant({ ...assistant, enableGenerateImage: false })
     }
+    if (isGenerateImageModel(model) && !assistant.enableGenerateImage) {
+      updateAssistant({ ...assistant, enableGenerateImage: true })
+    }
   }, [assistant, model, updateAssistant])
 
   const onMentionModel = (model: Model) => {
@@ -953,12 +951,6 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic, topic }) =
                 setInputValue={setText}
                 resizeTextArea={resizeTextArea}
               />
-              <GenerateImageButton
-                model={model}
-                assistant={assistant}
-                onEnableGenerateImage={onEnableGenerateImage}
-                ToolbarButton={ToolbarButton}
-              />
               <MentionModelsButton
                 ref={mentionModelsButtonRef}
                 mentionModels={mentionModels}