Jetpack AI: add general image generator style selector (#39917)

* import constants from ai-client. Add guessStyle and imageStyles to use-ai-image hook * allow style as image generation param, default to empty string as API endpoint expects it * pass down guessStyle and imageStyles to modal. Append style to generation call * add guess styling flow and style selector on modal * changelog * pass style param on regenerate and tryagain handlers * add todo comment on select disabled prop --------- Co-authored-by: Douglas Henri <[email protected]>
Automattic · Oct 28, 2024 · be9762c · be9762c
1 parent 291b09e
commit be9762c
Show file tree

Hide file tree

Showing 7 changed files with 202 additions and 40 deletions.
diff --git a/projects/js-packages/ai-client/changelog/add-jetpack-ai-general-image-style-selector b/projects/js-packages/ai-client/changelog/add-jetpack-ai-general-image-style-selector
@@ -0,0 +1,4 @@
+Significance: minor
+Type: added
+
+AI Client: export image generator hook constants
diff --git a/projects/js-packages/ai-client/src/index.ts b/projects/js-packages/ai-client/src/index.ts
@@ -16,6 +16,7 @@ export { default as useAudioTranscription } from './hooks/use-audio-transcriptio
 export { default as useTranscriptionPostProcessing } from './hooks/use-transcription-post-processing/index.js';
 export { default as useAudioValidation } from './hooks/use-audio-validation/index.js';
 export { default as useImageGenerator } from './hooks/use-image-generator/index.js';
+export * from './hooks/use-image-generator/constants.js';
 
 /*
  * Components: Icons

diff --git a/projects/plugins/jetpack/changelog/add-jetpack-ai-general-image-style-selector b/projects/plugins/jetpack/changelog/add-jetpack-ai-general-image-style-selector
@@ -0,0 +1,4 @@
+Significance: minor
+Type: other
+
+Jetpack AI: add styles dropdown on AI image generator modal
diff --git a/...extensions/plugins/ai-assistant-plugin/components/ai-image/components/ai-image-modal.scss b/...extensions/plugins/ai-assistant-plugin/components/ai-image/components/ai-image-modal.scss
@@ -4,7 +4,6 @@
 		margin-top: 32px;
 		flex-direction: column;
 		justify-content: center;
-		align-items: center;
 		gap: 8px;
 
 		.jetpack-ai-fair-usage-notice {

diff --git a/.../extensions/plugins/ai-assistant-plugin/components/ai-image/components/ai-image-modal.tsx b/.../extensions/plugins/ai-assistant-plugin/components/ai-image/components/ai-image-modal.tsx
@@ -1,11 +1,19 @@
 /**
  * External dependencies
  */
-import { AiModalPromptInput } from '@automattic/jetpack-ai-client';
-import { Button } from '@wordpress/components';
+import {
+	AiModalPromptInput,
+	IMAGE_STYLE_NONE,
+	IMAGE_STYLE_AUTO,
+	ImageStyleObject,
+	ImageStyle,
+} from '@automattic/jetpack-ai-client';
+import { useAnalytics } from '@automattic/jetpack-shared-extension-utils';
+import { Button, SelectControl } from '@wordpress/components';
 import { useCallback, useRef, useState, useEffect } from '@wordpress/element';
 import { __ } from '@wordpress/i18n';
 import { Icon, external } from '@wordpress/icons';
+import debugFactory from 'debug';
 /**
  * Internal dependencies
  */
@@ -17,6 +25,8 @@ import UsageCounter from './usage-counter';
 
 const FEATURED_IMAGE_UPGRADE_PROMPT_PLACEMENT = 'ai-image-generator';
 
+const debug = debugFactory( 'jetpack-ai:ai-image-modal' );
+
 export default function AiImageModal( {
 	title,
 	cost,
@@ -41,6 +51,8 @@ export default function AiImageModal( {
 	autoStart = false,
 	autoStartAction = null,
 	instructionsPlaceholder = null,
+	imageStyles = [],
+	onGuessStyle = null,
 }: {
 	title: string;
 	cost: number;
@@ -49,8 +61,8 @@ export default function AiImageModal( {
 	images: CarrouselImages;
 	currentIndex: number;
 	onClose: () => void;
-	onTryAgain: ( { userPrompt }: { userPrompt?: string } ) => void;
-	onGenerate: ( { userPrompt }: { userPrompt?: string } ) => void;
+	onTryAgain: ( { userPrompt, style }: { userPrompt?: string; style?: string } ) => void;
+	onGenerate: ( { userPrompt, style }: { userPrompt?: string; style?: string } ) => void;
 	generating: boolean;
 	notEnoughRequests: boolean;
 	requireUpgrade: boolean;
@@ -64,20 +76,50 @@ export default function AiImageModal( {
 	handleNextImage: () => void;
 	acceptButton: React.JSX.Element;
 	autoStart?: boolean;
-	autoStartAction?: ( { userPrompt }: { userPrompt?: string } ) => void;
+	autoStartAction?: ( { userPrompt, style }: { userPrompt?: string; style?: string } ) => void;
 	generateButtonLabel: string;
 	instructionsPlaceholder: string;
+	imageStyles?: Array< ImageStyleObject >;
+	onGuessStyle?: ( userPrompt: string ) => Promise< ImageStyle >;
 } ) {
+	const { tracks } = useAnalytics();
+	const { recordEvent: recordTracksEvent } = tracks;
 	const [ userPrompt, setUserPrompt ] = useState( '' );
 	const triggeredAutoGeneration = useRef( false );
+	const [ showStyleSelector, setShowStyleSelector ] = useState( false );
+	const [ style, setStyle ] = useState< ImageStyle >( null );
+	const [ styles, setStyles ] = useState< Array< ImageStyleObject > >( imageStyles || [] );
 
 	const handleTryAgain = useCallback( () => {
-		onTryAgain?.( { userPrompt } );
-	}, [ onTryAgain, userPrompt ] );
+		onTryAgain?.( { userPrompt, style } );
+	}, [ onTryAgain, userPrompt, style ] );
 
-	const handleGenerate = useCallback( () => {
-		onGenerate?.( { userPrompt } );
-	}, [ onGenerate, userPrompt ] );
+	const handleGenerate = useCallback( async () => {
+		if ( style === IMAGE_STYLE_AUTO ) {
+			recordTracksEvent( 'jetpack_ai_general_image_guess_style', {
+				context: 'block-editor',
+				tool: 'image',
+			} );
+			const guessedStyle = ( await onGuessStyle( userPrompt ) ) || IMAGE_STYLE_NONE;
+			setStyle( guessedStyle );
+			debug( 'guessed style', guessedStyle );
+			onGenerate?.( { userPrompt, style: guessedStyle } );
+		} else {
+			onGenerate?.( { userPrompt, style } );
+		}
+	}, [ onGenerate, userPrompt, style, onGuessStyle, recordTracksEvent ] );
+
+	const updateStyle = useCallback(
+		( imageStyle: ImageStyle ) => {
+			debug( 'change style', imageStyle );
+			setStyle( imageStyle );
+			recordTracksEvent( 'jetpack_ai_image_generator_switch_style', {
+				context: 'block-editor',
+				style: imageStyle,
+			} );
+		},
+		[ setStyle, recordTracksEvent ]
+	);
 
 	// Controllers
 	const instructionsDisabled = notEnoughRequests || generating || requireUpgrade;
@@ -99,11 +141,46 @@ export default function AiImageModal( {
 		}
 	}, [ placement, handleGenerate, autoStart, autoStartAction, userPrompt, open ] );
 
+	// initialize styles dropdown
+	useEffect( () => {
+		if ( imageStyles && imageStyles.length > 0 ) {
+			// Sort styles to have "None" and "Auto" first
+			setStyles(
+				[
+					imageStyles.find( ( { value } ) => value === IMAGE_STYLE_NONE ),
+					imageStyles.find( ( { value } ) => value === IMAGE_STYLE_AUTO ),
+					...imageStyles.filter(
+						( { value } ) => ! [ IMAGE_STYLE_NONE, IMAGE_STYLE_AUTO ].includes( value )
+					),
+				].filter( v => v ) // simplest way to get rid of empty values
+			);
+			setShowStyleSelector( true );
+			setStyle( IMAGE_STYLE_NONE );
+		}
+	}, [ imageStyles ] );
+
 	return (
 		<>
 			{ open && (
 				<AiAssistantModal handleClose={ onClose } title={ title }>
 					<div className="ai-image-modal__content">
+						{ showStyleSelector && (
+							<div style={ { display: 'flex', alignItems: 'center', gap: 16 } }>
+								<div style={ { fontWeight: 500, flexGrow: 1 } }>
+									{ __( 'Generate image', 'jetpack' ) }
+								</div>
+								<div>
+									<SelectControl
+										__nextHasNoMarginBottom
+										value={ style }
+										options={ styles }
+										onChange={ updateStyle }
+										// TODO: disable when necessary
+										// disabled={ isBusy || requireUpgrade }
+									/>
+								</div>
+							</div>
+						) }
 						<AiModalPromptInput
 							prompt={ userPrompt }
 							setPrompt={ setUserPrompt }

diff --git a/...pack/extensions/plugins/ai-assistant-plugin/components/ai-image/general-purpose-image.tsx b/...pack/extensions/plugins/ai-assistant-plugin/components/ai-image/general-purpose-image.tsx
@@ -5,6 +5,7 @@ import { useAnalytics } from '@automattic/jetpack-shared-extension-utils';
 import { Button } from '@wordpress/components';
 import { useCallback, useState } from '@wordpress/element';
 import { __, sprintf } from '@wordpress/i18n';
+import debugFactory from 'debug';
 /**
  * Internal dependencies
  */
@@ -30,6 +31,8 @@ type SetImageCallbackProps = {
 	url: string;
 };
 
+const debug = debugFactory( 'jetpack-ai:general-purpose-image' );
+
 export default function GeneralPurposeImage( {
 	placement,
 	onClose = () => {},
@@ -68,6 +71,8 @@ export default function GeneralPurposeImage( {
 		currentPointer,
 		images,
 		pointer,
+		imageStyles,
+		guessStyle,
 	} = useAiImage( {
 		cost: generalImageCost,
 		autoStart: false,
@@ -81,22 +86,27 @@ export default function GeneralPurposeImage( {
 	}, [ onClose ] );
 
 	const handleGenerate = useCallback(
-		( { userPrompt }: { userPrompt?: string } ) => {
+		async ( { userPrompt, style }: { userPrompt?: string; style?: string } ) => {
+			debug( 'handleGenerate', userPrompt, style );
+
 			// track the generate image event
 			recordEvent( 'jetpack_ai_general_image_generation_generate_image', {
 				placement,
 				model: generalImageActiveModel,
 				site_type: siteType,
+				style,
 			} );
-
-			processImageGeneration( { userPrompt, postContent, notEnoughRequests } ).catch( error => {
-				recordEvent( 'jetpack_ai_general_image_generation_error', {
-					placement,
-					error: error?.message,
-					model: generalImageActiveModel,
-					site_type: siteType,
-				} );
-			} );
+			processImageGeneration( { userPrompt, postContent, notEnoughRequests, style } ).catch(
+				error => {
+					recordEvent( 'jetpack_ai_general_image_generation_error', {
+						placement,
+						error: error?.message,
+						model: generalImageActiveModel,
+						site_type: siteType,
+						style,
+					} );
+				}
+			);
 		},
 		[
 			recordEvent,
@@ -110,23 +120,27 @@ export default function GeneralPurposeImage( {
 	);
 
 	const handleRegenerate = useCallback(
-		( { userPrompt }: { userPrompt?: string } ) => {
+		( { userPrompt, style }: { userPrompt?: string; style?: string } ) => {
+			debug( 'handleRegenerate', userPrompt );
 			// track the regenerate image event
 			recordEvent( 'jetpack_ai_general_image_generation_generate_another_image', {
 				placement,
 				model: generalImageActiveModel,
 				site_type: siteType,
+				style,
 			} );
 
 			setCurrent( crrt => crrt + 1 );
-			processImageGeneration( { userPrompt, postContent, notEnoughRequests } ).catch( error => {
-				recordEvent( 'jetpack_ai_general_image_generation_error', {
-					placement,
-					error: error?.message,
-					model: generalImageActiveModel,
-					site_type: siteType,
-				} );
-			} );
+			processImageGeneration( { userPrompt, postContent, notEnoughRequests, style } ).catch(
+				error => {
+					recordEvent( 'jetpack_ai_general_image_generation_error', {
+						placement,
+						error: error?.message,
+						model: generalImageActiveModel,
+						site_type: siteType,
+					} );
+				}
+			);
 		},
 		[
 			recordEvent,
@@ -141,22 +155,26 @@ export default function GeneralPurposeImage( {
 	);
 
 	const handleTryAgain = useCallback(
-		( { userPrompt }: { userPrompt?: string } ) => {
+		( { userPrompt, style }: { userPrompt?: string; style?: string } ) => {
+			debug( 'handleTryAgain', userPrompt );
 			// track the try again event
 			recordEvent( 'jetpack_ai_general_image_generation_try_again', {
 				placement,
 				model: generalImageActiveModel,
 				site_type: siteType,
+				style,
 			} );
 
-			processImageGeneration( { userPrompt, postContent, notEnoughRequests } ).catch( error => {
-				recordEvent( 'jetpack_ai_general_image_generation_error', {
-					placement,
-					error: error?.message,
-					model: generalImageActiveModel,
-					site_type: siteType,
-				} );
-			} );
+			processImageGeneration( { userPrompt, postContent, notEnoughRequests, style } ).catch(
+				error => {
+					recordEvent( 'jetpack_ai_general_image_generation_error', {
+						placement,
+						error: error?.message,
+						model: generalImageActiveModel,
+						site_type: siteType,
+					} );
+				}
+			);
 		},
 		[
 			recordEvent,
@@ -255,6 +273,8 @@ export default function GeneralPurposeImage( {
 			acceptButton={ acceptButton }
 			generateButtonLabel={ pointer?.current > 0 ? generateAgainText : generateText }
 			instructionsPlaceholder={ __( "Describe the image you'd like to create.", 'jetpack' ) }
+			imageStyles={ imageStyles }
+			onGuessStyle={ guessStyle }
 		/>
 	);
 }