diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 56ec929..a753e68 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -728,47 +728,33 @@
   return input_file
 
 
-def GetTargetFilesZipForCustomImagesUpdates(input_file, custom_images):
+def GetTargetFilesZipForCustomImagesUpdates(input_file, custom_images: dict):
   """Returns a target-files.zip for custom partitions update.
 
   This function modifies ab_partitions list with the desired custom partitions
   and puts the custom images into the target target-files.zip.
 
   Args:
-    input_file: The input target-files.zip filename.
+    input_file: The input target-files extracted directory
     custom_images: A map of custom partitions and custom images.
 
   Returns:
-    The filename of a target-files.zip which has renamed the custom images in
-    the IMAGES/ to their partition names.
+    The extracted dir of a target-files.zip which has renamed the custom images
+    in the IMAGES/ to their partition names.
   """
+  for custom_image in custom_images.values():
+    if not os.path.exists(os.path.join(input_file, "IMAGES", custom_image)):
+      raise ValueError("Specified custom image {} not found in target files {}, available images are {}",
+                       custom_image, input_file, os.listdir(os.path.join(input_file, "IMAGES")))
 
-  # First pass: use zip2zip to copy the target files contents, excluding
-  # the "custom" images that will be replaced.
-  target_file = common.MakeTempFile(prefix="targetfiles-", suffix=".zip")
-  cmd = ['zip2zip', '-i', input_file, '-o', target_file]
-
-  images = {}
   for custom_partition, custom_image in custom_images.items():
     default_custom_image = '{}.img'.format(custom_partition)
     if default_custom_image != custom_image:
-      src = 'IMAGES/' + custom_image
-      dst = 'IMAGES/' + default_custom_image
-      cmd.extend(['-x', dst])
-      images[dst] = src
+      src = os.path.join(input_file, 'IMAGES', custom_image)
+      dst = os.path.join(input_file, 'IMAGES', default_custom_image)
+      os.rename(src, dst)
 
-  common.RunAndCheckOutput(cmd)
-
-  # Second pass: write {custom_image}.img as {custom_partition}.img.
-  with zipfile.ZipFile(input_file, allowZip64=True) as input_zip:
-    with zipfile.ZipFile(target_file, 'a', allowZip64=True) as output_zip:
-      for dst, src in images.items():
-        data = input_zip.read(src)
-        logger.info("Update custom partition '%s'", dst)
-        common.ZipWriteStr(output_zip, dst, data)
-      output_zip.close()
-
-  return target_file
+  return input_file
 
 
 def GeneratePartitionTimestampFlags(partition_state):
@@ -1238,7 +1224,7 @@
       if len(words) == 2:
         if not words[1].isdigit():
           raise ValueError("Cannot parse value %r for option $COMPRESSION_LEVEL - only "
-                         "integers are allowed." % words[1])
+                           "integers are allowed." % words[1])
     elif o == "--security_patch_level":
       OPTIONS.security_patch_level = a
     elif o in ("--max_threads"):