diff --git a/delta_diff_generator.cc b/delta_diff_generator.cc
index 6279537..ee16454 100644
--- a/delta_diff_generator.cc
+++ b/delta_diff_generator.cc
@@ -58,15 +58,14 @@
 typedef map<const DeltaArchiveManifest_InstallOperation*,
             string> OperationNameMap;
 
+const size_t kRootFSPartitionSize = 1024 * 1024 * 1024;  // bytes
+const uint64_t kVersionNumber = 1;
+const uint64_t kFullUpdateChunkSize = 1024 * 1024;  // bytes
+
 namespace {
 const size_t kBlockSize = 4096;  // bytes
 const string kNonexistentPath = "";
 
-// TODO(adlr): switch from 1GiB to 2GiB when we no longer care about old
-// clients:
-const size_t kRootFSPartitionSize = 1 * 1024 * 1024 * 1024;  // bytes
-const uint64_t kVersionNumber = 1;
-const uint64_t kFullUpdateChunkSize = 1024 * 1024;  // bytes
 
 static const char* kInstallOperationTypes[] = {
   "REPLACE",
@@ -1387,6 +1386,7 @@
     const string& output_path,
     const string& private_key_path,
     off_t chunk_size,
+    size_t rootfs_partition_size,
     uint64_t* metadata_size) {
   TEST_AND_RETURN_FALSE(chunk_size == -1 || chunk_size % kBlockSize == 0);
   int old_image_block_count = 0, old_image_block_size = 0;
@@ -1402,6 +1402,15 @@
     LOG_IF(WARNING, old_image_block_count != new_image_block_count)
         << "Old and new images have different block counts.";
   }
+
+  // Sanity checks for the partition size.
+  TEST_AND_RETURN_FALSE(rootfs_partition_size % kBlockSize == 0);
+  size_t fs_size = static_cast<size_t>(new_image_block_size *
+                                       new_image_block_count);
+  LOG(INFO) << "Rootfs partition size: " << rootfs_partition_size;
+  LOG(INFO) << "Actual filesystem size: " << fs_size;
+  TEST_AND_RETURN_FALSE(rootfs_partition_size >= fs_size);
+
   // Sanity check kernel partition arg
   TEST_AND_RETURN_FALSE(utils::FileSize(new_kernel_part) >= 0);
 
@@ -1464,11 +1473,11 @@
                                                 &graph.back()));
 
       // Final scratch block (if there's space)
-      if (blocks.size() < (kRootFSPartitionSize / kBlockSize)) {
+      if (blocks.size() < (rootfs_partition_size / kBlockSize)) {
         scratch_vertex = graph.size();
         graph.resize(graph.size() + 1);
         CreateScratchNode(blocks.size(),
-                          (kRootFSPartitionSize / kBlockSize) - blocks.size(),
+                          (rootfs_partition_size / kBlockSize) - blocks.size(),
                           &graph.back());
       }
 
diff --git a/delta_diff_generator.h b/delta_diff_generator.h
index 5dfa364..c258a58 100644
--- a/delta_diff_generator.h
+++ b/delta_diff_generator.h
@@ -63,6 +63,7 @@
   // output_path is the filename where the delta update should be written.
   // If |chunk_size| is not -1, the delta payload is generated based on
   // |chunk_size| chunks rather than whole files.
+  // This method computes scratch space based on |rootfs_partition_size|.
   // Returns true on success. Also writes the size of the metadata into
   // |metadata_size|.
   static bool GenerateDeltaUpdateFile(const std::string& old_root,
@@ -74,6 +75,7 @@
                                       const std::string& output_path,
                                       const std::string& private_key_path,
                                       off_t chunk_size,
+                                      size_t rootfs_partition_size,
                                       uint64_t* metadata_size);
 
   // These functions are public so that the unit tests can access them:
@@ -265,6 +267,7 @@
 extern const char* const kBsdiffPath;
 extern const char* const kBspatchPath;
 extern const char* const kDeltaMagic;
+extern const size_t kRootFSPartitionSize;
 
 };  // namespace chromeos_update_engine
 
diff --git a/delta_performer_unittest.cc b/delta_performer_unittest.cc
index 57c16e7..8964a3c 100644
--- a/delta_performer_unittest.cc
+++ b/delta_performer_unittest.cc
@@ -423,6 +423,7 @@
             state->delta_path,
             private_key,
             chunk_size,
+            kRootFSPartitionSize,
             &state->metadata_size));
   }
 
diff --git a/generate_delta_main.cc b/generate_delta_main.cc
index bde08d4..7c8b54b 100644
--- a/generate_delta_main.cc
+++ b/generate_delta_main.cc
@@ -61,6 +61,9 @@
               "signature will be assigned a client version, starting from "
               "kSignatureOriginalVersion.");
 DEFINE_int32(chunk_size, -1, "Payload chunk size (-1 -- no limit/default)");
+DEFINE_int64(rootfs_partition_size,
+             chromeos_update_engine::kRootFSPartitionSize,
+             "RootFS partition size for the image once installed");
 
 // This file contains a simple program that takes an old path, a new path,
 // and an output file as arguments and the path to an output file and
@@ -274,6 +277,7 @@
                                                    FLAGS_out_file,
                                                    FLAGS_private_key,
                                                    FLAGS_chunk_size,
+                                                   FLAGS_rootfs_partition_size,
                                                    &metadata_size)) {
     return 1;
   }