feat: automatic registering of newer models and validate models + improve ui a bit

app.py

@@ -7,6 +7,8 @@ import numpy as np
 import pandas as pd
 import torch
 
+from hf_repo_utils import auto_register_kaloscope_variants, validate_models
+
 # Import inference implementations
 from inference_onnx import ONNXInference, softmax
 from inference_pytorch import PyTorchInference
@@ -25,27 +27,36 @@ except ImportError as e:
 # Define available models with their configuration
 # Format: "display_name": {
 #     "type": "onnx" or "pytorch",
-#     "path": "local path or repo:filename", # if left empty, will use repo id
-#     "repo_id": "huggingface repo id" (optional),
-#     "arch": "model architecture name"
+#     "path": "local path or repo:filename", # for local use - if left empty, will use repo id
+#     "repo_id": "huggingface repo id" (optional, but path cant be empty if repo_id empty),
+#     "subfolder": "subfolder in repo" (optional, if applicable),
+#     "arch": "model architecture name" - lsnet_xl_artist is expected for usual Kaloscope releases b, l, s exist but unused
 # }
 
 MODELS = {
-    "Kaloscope-onnx-opset18": {
+    "Kaloscope v1.1": {
+        "type": "pytorch",
+        "path": "",
+        "repo_id": "heathcliff01/Kaloscope",
+        "subfolder": "224-85.65",
+        "filename": "best_checkpoint.pth",
+        "arch": "lsnet_xl_artist",
+    },
+    "Kaloscope v1.0 ONNX": {
         "type": "onnx",
         "path": "",
         "repo_id": "DraconicDragon/Kaloscope-onnx-ema",
-        "filename": "lsnet_xl_artist-dynamo-opset18_merged.onnx",
+        "filename": "kaloscope_1-0.onnx",
         "arch": "lsnet_xl_artist",
     },
-    "Kaloscope-release ('best_checkpoint.pth')": {
+    "Kaloscope v1.0": {
         "type": "pytorch",
         "path": "",
         "repo_id": "heathcliff01/Kaloscope",
         "filename": "best_checkpoint.pth",
         "arch": "lsnet_xl_artist",
     },
-    "Kaloscope-release-ema": {
+    "Kaloscope v1.0 ema": {
         "type": "pytorch",
         "path": "",
         "repo_id": "DraconicDragon/Kaloscope-onnx-ema",
@@ -53,6 +64,8 @@ MODELS = {
         "arch": "lsnet_xl_artist",
     },
 }
+MODELS = validate_models(MODELS)
+auto_register_kaloscope_variants(MODELS)
 
 # Class mapping CSV configuration
 CSV_PATH = ""  # if left empty, will use repo's
@@ -111,7 +124,7 @@ def get_model_inference(model_name):
 
     if model_name not in model_cache:
         config = MODELS[model_name]
-        model_path = config["path"]
+        model_path = config.get("path") or ""
 
         print(f"Loading model: {model_name} ({config['filename']})")
 
@@ -119,13 +132,24 @@ def get_model_inference(model_name):
             raise gr.Error(f"Unknown model: {model_name}")
 
         # Check if local file exists, otherwise try to download
-        if not os.path.exists(model_path):
+        if not model_path or not os.path.exists(model_path):
             if "repo_id" in config and "filename" in config:
-                print(f"Model not found locally at {model_path}, attempting to download from Hugging Face...")
+                target_display = model_path or f"repo {config['repo_id']}"
+                print(f"Model not found locally at {target_display}, attempting to download from Hugging Face...")
                 try:
                     from huggingface_hub import hf_hub_download
 
-                    model_path = hf_hub_download(repo_id=config["repo_id"], filename=config["filename"])
+                    download_kwargs = {
+                        "repo_id": config["repo_id"],
+                        "filename": config["filename"],
+                    }
+                    if config.get("subfolder"):
+                        download_kwargs["subfolder"] = config["subfolder"]
+                    if config.get("revision"):
+                        download_kwargs["revision"] = config["revision"]
+
+                    model_path = hf_hub_download(**download_kwargs)
+                    config["path"] = model_path
                     print(f"Downloaded model to: {model_path}")
                 except Exception as e:
                     raise gr.Error(f"Could not load model from local path or Hugging Face: {e}")
@@ -394,7 +418,8 @@ with gr.Blocks(
                 model_selection = gr.Dropdown(
                     choices=[
                         (
-                            f"{name} | Repo: {MODELS[name].get('repo_id') or 'local'}",
+                            f"{name}",
+                            # f"{name} | Repo: {MODELS[name].get('repo_id') or 'local'}",
                             name,
                         )
                         for name in MODELS
@@ -408,13 +433,10 @@ with gr.Blocks(
                 time_display = gr.Markdown()  # populated after prediction
 
     gr.Markdown(
-        "Models sourced from [heathcliff01/Kaloscope](https://huggingface.co/heathcliff01/Kaloscope) (original release) "
-        + "and [DraconicDragon/Kaloscope-onnx-ema](https://huggingface.co/DraconicDragon/Kaloscope-onnx-ema) (ONNX converted and EMA weights)."
+        "Models sourced from [heathcliff01/Kaloscope](https://huggingface.co/heathcliff01/Kaloscope) (Original PyTorch releases) "
+        + "and [DraconicDragon/Kaloscope-onnx-ema](https://huggingface.co/DraconicDragon/Kaloscope-onnx-ema) (ONNX converted and EMA weights).  \n"
+        + "OpenVINO™ will be used to accelerate ONNX CPU inference with ONNX CPUExecutionProvider as fallback."
     )
-    gr.Markdown("OpenVINO™ will be used to accelerate CPU inference with ONNX CPUExecutionProvider as fallback.")
-    # gr.Markdown(
-    #    "ONNX models might output different scores compared to PyTorch possibly due to differences in numerical precision and operations (but they are still likely just as accurate overall)"
-    # )
 
     submit_btn.click(
         fn=predict,

hf_repo_utils.py

@@ -0,0 +1,133 @@
+"""Utility helpers for Kaloscope-related model management."""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Iterable, Tuple
+
+
+def _iter_repo_files(repo_info: Any) -> Iterable[Any]:
+    """Yield file metadata objects from a repository info response."""
+    siblings = getattr(repo_info, "siblings", None)
+    if not siblings:
+        return []
+    return siblings
+
+
+def auto_register_kaloscope_variants(models: Dict[str, Dict[str, Any]]) -> None:
+    """Discover Kaloscope checkpoints in subfolders and append them to the model list."""
+    try:
+        from huggingface_hub import HfApi
+    except ImportError:
+        print("huggingface_hub is not available; skipping Kaloscope auto-discovery.")
+        return
+
+    repo_id = "heathcliff01/Kaloscope"
+    api = HfApi()
+
+    try:
+        repo_info = api.model_info(repo_id=repo_id)
+    except Exception as exc:
+        print(f"Failed to query Hugging Face repo info: {exc}")
+        return
+
+    existing_entries: set[Tuple[str, str, str]] = {
+        (
+            config.get("repo_id", ""),
+            config.get("subfolder", ""),
+            config.get("filename", ""),
+        )
+        for config in models.values()
+        if config.get("repo_id") and config.get("filename")
+    }
+
+    discovered: list[Tuple[str, Dict[str, Any], Tuple[str, str, str]]] = []
+
+    for sibling in _iter_repo_files(repo_info):
+        path = getattr(sibling, "rfilename", None) or getattr(sibling, "path", None)
+        if not path or not path.endswith(".pth"):
+            continue
+
+        subfolder = os.path.dirname(path).strip("/")
+        filename = os.path.basename(path)
+
+        key = (repo_id, subfolder, filename)
+        if key in existing_entries:
+            continue
+
+        # Use repo_id/folder format if in subfolder, otherwise repo_id/filename
+        repo_name = repo_id.split("/")[-1]
+        if subfolder:
+            display_name = f"{repo_name}/{subfolder.split('/')[-1]}"
+        else:
+            display_name = f"{repo_name}/{os.path.splitext(filename)[0]}"
+
+        config: Dict[str, Any] = {
+            "type": "pytorch",
+            "path": "",
+            "repo_id": repo_id,
+            "filename": filename,
+            "arch": "lsnet_xl_artist",
+        }
+
+        if subfolder:
+            config["subfolder"] = subfolder
+
+        discovered.append((display_name, config, key))
+
+    # Sort alphabetically by display name
+    discovered.sort(key=lambda item: item[0])
+
+    for display_name, config, key in discovered:
+        candidate_name = display_name
+        suffix = 2
+        while candidate_name in models:
+            candidate_name = f"{display_name} #{suffix}"
+            suffix += 1
+        models[candidate_name] = config
+        existing_entries.add(key)
+        print(f"Auto-registered: {candidate_name}")
+
+
+def validate_models(models: Dict[str, Dict[str, Any]]):
+    """
+    Validate models at startup and remove invalid entries.
+    Returns list of valid model names.
+    """
+    valid_models = {}
+
+    for model_name, config in models.items():
+        model_path = config.get("path") or ""
+
+        # If local path exists, it's valid
+        if model_path and os.path.exists(model_path):
+            valid_models[model_name] = config
+            continue
+
+        # Check if it can be downloaded from HF
+        if "repo_id" in config and "filename" in config:
+            try:
+                from huggingface_hub import file_exists
+
+                # Build kwargs for checking file existence
+                check_kwargs = {
+                    "repo_id": config["repo_id"],
+                    "filename": config["filename"],
+                    "repo_type": "model",
+                }
+                if config.get("subfolder"):
+                    check_kwargs["filename"] = f"{config['subfolder']}/{config['filename']}"
+                if config.get("revision"):
+                    check_kwargs["revision"] = config["revision"]
+
+                # Check if file exists on HF
+                if file_exists(**check_kwargs):
+                    valid_models[model_name] = config
+                else:
+                    print(f"Skipping {model_name}: file not found on Hugging Face")
+            except Exception as e:
+                print(f"Skipping {model_name}: validation error - {e}")
+        else:
+            print(f"Skipping {model_name}: no valid path or repo configuration")
+
+    return valid_models