From 37eff915d6dae6f0782b41b5a5b2ad5b50b0a485 Mon Sep 17 00:00:00 2001
From: Alberto Cano <34340962+canoalberto@users.noreply.github.com>
Date: Thu, 4 Jan 2024 04:27:26 +0100
Subject: [PATCH] Use --disk-cache-dir for all caches

---
 modules/chat.py           | 19 ++++++++++---------
 modules/html_generator.py |  7 ++++---
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/modules/chat.py b/modules/chat.py
index dccd20b7..f5488e3a 100644
--- a/modules/chat.py
+++ b/modules/chat.py
@@ -560,17 +560,17 @@ def replace_character_names(text, name1, name2):
 
 
 def generate_pfp_cache(character):
-    cache_folder = Path("cache")
+    cache_folder = Path(shared.args.disk_cache_dir)
     if not cache_folder.exists():
         cache_folder.mkdir()
 
     for path in [Path(f"characters/{character}.{extension}") for extension in ['png', 'jpg', 'jpeg']]:
         if path.exists():
             original_img = Image.open(path)
-            original_img.save(Path('cache/pfp_character.png'), format='PNG')
+            original_img.save(Path('{cache_folder}/pfp_character.png'), format='PNG')
 
             thumb = make_thumbnail(original_img)
-            thumb.save(Path('cache/pfp_character_thumb.png'), format='PNG')
+            thumb.save(Path('{cache_folder}/pfp_character_thumb.png'), format='PNG')
 
             return thumb
 
@@ -594,8 +594,9 @@ def load_character(character, name1, name2):
 
     file_contents = open(filepath, 'r', encoding='utf-8').read()
     data = json.loads(file_contents) if extension == "json" else yaml.safe_load(file_contents)
+    cache_folder = Path(shared.args.disk_cache_dir)
 
-    for path in [Path("cache/pfp_character.png"), Path("cache/pfp_character_thumb.png")]:
+    for path in [Path("{cache_folder}/pfp_character.png"), Path("{cache_folder}/pfp_character_thumb.png")]:
         if path.exists():
             path.unlink()
 
@@ -713,17 +714,17 @@ def check_tavern_character(img):
 
 
 def upload_your_profile_picture(img):
-    cache_folder = Path("cache")
+    cache_folder = Path(shared.args.disk_cache_dir)
     if not cache_folder.exists():
         cache_folder.mkdir()
 
     if img is None:
-        if Path("cache/pfp_me.png").exists():
-            Path("cache/pfp_me.png").unlink()
+        if Path("{cache_folder}/pfp_me.png").exists():
+            Path("{cache_folder}/pfp_me.png").unlink()
     else:
         img = make_thumbnail(img)
-        img.save(Path('cache/pfp_me.png'))
-        logger.info('Profile picture saved to "cache/pfp_me.png"')
+        img.save(Path('{cache_folder}/pfp_me.png'))
+        logger.info('Profile picture saved to "{cache_folder}/pfp_me.png"')
 
 
 def generate_character_yaml(name, greeting, context):
diff --git a/modules/html_generator.py b/modules/html_generator.py
index 789cf639..71a61bd0 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -8,6 +8,7 @@ import markdown
 from PIL import Image, ImageOps
 
 from modules.utils import get_available_chat_styles
+from modules import shared
 
 # This is to store the paths to the thumbnails of the profile pictures
 image_cache = {}
@@ -170,7 +171,7 @@ def make_thumbnail(image):
 
 
 def get_image_cache(path):
-    cache_folder = Path("cache")
+    cache_folder = Path(shared.args.disk_cache_dir)
     if not cache_folder.exists():
         cache_folder.mkdir()
 
@@ -178,8 +179,8 @@ def get_image_cache(path):
     if (path in image_cache and mtime != image_cache[path][0]) or (path not in image_cache):
         img = make_thumbnail(Image.open(path))
 
-        old_p = Path(f'cache/{path.name}_cache.png')
-        p = Path(f'cache/cache_{path.name}.png')
+        old_p = Path(f'{cache_folder}/{path.name}_cache.png')
+        p = Path(f'{cache_folder}/cache_{path.name}.png')
         if old_p.exists():
             old_p.rename(p)