From 9ff6271912088587c05d003f88fe44a959611294 Mon Sep 17 00:00:00 2001
From: Siddharth Thevaril <siddharth.thevaril@gmail.com>
Date: Fri, 15 Mar 2024 21:51:25 +0530
Subject: [PATCH 1/2] feat: add 'Generate image' button to media upload
 components

---
 .eslintrc.json                                |  3 +-
 .../Classifai/Features/ImageGeneration.php    |  8 +++
 src/js/extend-image-block-generate-image.js   | 58 +++++++++++++++++++
 webpack.config.js                             |  1 +
 4 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 src/js/extend-image-block-generate-image.js

diff --git a/.eslintrc.json b/.eslintrc.json
index b73aa306d..0641969fa 100644
--- a/.eslintrc.json
+++ b/.eslintrc.json
@@ -19,7 +19,8 @@
 		"_": "readonly",
 		"File": "readonly",
 		"Headers": "readonly",
-		"requestAnimationFrame": "readonly"
+		"requestAnimationFrame": "readonly",
+		"React": "readonly"
 	},
 	"extends": ["plugin:@wordpress/eslint-plugin/recommended"],
 	"ignorePatterns": ["*.json"]
diff --git a/includes/Classifai/Features/ImageGeneration.php b/includes/Classifai/Features/ImageGeneration.php
index e2412957e..56df86aef 100644
--- a/includes/Classifai/Features/ImageGeneration.php
+++ b/includes/Classifai/Features/ImageGeneration.php
@@ -209,6 +209,14 @@ public function enqueue_admin_scripts( string $hook_suffix = '' ) {
 			true
 		);
 
+		wp_enqueue_script(
+			'classifai-extend-image-blocks',
+			CLASSIFAI_PLUGIN_URL . 'dist/extend-image-blocks.js',
+			get_asset_info( 'extend-image-blocks', 'dependencies' ),
+			get_asset_info( 'extend-image-blocks', 'version' ),
+			true
+		);
+
 		/**
 		 * Filter the default attribution added to generated images.
 		 *
diff --git a/src/js/extend-image-block-generate-image.js b/src/js/extend-image-block-generate-image.js
new file mode 100644
index 000000000..0778f63f4
--- /dev/null
+++ b/src/js/extend-image-block-generate-image.js
@@ -0,0 +1,58 @@
+import { Button } from '@wordpress/components';
+import { _nx } from '@wordpress/i18n';
+import { useBlockProps } from '@wordpress/block-editor';
+
+/**
+ * Adds a `Generate image` button to the media-related blocks.
+ * @see {@link https://github.com/10up/classifai/issues/724}
+ *
+ * @param {React.ReactNode} Component The Wrapped Media upload component.
+ * @return {React.ReactNode} The transformed Media upload component.
+ */
+function addImageGenerationLink( Component ) {
+	return function ( props ) {
+		const { render, mode, ...rest } = props;
+		let blockProps;
+
+		try {
+			blockProps = useBlockProps();
+		} catch ( e ) {
+			return <Component { ...props } />;
+		}
+
+		const { 'data-type': blockName } = blockProps;
+
+		let isSingle = 1;
+
+		if ( blockName && 'core/gallery' === blockName ) {
+			isSingle = Infinity;
+		}
+
+		return (
+			<>
+				<Component
+					{ ...rest }
+					mode="generate"
+					render={ ( { open } ) => (
+						<Button variant="tertiary" onClick={ open }>
+							{ _nx(
+								'Generate image',
+								'Generate images',
+								isSingle,
+								'Image or gallery upload',
+								'classifai'
+							) }
+						</Button>
+					) }
+				/>
+				<Component { ...props } />
+			</>
+		);
+	};
+}
+
+wp.hooks.addFilter(
+	'editor.MediaUpload',
+	'classifai/image-generation-link',
+	addImageGenerationLink
+);
diff --git a/webpack.config.js b/webpack.config.js
index 4d6c4fb8c..6acba9b7c 100644
--- a/webpack.config.js
+++ b/webpack.config.js
@@ -42,6 +42,7 @@ module.exports = {
 		'generate-image-media-upload': [
 			'./src/js/media-modal/views/generate-image-media-upload.js',
 		],
+		'extend-image-blocks': './src/js/extend-image-block-generate-image.js',
 	},
 	module: {
 		rules: [

From 754bd949211a233bdf7737d9bfc2559bf310756a Mon Sep 17 00:00:00 2001
From: Siddharth Thevaril <siddharth.thevaril@gmail.com>
Date: Fri, 29 Mar 2024 20:12:40 +0530
Subject: [PATCH 2/2] specify allowed blocks

---
 src/js/extend-image-block-generate-image.js | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/js/extend-image-block-generate-image.js b/src/js/extend-image-block-generate-image.js
index 0778f63f4..829e3db45 100644
--- 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}
@@ -11,7 +18,7 @@ import { useBlockProps } from '@wordpress/block-editor';
  */
 function addImageGenerationLink( Component ) {
 	return function ( props ) {
-		const { render, mode, ...rest } = props;
+		const { render, ...rest } = props;
 		let blockProps;
 
 		try {
@@ -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 ) {