offload: fixes for 4090 batch size

bghira · Jul 16, 2023 · c733e0a · c733e0a
1 parent 1d04d0e
commit c733e0a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/discord_tron_client/classes/app_config.py b/discord_tron_client/classes/app_config.py
@@ -282,4 +282,8 @@ def bark_subsystem_type(self):
     def enable_compel(self):
         return self.config.get("use_compel_prompt_weighting", True)
     def enable_compile(self):
-        return self.config.get('enable_torch_compile', True)
+        return self.config.get('enable_torch_compile', True)
+    def enable_cpu_offload(self):
+        return self.config.get('enable_cpu_offload', True)
+    def maximum_batch_size(self):
+        return max(self.config.get('maximum_batch_size', 4), 1)
diff --git a/discord_tron_client/classes/image_manipulation/pipeline.py b/discord_tron_client/classes/image_manipulation/pipeline.py
@@ -203,11 +203,7 @@ def _run_pipeline(
     ):
         original_stderr = sys.stderr
         sys.stderr = self.tqdm_capture
-        batch_size = 4
-        if hardware.should_offload():
-            batch_size = 2
-        if hardware.should_sequential_offload():
-            batch_size = 1
+        batch_size = self.config.maximum_batch_size()
         try:
             alt_weight_algorithm = user_config.get("alt_weight_algorithm", False)
             use_latent_result = user_config.get('latent_refiner', True)