ai / includes / Abilities / Image / Generate_Image.php

Generate_Image.php in AI 0.2.0, at includes/Abilities/Image/Generate_Image.php

144 lines 3.0 KB

Raw Download Zip

1	<?php
2	/**
3	* Image generation WordPress Ability implementation.
4	*
5	* @package WordPress\AI
6	*/
7
8	declare( strict_types=1 );
9
10	namespace WordPress\AI\Abilities\Image;
11
12	use WP_Error;
13	use WordPress\AI\Abstracts\Abstract_Ability;
14	use WordPress\AI_Client\AI_Client;
15	use WordPress\AiClient\Files\Enums\FileTypeEnum;
16
17	use function WordPress\AI\get_preferred_image_models;
18
19	/**
20	* Image generation WordPress Ability.
21	*
22	* @since 0.2.0
23	*/
24	class Generate_Image extends Abstract_Ability {
25
26	/**
27	* {@inheritDoc}
28	*
29	* @since 0.2.0
30	*/
31	protected function input_schema(): array {
32	return array(
33	'type' => 'object',
34	'properties' => array(
35	'prompt' => array(
36	'type' => 'string',
37	'sanitize_callback' => 'sanitize_text_field',
38	'description' => esc_html__( 'Prompt used to generate an image.', 'ai' ),
39	),
40	),
41	'required' => array( 'prompt' ),
42	);
43	}
44
45	/**
46	* {@inheritDoc}
47	*
48	* @since 0.2.0
49	*/
50	protected function output_schema(): array {
51	return array(
52	'type' => 'string',
53	'description' => esc_html__( 'The base64 encoded image data.', 'ai' ),
54	);
55	}
56
57	/**
58	* {@inheritDoc}
59	*
60	* @since 0.2.0
61	*/
62	protected function execute_callback( $input ) {
63	// Generate the image.
64	$result = $this->generate_image( $input['prompt'] );
65
66	// If we have an error, return it.
67	if ( is_wp_error( $result ) ) {
68	return $result;
69	}
70
71	// If we have no results, return an error.
72	if ( empty( $result ) ) {
73	return new WP_Error(
74	'no_results',
75	esc_html__( 'No image was generated.', 'ai' )
76	);
77	}
78
79	// Return the image data in the format the Ability expects.
80	return sanitize_text_field( trim( $result ) );
81	}
82
83	/**
84	* {@inheritDoc}
85	*
86	* @since 0.2.0
87	*/
88	protected function permission_callback( $args ) {
89	// Ensure the user has permission to upload files.
90	if ( ! current_user_can( 'upload_files' ) ) {
91	return new WP_Error(
92	'insufficient_capabilities',
93	esc_html__( 'You do not have permission to generate images.', 'ai' )
94	);
95	}
96
97	return true;
98	}
99
100	/**
101	* {@inheritDoc}
102	*
103	* @since 0.2.0
104	*/
105	protected function meta(): array {
106	return array(
107	'show_in_rest' => true,
108	);
109	}
110
111	/**
112	* Generates an image from the given prompt.
113	*
114	* @since 0.2.0
115	*
116	* @param string $prompt The prompt to generate an image from.
117	* @return string\|\WP_Error The generated image data, or a WP_Error if there was an error.
118	*/
119	protected function generate_image( string $prompt ) { // phpcs:ignore Generic.NamingConventions.ConstructorName.OldStyle
120	// Generate the image using the AI client.
121	$file = AI_Client::prompt_with_wp_error( $prompt )
122	->as_output_file_type( FileTypeEnum::inline() )
123	->using_model_preference( ...get_preferred_image_models() )
124	->generate_image();
125
126	// If we have an error, return it.
127	if ( is_wp_error( $file ) ) {
128	return $file;
129	}
130
131	// Return the base64 encoded image data.
132	$data = $file->getBase64Data();
133
134	if ( empty( $data ) ) {
135	return new WP_Error(
136	'no_image_data',
137	esc_html__( 'No image data was generated.', 'ai' )
138	);
139	}
140
141	return $data;
142	}
143	}
144