Ethanfel
11beba1c47
Remove OmniVoice = None fallback in nodes/loader.py so missing omnivoice gives a clear ImportError instead of a confusing AttributeError. Restore __init__.py to clean form without the try/except that silently swallowed real import errors. Add omnivoice mock to conftest.py and register a pytest plugin that prevents pytest from treating the project root as a Package node (which would try to import __init__.py outside a package context and fail on the relative import). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
import os
|
|
import torch
|
|
|
|
from omnivoice import OmniVoice
|
|
|
|
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:
|
|
@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,)
|