Merge pull request #1429 from bghira/feature/disable-text-enc-offload-at-startup

disable offload at startup by default, allow opt-in via --offload_during_startup

Author: bghira

helpers/configuration/cmd_args.py

@@ -1268,6 +1268,14 @@ def get_argument_parser():
             " remote repository specified with --pretrained_model_name_or_path."
         ),
     )
+    parser.add_argument(
+        "--offload_during_startup",
+        action="store_true",
+        help=(
+            "When set, text encoders, the VAE, or other models will be moved to and from the CPU as needed, which can slow"
+            " down startup, but saves VRAM. This is useful for video models or high-resolution pre-caching of latent embeds."
+        ),
+    )
     parser.add_argument(
         "--offload_param_path",
         type=str,

helpers/data_backend/factory.py

@@ -455,7 +455,7 @@ def configure_parquet_database(backend: dict, args, data_backend: BaseDataBacken
 
 def move_text_encoders(args, text_encoders: list, target_device: str):
     """Move text encoders to the target device."""
-    if text_encoders is None:
+    if text_encoders is None or not args.offload_during_startup:
         return
     # we'll move text encoder only if their precision arg is no_change
     # otherwise, we assume the user has already moved them to the correct device due to quantisation.