diff --git a/tools/releasetools/non_ab_ota.py b/tools/releasetools/non_ab_ota.py
index c4fd809..667891c 100644
--- a/tools/releasetools/non_ab_ota.py
+++ b/tools/releasetools/non_ab_ota.py
@@ -23,6 +23,7 @@
 from check_target_files_vintf import CheckVintfIfTrebleEnabled, HasPartition
 from common import OPTIONS
 from ota_utils import UNZIP_PATTERN, FinalizeMetadata, GetPackageMetadata, PropertyFiles
+import subprocess
 
 logger = logging.getLogger(__name__)
 
@@ -277,7 +278,8 @@
   needed_property_files = (
       NonAbOtaPropertyFiles(),
   )
-  FinalizeMetadata(metadata, staging_file, output_file, needed_property_files, package_key=OPTIONS.package_key)
+  FinalizeMetadata(metadata, staging_file, output_file,
+                   needed_property_files, package_key=OPTIONS.package_key)
 
 
 def WriteBlockIncrementalOTAPackage(target_zip, source_zip, output_file):
@@ -532,7 +534,8 @@
   needed_property_files = (
       NonAbOtaPropertyFiles(),
   )
-  FinalizeMetadata(metadata, staging_file, output_file, needed_property_files, package_key=OPTIONS.package_key)
+  FinalizeMetadata(metadata, staging_file, output_file,
+                   needed_property_files, package_key=OPTIONS.package_key)
 
 
 def GenerateNonAbOtaPackage(target_file, output_file, source_file=None):
@@ -555,8 +558,18 @@
   if OPTIONS.extracted_input is not None:
     OPTIONS.input_tmp = OPTIONS.extracted_input
   else:
-    logger.info("unzipping target target-files...")
-    OPTIONS.input_tmp = common.UnzipTemp(target_file, UNZIP_PATTERN)
+    if not os.path.isdir(target_file):
+      logger.info("unzipping target target-files...")
+      OPTIONS.input_tmp = common.UnzipTemp(target_file, UNZIP_PATTERN)
+    else:
+      OPTIONS.input_tmp = target_file
+      tmpfile = common.MakeTempFile(suffix=".zip")
+      os.unlink(tmpfile)
+      common.RunAndCheckOutput(
+          ["zip", tmpfile, "-r", ".", "-0"], cwd=target_file)
+      assert zipfile.is_zipfile(tmpfile)
+      target_file = tmpfile
+
   OPTIONS.target_tmp = OPTIONS.input_tmp
 
   # If the caller explicitly specified the device-specific extensions path via
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 1a4a895..04ef5ef 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -273,7 +273,6 @@
                        PayloadGenerator, SECURITY_PATCH_LEVEL_PROP_NAME, ExtractTargetFiles, CopyTargetFilesDir)
 from common import DoesInputFileContain, IsSparseImage
 import target_files_diff
-from check_target_files_vintf import CheckVintfIfTrebleEnabled
 from non_ab_ota import GenerateNonAbOtaPackage
 from payload_signer import PayloadSigner
 
@@ -954,6 +953,7 @@
   target_info.info_dict['ab_partitions'] = common.ReadFromInputFile(target_file,
                                                                     AB_PARTITIONS).strip().split("\n")
 
+  from check_target_files_vintf import CheckVintfIfTrebleEnabled
   CheckVintfIfTrebleEnabled(target_file, target_info)
 
   # Metadata to comply with Android OTA package format.