diff --git a/src/efifile.c b/src/efifile.c
index 4e75d3f..a7a9cf3 100644
--- a/src/efifile.c
+++ b/src/efifile.c
@@ -145,6 +145,7 @@ void efi_extract ( EFI_HANDLE handle ) {
 		CHAR16 name[ VDISK_NAME_LEN + 1 /* WNUL */ ];
 	} __attribute__ (( packed )) info;
 	char name[ VDISK_NAME_LEN + 1 /* NUL */ ];
+	struct vdisk_file *wim = NULL;
 	struct vdisk_file *vfile;
 	EFI_FILE_PROTOCOL *root;
 	EFI_FILE_PROTOCOL *file;
@@ -212,17 +213,22 @@ void efi_extract ( EFI_HANDLE handle ) {
 		} else if ( wcscasecmp ( ( wname + ( wcslen ( wname ) - 4 ) ),
 					 L".wim" ) == 0 ) {
 			DBG ( "...found WIM file %ls\n", wname );
-			vdisk_patch_file ( vfile, patch_wim );
-			if ( ( ! bootmgfw ) &&
-			     ( bootmgfw = wim_add_file ( vfile, cmdline_index,
-							 bootmgfw_path,
-							 efi_bootarch() ) ) ) {
-				DBG ( "...extracted %ls\n", bootmgfw_path );
-			}
-			wim_add_files ( vfile, cmdline_index, efi_wim_paths );
+			wim = vfile;
 		}
 	}
 
+	/* Process WIM image */
+	if ( wim ) {
+		vdisk_patch_file ( wim, patch_wim );
+		if ( ( ! bootmgfw ) &&
+		     ( bootmgfw = wim_add_file ( wim, cmdline_index,
+						 bootmgfw_path,
+						 efi_bootarch() ) ) ) {
+			DBG ( "...extracted %ls\n", bootmgfw_path );
+		}
+		wim_add_files ( wim, cmdline_index, efi_wim_paths );
+	}
+
 	/* Check that we have a boot file */
 	if ( ! bootmgfw ) {
 		die ( "FATAL: no %ls or bootmgfw.efi found\n",
diff --git a/src/main.c b/src/main.c
index 67fc859..f7165b2 100644
--- a/src/main.c
+++ b/src/main.c
@@ -77,6 +77,9 @@ static const wchar_t *wim_paths[] = {
 /** bootmgr.exe file */
 static struct vdisk_file *bootmgr;
 
+/** WIM image file */
+static struct vdisk_file *bootwim;
+
 /** Minimal length of embedded bootmgr.exe */
 #define BOOTMGR_MIN_LEN 16384
 
@@ -364,14 +367,7 @@ static int add_file ( const char *name, void *data, size_t len ) {
 	} else if ( strcasecmp ( ( name + strlen ( name ) - 4 ),
 				 ".wim" ) == 0 ) {
 		DBG ( "...found WIM file %s\n", name );
-		vdisk_patch_file ( file, patch_wim );
-		if ( ( ! bootmgr ) &&
-		     ( bootmgr = wim_add_file ( file, cmdline_index,
-						bootmgr_path,
-						L"bootmgr.exe" ) ) ) {
-			DBG ( "...extracted bootmgr.exe\n" );
-		}
-		wim_add_files ( file, cmdline_index, wim_paths );
+		bootwim = file;
 	}
 
 	return 0;
@@ -447,6 +443,18 @@ int main ( void ) {
 	if ( cpio_extract ( initrd, initrd_len, add_file ) != 0 )
 		die ( "FATAL: could not extract initrd files\n" );
 
+	/* Process WIM image */
+	if ( bootwim ) {
+		vdisk_patch_file ( bootwim, patch_wim );
+		if ( ( ! bootmgr ) &&
+		     ( bootmgr = wim_add_file ( bootwim, cmdline_index,
+						bootmgr_path,
+						L"bootmgr.exe" ) ) ) {
+			DBG ( "...extracted bootmgr.exe\n" );
+		}
+		wim_add_files ( bootwim, cmdline_index, wim_paths );
+	}
+
 	/* Add INT 13 drive */
 	callback.drive = initialise_int13();
 
diff --git a/test/testwimboot b/test/testwimboot
index 642ee53..638abf8 100755
--- a/test/testwimboot
+++ b/test/testwimboot
@@ -59,6 +59,7 @@ def ipxe_script(uuid, version, arch, bootmgr, bcd, bootsdi, bootargs):
     #!ipxe
     kernel ../wimboot {bootargs}
     initrd -n qr.txt qr-{uuid}.txt qr.txt
+    initrd ../images/{version}/{arch}/sources/boot.wim boot.wim
     initrd ../winpeshl.ini winpeshl.ini
     """).lstrip()
     if bootmgr:
@@ -74,7 +75,6 @@ def ipxe_script(uuid, version, arch, bootmgr, bcd, bootsdi, bootargs):
         initrd ../images/{version}/{arch}/boot/boot.sdi boot.sdi
         """).lstrip()
     script += textwrap.dedent(f"""
-    initrd ../images/{version}/{arch}/sources/boot.wim boot.wim
     boot
     """).lstrip()
     return script