From a9b2b32c5a04e82968e022b91eceb8a9b266f4a8 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Sat, 22 Mar 2025 22:45:45 +0800
Subject: [PATCH] feat: add GenerateImageButton component to Inputbar for image
 generation functionality

---
 .../home/Inputbar/GenerateImageButton.tsx     | 36 +++++++++++++++++++
 .../src/pages/home/Inputbar/Inputbar.tsx      | 22 ++++--------
 2 files changed, 43 insertions(+), 15 deletions(-)
 create mode 100644 src/renderer/src/pages/home/Inputbar/GenerateImageButton.tsx

diff --git a/src/renderer/src/pages/home/Inputbar/GenerateImageButton.tsx b/src/renderer/src/pages/home/Inputbar/GenerateImageButton.tsx
new file mode 100644
index 00000000..41b6ccc5
--- /dev/null
+++ b/src/renderer/src/pages/home/Inputbar/GenerateImageButton.tsx
@@ -0,0 +1,36 @@
+import { PictureOutlined } from '@ant-design/icons'
+import { isGenerateImageModel } from '@renderer/config/models'
+import { Assistant, Model } from '@renderer/types'
+import { Tooltip } from 'antd'
+import { FC } from 'react'
+import { useTranslation } from 'react-i18next'
+
+interface Props {
+  assistant: Assistant
+  model: Model
+  ToolbarButton: any
+  onEnableGenerateImage: () => void
+}
+
+const GenerateImageButton: FC<Props> = ({ model, ToolbarButton, assistant, onEnableGenerateImage }) => {
+  const { t } = useTranslation()
+
+  if (!isGenerateImageModel(model)) {
+    return null
+  }
+
+  return (
+    <Tooltip
+      placement="top"
+      title={
+        isGenerateImageModel(model) ? t('chat.input.generate_image') : t('chat.input.generate_image_not_supported')
+      }
+      arrow>
+      <ToolbarButton type="text" disabled={!isGenerateImageModel(model)} onClick={onEnableGenerateImage}>
+        <PictureOutlined style={{ color: assistant.enableGenerateImage ? 'var(--color-link)' : 'var(--color-icon)' }} />
+      </ToolbarButton>
+    </Tooltip>
+  )
+}
+
+export default GenerateImageButton
diff --git a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
index d60ccb49..1dadb054 100644
--- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
+++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
@@ -7,7 +7,6 @@ import {
   GlobalOutlined,
   HolderOutlined,
   PauseCircleOutlined,
-  PictureOutlined,
   QuestionCircleOutlined
 } from '@ant-design/icons'
 import TranslateButton from '@renderer/components/TranslateButton'
@@ -45,6 +44,7 @@ import styled from 'styled-components'
 import NarrowLayout from '../Messages/NarrowLayout'
 import AttachmentButton from './AttachmentButton'
 import AttachmentPreview from './AttachmentPreview'
+import GenerateImageButton from './GenerateImageButton'
 import KnowledgeBaseButton from './KnowledgeBaseButton'
 import MCPToolsButton from './MCPToolsButton'
 import MentionModelsButton from './MentionModelsButton'
@@ -745,20 +745,12 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic, topic }) =
                   ToolbarButton={ToolbarButton}
                 />
               )}
-              <Tooltip
-                placement="top"
-                title={
-                  !isGenerateImageModel(model)
-                    ? t('chat.input.generate_image_not_supported')
-                    : t('chat.input.generate_image')
-                }
-                arrow>
-                <ToolbarButton type="text" disabled={!isGenerateImageModel(model)} onClick={onEnableGenerateImage}>
-                  <PictureOutlined
-                    style={{ color: assistant.enableGenerateImage ? 'var(--color-link)' : 'var(--color-icon)' }}
-                  />
-                </ToolbarButton>
-              </Tooltip>
+              <GenerateImageButton
+                model={model}
+                assistant={assistant}
+                onEnableGenerateImage={onEnableGenerateImage}
+                ToolbarButton={ToolbarButton}
+              />
               <MentionModelsButton
                 mentionModels={mentionModels}
                 onMentionModel={(model) => onMentionModel(model, mentionFromKeyboard)}