From bab643ee3d8cfbd07670085bae03f4ab1676392c Mon Sep 17 00:00:00 2001
From: hayden <hayden.mike.zhang@gmail.com>
Date: Sat, 2 Nov 2024 19:47:17 +0800
Subject: [PATCH] feat: Add web version check

---
 __init__.py    |  3 +++
 py/utils.py    | 68 +++++++++++++++++++++++++++++++++++++++++++++++++-
 vite.config.ts | 23 ++++++++++++++++-
 3 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/__init__.py b/__init__.py
index 7656106..87cbb38 100644
--- a/__init__.py
+++ b/__init__.py
@@ -8,6 +8,9 @@ from .py import utils
 config.extension_uri = os.path.dirname(__file__)
 utils.resolve_model_base_paths()
 
+version = utils.get_current_version()
+utils.download_web_distribution(version)
+
 
 import logging
 from aiohttp import web
diff --git a/py/utils.py b/py/utils.py
index f560b94..90c6515 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -1,13 +1,79 @@
 import os
-import comfy.utils
 import json
+import yaml
+import shutil
+import tarfile
 import logging
+import requests
+import configparser
+
+import comfy.utils
 import folder_paths
+
 from aiohttp import web
 from typing import Any
 from . import config
 
 
+def get_current_version():
+    try:
+        pyproject_path = os.path.join(config.extension_uri, "pyproject.toml")
+        config_parser = configparser.ConfigParser()
+        config_parser.read(pyproject_path)
+        version = config_parser.get("project", "version")
+        return version.strip("'\"")
+    except:
+        return "0.0.0"
+
+
+def download_web_distribution(version: str):
+    web_path = os.path.join(config.extension_uri, "web")
+    dev_web_file = os.path.join(web_path, "manager-dev.js")
+    if os.path.exists(dev_web_file):
+        return
+
+    web_version = "0.0.0"
+    version_file = os.path.join(web_path, "version.yaml")
+    if os.path.exists(version_file):
+        with open(version_file, "r") as f:
+            version_content = yaml.safe_load(f)
+            web_version = version_content.get("version", web_version)
+
+    if version == web_version:
+        return
+
+    try:
+        logging.info(f"current version {version}, web version {web_version}")
+        logging.info("Downloading web distribution...")
+        download_url = f"https://github.com/hayden-fr/ComfyUI-Model-Manager/releases/download/v{version}/dist.tar.gz"
+        response = requests.get(download_url, stream=True)
+        response.raise_for_status()
+
+        temp_file = os.path.join(config.extension_uri, "temp.tar.gz")
+        with open(temp_file, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+
+        if os.path.exists(web_path):
+            shutil.rmtree(web_path)
+
+        logging.info("Extracting web distribution...")
+        with tarfile.open(temp_file, "r:gz") as tar:
+            members = [
+                member for member in tar.getmembers() if member.name.startswith("web/")
+            ]
+            tar.extractall(path=config.extension_uri, members=members)
+
+        os.remove(temp_file)
+        logging.info("Web distribution downloaded successfully.")
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Failed to download web distribution: {e}")
+    except tarfile.TarError as e:
+        logging.error(f"Failed to extract web distribution: {e}")
+    except Exception as e:
+        logging.error(f"An unexpected error occurred: {e}")
+
+
 def resolve_model_base_paths():
     folders = list(folder_paths.folder_names_and_paths.keys())
     config.model_base_paths = {}
diff --git a/vite.config.ts b/vite.config.ts
index 6462473..c912982 100644
--- a/vite.config.ts
+++ b/vite.config.ts
@@ -86,8 +86,29 @@ function dev(): Plugin {
   }
 }
 
+function createWebVersion(): Plugin {
+  return {
+    name: 'vite-plugin-web-version',
+    apply: 'build',
+    enforce: 'post',
+    writeBundle() {
+      const pyProjectContent = fs.readFileSync('pyproject.toml', 'utf8')
+      const [, version] = pyProjectContent.match(/version = "(.*)"/) ?? []
+
+      const metadata = [
+        `version: ${version}`,
+        `build_time: ${new Date().toISOString()}`,
+        '',
+      ].join('\n')
+
+      const metadataFilePath = path.join(__dirname, 'web', 'version.yaml')
+      fs.writeFileSync(metadataFilePath, metadata, 'utf-8')
+    },
+  }
+}
+
 export default defineConfig({
-  plugins: [vue(), css(), output(), dev()],
+  plugins: [vue(), css(), output(), dev(), createWebVersion()],
 
   build: {
     outDir: 'web',