diff --git a/src/components/graph/SelectionToolbox.vue b/src/components/graph/SelectionToolbox.vue
index fdb78e59b..5caa3bee7 100644
--- a/src/components/graph/SelectionToolbox.vue
+++ b/src/components/graph/SelectionToolbox.vue
@@ -30,6 +30,7 @@ import { computed } from 'vue'
 import BypassButton from '@/components/graph/selectionToolbox/BypassButton.vue'
 import ColorPickerButton from '@/components/graph/selectionToolbox/ColorPickerButton.vue'
 import DeleteButton from '@/components/graph/selectionToolbox/DeleteButton.vue'
+import EditModelButton from '@/components/graph/selectionToolbox/EditModelButton.vue'
 import ExecuteButton from '@/components/graph/selectionToolbox/ExecuteButton.vue'
 import ExtensionCommandButton from '@/components/graph/selectionToolbox/ExtensionCommandButton.vue'
 import HelpButton from '@/components/graph/selectionToolbox/HelpButton.vue'
diff --git a/src/components/graph/selectionToolbox/EditModelButton.vue b/src/components/graph/selectionToolbox/EditModelButton.vue
index 702b18394..87223d219 100644
--- a/src/components/graph/selectionToolbox/EditModelButton.vue
+++ b/src/components/graph/selectionToolbox/EditModelButton.vue
@@ -1,6 +1,6 @@
 <template>
   <Button
-    v-show="canvasStore.nodeSelected || canvasStore.groupSelected"
+    v-show="isImageOutputSelected"
     v-tooltip.top="{
       value: t('commands.Comfy_Canvas_AddEditModelStep.label'),
       showDelay: 1000
@@ -14,12 +14,24 @@
 
 <script setup lang="ts">
 import Button from 'primevue/button'
+import { computed } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import { useCommandStore } from '@/stores/commandStore'
 import { useCanvasStore } from '@/stores/graphStore'
+import { isLGraphNode } from '@/utils/litegraphUtil'
 
 const { t } = useI18n()
 const commandStore = useCommandStore()
 const canvasStore = useCanvasStore()
+
+const isImageOutputOrEditModelNode = (node: unknown) =>
+  isLGraphNode(node) &&
+  (node.images?.length || node.type === 'workflow>FLUX.1 Kontext Image Edit')
+
+const isImageOutputSelected = computed(
+  () =>
+    canvasStore.selectedItems.length === 1 &&
+    isImageOutputOrEditModelNode(canvasStore.selectedItems[0])
+)
 </script>