diff --git a/__init__.py b/__init__.py index 8a2ff3d..4272a71 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,6 @@ """ComfyUI-MOSS-TTS — Custom nodes for MOSS-TTS 1.7B zero-shot voice cloning.""" -from nodes import ( +from .nodes import ( MOSSTTSLoader, MOSSTTSVoiceReference, MOSSTTSGenerate, diff --git a/lib/audio.py b/lib/audio.py index 2fe8db7..68a1e27 100644 --- a/lib/audio.py +++ b/lib/audio.py @@ -26,9 +26,13 @@ def audio_dict_to_wav(audio_dict: dict, dest_dir: str) -> str: def tts_output_to_audio_dict(audio: torch.Tensor, sample_rate: int) -> dict: - """Pack a 1D MOSS-TTS audio tensor into ComfyUI AUDIO format [1, 1, T].""" + """Pack a 1D MOSS-TTS audio tensor into ComfyUI AUDIO format [1, 1, T]. + + Always returns the tensor on CPU so downstream nodes don't need to handle GPU tensors. + """ if audio.dim() != 1: raise ValueError(f"Expected 1D audio tensor, got shape {tuple(audio.shape)}") + audio = audio.detach().cpu() return { "waveform": audio.unsqueeze(0).unsqueeze(0), "sample_rate": int(sample_rate), diff --git a/lib/model_state.py b/lib/model_state.py index 68407aa..3874ed6 100644 --- a/lib/model_state.py +++ b/lib/model_state.py @@ -41,6 +41,13 @@ def resolve_dtype(requested: str, device: str) -> torch.dtype: }[requested] +def _resolved_dtype_str(requested: str, device: str) -> str: + """Return the canonical dtype string ('fp16'/'bf16'/'fp32') for cache keying.""" + if requested == "auto": + return "fp16" if device == "cuda" else "fp32" + return requested + + def resolve_attn_impl(requested: str, device: str) -> str: if requested == "auto": return "sdpa" if device == "cuda" else "eager" @@ -74,7 +81,7 @@ def get_or_load( resolved_dtype = resolve_dtype(dtype_str, resolved_device) resolved_attn = resolve_attn_impl(attn_impl, resolved_device) - key = (model_id, resolved_device, dtype_str, resolved_attn) + key = (model_id, resolved_device, _resolved_dtype_str(dtype_str, resolved_device), resolved_attn) if key in _MODEL_CACHE: entry = _MODEL_CACHE[key] entry["keep_loaded"] = keep_loaded @@ -135,3 +142,16 @@ def cleanup_all() -> str: free, total = torch.cuda.mem_get_info() return f"Cleared {n} cached model(s). Free VRAM: {free / 1024**3:.2f} GB / {total / 1024**3:.2f} GB" return f"Cleared {n} cached model(s)." + + +def cleanup_one(key: tuple) -> bool: + """Remove a single cached entry by key. Returns True if found and dropped.""" + if key not in _MODEL_CACHE: + return False + entry = _MODEL_CACHE.pop(key) + entry["model"] = None + entry["processor"] = None + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return True diff --git a/lib/voice.py b/lib/voice.py index 163fac4..13f61b6 100644 --- a/lib/voice.py +++ b/lib/voice.py @@ -4,7 +4,7 @@ import logging import os from typing import Optional -from lib.audio import audio_dict_to_wav +from .audio import audio_dict_to_wav log = logging.getLogger(__name__) diff --git a/nodes/__init__.py b/nodes/__init__.py index 18c9951..ade5f45 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -1,7 +1,7 @@ -from nodes.loader import MOSSTTSLoader -from nodes.voice_reference import MOSSTTSVoiceReference -from nodes.generate import MOSSTTSGenerate -from nodes.unload import MOSSTTSUnload +from .loader import MOSSTTSLoader +from .voice_reference import MOSSTTSVoiceReference +from .generate import MOSSTTSGenerate +from .unload import MOSSTTSUnload __all__ = [ "MOSSTTSLoader", diff --git a/nodes/generate.py b/nodes/generate.py index 3660335..42f3e4b 100644 --- a/nodes/generate.py +++ b/nodes/generate.py @@ -6,11 +6,11 @@ import time import torch import torchaudio -from lib.audio import tts_output_to_audio_dict -from lib.delay_config import DelayGenerationConfig -from lib.generation import build_layers_config, format_status -from lib.model_state import cleanup_all -from lib.presets import PRESET_NAMES, resolve_preset +from ..lib.audio import tts_output_to_audio_dict +from ..lib.delay_config import DelayGenerationConfig +from ..lib.generation import build_layers_config, format_status +from ..lib.model_state import cleanup_all +from ..lib.presets import PRESET_NAMES, resolve_preset _DEFAULT_VOICE = {"path": None, "source": "none"} diff --git a/nodes/loader.py b/nodes/loader.py index 9357b2b..42099c4 100644 --- a/nodes/loader.py +++ b/nodes/loader.py @@ -1,6 +1,6 @@ """ComfyUI node: MOSSTTSLoader.""" -from lib.model_state import get_or_load +from ..lib.model_state import get_or_load class MOSSTTSLoader: diff --git a/nodes/unload.py b/nodes/unload.py index 52213c6..8c65989 100644 --- a/nodes/unload.py +++ b/nodes/unload.py @@ -1,6 +1,6 @@ """ComfyUI node: MOSSTTSUnload.""" -from lib.model_state import cleanup_all +from ..lib.model_state import cleanup_all class MOSSTTSUnload: diff --git a/nodes/voice_reference.py b/nodes/voice_reference.py index f32d20a..ff4ab9c 100644 --- a/nodes/voice_reference.py +++ b/nodes/voice_reference.py @@ -2,7 +2,7 @@ import os -from lib.voice import build_voice_reference +from ..lib.voice import build_voice_reference def _comfy_temp_dir() -> str: diff --git a/pyproject.toml b/pyproject.toml index e895da5..f60b6c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ dev = ["pytest>=7", "pytest-mock"] [project.urls] -Homepage = "https://github.com/OpenMOSS/MOSS-TTS" +Homepage = "https://github.com/Eutectico/ComfyUI-MOSS-TTS" [tool.comfy] PublisherId = "" diff --git a/tests/test_model_state.py b/tests/test_model_state.py index 81065f2..4c4c05f 100644 --- a/tests/test_model_state.py +++ b/tests/test_model_state.py @@ -92,3 +92,16 @@ def test_resolve_attn_impl_flash_falls_back_when_unavailable(monkeypatch): monkeypatch.setattr(builtins, "__import__", fake_import) assert model_state.resolve_attn_impl("flash_attention_2", "cuda") == "sdpa" + + +def test_auto_dtype_collapses_with_explicit_on_same_device(): + """auto and the matching explicit dtype should hit the same cache entry.""" + auto_model_p, auto_proc_p = _patches() + with auto_model_p as MockModel, auto_proc_p as MockProc, \ + patch("lib.model_state.torch.cuda.is_available", return_value=False): + MockModel.from_pretrained.return_value = MagicMock() + MockProc.from_pretrained.return_value = MagicMock() + e1 = model_state.get_or_load("x", "cpu", "auto", "eager") + e2 = model_state.get_or_load("x", "cpu", "fp32", "eager") + assert e1 is e2 + assert MockModel.from_pretrained.call_count == 1