diff --git a/modules/ai/assets/js/editor/utils/screenshot.js b/modules/ai/assets/js/editor/utils/screenshot.js
index b486593c4b4d..04d0053bd840 100644
--- a/modules/ai/assets/js/editor/utils/screenshot.js
+++ b/modules/ai/assets/js/editor/utils/screenshot.js
@@ -20,6 +20,10 @@ export const takeScreenshot = async ( template ) => {
 	// Wait for the container to render.
 	await waitForContainer( container.id );
 
+	if ( template.elements.length ) {
+		await Promise.all( template.elements.map( ( child ) => waitForContainer( child.id ) ) );
+	}
+
 	let screenshot;
 
 	try {