Refactor code to improve initialization process

Author: yoland68

comfy/__main__.py

@@ -1,4 +1,5 @@
-from comfy.cmdline import app
+from comfy.cmdline import init, app
 
 if __name__ == '__main__':  # pragma: nocover
+    init()
     app()

comfy/cmdline.py

@@ -2,28 +2,75 @@
 from typing_extensions import Annotated
 from comfy.command.models import models
 from rich import print
+import os
+import subprocess
 
 
 from comfy.command import custom_nodes
 from comfy.command import install as install_inner
 from comfy.command import run as run_inner
 from comfy import constants
 from comfy.env_checker import EnvChecker
+from comfy.meta_data import MetadataManager
+from comfy import env_checker
+from rich.console import Console
+import time
 
 app = typer.Typer()
 
+
+def init():
+    _env_checker = EnvChecker()
+    metadata_manager = MetadataManager()
+    start_time = time.time()
+    metadata_manager.scan_dir()
+    end_time = time.time()
+    
+    print(f"scan_dir took {end_time - start_time:.2f} seconds to run")
+
+
 @app.callback(invoke_without_command=True)
 def no_command(ctx: typer.Context):
     if ctx.invoked_subcommand is None:
         print(ctx.get_help())
         ctx.exit()
 
+
 @app.command(help="Download and install ComfyUI")
 def install(
-    url: str = constants.COMFY_GITHUB_URL,
-    workspace: Annotated[str, typer.Option(help="Path to the output directory.")] = None,
-    ):
-    install_inner.execute(url)
+    url: Annotated[
+        str,
+        typer.Option(show_default=False)
+    ] = constants.COMFY_GITHUB_URL,
+    workspace: Annotated[
+        str,
+        typer.Option(
+            show_default=False,
+            help="Path to ComfyUI workspace")
+    ] = None,
+):
+    checker = EnvChecker()
+    if checker.python_version.major < 3:
+        print(
+            "[bold red]Python version 3.6 or higher is required to run ComfyUI.[/bold red]"
+        )
+        print(
+            f"You are currently using Python version {env_checker.format_python_version(checker.python_version)}."
+        )
+    if checker.currently_in_comfy_repo:
+        console = Console()
+        # TODO: warn user that you are teh
+        
+
+    install_inner.execute(url, workspace)
+
+
+def update(self):
+    print(f"Updating ComfyUI in {self.workspace}...")
+    os.chdir(self.workspace)
+    subprocess.run(["git", "pull"])
+    subprocess.run(["pip", "install", "-r", "requirements.txt"])
+
 
 @app.command(help="Run workflow file")
 def run(

comfy/command/install.py

@@ -4,7 +4,7 @@
 from rich import print
 from comfy import env_checker
 
-def execute(url: str, *args, **kwargs):
+def execute(url: str, comfy_workspace: str, *args, **kwargs):
     print(f"Installing from {url}")
 
     checker = env_checker.EnvChecker()

comfy/constants.py

@@ -7,4 +7,8 @@
     "git@github.com:drip-art/comfy.git",
     "https://github.com/comfyanonymous/ComfyUI.git",
     "https://github.com/drip-art/ComfyUI.git",
-]
+]
+
+# Referencing supported pt extension from ComfyUI
+# https://github.com/comfyanonymous/ComfyUI/blob/a88b0ebc2d2f933c94e42aa689c42e836eedaf3c/folder_paths.py#L5
+SUPPORTED_PT_EXTENSIONS = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors'])

comfy/env_checker.py

@@ -30,8 +30,8 @@ def format_python_version(version_info):
         str: The formatted Python version string.
     """
     if version_info.major == 3 and version_info.minor > 8:
-        return "{}.{}".format(version_info.major, version_info.minor)
-    return "[bold red]{}.{}[/bold red]".format(version_info.major, version_info.minor)
+        return "{}.{}.{}".format(version_info.major, version_info.minor, version_info.micro)
+    return "[bold red]{}.{}.{}[/bold red]".format(version_info.major, version_info.minor, version_info.micro)
 
 
 def check_comfy_server_running():
@@ -71,8 +71,9 @@ class EnvChecker(object):
     def __init__(self):
         self.virtualenv_path = None
         self.conda_env = None
-        self.python_version = None
+        self.python_version: None = None
         self.currently_in_comfy_repo = False
+        self.comfy_repo = None
         self.check()
 
     def check(self):
@@ -93,9 +94,12 @@ def check(self):
             self.currently_in_comfy_repo = (
                 repo.remotes.origin.url in constants.COMFY_ORIGIN_URL_CHOICES
             )
+            if self.currently_in_comfy_repo:
+                self.comfy_repo = repo
         except git.exc.InvalidGitRepositoryError:
             self.currently_in_comfy_repo = False
 
+
     def print(self):
         table = Table(":laptop_computer: Environment", "Value")
         table.add_row("Python Version", format_python_version(sys.version_info))

comfy/meta_data.py

@@ -1,7 +1,11 @@
 import os
 import yaml
+import concurrent.futures
+from pathlib import Path
+import os
 
 from comfy.utils import singleton
+from comfy import constants
 
 
 @singleton
@@ -14,6 +18,35 @@ def __init__(self):
         self.metadata_file = None
         self.metadata = {}
 
+    
+    def scan_dir(self):
+        config_files = []
+        for root, dirs, files in os.walk("."):
+            for file in files:
+                if file.endswith(constants.SUPPORTED_PT_EXTENSIONS):
+                    config_files.append(os.path.join(root, file))
+        return config_files
+
+
+    def scan_dir_concur(self):
+        base_path = Path(".")
+        model_files = []
+
+        # Function to check if the file is config.json
+        def check_file(path):
+            if path.name.endswith(constants.SUPPORTED_PT_EXTENSIONS):
+                return str(path)
+
+        # Use ThreadPoolExecutor to manage concurrency
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = [executor.submit(check_file, p) for p in base_path.rglob('*')]
+            for future in concurrent.futures.as_completed(futures):
+                if future.result():
+                    model_files.append(future.result())
+
+        return model_files
+        
+
     def load_metadata(self):
         if os.path.exists(self.metadata_file):
             with open(self.metadata_file, "r", encoding="utf-8") as file: