|
import json |
|
import os |
|
import platform |
|
import time |
|
|
|
from pathlib import Path |
|
|
|
import folder_paths |
|
import nodes |
|
|
|
|
|
class Loader: |
|
def __init__(self): |
|
pass |
|
|
|
__ROOT_PATH = os.path.dirname(os.path.abspath(__file__)) |
|
__TEMPLATE_PATH = os.path.join(__ROOT_PATH, "resources/template.json") |
|
__TIMESTAMP_PATH = os.path.join(__ROOT_PATH, "resources/timestamp.json") |
|
__CONFIG_PATH = os.path.join(__ROOT_PATH, "config.json") |
|
__GIT_PATH = Path(os.path.join(__ROOT_PATH, ".git")) |
|
|
|
__DAY_SECONDS = 24 * 60 * 60 |
|
__WEEK_SECONDS = 7 * __DAY_SECONDS |
|
__MONTH_SECONDS = 30 * __DAY_SECONDS |
|
|
|
def __log(self, text): |
|
print("\033[92m[Allor]\033[0m: " + text) |
|
|
|
def __error(self, text): |
|
print("\033[91m[Allor]\033[0m: " + text) |
|
|
|
def __notification(self, text): |
|
print("\033[94m[Allor]\033[0m: " + text) |
|
|
|
def __new_line(self): |
|
print() |
|
|
|
def __create_config(self): |
|
with open(self.__CONFIG_PATH, "w", encoding="utf-8") as f: |
|
json.dump(self.__template(), f, ensure_ascii=False, indent=4) |
|
|
|
def __create_timestamp(self): |
|
with open(self.__TIMESTAMP_PATH, "w", encoding="utf-8") as f: |
|
json.dump({"timestamp": 0}, f, ensure_ascii=False, indent=4) |
|
|
|
def __get_template(self): |
|
with open(self.__TEMPLATE_PATH, "r") as f: |
|
template = json.load(f) |
|
|
|
if "__comment" in template: |
|
del template["__comment"] |
|
|
|
return template |
|
|
|
def __get_config(self): |
|
with open(self.__CONFIG_PATH, "r") as f: |
|
return json.load(f) |
|
|
|
def __get_timestamp(self): |
|
with open(self.__TIMESTAMP_PATH, "r") as f: |
|
return json.load(f) |
|
|
|
def __update_config(self, template, source): |
|
def update_source(__template, __source): |
|
for k, v in __template.items(): |
|
if k not in __source: |
|
if isinstance(v, dict): |
|
__source[k] = {} |
|
else: |
|
__source[k] = v |
|
|
|
if isinstance(v, dict): |
|
__source[k] = update_source(v, __source[k]) |
|
|
|
return __source |
|
|
|
def delete_keys(__template, __source): |
|
keys_to_delete = [k for k in __source if k not in __template] |
|
|
|
for k in keys_to_delete: |
|
del __source[k] |
|
|
|
return __source |
|
|
|
def sync_order(__template, __source): |
|
new_source = {} |
|
|
|
for key in __template: |
|
if key in __source: |
|
if isinstance(__template[key], dict): |
|
new_source[key] = sync_order(__template[key], __source[key]) |
|
else: |
|
new_source[key] = __source[key] |
|
|
|
return new_source |
|
|
|
source = update_source(template, source) |
|
source = delete_keys(template, source) |
|
source = sync_order(template, source) |
|
|
|
with open(self.__CONFIG_PATH, "w", encoding="utf-8") as f: |
|
json.dump(source, f, ensure_ascii=False, indent=4) |
|
|
|
def __update_timestamp(self): |
|
with open(self.__TIMESTAMP_PATH, "w", encoding="utf-8") as f: |
|
json.dump({"timestamp": time.time()}, f, ensure_ascii=False, indent=4) |
|
|
|
__template = __get_template |
|
__config = __get_config |
|
__timestamp = __get_timestamp |
|
|
|
def __get_fonts_folder_path(self): |
|
system = platform.system() |
|
user_home = os.path.expanduser('~') |
|
|
|
config_font_path = os.path.join(folder_paths.base_path, *self.__config()["fonts"]["folder_path"].replace("\\", "/").split("/")) |
|
|
|
if not os.path.exists(config_font_path): |
|
os.makedirs(config_font_path, exist_ok=True) |
|
|
|
paths = [config_font_path] |
|
|
|
if self.__config()["fonts"]["system_fonts"]: |
|
if system == "Windows": |
|
paths.append(os.path.join(os.environ["WINDIR"], "Fonts")) |
|
elif system == "Darwin": |
|
paths.append(os.path.join("/Library", "Fonts")) |
|
elif system == "Linux": |
|
paths.append(os.path.join("/usr", "share", "fonts")) |
|
paths.append(os.path.join("/usr", "local", "share", "fonts")) |
|
|
|
if self.__config()["fonts"]["user_fonts"]: |
|
if system == "Darwin": |
|
paths.append(os.path.join(user_home, "Library", "Fonts")) |
|
elif system == "Linux": |
|
paths.append(os.path.join(user_home, ".fonts")) |
|
|
|
return [path for path in paths if os.path.exists(path)] |
|
|
|
def __get_keys(self, json_obj, prefix=''): |
|
keys = [] |
|
|
|
for k, v in json_obj.items(): |
|
if isinstance(v, dict): |
|
keys.extend(self.__get_keys(v, prefix + k + '.')) |
|
else: |
|
keys.append(prefix + k) |
|
|
|
return set(keys) |
|
|
|
def __check_json_keys(self, json1, json2): |
|
keys1 = self.__get_keys(json1) |
|
keys2 = self.__get_keys(json2) |
|
|
|
return keys1 == keys2 |
|
|
|
def setup_config(self): |
|
if not os.path.exists(self.__CONFIG_PATH): |
|
self.__log("Creating config.json") |
|
self.__create_config() |
|
else: |
|
if not self.__check_json_keys(self.__template(), self.__config()): |
|
self.__log("Updating config.json") |
|
self.__update_config(self.__template(), self.__config()) |
|
|
|
def setup_timestamp(self): |
|
if not os.path.exists(self.__TIMESTAMP_PATH): |
|
self.__log("Creating timestamp.json") |
|
self.__create_timestamp() |
|
|
|
def check_updates(self): |
|
branch_name = self.__config()["updates"]["branch_name"] |
|
update_frequency = self.__config()["updates"]["update_frequency"].lower() |
|
valid_frequencies = ["always", "day", "week", "month", "never"] |
|
time_difference = time.time() - self.__timestamp()["timestamp"] |
|
|
|
if update_frequency == valid_frequencies[0]: |
|
it_is_time_for_update = True |
|
elif update_frequency == valid_frequencies[1]: |
|
it_is_time_for_update = time_difference >= self.__DAY_SECONDS |
|
elif update_frequency == valid_frequencies[2]: |
|
it_is_time_for_update = time_difference >= self.__WEEK_SECONDS |
|
elif update_frequency == valid_frequencies[3]: |
|
it_is_time_for_update = time_difference >= self.__MONTH_SECONDS |
|
elif update_frequency == valid_frequencies[4]: |
|
it_is_time_for_update = False |
|
else: |
|
self.__error(f"Unknown update frequency - {update_frequency}, available: {valid_frequencies}") |
|
|
|
return |
|
|
|
if it_is_time_for_update: |
|
if not (self.__GIT_PATH.exists() or self.__GIT_PATH.is_dir()): |
|
self.__error("Root directory of Allor is not a git repository. Update canceled.") |
|
|
|
return |
|
|
|
try: |
|
import git |
|
|
|
from git import Repo |
|
from git import GitCommandError |
|
|
|
|
|
repo = Repo(self.__ROOT_PATH, odbt=git.db.GitDB) |
|
current_commit = repo.head.commit.hexsha |
|
|
|
repo.remotes.origin.fetch() |
|
|
|
latest_commit = getattr(repo.remotes.origin.refs, branch_name).commit.hexsha |
|
|
|
if current_commit == latest_commit: |
|
if self.__config()["updates"]["notify_if_no_new_updates"]: |
|
self.__notification("No new updates.") |
|
else: |
|
if self.__config()["updates"]["notify_if_has_new_updates"]: |
|
self.__notification("New updates are available.") |
|
|
|
if self.__config()["updates"]["auto_update"]: |
|
update_mode = self.__config()["updates"]["update_mode"].lower() |
|
valid_modes = ["soft", "hard"] |
|
|
|
if repo.active_branch.name != branch_name: |
|
try: |
|
repo.git.checkout(branch_name) |
|
except GitCommandError: |
|
self.__error(f"An error occurred while switching to the branch {branch_name}.") |
|
|
|
return |
|
|
|
if update_mode == "soft": |
|
try: |
|
repo.git.pull() |
|
except GitCommandError: |
|
self.__error("An error occurred during the update. " |
|
"It is recommended to use \"hard\" update mode. " |
|
"But be careful, it erases all personal changes from Allor repository.") |
|
|
|
elif update_mode == "hard": |
|
repo.git.reset('--hard', 'origin/' + branch_name) |
|
else: |
|
self.__error(f"Unknown update mode - {update_mode}, available: {valid_modes}") |
|
|
|
return |
|
|
|
self.__notification("Update complete.") |
|
|
|
self.__update_timestamp() |
|
|
|
except ImportError: |
|
self.__error("GitPython is not installed.") |
|
|
|
def setup_rembg(self): |
|
os.environ["U2NET_HOME"] = folder_paths.models_dir + "/onnx" |
|
|
|
def setup_paths(self): |
|
fonts_folder_path = self.__get_fonts_folder_path() |
|
|
|
folder_paths.folder_names_and_paths["onnx"] = ([os.path.join(folder_paths.models_dir, "onnx")], {".onnx"}) |
|
folder_paths.folder_names_and_paths["fonts"] = (fonts_folder_path, {".otf", ".ttf"}) |
|
|
|
def setup_override(self): |
|
override_nodes_len = 0 |
|
|
|
def override(function): |
|
start_len = nodes.NODE_CLASS_MAPPINGS.__len__() |
|
|
|
nodes.NODE_CLASS_MAPPINGS = dict( |
|
filter(function, nodes.NODE_CLASS_MAPPINGS.items()) |
|
) |
|
|
|
return start_len - nodes.NODE_CLASS_MAPPINGS.__len__() |
|
|
|
if self.__config()["override"]["postprocessing"]: |
|
override_nodes_len += override(lambda item: not item[1].CATEGORY.startswith("image/postprocessing")) |
|
|
|
if self.__config()["override"]["transform"]: |
|
override_nodes_len += override(lambda item: not item[0] == "ImageScale" and not item[0] == "ImageScaleBy" and not item[0] == "ImageInvert") |
|
|
|
if self.__config()["override"]["debug"]: |
|
nodes.VAEDecodeTiled.CATEGORY = "latent" |
|
nodes.VAEEncodeTiled.CATEGORY = "latent" |
|
|
|
override_nodes_len += override(lambda item: not item[1].CATEGORY.startswith("_for_testing")) |
|
|
|
self.__log(str(override_nodes_len) + " nodes were overridden.") |
|
|
|
def get_modules(self): |
|
modules = dict() |
|
|
|
if self.__config()["modules"]["AlphaChanel"]: |
|
from .modules import AlphaChanel |
|
modules.update(AlphaChanel.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["Clamp"]: |
|
from .modules import Clamp |
|
modules.update(Clamp.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageBatch"]: |
|
from .modules import ImageBatch |
|
modules.update(ImageBatch.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageComposite"]: |
|
from .modules import ImageComposite |
|
modules.update(ImageComposite.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageContainer"]: |
|
from .modules import ImageContainer |
|
modules.update(ImageContainer.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageDraw"]: |
|
from .modules import ImageDraw |
|
modules.update(ImageDraw.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageEffects"]: |
|
from .modules import ImageEffects |
|
modules.update(ImageEffects.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageFilter"]: |
|
from .modules import ImageFilter |
|
modules.update(ImageFilter.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageNoise"]: |
|
from .modules import ImageNoise |
|
modules.update(ImageNoise.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageSegmentation"]: |
|
from .modules import ImageSegmentation |
|
modules.update(ImageSegmentation.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageText"]: |
|
from .modules import ImageText |
|
modules.update(ImageText.NODE_CLASS_MAPPINGS) |
|
|
|
if self.__config()["modules"]["ImageTransform"]: |
|
from .modules import ImageTransform |
|
modules.update(ImageTransform.NODE_CLASS_MAPPINGS) |
|
|
|
modules_len = dict( |
|
filter( |
|
lambda item: item[1], |
|
self.__config()["modules"].items() |
|
) |
|
).__len__() |
|
|
|
nodes_len = modules.__len__() |
|
|
|
self.__log(str(modules_len) + " modules were enabled.") |
|
self.__log(str(nodes_len) + " nodes were loaded.") |
|
|
|
return modules |
|
|
