| 1 |
import { GenerateForm } from '@draft/components/image-generation/GenerateForm'; |
| 2 |
import { ImagePreview } from '@draft/components/image-generation/ImagePreview'; |
| 3 |
import { useRouter } from '@draft/hooks/useRouter'; |
| 4 |
import { pageState } from '@draft/state/factory'; |
| 5 |
import { backArrow } from '@draft/svg/BackArrow'; |
| 6 |
import { generateImage } from '@shared/api/DataApi'; |
| 7 |
import { useImageGenerationStore } from '@shared/state/generate-images'; |
| 8 |
import { |
| 9 |
BaseControl, |
| 10 |
__experimentalHeading as Heading, |
| 11 |
Panel, |
| 12 |
PanelBody, |
| 13 |
} from '@wordpress/components'; |
| 14 |
import { useEffect, useRef, useState } from '@wordpress/element'; |
| 15 |
import { __ } from '@wordpress/i18n'; |
| 16 |
|
| 17 |
const usePageState = pageState('AI Image', (set) => ({ |
| 18 |
imageDetails: { src: '', id: undefined }, |
| 19 |
setImageDetails: (newState) => { |
| 20 |
set((state) => ({ ...state, imageDetails: newState })); |
| 21 |
}, |
| 22 |
})); |
| 23 |
|
| 24 |
export const GenerateImage = () => { |
| 25 |
const { |
| 26 |
imageCredits: curCredits, |
| 27 |
updateImageCredits, |
| 28 |
subtractOneCredit, |
| 29 |
aiImageOptions, |
| 30 |
} = useImageGenerationStore(); |
| 31 |
const [isGenerating, setIsGenerating] = useState(false); |
| 32 |
const [errorMessage, setErrorMessage] = useState(''); |
| 33 |
const abortController = useRef(null); |
| 34 |
const noCredits = curCredits.remaining === 0; |
| 35 |
const { imageDetails, setImageDetails } = usePageState(); |
| 36 |
const { goBack } = useRouter(); |
| 37 |
|
| 38 |
const clearImageResponse = () => setImageDetails({ src: '', id: undefined }); |
| 39 |
const handleSubmit = async (event) => { |
| 40 |
event.preventDefault(); |
| 41 |
setErrorMessage(''); |
| 42 |
if (noCredits) { |
| 43 |
abortController.current?.abort(); |
| 44 |
return; |
| 45 |
} |
| 46 |
|
| 47 |
try { |
| 48 |
setIsGenerating(true); |
| 49 |
subtractOneCredit(); |
| 50 |
abortController.current = new AbortController(); |
| 51 |
const { imageCredits, images, id } = await generateImage( |
| 52 |
aiImageOptions, |
| 53 |
abortController.current.signal, |
| 54 |
); |
| 55 |
updateImageCredits(imageCredits); |
| 56 |
setImageDetails({ src: images[0].url, id }); |
| 57 |
} catch (error) { |
| 58 |
// If the request was aborted (canceled), don't show an error |
| 59 |
if (error?.code === 20) return; |
| 60 |
// If we didn't get back any credit info, it was a server error |
| 61 |
if (!error?.imageCredits) { |
| 62 |
// Pause to prefent flickering |
| 63 |
await new Promise((resolve) => setTimeout(resolve, 1000)); |
| 64 |
setErrorMessage(error.message); |
| 65 |
// Add back the credit we subtracted |
| 66 |
updateImageCredits({ remaining: curCredits.remaining }); |
| 67 |
return; |
| 68 |
} |
| 69 |
// Probably out of credits here |
| 70 |
updateImageCredits(error.imageCredits); |
| 71 |
setErrorMessage(error.message); |
| 72 |
} finally { |
| 73 |
setIsGenerating(false); |
| 74 |
} |
| 75 |
}; |
| 76 |
|
| 77 |
useEffect(() => { |
| 78 |
if (imageDetails.src || isGenerating) return; |
| 79 |
// refocus when image is removed |
| 80 |
document.getElementById('draft-ai-image-textarea')?.focus(); |
| 81 |
}, [imageDetails.src, isGenerating]); |
| 82 |
|
| 83 |
return ( |
| 84 |
<Panel> |
| 85 |
<div className="flex h-12 items-center gap-1 pl-1"> |
| 86 |
<button |
| 87 |
className="h-9 w-9 border-0 bg-transparent" |
| 88 |
onClick={goBack} |
| 89 |
type="button" |
| 90 |
aria-label={__('Go Back', 'extendify-local')} |
| 91 |
> |
| 92 |
{backArrow} |
| 93 |
</button> |
| 94 |
<Heading className="mb-0"> |
| 95 |
{__('AI Image Generator', 'extendify-local')} |
| 96 |
</Heading> |
| 97 |
</div> |
| 98 |
<PanelBody className="-mt-4 border-t-0"> |
| 99 |
<BaseControl label={__('Image Description', 'extendify-local')}> |
| 100 |
<ImagePreview |
| 101 |
prompt={aiImageOptions.prompt} |
| 102 |
size={aiImageOptions.size} |
| 103 |
isGenerating={isGenerating} |
| 104 |
id={imageDetails?.id} |
| 105 |
src={imageDetails?.src} |
| 106 |
clearImageResponse={clearImageResponse} |
| 107 |
/> |
| 108 |
{imageDetails.src ? null : ( |
| 109 |
<form onSubmit={handleSubmit} className="flex flex-col gap-5"> |
| 110 |
<GenerateForm |
| 111 |
isGenerating={isGenerating} |
| 112 |
errorMessage={errorMessage} |
| 113 |
/> |
| 114 |
</form> |
| 115 |
)} |
| 116 |
</BaseControl> |
| 117 |
</PanelBody> |
| 118 |
</Panel> |
| 119 |
); |
| 120 |
}; |
| 121 |
|
| 122 |
export const routes = [ |
| 123 |
{ |
| 124 |
slug: 'ai-image', |
| 125 |
title: __('AI Image', 'extendify-local'), |
| 126 |
component: GenerateImage, |
| 127 |
}, |
| 128 |
]; |
| 129 |
|
