diff --git a/filesystem/android_device.go b/filesystem/android_device.go
index 86f853c..fe30a26 100644
--- a/filesystem/android_device.go
+++ b/filesystem/android_device.go
@@ -208,17 +208,6 @@
 		toCopy = append(toCopy, targetFilesZipCopy{a.partitionProps.Recovery_partition_name, "VENDOR_BOOT/RAMDISK"})
 	}
 
-	// Create an IMAGES/ subdirectory
-	builder.Command().Textf("mkdir -p %s/IMAGES/", targetFilesDir.String())
-	if a.deviceProps.Bootloader != nil {
-		builder.Command().Textf("cp ").Input(android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Bootloader))).Textf(" %s/IMAGES/bootloader", targetFilesDir.String())
-	}
-	// Copy the vbmeta img files to IMAGES/
-	for _, vbmetaPartition := range a.partitionProps.Vbmeta_partitions {
-		vbmetaInfo, _ := android.OtherModuleProvider(ctx, ctx.GetDirectDepWithTag(vbmetaPartition, filesystemDepTag), vbmetaPartitionProvider)
-		builder.Command().Textf("cp ").Input(vbmetaInfo.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
-	}
-
 	for _, zipCopy := range toCopy {
 		if zipCopy.srcModule == nil {
 			continue
@@ -239,7 +228,6 @@
 			// Create the ROOT partition in target_files.zip
 			builder.Command().Textf("rsync --links --exclude=system/* %s/ -r %s/ROOT", fsInfo.RootDir, targetFilesDir.String())
 		}
-		builder.Command().Textf("cp %s %s/IMAGES/", fsInfo.Output, targetFilesDir.String())
 	}
 	// Copy cmdline, kernel etc. files of boot images
 	if a.partitionProps.Vendor_boot_partition_name != nil {
@@ -276,6 +264,8 @@
 		builder.Command().Textf("cp %s %s/OTA/android-info.txt", android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Android_info)), targetFilesDir)
 	}
 
+	a.copyImagesToTargetZip(ctx, builder, targetFilesDir)
+
 	builder.Command().
 		BuiltTool("soong_zip").
 		Text("-d").
@@ -286,6 +276,26 @@
 	builder.Build("target_files_"+ctx.ModuleName(), "Build target_files.zip")
 }
 
+func (a *androidDevice) copyImagesToTargetZip(ctx android.ModuleContext, builder *android.RuleBuilder, targetFilesDir android.WritablePath) {
+	// Create an IMAGES/ subdirectory
+	builder.Command().Textf("mkdir -p %s/IMAGES", targetFilesDir.String())
+	if a.deviceProps.Bootloader != nil {
+		builder.Command().Textf("cp ").Input(android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Bootloader))).Textf(" %s/IMAGES/bootloader", targetFilesDir.String())
+	}
+	// Copy the filesystem ,boot and vbmeta img files to IMAGES/
+	ctx.VisitDirectDepsProxyWithTag(filesystemDepTag, func(child android.ModuleProxy) {
+		if info, ok := android.OtherModuleProvider(ctx, child, FilesystemProvider); ok {
+			builder.Command().Textf("cp ").Input(info.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
+		} else if info, ok := android.OtherModuleProvider(ctx, child, BootimgInfoProvider); ok {
+			builder.Command().Textf("cp ").Input(info.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
+		} else if info, ok := android.OtherModuleProvider(ctx, child, vbmetaPartitionProvider); ok {
+			builder.Command().Textf("cp ").Input(info.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
+		} else {
+			ctx.ModuleErrorf("Module %s does not provide an .img file output for target_files.zip", child.Name())
+		}
+	})
+}
+
 func (a *androidDevice) getFilesystemInfo(ctx android.ModuleContext, depName string) FilesystemInfo {
 	fsMod := ctx.GetDirectDepWithTag(depName, filesystemDepTag)
 	fsInfo, ok := android.OtherModuleProvider(ctx, fsMod, FilesystemProvider)
diff --git a/filesystem/bootimg.go b/filesystem/bootimg.go
index 0a3a177..803e981 100644
--- a/filesystem/bootimg.go
+++ b/filesystem/bootimg.go
@@ -235,6 +235,7 @@
 		Kernel:     b.getKernelPath(ctx),
 		Dtb:        b.getDtbPath(ctx),
 		Bootconfig: b.getBootconfigPath(ctx),
+		Output:     output,
 	})
 }
 
@@ -245,6 +246,7 @@
 	Kernel     android.Path
 	Dtb        android.Path
 	Bootconfig android.Path
+	Output     android.Path
 }
 
 func (b *bootimg) getKernelPath(ctx android.ModuleContext) android.Path {