From f6e5d420a182aa9eb8bbfacb83bd89923933a8a6 Mon Sep 17 00:00:00 2001
From: Federico Volovnik-Jose <fedejose@gmail.com>
Date: Mon, 27 Apr 2026 15:04:44 +0200
Subject: [PATCH] feat(loader): expose audio_tokenizer_device option for
 memory-tight GPUs

Add a new audio_tokenizer_device input on MOSSTTSLoader (auto/cuda/cpu).
auto mirrors the main model's device (previous behaviour). cpu keeps the
audio tokenizer on host RAM, saving ~1-3 GB of VRAM at the cost of some
CPU latency at the start (encoding the reference clip) and end (decoding
the generated codes) of each generation.

Useful on 16-20 GB GPUs where MOSS-TTS' main model + activations push
right against the ceiling. Including the audio_tokenizer_device in the
cache key prevents stale entries when a user toggles the setting.
---
 lib/model_state.py | 23 ++++++++++++++++++++---
 nodes/loader.py    | 11 ++++++++---
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/lib/model_state.py b/lib/model_state.py
index 0ec4701..715f83b 100644
--- a/lib/model_state.py
+++ b/lib/model_state.py
@@ -267,14 +267,31 @@ def get_or_load(
     dtype_str: str,
     attn_impl: str,
     keep_loaded: bool = True,
+    audio_tokenizer_device: str = "auto",
 ) -> dict:
-    """Load or fetch the cached MOSS-TTS model + processor."""
+    """Load or fetch the cached MOSS-TTS model + processor.
+
+    audio_tokenizer_device: where to place the audio_tokenizer (a separate
+        small model used to encode reference audio and decode generated codes).
+        "auto" mirrors the main model's device. "cpu" keeps it off the GPU,
+        saving ~1-3 GB of VRAM at the cost of some latency at the start and
+        end of each generation. Useful on memory-tight GPUs.
+    """
     _ensure_atexit()
     resolved_device = resolve_device(device)
     resolved_dtype = resolve_dtype(dtype_str, resolved_device)
     resolved_attn = resolve_attn_impl(attn_impl, resolved_device)
+    resolved_at_device = (
+        resolved_device if audio_tokenizer_device == "auto" else audio_tokenizer_device
+    )
 
-    key = (model_id, resolved_device, _resolved_dtype_str(dtype_str, resolved_device), resolved_attn)
+    key = (
+        model_id,
+        resolved_device,
+        _resolved_dtype_str(dtype_str, resolved_device),
+        resolved_attn,
+        resolved_at_device,
+    )
     if key in _MODEL_CACHE:
         entry = _MODEL_CACHE[key]
         entry["keep_loaded"] = keep_loaded
@@ -288,7 +305,7 @@ def get_or_load(
     processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
     _fix_moss_model_config_token_ids(processor)
     if hasattr(processor, "audio_tokenizer") and processor.audio_tokenizer is not None:
-        processor.audio_tokenizer = processor.audio_tokenizer.to(resolved_device)
+        processor.audio_tokenizer = processor.audio_tokenizer.to(resolved_at_device)
 
     if resolved_device == "cuda":
         torch.cuda.empty_cache()
diff --git a/nodes/loader.py b/nodes/loader.py
index 1d20af0..c7375bf 100644
--- a/nodes/loader.py
+++ b/nodes/loader.py
@@ -21,6 +21,10 @@ class MOSSTTSLoader:
                     {"default": "sdpa"},
                 ),
                 "keep_loaded": ("BOOLEAN", {"default": True}),
+                "audio_tokenizer_device": (
+                    ["auto", "cuda", "cpu"],
+                    {"default": "auto"},
+                ),
             }
         }
 
@@ -30,21 +34,22 @@ class MOSSTTSLoader:
     CATEGORY = "MOSS-TTS/loaders"
 
     @classmethod
-    def IS_CHANGED(cls, model_id, device, dtype, attn_implementation, keep_loaded):
+    def IS_CHANGED(cls, model_id, device, dtype, attn_implementation, keep_loaded, audio_tokenizer_device):
         # When keep_loaded is False, the previous run's Generate node freed the
         # model from the cache. ComfyUI's per-node output cache would otherwise
         # hand the now-empty handle to Generate again. Returning a fresh value
         # each call forces ComfyUI to re-execute load() and reload the model.
         if not keep_loaded:
             return time.time()
-        return f"{model_id}|{device}|{dtype}|{attn_implementation}|keep"
+        return f"{model_id}|{device}|{dtype}|{attn_implementation}|{audio_tokenizer_device}|keep"
 
-    def load(self, model_id, device, dtype, attn_implementation, keep_loaded):
+    def load(self, model_id, device, dtype, attn_implementation, keep_loaded, audio_tokenizer_device):
         entry = get_or_load(
             model_id=model_id,
             device=device,
             dtype_str=dtype,
             attn_impl=attn_implementation,
             keep_loaded=keep_loaded,
+            audio_tokenizer_device=audio_tokenizer_device,
         )
         return (entry,)