Extended Dreambooth: Train TI tokens separately

1 files changed, 7 insertions, 0 deletions

diff --git a/data/csv.py b/data/csv.py
index 2b1e202..002fdd2 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -15,6 +15,9 @@ from data.keywords import prompt_to_keywords, keywords_to_prompt
 from models.clip.util import unify_input_ids
 
 
+cache = {}
+
+
 interpolations = {
     "linear": transforms.InterpolationMode.NEAREST,
     "bilinear": transforms.InterpolationMode.BILINEAR,
@@ -24,9 +27,13 @@ interpolations = {
 
 
 def get_image(path):
+    if path in cache:
+        return cache[path]
+
     image = Image.open(path)
     if not image.mode == "RGB":
         image = image.convert("RGB")
+    cache[path] = image
     return image