Add nightly and stable difference on torch install

Comfy-Org · Jul 11, 2024 · 0478fbd · 0478fbd
1 parent 2af1b4c
commit 0478fbd
Showing 1 changed file with 47 additions and 8 deletions.
diff --git a/action.yml b/action.yml
@@ -10,13 +10,17 @@ inputs:
     required: false
     default: "3.9"
   cuda_version:
-    description: "CUDA Version. Will be used in the conda install command for pytorch. 11.8 or 12.1. Only valid when os is linux or windows."
+    description: "CUDA Version. Will be used in the conda install command for pytorch."
     required: false
     default: "12.1"
   torch_version:
-    description: "Pytorch Version. Will be used in the conda install command for pytorch. 1.10.0 or 1.11.0. Only valid when os is linux or windows."
+    description: "Pytorch Version type, 'stable', 'nightly', 'specific'"
     required: false
     default: 'stable'
+  specific_torch_install:
+    description: "Command for installing specific torch version. E.g. `conda install bla bla`"
+    required: false
+    default: ''
   models-json:
     description: 'JSON string containing models and their download URLs. The models will be downloaded into the exact directory relative to /ComfyUI/models/. eg { "model_name": { url: "https://example.com/model.pth", "directory": "checkpoints" } }'
     required: false
@@ -100,17 +104,36 @@ runs:
         python-version: ${{ inputs.python_version }}
 
     - name: '[Unix-Mac-Only] Install Pytorch nightly'
-      if: ${{ inputs.os == 'macos' }}
+      if: ${{ inputs.os == 'macos' && inputs.torch_version == 'nightly' }}
+      shell: bash -el {0}
+      run: |
+        conda install pytorch-nightly::pytorch torchvision torchaudio -c pytorch-nightly
+
+    - name: '[Unix-Mac-Only] Install Pytorch stable'
+      if: ${{ inputs.os == 'macos' && inputs.torch_version == 'stable' }}
       shell: bash -el {0}
       run: |
-        conda install pytorch torchvision torchaudio -c pytorch-nightly
+        conda install pytorch::pytorch torchvision torchaudio -c pytorch
+    
+    - name: '[Unix-Linux-Only] Install Pytorch nightly'
+      if: ${{ inputs.os == 'linux' && inputs.torch_version == 'nightly' }}
+      shell: bash -el {0}
+      run: |
+        conda install pytorch torchvision torchaudio pytorch-cuda=${{ inputs.cuda_version }} -c pytorch-nightly -c nvidia
 
-    - name: '[Unix-Linux-Only] Install Pytorch'
-      if: ${{ inputs.os == 'linux' }}
+    - name: '[Unix-Linux-Only] Install Pytorch stable'
+      if: ${{ inputs.os == 'linux' && inputs.torch_version == 'stable' }}
       shell: bash -el {0}
       run: |
         conda install pytorch torchvision torchaudio pytorch-cuda=${{ inputs.cuda_version}} -c pytorch -c nvidia
 
+    - name: '[Unix] Install Pytorch specific version'
+      if: ${{ inputs.os == 'macos' && inputs.torch_version == 'specific' }}
+      shell: bash -el {0}
+      run: |
+        echo "Run specific torch install command"
+        ${{ inputs.specific_torch_install }}
+
     - name: '[Unix] Install dependencies'
       if: ${{ inputs.os != 'windows' }}
       shell: bash -el {0}
@@ -265,13 +288,29 @@ runs:
         python-version: ${{ inputs.python_version }}
       continue-on-error: true
 
-    - name: '[Win-Only] Install Pytorch'
-      if: ${{ inputs.os == 'windows' }}
+    - name: '[Win-Only] Install Pytorch nightly'
+      if: ${{ inputs.os == 'windows' && inputs.torch_version == 'nightly'}}
+      shell: powershell
+      run: |
+        conda activate gha-comfyui-${{ inputs.python_version }}-${{ inputs.torch_version }}
+        Get-Command pip
+        pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
+
+    - name: '[Win-Only] Install Pytorch stable'
+      if: ${{ inputs.os == 'windows' && inputs.torch_version =='stable' }}
       shell: powershell
       run: |
         conda activate gha-comfyui-${{ inputs.python_version }}-${{ inputs.torch_version }}
         Get-Command pip
         pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+    
+    - name: '[Win] Install Pytorch specific version'
+      if: ${{ inputs.os == 'windows' && inputs.torch_version == 'specific' }}
+      shell: powershell
+      run: |
+        conda activate gha-comfyui-${{ inputs.python_version }}-${{ inputs.torch_version }}
+        Write-Output "Run specific torch install command"
+        ${{ inputs.specific_torch_install }}
 
     - name: '[Win] Install dependencies'
       if: ${{ inputs.os == 'windows' }}