diff --git a/src/app.tsx b/src/app.tsx
index d75a80c..34e385f 100644
--- a/src/app.tsx
+++ b/src/app.tsx
@@ -262,7 +262,9 @@ export function App() {
         chatStore.tts_key,
         chatStore.tts_speed,
         chatStore.tts_speed_enabled,
-        chatStore.toolsString
+        chatStore.toolsString,
+        chatStore.image_gen_api,
+        chatStore.image_gen_key
       )
     );
     setSelectedChatIndex(newKey as number);
diff --git a/src/chatbox.tsx b/src/chatbox.tsx
index 3137d57..20cceef 100644
--- a/src/chatbox.tsx
+++ b/src/chatbox.tsx
@@ -675,7 +675,8 @@ export default function ChatBOX(props: {
       )}
 
       <div className="flex justify-between">
-        {chatStore.model.match("vision") && (
+        {(chatStore.model.match("vision") ||
+          (chatStore.image_gen_api && chatStore.image_gen_key)) && (
           <button
             className="disabled:line-through disabled:bg-slate-500 rounded m-1 p-1 border-2 bg-cyan-400 hover:bg-cyan-600"
             disabled={showGenerating || !chatStore.apiKey}