From 840ac77c9d35682a7227c52eb80b0b3006adcc49 Mon Sep 17 00:00:00 2001
From: Nourepide <nourepide@gmail.com>
Date: Tue, 14 Nov 2023 01:14:59 +0300
Subject: [PATCH] Add auto updates.

---
 Loader.py              | 36 ++++++++++++++++++++++++++++++++++++
 __init__.py            |  1 +
 requirements.txt       |  3 +++
 template/template.json |  6 ++++++
 4 files changed, 46 insertions(+)

diff --git a/Loader.py b/Loader.py
index 7178911..3bf1197 100644
--- a/Loader.py
+++ b/Loader.py
@@ -16,6 +16,12 @@ class Loader:
     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 __get_template(self):
         with open(self.__template_path, "r") as f:
             template = json.load(f)
@@ -84,6 +90,36 @@ class Loader:
                 self.__log("Updating config.json")
                 self.__update_config(self.__template(), self.__config())
 
+    def check_updates(self):
+        if self.__config()["updates"]["check_updates"]:
+            try:
+                import git
+            except ImportError:
+                self.__error("GitPython is not installed.")
+                return
+
+        from git import Repo
+
+        # noinspection PyTypeChecker, PyUnboundLocalVariable
+        repo = Repo(self.__root_path, odbt=git.db.GitDB)
+        current_commit = repo.head.commit.hexsha
+
+        repo.remotes.origin.fetch()
+
+        # noinspection PyUnresolvedReferences
+        latest_commit = repo.head.ref.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"]:
+                repo.remotes.origin.pull()
+                self.__notification("Update complete.")
+
     def setup_rembg(self):
         os.environ["U2NET_HOME"] = folder_paths.models_dir + "/onnx"
 
diff --git a/__init__.py b/__init__.py
index d5f1cdb..f012c12 100644
--- a/__init__.py
+++ b/__init__.py
@@ -3,6 +3,7 @@ from .Loader import Loader
 loader = Loader()
 
 loader.setup_config()
+loader.check_updates()
 loader.setup_rembg()
 loader.setup_paths()
 loader.setup_override()
diff --git a/requirements.txt b/requirements.txt
index 095a920..28e2be3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,9 @@
 rembg
 onnx
 
+# Optional for auto updates
+gitpython
+
 # If you want you can install onnxruntime with GPU support
 # onnxruntime[gpu]
 
diff --git a/template/template.json b/template/template.json
index 3ee9fb7..c7b44b6 100644
--- a/template/template.json
+++ b/template/template.json
@@ -22,5 +22,11 @@
     "postprocessing": false,
     "transform": false,
     "debug": false
+  },
+  "updates": {
+    "check_updates": true,
+    "notify_if_has_new_updates": true,
+    "notify_if_no_new_updates": true,
+    "auto_update": true
   }
 }