diff --git a/index.js b/index.js
index 2de357e..e545da1 100644
--- a/index.js
+++ b/index.js
@@ -3,9 +3,31 @@ const pathToBrowserBuild = require.resolve('happo-e2e/browser.build.js');
 
 const controller = new Controller();
 
+async function lazyLoadBrowserBundle(page) {
+  if (
+    await page.evaluate(() => typeof window.happoTakeDOMSnapshot === 'undefined')
+  ) {
+    await page.addScriptTag({ path: pathToBrowserBuild });
+
+    // Add timeout check for happoTakeDOMSnapshot
+    try {
+      await page.waitForFunction(
+        () => typeof window.happoTakeDOMSnapshot !== 'undefined',
+        { timeout: 10000 },
+      );
+    } catch (error) {
+      throw new Error('Timed out waiting for happoTakeDOMSnapshot to be defined');
+    }
+  }
+}
+
 module.exports = {
-  async init(contextOrPage) {
-    await contextOrPage.addInitScript({ path: pathToBrowserBuild });
+  async init(pageOrContext) {
+    if (pageOrContext) {
+      console.warn(
+        '[HAPPO] You no longer need to pass a page or context to happoPlaywright.init()',
+      );
+    }
     await controller.init();
   },
 
@@ -30,6 +52,8 @@ module.exports = {
       );
     }
 
+    await lazyLoadBrowserBundle(page);
+
     const elementHandle = handleOrLocator.elementHandle
       ? await handleOrLocator.elementHandle()
       : handleOrLocator;
diff --git a/tests/app.spec.js b/tests/app.spec.js
index c37477e..82e195c 100644
--- a/tests/app.spec.js
+++ b/tests/app.spec.js
@@ -1,11 +1,11 @@
 const { test, expect } = require('@playwright/test');
 const happoPlaywright = require('../');
 
-test.beforeEach(async ({ context }) => {
-  await happoPlaywright.init(context);
+test.beforeAll(async () => {
+  await happoPlaywright.init();
 });
 
-test.afterEach(async () => {
+test.afterAll(async () => {
   await happoPlaywright.finish();
 });