feat: add OmniVoiceModelLoader node

Implements OmniVoiceModelLoader with INPUT_TYPES, RETURN_TYPES, and load_model supporting both HuggingFace auto-download and local path sources. Adds TDD test suite and pytest infrastructure (conftest.py, pytest.ini) to enable testing outside ComfyUI without omnivoice installed. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-05 08:52:26 +02:00
parent 0ed43a83ca
commit 069169485d
6 changed files with 144 additions and 10 deletions
@@ -1,2 +1,63 @@
+import os
+import torch
+
+try:
+    from omnivoice import OmniVoice
+except ImportError:
+    OmniVoice = None  # type: ignore[assignment,misc]
+
+try:
+    import folder_paths
+    CACHE_DIR = os.path.join(folder_paths.models_dir, "omnivoice")
+except ImportError:
+    CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "omnivoice")
+
+
+DTYPE_MAP = {
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "float32": torch.float32,
+}
+
+
 class OmniVoiceModelLoader:
-    pass
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_source": (
+                    ["Auto-download (HuggingFace)", "Local path"],
+                    {"default": "Auto-download (HuggingFace)"},
+                ),
+                "device": (
+                    ["cuda:0", "cuda:1", "cpu"],
+                    {"default": "cuda:0"},
+                ),
+                "dtype": (
+                    ["float16", "bfloat16", "float32"],
+                    {"default": "float16"},
+                ),
+            },
+            "optional": {
+                "local_path": ("STRING", {"default": ""}),
+            },
+        }
+
+    RETURN_TYPES = ("OMNIVOICE_MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_model"
+    CATEGORY = "OmniVoice"
+
+    def load_model(self, model_source, device, dtype, local_path=""):
+        if model_source == "Local path" and local_path:
+            source = local_path
+        else:
+            source = "k2-fsa/OmniVoice"
+
+        model = OmniVoice.from_pretrained(
+            source,
+            device_map=device,
+            dtype=DTYPE_MAP[dtype],
+            cache_dir=CACHE_DIR,
+        )
+        return (model,)