Allow override of using safetensors

bghira · Aug 23, 2023 · ebc1db2 · ebc1db2
1 parent ab3ed84
commit ebc1db2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/discord_tron_client/classes/app_config.py b/discord_tron_client/classes/app_config.py
@@ -295,4 +295,6 @@ def enable_offload(self):
     def enable_sequential_offload(self):
         return self.get_config_value('enable_sequential_offload', False)
     def maximum_batch_size(self):
-        return max(self.get_config_value('maximum_batch_size', 4), 1)
+        return max(self.get_config_value('maximum_batch_size', 4), 1)
+    def use_safetensors(self):
+        return self.get_config_value('use_safetensors', True)
diff --git a/discord_tron_client/classes/image_manipulation/pipeline.py b/discord_tron_client/classes/image_manipulation/pipeline.py
@@ -108,6 +108,7 @@ def _prepare_pipe(
             prompt_variation=variation,
             promptless_variation=promptless_variation,
             upscaler=upscaler,
+            use_safetensors=self.config.use_safetensors(),
         )
         logging.info("Copied pipe to the local context")
         return pipe
@@ -523,6 +524,8 @@ def _controlnet_pipeline(
         #     prompt_embed, negative_embed = controlnet_prompt_manager.process_long_prompt(
         #         positive_prompt=prompt, negative_prompt=negative_prompt
         #     )
+        pipe.vae.disable_tiling()
+        pipe.vae.disable_slicing()
         new_image = pipe(
             # prompt_embeds=prompt_embed,
             # negative_prompt_embeds=negative_embed,