diff --git a/releasetools.py b/releasetools.py
index 2eb5208..78692ea 100644
--- a/releasetools.py
+++ b/releasetools.py
@@ -13,25 +13,25 @@
 import common
 
 def FullOTA_InstallEnd(info):
-  OTA_InstallEnd(info)
+  OTA_InstallEnd(info, info.input_zip)
   return
 
 def IncrementalOTA_InstallEnd(info):
-  OTA_InstallEnd(info)
+  OTA_InstallEnd(info, info.target_zip)
   return
 
-def AddImage(info, basename, dest):
+def AddImage(info, input_zip, basename, dest):
   path = "IMAGES/" + basename
-  if path not in info.input_zip.namelist():
+  if path not in input_zip.namelist():
     return
 
-  data = info.input_zip.read(path)
+  data = input_zip.read(path)
   common.ZipWriteStr(info.output_zip, basename, data)
   info.script.AppendExtra('package_extract_file("%s", "%s");' % (basename, dest))
 
-def OTA_InstallEnd(info):
+def OTA_InstallEnd(info, input_zip):
   info.script.Print("Patching firmware images...")
-  AddImage(info, "dtbo.img", "/dev/block/bootdevice/by-name/dtbo")
-  AddImage(info, "vbmeta.img", "/dev/block/bootdevice/by-name/vbmeta")
-  AddImage(info, "vbmeta_system.img", "/dev/block/bootdevice/by-name/vbmeta_system")
+  AddImage(info, input_zip, "dtbo.img", "/dev/block/bootdevice/by-name/dtbo")
+  AddImage(info, input_zip, "vbmeta.img", "/dev/block/bootdevice/by-name/vbmeta")
+  AddImage(info, input_zip, "vbmeta_system.img", "/dev/block/bootdevice/by-name/vbmeta_system")
   return