File size: 2,036 Bytes
19b3da3
 
 
 
 
 
 
 
 
 
 
a3d6c18
 
19b3da3
 
 
 
1bc457e
 
19b3da3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import json
import os
import re

from internals.data.dataAccessor import getCharacters
from internals.util.config import root_dir


class Avatar:
    __avatars = {}

    def load_local(self, model_dir):
        self.__find_available_characters(model_dir)
        if len(self.__avatars.items()) > 0:
            print("Local characters", self.__avatars)

    def fetch_from_network(self, model_id: int):
        if not model_id:
            return
        characters = getCharacters(str(model_id))
        if characters is not None:
            for character in characters:
                item = {
                    "avatarName": str(character["title"]).lower(),
                    "codename": character["tag"],
                    "extraPrompt": character["extraData"]["extraPrompt"],
                }
                self.__avatars[item["avatarName"]] = item

    def add_code_names(self, prompt):
        array_of_objects = self.__avatars.values()

        for obj in array_of_objects:
            prompt = (
                re.sub(
                    r"\b" + obj["avatarName"] + r"\b",
                    obj["extraPrompt"],
                    prompt,
                    flags=re.IGNORECASE,
                )
                + " "
            )
        print(prompt)
        return prompt

    def __find_available_characters(self, path: str):
        if os.path.exists(path + "/characters.json"):
            print(path)
            try:
                print("Loading characters")
                with open(path + "/characters.json") as f:
                    data = json.load(f)
                    print("Characters: ", data)
                    if "avatarName" in data[0]:
                        for item in data:
                            self.__avatars[item["avatarName"]] = item
                        print("Avatars", self.__avatars)
                    else:
                        print("Invalid characters.json file")
            except Exception as e:
                print("Error Loading characters", e)