specify allowed blocks

10up · Mar 29, 2024 · 7d3e318 · 7d3e318
1 parent 9ff6271
commit 7d3e318
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/src/js/extend-image-block-generate-image.js b/src/js/extend-image-block-generate-image.js
@@ -2,6 +2,13 @@ import { Button } from '@wordpress/components';
 import { _nx } from '@wordpress/i18n';
 import { useBlockProps } from '@wordpress/block-editor';
 
+const allowedCoreBlocks = [
+	'core/image',
+	'core/gallery',
+	'core/media-text',
+	'core/cover',
+];
+
 /**
  * Adds a `Generate image` button to the media-related blocks.
  * @see {@link https://github.com/10up/classifai/issues/724}
@@ -22,6 +29,10 @@ function addImageGenerationLink( Component ) {
 
 		const { 'data-type': blockName } = blockProps;
 
+		if ( ! allowedCoreBlocks.includes( blockName ) ) {
+			return <Component { ...props } />;
+		}
+
 		let isSingle = 1;
 
 		if ( blockName && 'core/gallery' === blockName ) {