feat: full models list in models.yaml

.github/scripts/scope.js

@@ -28,7 +28,8 @@ module.exports = ({core}) => {
     if (!m || !n) return false;
     return m["name"] === n["name"] &&
       m["source"] === n["source"] &&
-      m["hf-name"] === n["hf-name"];
+      m["hf-name"] === n["hf-name"] &&
+      m["runner"] === n["runner"];
   }
 
   for (const key of keys) {
@@ -45,6 +46,11 @@ module.exports = ({core}) => {
     // !n: deleted, which we ignore
   }
 
+  // Potential shape change:
+  // Instead of this binary choice, output all the models
+  // And annotate each model with which tasks should run for it
+  // (like export, benchmark, upload, etc)
+  // That means every model will make it to the index job through the same path
   core.setOutput('to_export', JSON.stringify(to_export));
   core.setOutput('unchanged', JSON.stringify(unchanged));
 }

.github/workflows/export.yaml

@@ -1,7 +1,7 @@
 on:
   workflow_call:
     secrets:
-      HF_AUTH_TOKEN:
+      HF_TOKEN:
         required: true
     inputs:
       model-name:
@@ -10,6 +10,9 @@ on:
       model-source:
         required: true
         type: string
+      runner:
+        required: true
+        type: string
       hf-name:
         required: false
         type: string
@@ -19,15 +22,18 @@ on:
 
 jobs:
   export:
-    runs-on: ubuntu-latest
+    runs-on: ${{ inputs.runner }}
     steps:
       - name: Checkout
         uses: actions/checkout@v4
 
       - uses: astral-sh/setup-uv@v5
-      - run: uv sync # TODO: cache uv env (does setup-uv do that already automatically?)
+      - run: uv sync
 
       - run: uv run immich_model_exporter export "${{ inputs.model-name }}" "${{ inputs.model-source }}" --hf-model-name "${{ inputs.hf-name }}"
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+
       - uses: actions/upload-artifact@v4
         with:
           name: ${{ inputs.model-name }}
@@ -38,14 +44,14 @@ jobs:
 
   upload:
     if: ${{ inputs.upload }}
-    runs-on: ubuntu-latest
+    runs-on: ${{ inputs.runner }}
     needs: export
     steps:
       - name: Checkout
         uses: actions/checkout@v4
 
       - uses: astral-sh/setup-uv@v5
-      - run: uv sync # TODO: cache uv env (does setup-uv do that already automatically?)
+      - run: uv sync
 
       - uses: actions/download-artifact@v4
         with:
@@ -55,4 +61,4 @@ jobs:
 
       - run: uv run immich_model_exporter upload "${{ inputs.model-name }}" --hf-model-name "${{ inputs.hf-name }}" --hf-organization immich-testing
         env:
-          HF_AUTH_TOKEN: ${{ secrets.HF_AUTH_TOKEN }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}

.github/workflows/push.yaml

@@ -66,6 +66,7 @@ jobs:
             script({core})
 
 
+  # TODO: Explicitly skip this if nothing to do
   export:
     uses: ./.github/workflows/export.yaml
     needs: configure
@@ -80,4 +81,5 @@ jobs:
       model-name: ${{ matrix.name }}
       model-source: ${{ matrix.source }}
       hf-name: ${{ matrix.hf-name }}
+      runner: ${{ matrix.runner || 'ubuntu-latest'}}
       upload: ${{ inputs.force || github.event_name == 'release'}}

immich_model_exporter/__init__.py

@@ -147,7 +147,6 @@ def upload(
     input_dir: Path = Path("models"),
     hf_model_name: str | None = None,
     hf_organization: str = "immich-app",
-    hf_auth_token: Annotated[str | None, typer.Option(envvar="HF_AUTH_TOKEN")] = None,
 ) -> None:
     from huggingface_hub import create_repo, upload_folder
 
