stable-diffusion-webui

1
import json
2
import os
3
import os.path
4
import threading
5
import time
6

7
from modules.paths import data_path, script_path
8

9
cache_filename = os.environ.get('SD_WEBUI_CACHE_FILE', os.path.join(data_path, "cache.json"))
10
cache_data = None
11
cache_lock = threading.Lock()
12

13
dump_cache_after = None
14
dump_cache_thread = None
15

16

17
def dump_cache():
18
    """
19
    Marks cache for writing to disk. 5 seconds after no one else flags the cache for writing, it is written.
20
    """
21

22
    global dump_cache_after
23
    global dump_cache_thread
24

25
    def thread_func():
26
        global dump_cache_after
27
        global dump_cache_thread
28

29
        while dump_cache_after is not None and time.time() < dump_cache_after:
30
            time.sleep(1)
31

32
        with cache_lock:
33
            cache_filename_tmp = cache_filename + "-"
34
            with open(cache_filename_tmp, "w", encoding="utf8") as file:
35
                json.dump(cache_data, file, indent=4, ensure_ascii=False)
36

37
            os.replace(cache_filename_tmp, cache_filename)
38

39
            dump_cache_after = None
40
            dump_cache_thread = None
41

42
    with cache_lock:
43
        dump_cache_after = time.time() + 5
44
        if dump_cache_thread is None:
45
            dump_cache_thread = threading.Thread(name='cache-writer', target=thread_func)
46
            dump_cache_thread.start()
47

48

49
def cache(subsection):
50
    """
51
    Retrieves or initializes a cache for a specific subsection.
52

53
    Parameters:
54
        subsection (str): The subsection identifier for the cache.
55

56
    Returns:
57
        dict: The cache data for the specified subsection.
58
    """
59

60
    global cache_data
61

62
    if cache_data is None:
63
        with cache_lock:
64
            if cache_data is None:
65
                try:
66
                    with open(cache_filename, "r", encoding="utf8") as file:
67
                        cache_data = json.load(file)
68
                except FileNotFoundError:
69
                    cache_data = {}
70
                except Exception:
71
                    os.replace(cache_filename, os.path.join(script_path, "tmp", "cache.json"))
72
                    print('[ERROR] issue occurred while trying to read cache.json, move current cache to tmp/cache.json and create new cache')
73
                    cache_data = {}
74

75
    s = cache_data.get(subsection, {})
76
    cache_data[subsection] = s
77

78
    return s
79

80

81
def cached_data_for_file(subsection, title, filename, func):
82
    """
83
    Retrieves or generates data for a specific file, using a caching mechanism.
84

85
    Parameters:
86
        subsection (str): The subsection of the cache to use.
87
        title (str): The title of the data entry in the subsection of the cache.
88
        filename (str): The path to the file to be checked for modifications.
89
        func (callable): A function that generates the data if it is not available in the cache.
90

91
    Returns:
92
        dict or None: The cached or generated data, or None if data generation fails.
93

94
    The `cached_data_for_file` function implements a caching mechanism for data stored in files.
95
    It checks if the data associated with the given `title` is present in the cache and compares the
96
    modification time of the file with the cached modification time. If the file has been modified,
97
    the cache is considered invalid and the data is regenerated using the provided `func`.
98
    Otherwise, the cached data is returned.
99

100
    If the data generation fails, None is returned to indicate the failure. Otherwise, the generated
101
    or cached data is returned as a dictionary.
102
    """
103

104
    existing_cache = cache(subsection)
105
    ondisk_mtime = os.path.getmtime(filename)
106

107
    entry = existing_cache.get(title)
108
    if entry:
109
        cached_mtime = entry.get("mtime", 0)
110
        if ondisk_mtime > cached_mtime:
111
            entry = None
112

113
    if not entry or 'value' not in entry:
114
        value = func()
115
        if value is None:
116
            return None
117

118
        entry = {'mtime': ondisk_mtime, 'value': value}
119
        existing_cache[title] = entry
120

121
        dump_cache()
122

123
    return entry['value']
124