From 0ea27011d58e3c2fcb0a37f91313ecaedef9f982 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sun, 4 Aug 2024 11:07:19 -0600
Subject: [PATCH] Bug fix

---
 build_and_push_docker.yaml        | 8 ++++++++
 toolkit/stable_diffusion_model.py | 4 ++--
 2 files changed, 10 insertions(+), 2 deletions(-)
 create mode 100644 build_and_push_docker.yaml

diff --git a/build_and_push_docker.yaml b/build_and_push_docker.yaml
new file mode 100644
index 00000000..bbda385a
--- /dev/null
+++ b/build_and_push_docker.yaml
@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+echo "Docker builds from the repo, not this dir. Make sure changes are pushed to the repo."
+# wait 2 seconds
+sleep 2
+docker build --build-arg CACHEBUST=$(date +%s) -t aitoolkit:latest -f docker/Dockerfile .
+docker tag aitoolkit:latest ostris/aitoolkit:latest
+docker push ostris/aitoolkit:latest
\ No newline at end of file
diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py
index f48ecdbb..0f6d23c2 100644
--- a/toolkit/stable_diffusion_model.py
+++ b/toolkit/stable_diffusion_model.py
@@ -468,8 +468,8 @@ class StableDiffusion:
                 transformer_path,
                 subfolder=subfolder,
                 torch_dtype=dtype,
-                low_cpu_mem_usage=False,
-                device_map=None
+                # low_cpu_mem_usage=False,
+                # device_map=None
             )
             transformer.to(self.device_torch, dtype=dtype)
             flush()