stable-diffusion-webui

1
import json
2
import os
3
import sys
4

5
import platform
6
import hashlib
7
import pkg_resources
8
import psutil
9
import re
10

11
import launch
12
from modules import paths_internal, timer, shared, extensions, errors
13

14
checksum_token = "DontStealMyGamePlz__WINNERS_DONT_USE_DRUGS__DONT_COPY_THAT_FLOPPY"
15
environment_whitelist = {
16
    "GIT",
17
    "INDEX_URL",
18
    "WEBUI_LAUNCH_LIVE_OUTPUT",
19
    "GRADIO_ANALYTICS_ENABLED",
20
    "PYTHONPATH",
21
    "TORCH_INDEX_URL",
22
    "TORCH_COMMAND",
23
    "REQS_FILE",
24
    "XFORMERS_PACKAGE",
25
    "CLIP_PACKAGE",
26
    "OPENCLIP_PACKAGE",
27
    "ASSETS_REPO",
28
    "STABLE_DIFFUSION_REPO",
29
    "K_DIFFUSION_REPO",
30
    "BLIP_REPO",
31
    "ASSETS_COMMIT_HASH",
32
    "STABLE_DIFFUSION_COMMIT_HASH",
33
    "K_DIFFUSION_COMMIT_HASH",
34
    "BLIP_COMMIT_HASH",
35
    "COMMANDLINE_ARGS",
36
    "IGNORE_CMD_ARGS_ERRORS",
37
}
38

39

40
def pretty_bytes(num, suffix="B"):
41
    for unit in ["", "K", "M", "G", "T", "P", "E", "Z", "Y"]:
42
        if abs(num) < 1024 or unit == 'Y':
43
            return f"{num:.0f}{unit}{suffix}"
44
        num /= 1024
45

46

47
def get():
48
    res = get_dict()
49

50
    text = json.dumps(res, ensure_ascii=False, indent=4)
51

52
    h = hashlib.sha256(text.encode("utf8"))
53
    text = text.replace(checksum_token, h.hexdigest())
54

55
    return text
56

57

58
re_checksum = re.compile(r'"Checksum": "([0-9a-fA-F]{64})"')
59

60

61
def check(x):
62
    m = re.search(re_checksum, x)
63
    if not m:
64
        return False
65

66
    replaced = re.sub(re_checksum, f'"Checksum": "{checksum_token}"', x)
67

68
    h = hashlib.sha256(replaced.encode("utf8"))
69
    return h.hexdigest() == m.group(1)
70

71

72
def get_dict():
73
    ram = psutil.virtual_memory()
74

75
    res = {
76
        "Platform": platform.platform(),
77
        "Python": platform.python_version(),
78
        "Version": launch.git_tag(),
79
        "Commit": launch.commit_hash(),
80
        "Script path": paths_internal.script_path,
81
        "Data path": paths_internal.data_path,
82
        "Extensions dir": paths_internal.extensions_dir,
83
        "Checksum": checksum_token,
84
        "Commandline": get_argv(),
85
        "Torch env info": get_torch_sysinfo(),
86
        "Exceptions": errors.get_exceptions(),
87
        "CPU": {
88
            "model": platform.processor(),
89
            "count logical": psutil.cpu_count(logical=True),
90
            "count physical": psutil.cpu_count(logical=False),
91
        },
92
        "RAM": {
93
            x: pretty_bytes(getattr(ram, x, 0)) for x in ["total", "used", "free", "active", "inactive", "buffers", "cached", "shared"] if getattr(ram, x, 0) != 0
94
        },
95
        "Extensions": get_extensions(enabled=True),
96
        "Inactive extensions": get_extensions(enabled=False),
97
        "Environment": get_environment(),
98
        "Config": get_config(),
99
        "Startup": timer.startup_record,
100
        "Packages": sorted([f"{pkg.key}=={pkg.version}" for pkg in pkg_resources.working_set]),
101
    }
102

103
    return res
104

105

106
def get_environment():
107
    return {k: os.environ[k] for k in sorted(os.environ) if k in environment_whitelist}
108

109

110
def get_argv():
111
    res = []
112

113
    for v in sys.argv:
114
        if shared.cmd_opts.gradio_auth and shared.cmd_opts.gradio_auth == v:
115
            res.append("<hidden>")
116
            continue
117

118
        if shared.cmd_opts.api_auth and shared.cmd_opts.api_auth == v:
119
            res.append("<hidden>")
120
            continue
121

122
        res.append(v)
123

124
    return res
125

126
re_newline = re.compile(r"\r*\n")
127

128

129
def get_torch_sysinfo():
130
    try:
131
        import torch.utils.collect_env
132
        info = torch.utils.collect_env.get_env_info()._asdict()
133

134
        return {k: re.split(re_newline, str(v)) if "\n" in str(v) else v for k, v in info.items()}
135
    except Exception as e:
136
        return str(e)
137

138

139
def get_extensions(*, enabled):
140

141
    try:
142
        def to_json(x: extensions.Extension):
143
            return {
144
                "name": x.name,
145
                "path": x.path,
146
                "version": x.version,
147
                "branch": x.branch,
148
                "remote": x.remote,
149
            }
150

151
        return [to_json(x) for x in extensions.extensions if not x.is_builtin and x.enabled == enabled]
152
    except Exception as e:
153
        return str(e)
154

155

156
def get_config():
157
    try:
158
        return shared.opts.data
159
    except Exception as e:
160
        return str(e)
161