From 293b96c816561ca76817dc415990badba2e1ded4 Mon Sep 17 00:00:00 2001
From: Max Klein <telamonian@users.noreply.github.com>
Date: Thu, 1 Aug 2024 21:53:38 -0400
Subject: [PATCH] Allowing picking of specific branch in
 `--url`/`--manager-url` options for `comfy install` command (#147)

* add support for choosing specific branch to `--manager-url` option

* update help for `--manager-url` to include how to pick branch

* add support for picking specific branch to `--url` option

* fix type hinting of `execute` signature

* small cleanup of url/branch splitting code
---
 comfy_cli/cmdline.py         | 14 ++++++++++++--
 comfy_cli/command/install.py | 24 +++++++++++++++++++-----
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py
index 6949240..32c3e39 100644
--- a/comfy_cli/cmdline.py
+++ b/comfy_cli/cmdline.py
@@ -140,9 +140,19 @@ def entry(
 @app.command(help="Download and install ComfyUI and ComfyUI-Manager")
 @tracking.track_command()
 def install(
-    url: Annotated[str, typer.Option(show_default=False)] = constants.COMFY_GITHUB_URL,
+    url: Annotated[
+        str,
+        typer.Option(
+            show_default=False,
+            help="url or local path pointing to the ComfyUI core git repo to be installed. A specific branch can optionally be specified using a setuptools-like syntax, eg https://foo.git@bar",
+        ),
+    ] = constants.COMFY_GITHUB_URL,
     manager_url: Annotated[
-        str, typer.Option(show_default=False)
+        str,
+        typer.Option(
+            show_default=False,
+            help="url or local path pointing to the ComfyUI-Manager git repo to be installed. A specific branch can optionally be specified using a setuptools-like syntax, eg https://foo.git@bar",
+        ),
     ] = constants.COMFY_MANAGER_GITHUB_URL,
     restore: Annotated[
         bool,
diff --git a/comfy_cli/command/install.py b/comfy_cli/command/install.py
index 54592df..553046e 100644
--- a/comfy_cli/command/install.py
+++ b/comfy_cli/command/install.py
@@ -1,10 +1,10 @@
+from rich import print
 import os
 import platform
 import subprocess
 import sys
-
 import typer
-from rich import print
+from typing import Optional
 
 from comfy_cli import constants, ui, utils
 from comfy_cli.command.custom_nodes.command import update_node_id_cache
@@ -163,7 +163,7 @@ def execute(
     comfy_path: str,
     restore: bool,
     skip_manager: bool,
-    commit=None,
+    commit: Optional[str] = None,
     gpu: constants.GPU_OPTION = None,
     cuda_version: constants.CUDAVersion = constants.CUDAVersion.v12_1,
     plat: constants.OS = None,
@@ -189,7 +189,14 @@ def execute(
         os.makedirs(parent_path, exist_ok=True)
 
     if not os.path.exists(repo_dir):
-        subprocess.run(["git", "clone", url, repo_dir], check=True)
+        if "@" in url:
+            # clone specific branch
+            url, branch = url.rsplit("@", 1)
+
+            subprocess.run(["git", "clone", "-b", branch, url, repo_dir], check=True)
+        else:
+            subprocess.run(["git", "clone", url, repo_dir], check=True)
+
     elif not check_comfy_repo(repo_dir)[0]:
         print(
             f"[bold red]'{repo_dir}' already exists. But it is an invalid ComfyUI repository. Remove it and retry.[/bold red]"
@@ -226,7 +233,14 @@ def execute(
         else:
             print("\nInstalling ComfyUI-Manager..")
 
-            subprocess.run(["git", "clone", manager_url, manager_repo_dir], check=True)
+            if "@" in manager_url:
+                # clone specific branch
+                manager_url, manager_branch = manager_url.rsplit("@", 1)
+
+                subprocess.run(["git", "clone", "-b", manager_branch, manager_url, manager_repo_dir], check=True)
+            else:
+                subprocess.run(["git", "clone", manager_url, manager_repo_dir], check=True)
+
             install_manager_dependencies(repo_dir)
 
         update_node_id_cache()