Source: Features/ImageGeneration.php

<?php

namespace Classifai\Features;

use Classifai\Services\ImageProcessing;
use Classifai\Providers\OpenAI\DallE;
use WP_REST_Server;
use WP_REST_Request;
use WP_Error;

use function Classifai\get_asset_info;
use function Classifai\render_disable_feature_link;

/**
 * Class ImageGeneration
 */
class ImageGeneration extends Feature {
	/**
	 * ID of the current feature.
	 *
	 * @var string
	 */
	const ID = 'feature_image_generation';

	/**
	 * Constructor.
	 */
	public function __construct() {
		$this->label = __( 'Image Generation', 'classifai' );

		// Contains all providers that are registered to the service.
		$this->provider_instances = $this->get_provider_instances( ImageProcessing::get_service_providers() );

		// Contains just the providers this feature supports.
		$this->supported_providers = [
			DallE::ID => __( 'OpenAI DALL·E 3', 'classifai' ),
		];
	}

	/**
	 * Set up necessary hooks.
	 *
	 * We utilize this so we can register the REST route.
	 */
	public function setup() {
		parent::setup();
		add_action( 'rest_api_init', [ $this, 'register_endpoints' ] );
	}

	/**
	 * Set up necessary hooks.
	 */
	public function feature_setup() {
		add_action( 'admin_menu', [ $this, 'register_generate_media_page' ], 0 );
		add_action( 'admin_enqueue_scripts', [ $this, 'enqueue_admin_scripts' ] );
		add_action( 'print_media_templates', [ $this, 'print_media_templates' ] );
	}

	/**
	 * Register any needed endpoints.
	 */
	public function register_endpoints() {
		$route = 'generate-image';

		/**
		 * Filter the arguments for the REST route.
		 *
		 * This allows for adding or modifying the arguments for the route.
		 * The filter name is dynamic and based on the route.
		 * Example: classifai_feature_image_generation_rest_route_generate-image_args
		 *
		 * @since 3.0.0
		 * @hook classifai_{feature}_rest_route_{route}_args
		 *
		 * @param {array} $args Array of arguments for the REST route.
		 *
		 * @return {array} Modified array of arguments.
		 */
		$args = apply_filters(
			'classifai_' . static::ID . '_rest_route_' . $route . '_args',
			[
				'methods'             => WP_REST_Server::READABLE,
				'callback'            => [ $this, 'rest_endpoint_callback' ],
				'args'                => [
					'prompt' => [
						'required'          => true,
						'type'              => 'string',
						'sanitize_callback' => 'sanitize_text_field',
						'validate_callback' => 'rest_validate_request_arg',
						'description'       => esc_html__( 'Prompt used to generate an image', 'classifai' ),
					],
				],
				'permission_callback' => [ $this, 'generate_image_permissions_check' ],
			]
		);

		register_rest_route(
			'classifai/v1',
			$route,
			$args
		);
	}

	/**
	 * Check if a given request has access to generate an image.
	 *
	 * This check ensures we have a valid user with proper capabilities
	 * making the request, that we are properly authenticated with OpenAI
	 * and that image generation is turned on.
	 *
	 * @return WP_Error|bool
	 */
	public function generate_image_permissions_check() {
		// Ensure the feature is enabled. Also runs a user check.
		if ( ! $this->is_feature_enabled() ) {
			return new WP_Error( 'not_enabled', esc_html__( 'Image generation not currently enabled.', 'classifai' ) );
		}

		return true;
	}

	/**
	 * Generic request handler for all our custom routes.
	 *
	 * @param WP_REST_Request $request The full request object.
	 * @return \WP_REST_Response
	 */
	public function rest_endpoint_callback( WP_REST_Request $request ) {
		$route = $request->get_route();

		if ( strpos( $route, '/classifai/v1/generate-image' ) === 0 ) {
			return rest_ensure_response(
				$this->run(
					$request->get_param( 'prompt' ),
					'image_gen',
					$request->get_params(),
				)
			);
		}

		return parent::rest_endpoint_callback( $request );
	}

	/**
	 * Registers a Media > Generate Image submenu.
	 */
	public function register_generate_media_page() {
		if ( ! $this->is_feature_enabled() ) {
			return;
		}

		$settings         = $this->get_settings();
		$provider_id      = $settings['provider'];
		$number_of_images = absint( $settings[ $provider_id ]['number_of_images'] );

		add_submenu_page(
			'upload.php',
			$number_of_images > 1 ? esc_html__( 'Generate Images', 'classifai' ) : esc_html__( 'Generate Image', 'classifai' ),
			$number_of_images > 1 ? esc_html__( 'Generate Images', 'classifai' ) : esc_html__( 'Generate Image', 'classifai' ),
			'upload_files',
			esc_url( admin_url( 'upload.php?action=classifai-generate-image&mode=grid' ) ),
			''
		);
	}

