fix: convert to relative imports + cache key/cpu/cleanup_one fixes

- Use relative imports across nodes/ and lib/ so ComfyUI's loader resolves the
  local subpackages instead of colliding with ComfyUI's own top-level nodes module.
- tts_output_to_audio_dict now returns CPU tensors so downstream nodes don't see CUDA tensors.
- model_state cache key uses resolved dtype string, so dtype=auto and the matching
  explicit choice share one cache slot (no spurious 13 GB reload).
- Add cleanup_one(key) helper as called for in spec section 4.
- Fix pyproject Homepage URL to point at this repo.

__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,

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),

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

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__)
 

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",

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"}

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:

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:

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:

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 = ""

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