Add HF and Civitai token tests

ai-dock · Aug 7, 2024 · 67b9f32 · 67b9f32
1 parent f1435d9
commit 67b9f32
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 4 deletions.
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -18,7 +18,7 @@ jobs:
         build:
           # Undeclared SHA tags with latest commit from master branch
           # Only building periodic sha tagged images
-          - {latest: "false", sha: "de17a97", python: "3.10", pytorch: "2.3.1"}
+          - {latest: "true", sha: "de17a97", python: "3.10", pytorch: "2.3.1"}
     steps:
       -
         name: Free Space
@@ -102,7 +102,7 @@ jobs:
         build:
           # Undeclared SHA tags with latest commit from master branch
           # Only building periodic sha tagged images
-          - {latest: "false", sha: "de17a97", python: "3.10", pytorch: "2.3.1", cuda: "12.1.1-base"}
+          - {latest: "true", sha: "de17a97", python: "3.10", pytorch: "2.3.1", cuda: "12.1.1-base"}
     steps:
       -
         name: Free Space
@@ -189,7 +189,7 @@ jobs:
         build:
           # Undeclared SHA tags with latest commit from master branch
           # Only building periodic sha tagged images
-          - {latest: "false", sha: "de17a97", python: "3.10", pytorch: "2.3.1", rocm: "6.0-runtime"}
+          - {latest: "true", sha: "de17a97", python: "3.10", pytorch: "2.3.1", rocm: "6.0-runtime"}
     steps:
       -
         name: Free Space

diff --git a/config/provisioning/default.sh b/config/provisioning/default.sh
@@ -6,6 +6,8 @@
 
 # Packages are installed after nodes so we can fix them...
 
+#DEFAULT_WORKFLOW="https://..."
+
 APT_PACKAGES=(
     #"package-1"
     #"package-2"
@@ -147,6 +149,14 @@ function provisioning_get_nodes() {
     done
 }
 
+function provisioning_get_default_workflow() {
+    if [[ -n $DEFAULT_WORKFLOW ]]; then
+        workflow_json=$(curl -s "$DEFAULT_WORKFLOW")
+        if [[ -n $workflow_json ]]; then
+            echo "export const defaultGraph = $workflow_json;" > /opt/ComfyUI/web/scripts/defaultGraph.js
+        fi
+    fi
+}
 
 function provisioning_get_models() {
     if [[ -z $2 ]]; then return 1; fi
@@ -174,6 +184,38 @@ function provisioning_print_end() {
     printf "\nProvisioning complete:  Web UI will start now\n\n"
 }
 
+function provisioning_has_valid_hf_token() {
+    [[ -n "$HF_TOKEN" ]] || return 1
+    url="https://huggingface.co/api/whoami-v2"
+
+    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET "$url" \
+        -H "Authorization: Bearer $HF_TOKEN" \
+        -H "Content-Type: application/json")
+
+    # Check if the token is valid
+    if [ "$response" -eq 200 ]; then
+        return 0
+    else
+        return 1
+    fi
+}
+
+function provisioning_has_valid_civitai_token() {
+    [[ -n "$CIVITAI_TOKEN" ]] || return 1
+    url="https://civitai.com/api/v1/models?hidden=1&limit=1"
+
+    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET "$url" \
+        -H "Authorization: Bearer $HF_TOKEN" \
+        -H "Content-Type: application/json")
+
+    # Check if the token is valid
+    if [ "$response" -eq 200 ]; then
+        return 0
+    else
+        return 1
+    fi
+}
+
 # Download from $1 URL to $2 file path
 function provisioning_download() {
     if [[ -n $HF_TOKEN && $1 =~ ^https://([a-zA-Z0-9_-]+\.)?huggingface\.co(/|$|\?) ]]; then

diff --git a/config/provisioning/flux.sh b/config/provisioning/flux.sh
@@ -36,7 +36,7 @@ UNET_MODELS=(
 VAE_MODELS=(
 )
 
-if [[ -n $HF_TOKEN ]]; then
+if provisioning_has_valid_hf_token; then
     UNET_MODELS+=("https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors")
     VAE_MODELS+=("https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/ae.safetensors")
 else
@@ -175,6 +175,38 @@ function provisioning_print_end() {
     printf "\nProvisioning complete:  Web UI will start now\n\n"
 }
 
+function provisioning_has_valid_hf_token() {
+    [[ -n "$HF_TOKEN" ]] || return 1
+    url="https://huggingface.co/api/whoami-v2"
+
+    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET "$url" \
+        -H "Authorization: Bearer $HF_TOKEN" \
+        -H "Content-Type: application/json")
+
+    # Check if the token is valid
+    if [ "$response" -eq 200 ]; then
+        return 0
+    else
+        return 1
+    fi
+}
+
+function provisioning_has_valid_civitai_token() {
+    [[ -n "$CIVITAI_TOKEN" ]] || return 1
+    url="https://civitai.com/api/v1/models?hidden=1&limit=1"
+
+    response=$(curl -o /dev/null -s -w "%{http_code}" -X GET "$url" \
+        -H "Authorization: Bearer $HF_TOKEN" \
+        -H "Content-Type: application/json")
+
+    # Check if the token is valid
+    if [ "$response" -eq 200 ]; then
+        return 0
+    else
+        return 1
+    fi
+}
+
 # Download from $1 URL to $2 file path
 function provisioning_download() {
     if [[ -n $HF_TOKEN && $1 =~ ^https://([a-zA-Z0-9_-]+\.)?huggingface\.co(/|$|\?) ]]; then