	/**
	 * Enqueue the admin scripts.
	 *
	 * @since 2.4.0 Use get_asset_info to get the asset version and dependencies.
	 *
	 * @param string $hook_suffix The current admin page.
	 */
	public function enqueue_admin_scripts( string $hook_suffix = '' ) {
		if ( 'post.php' !== $hook_suffix && 'post-new.php' !== $hook_suffix && 'upload.php' !== $hook_suffix ) {
			return;
		}

		if ( ! $this->is_feature_enabled() ) {
			return;
		}

		$settings         = $this->get_settings();
		$provider_id      = $settings['provider'];
		$number_of_images = absint( $settings[ $provider_id ]['number_of_images'] );

		wp_enqueue_media();

		wp_enqueue_style(
			'classifai-image-processing-style',
			CLASSIFAI_PLUGIN_URL . 'dist/media-modal.css',
			[],
			get_asset_info( 'media-modal', 'version' ),
			'all'
		);

		wp_enqueue_script(
			'classifai-generate-images',
			CLASSIFAI_PLUGIN_URL . 'dist/media-modal.js',
			array_merge( get_asset_info( 'media-modal', 'dependencies' ), array( 'jquery', 'wp-api' ) ),
			get_asset_info( 'media-modal', 'version' ),
			true
		);

		wp_enqueue_script(
			'classifai-inserter-media-category',
			CLASSIFAI_PLUGIN_URL . 'dist/inserter-media-category.js',
			get_asset_info( 'inserter-media-category', 'dependencies' ),
			get_asset_info( 'inserter-media-category', 'version' ),
			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.
		 *
		 * @since 2.1.0
		 * @hook classifai_dalle_caption
		 *
		 * @param {string} $caption Attribution to be added as a caption to the image.
		 *
		 * @return {string} Caption.
		 */
		$caption = apply_filters(
			'classifai_dalle_caption',
			sprintf(
				/* translators: %1$s is replaced with the OpenAI DALL·E URL */
				esc_html__( 'Image generated by <a href="%s">OpenAI\'s DALL·E</a>', 'classifai' ),
				'https://openai.com/research/dall-e'
			)
		);

		wp_localize_script(
			'classifai-generate-images',
			'classifaiDalleData',
			[
				'endpoint'   => 'classifai/v1/generate-image',
				'tabText'    => $number_of_images > 1 ? esc_html__( 'Generate images', 'classifai' ) : esc_html__( 'Generate image', 'classifai' ),
				'errorText'  => esc_html__( 'Something went wrong. No results found', 'classifai' ),
				'buttonText' => esc_html__( 'Select image', 'classifai' ),
				'caption'    => $caption,
			]
		);

		if ( 'upload.php' === $hook_suffix ) {
			$action = isset( $_GET['action'] ) ? sanitize_key( wp_unslash( $_GET['action'] ) ) : ''; // phpcs:ignore WordPress.Security.NonceVerification.Recommended

			if ( 'classifai-generate-image' === $action ) {
				wp_enqueue_script(
					'classifai-generate-images-media-upload',
					CLASSIFAI_PLUGIN_URL . 'dist/generate-image-media-upload.js',
					array_merge( get_asset_info( 'generate-image-media-upload', 'dependencies' ), array( 'jquery' ) ),
					get_asset_info( 'classifai-generate-images-media-upload', 'version' ),
					true
				);

				wp_localize_script(
					'classifai-generate-images-media-upload',
					'classifaiGenerateImages',
					[
						'upload_url' => esc_url( admin_url( 'upload.php' ) ),
					]
				);
			}
		}
	}

	/**
	 * Print the templates we need for our media modal integration.
	 */
	public function print_media_templates() {
		if ( ! $this->is_feature_enabled() ) {
			return;
		}

		$settings          = $this->get_settings();
		$provider_id       = $settings['provider'];
		$number_of_images  = absint( $settings[ $provider_id ]['number_of_images'] );
		$provider_instance = $this->get_feature_provider_instance( $provider_id );
		?>

		<?php // Template for the Generate images tab content. Includes prompt input. ?>
		<script type="text/html" id="tmpl-dalle-prompt">
			<div class="prompt-view">
				<p>
					<?php
					if ( $number_of_images > 1 ) {
						esc_html_e( 'Enter a prompt below to generate images.', 'classifai' );
					} else {
						esc_html_e( 'Enter a prompt below to generate an image.', 'classifai' );
					}
					?>
				</p>
				<p>
					<?php
					if ( $number_of_images > 1 ) {
						esc_html_e( 'Once images are generated, choose one or more of those to import into your Media Library and then choose one image to insert.', 'classifai' );
					} else {
						esc_html_e( 'Once an image is generated, you can import it into your Media Library and then select to insert.', 'classifai' );
					}
					?>
				</p>
				<textarea class="prompt" placeholder="<?php esc_attr_e( 'Enter prompt', 'classifai' ); ?>" rows="4" maxlength="<?php echo absint( $provider_instance->max_prompt_chars ); ?>"></textarea>
				<button type="button" class="button button-secondary button-large button-generate">
					<?php
					if ( $number_of_images > 1 ) {
						esc_html_e( 'Generate images', 'classifai' );
					} else {
						esc_html_e( 'Generate image', 'classifai' );
					}
					?>
				</button>
				<span class="error"></span>
			</div>
			<div class="generated-images">
				<h2 class="prompt-text hidden">
					<?php
					if ( $number_of_images > 1 ) {
						esc_html_e( 'Images generated from prompt:', 'classifai' );
					} else {
						esc_html_e( 'Image generated from prompt:', 'classifai' );
					}
					?>
					<span></span>
				</h2>
				<span class="spinner"></span>
				<ul></ul>
				<p>
					<?php render_disable_feature_link( 'feature_image_generation' ); ?>
				</p>
			</div>
		</script>

		<?php
		// Template for a single generated image.
		/* phpcs:disable WordPressVIPMinimum.Security.Mustache.OutputNotation */
		?>
		<script type="text/html" id="tmpl-dalle-image">
			<div class="generated-image">
				<img src="data:image/png;base64,{{{ data.url }}}" />
				<button type="button" class="components-button button-secondary button-import"><?php esc_html_e( 'Import into Media Library', 'classifai' ); ?></button>
				<button type="button" class="components-button is-tertiary button-import-insert"><?php esc_html_e( 'Import and Insert', 'classifai' ); ?></button>
				<span class="spinner"></span>
				<span class="error"></span>
			</div>
		</script>
		<?php
		/* phpcs:enable WordPressVIPMinimum.Security.Mustache.OutputNotation */
	}

	/**
	 * Assigns user roles to the $roles array.
	 */
	public function setup_roles() {
		$default_settings = $this->get_default_settings();

		// Get all roles that have the upload_files cap.
		$roles = get_editable_roles() ?? [];
		$roles = array_filter(
			$roles,
			function ( $role ) {
				return isset( $role['capabilities'], $role['capabilities']['upload_files'] ) && $role['capabilities']['upload_files'];
			}
		);
		$roles = array_combine( array_keys( $roles ), array_column( $roles, 'name' ) );

		/**
		 * Filter the allowed WordPress roles for image generation.
		 *
		 * @since 2.3.0
		 * @hook classifai_feature_image_generation_roles
		 *
		 * @param {array} $roles            Array of arrays containing role information.
		 * @param {array} $default_settings Default setting values.
		 *
		 * @return {array} Roles array.
		 */
		$this->roles = apply_filters( 'classifai_' . static::ID . '_roles', $roles, $default_settings );
	}

	/**
	 * Get the description for the enable field.
	 *
	 * @return string
	 */
	public function get_enable_description(): string {
		return esc_html__( 'When enabled, a new Generate images tab will be shown in the media upload flow, allowing you to generate and import images.', 'classifai' );
	}

	/**
	 * Returns true if the feature meets all the criteria to be enabled.
	 *
	 * @return bool
	 */
	public function is_feature_enabled(): bool {
		$settings           = $this->get_settings();
		$is_feature_enabled = parent::is_feature_enabled() && current_user_can( 'upload_files' );

		/** This filter is documented in includes/Classifai/Features/Feature.php */
		return apply_filters( 'classifai_' . static::ID . '_is_feature_enabled', $is_feature_enabled, $settings );
	}

	/**
	 * Returns the default settings for the feature.
	 *
	 * @return array
	 */
	public function get_feature_default_settings(): array {
		return [
			'provider' => DallE::ID,
		];
	}

	/**
	 * Generates feature setting data required for migration from
	 * ClassifAI < 3.0.0 to 3.0.0
	 *
	 * @return array
	 */
	public function migrate_settings() {
		$old_settings = get_option( 'classifai_openai_dalle', array() );
		$new_settings = $this->get_default_settings();

		$new_settings['provider'] = 'openai_dalle';

		if ( isset( $old_settings['enable_image_gen'] ) ) {
			$new_settings['status'] = $old_settings['enable_image_gen'];
		}

		if ( isset( $old_settings['number'] ) ) {
			$new_settings['openai_dalle']['number_of_images'] = $old_settings['number'];
		}

		if ( isset( $old_settings['size'] ) ) {
			$new_settings['openai_dalle']['image_size'] = $old_settings['size'];
		}

		if ( isset( $old_settings['api_key'] ) ) {
			$new_settings['openai_dalle']['api_key'] = $old_settings['api_key'];
		}

		if ( isset( $old_settings['authenticated'] ) ) {
			$new_settings['openai_dalle']['authenticated'] = $old_settings['authenticated'];
		}

		if ( isset( $old_settings['image_generation_roles'] ) ) {
			$new_settings['roles'] = $old_settings['image_generation_roles'];
		}

		if ( isset( $old_settings['image_generation_users'] ) ) {
			$new_settings['users'] = $old_settings['image_generation_users'];
		}

		if ( isset( $old_settings['image_generation_user_based_opt_out'] ) ) {
			$new_settings['user_based_opt_out'] = $old_settings['image_generation_user_based_opt_out'];
		}

		return $new_settings;
	}
}