Comfy-Org · robinjhuang · Nov 12, 2024 · Nov 7, 2024 · Nov 7, 2024 · Nov 7, 2024
diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py
@@ -242,6 +242,10 @@ def install(
             help="Use new fast dependency installer",
         ),
     ] = False,
+    manager_commit: Annotated[
+        Optional[str],
+        typer.Option(help="Specify commit hash for ComfyUI-Manager"),
+    ] = None,
 ):
     check_for_updates()
     checker = EnvChecker()
@@ -279,6 +283,7 @@ def install(
             skip_torch_or_directml=skip_torch_or_directml,
             skip_requirement=skip_requirement,
             fast_deps=fast_deps,
+            manager_commit=manager_commit,
         )
         rprint(f"ComfyUI is installed at: {comfy_path}")
         return None
@@ -346,6 +351,7 @@ def install(
         skip_torch_or_directml=skip_torch_or_directml,
         skip_requirement=skip_requirement,
         fast_deps=fast_deps,
+        manager_commit=manager_commit,
     )
 
     rprint(f"ComfyUI is installed at: {comfy_path}")

diff --git a/comfy_cli/command/install.py b/comfy_cli/command/install.py
@@ -157,6 +157,7 @@ def execute(
     skip_manager: bool,
     version: str,
     commit: Optional[str] = None,
+    manager_commit: Optional[str] = None,
     gpu: constants.GPU_OPTION = None,
     cuda_version: constants.CUDAVersion = constants.CUDAVersion.v12_1,
     plat: constants.OS = None,
@@ -188,13 +189,17 @@ def execute(
         clone_comfyui(url=url, repo_dir=repo_dir)
 
     if version != "nightly":
-        checkout_stable_comfyui(version=version, repo_dir=repo_dir)
+        try:
+            checkout_stable_comfyui(version=version, repo_dir=repo_dir)
+        except GitHubRateLimitError as e:
+            rprint(f"[bold red]Error checking out ComfyUI version: {e}[/bold red]")
+            sys.exit(1)
 
     elif not check_comfy_repo(repo_dir)[0]:
         rprint(
             f"[bold red]'{repo_dir}' already exists. But it is an invalid ComfyUI repository. Remove it and retry.[/bold red]"
         )
-        exit(-1)
+        sys.exit(-1)
 
     # checkout specified commit
     if commit is not None:
@@ -234,6 +239,8 @@ def execute(
                 )
             else:
                 subprocess.run(["git", "clone", manager_url, manager_repo_dir], check=True)
+                if manager_commit is not None:
+                    subprocess.run(["git", "checkout", manager_commit], check=True, cwd=manager_repo_dir)
 
             if not fast_deps:
                 pip_install_manager_dependencies(repo_dir)
@@ -281,12 +288,39 @@ def validate_version(version: str) -> Optional[str]:
         ) from exc
 
 
+class GitHubRateLimitError(Exception):
+    """Raised when GitHub API rate limit is exceeded"""
+
+
 def fetch_github_releases(repo_owner: str, repo_name: str) -> List[Dict[str, str]]:
     """
     Fetch the list of releases from the GitHub API.
+    Handles rate limiting by logging the wait time.
     """
     url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/releases"
-    response = requests.get(url)
+
+    headers = {}
+    if github_token := os.getenv("GITHUB_TOKEN"):
+        headers["Authorization"] = f"Bearer {github_token}"
+
+    response = requests.get(url, headers=headers, timeout=5)
+
+    # Handle rate limiting
+    if response.status_code in (403, 429):
+        # Check rate limit headers
+        remaining = int(response.headers.get("x-ratelimit-remaining", 0))
+        if remaining == 0:
+            reset_time = int(response.headers.get("x-ratelimit-reset", 0))
+            message = f"Primary rate limit from Github exceeded! Please retry after: {reset_time})"
+            raise GitHubRateLimitError(message)
+
+        if "retry-after" in response.headers:
+            wait_seconds = int(response.headers["retry-after"])
+            message = f"Rate limit from Github exceeded! Please wait {wait_seconds} seconds before retrying."
+            rprint(f"[yellow]{message}[/yellow]")
+            raise GitHubRateLimitError(message)
+
+    response.raise_for_status()
     return response.json()
 
 

diff --git a/tests/comfy_cli/test_install.py b/tests/comfy_cli/test_install.py
@@ -53,7 +53,7 @@ def test_fetch_releases_su
5E4B
ccess(mock_get):
     assert len(releases) == 2
     assert releases[0]["tag_name"] == "v1.0.0"
     assert releases[1]["tag_name"] == "v1.1.0"
-    mock_get.assert_called_once_with("https://api.github.com/repos/owner/repo/releases")
+    mock_get.assert_called_once_with("https://api.github.com/repos/owner/repo/releases", headers={}, timeout=5)
 
 
 @patch("requests.get")