@@ -158,14 +157,13 @@ def upload(
 
     @retry(stop=stop_after_attempt(5), wait=wait_fixed(5))
     def upload_model() -> None:
-        create_repo(repo_id, exist_ok=True, token=hf_auth_token)
+        create_repo(repo_id, exist_ok=True)
         upload_folder(
             repo_id=repo_id,
             folder_path=model_dir,
             # remote repo files to be deleted before uploading
             # deletion is in the same commit as the upload, so it's atomic
             delete_patterns=DELETE_PATTERNS,
-            token=hf_auth_token,
         )
 
     upload_model()

immich_model_exporter/run.py

@@ -1,92 +1,29 @@
 import subprocess
 from pathlib import Path
+from typing import NamedTuple
 
+import yaml
 from exporters.constants import ModelSource
 
-from immich_model_exporter import clean_name
 from immich_model_exporter.exporters.constants import SOURCE_TO_TASK
 
-mclip = [
-    "M-CLIP/LABSE-Vit-L-14",
-    "M-CLIP/XLM-Roberta-Large-Vit-B-16Plus",
-    "M-CLIP/XLM-Roberta-Large-Vit-B-32",
-    "M-CLIP/XLM-Roberta-Large-Vit-L-14",
-]
-
-openclip = [
-    "RN101__openai",
-    "RN101__yfcc15m",
-    "RN50__cc12m",
-    "RN50__openai",
-    "RN50__yfcc15m",
-    "RN50x16__openai",
-    "RN50x4__openai",
-    "RN50x64__openai",
-    "ViT-B-16-SigLIP-256__webli",
-    "ViT-B-16-SigLIP-384__webli",
-    "ViT-B-16-SigLIP-512__webli",
-    "ViT-B-16-SigLIP-i18n-256__webli",
-    "ViT-B-16-SigLIP2__webli",
-    "ViT-B-16-SigLIP__webli",
-    "ViT-B-16-plus-240__laion400m_e31",
-    "ViT-B-16-plus-240__laion400m_e32",
-    "ViT-B-16__laion400m_e31",
-    "ViT-B-16__laion400m_e32",
-    "ViT-B-16__openai",
-    "ViT-B-32-SigLIP2-256__webli",
-    "ViT-B-32__laion2b-s34b-b79k",
-    "ViT-B-32__laion2b_e16",
-    "ViT-B-32__laion400m_e31",
-    "ViT-B-32__laion400m_e32",
-    "ViT-B-32__openai",
-    "ViT-H-14-378-quickgelu__dfn5b",
-    "ViT-H-14-quickgelu__dfn5b",
-    "ViT-H-14__laion2b-s32b-b79k",
-    "ViT-L-14-336__openai",
-    "ViT-L-14-quickgelu__dfn2b",
-    "ViT-L-14__laion2b-s32b-b82k",
-    "ViT-L-14__laion400m_e31",
-    "ViT-L-14__laion400m_e32",
-    "ViT-L-14__openai",
-    "ViT-L-16-SigLIP-256__webli",
-    "ViT-L-16-SigLIP-384__webli",
-    "ViT-L-16-SigLIP2-256__webli",
-    "ViT-L-16-SigLIP2-384__webli",
-    "ViT-L-16-SigLIP2-512__webli",
-    "ViT-SO400M-14-SigLIP-384__webli",
-    "ViT-SO400M-14-SigLIP2-378__webli",
-    "ViT-SO400M-14-SigLIP2__webli",
-    "ViT-SO400M-16-SigLIP2-256__webli",
-    "ViT-SO400M-16-SigLIP2-384__webli",
-    "ViT-SO400M-16-SigLIP2-512__webli",
-    "ViT-gopt-16-SigLIP2-256__webli",
-    "ViT-gopt-16-SigLIP2-384__webli",
-    "nllb-clip-base-siglip__mrl",
-    "nllb-clip-base-siglip__v1",
-    "nllb-clip-large-siglip__mrl",
-    "nllb-clip-large-siglip__v1",
-    "xlm-roberta-base-ViT-B-32__laion5b_s13b_b90k",
-    "xlm-roberta-large-ViT-H-14__frozen_laion5b_s13b_b90k",
-]
-
-insightface = [
-    "antelopev2",
-    "buffalo_l",
-    "buffalo_m",
-    "buffalo_s",
-]
-
-
-def export_models(models: list[str], source: ModelSource) -> None:
+
+class ModelSpec(NamedTuple):
+    name: str
+    hf_name: str
+    source: ModelSource
+
+
+def export_models(models: list[ModelSpec]) -> None:
     profiling_dir = Path("profiling")
     profiling_dir.mkdir(exist_ok=True)
     for model in models:
         try:
-            model_dir = f"models/{clean_name(model)}"
-            task = SOURCE_TO_TASK[source]
+            model_dir = f"models/{model.name}"
+            task = SOURCE_TO_TASK[model.source]
 
-            print(f"Processing model {model}")
-            subprocess.check_call(["python", "-m", "immich_model_exporter", "export", model, source])
+            print(f"Processing model {model.name}")
+            subprocess.check_call(["python", "-m", "immich_model_exporter", "export", model.hf_name, model.source])
             subprocess.check_call(
                 [
                     "python",
@@ -104,10 +41,25 @@ def export_models(models: list[str], source: ModelSource) -> None:
             print(f"Failed to export model {model}: {e}")
 
 
+def read_models() -> list[ModelSpec]:
+    with open("../models.yaml") as f:
+        y = yaml.safe_load(f)
+    models = []
+    for m in y["models"]:
+        name = m["name"]
+        source = ModelSource(m["source"])
+        hfName = m["hf-name"] if "hf-name" in m else name
+        model = ModelSpec(name, hfName, source)
+        models.append(model)
+    return models
+
+
 if __name__ == "__main__":
-    export_models(mclip, ModelSource.MCLIP)
-    export_models(openclip, ModelSource.OPENCLIP)
-    export_models(insightface, ModelSource.INSIGHTFACE)
+    models = read_models()
+
+    export_models(models)
+
+    openclip_names = [m.hf_name.replace("__", ",") for m in models if m.source == ModelSource.OPENCLIP]
 
     Path("results").mkdir(exist_ok=True)
     dataset_root = Path("datasets")
@@ -119,7 +71,7 @@ def export_models(models: list[str], source: ModelSource) -> None:
             "clip_benchmark",
             "eval",
             "--pretrained_model",
-            *[name.replace("__", ",") for name in openclip],
+            *openclip_names,
             "--task",
             "zeroshot_retrieval",
             "--dataset",
@@ -181,7 +133,7 @@ def export_models(models: list[str], source: ModelSource) -> None:
             "clip_benchmark",
             "eval",
             "--pretrained_model",
-            *[name.replace("__", ",") for name in openclip],
+            *openclip_names,
             "--task",
             "zeroshot_retrieval",
             "--dataset",
@@ -219,7 +171,7 @@ def export_models(models: list[str], source: ModelSource) -> None:
             "clip_benchmark",
             "eval",
             "--pretrained_model",
-            *[name.replace("__", ",") for name in openclip],
+            *openclip_names,
             "--task",
             "zeroshot_retrieval",
             "--dataset",