A full fp16 cast on the audio_tokenizer saved ~3 GB VRAM but produced
silent audio: fp16 lacks the dynamic range LayerNorm/RMSNorm need to
keep their normalization statistics meaningful, and the decoder
collapsed to near-zero output.
Now: cast all weights to the resolved dtype, then explicitly cast
LayerNorm/GroupNorm/RMSNorm back to fp32. A forward pre-hook on every
leaf-ish module aligns incoming float tensors with each module's own
param dtype, so:
- fp32 wav input -> fp16 Linear: hook casts to fp16
- fp16 hidden state -> fp32 LayerNorm: hook casts to fp32
- fp32 norm output -> fp16 next layer: that layer's hook casts to fp16
Tested with a tiny mixed Linear/Conv1d/LayerNorm stand-in: fp32 inputs
flow through fp16 Linear, fp16 hidden states reach fp32 norm cleanly.
End-to-end audio quality has to be verified on the real model.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Adds bitsandbytes 8-bit quantization on the main LM as an alternative
memory-saving path. The audio_tokenizer (where fp16 broke output
quality) stays fp32 and is unaffected.
Wiring:
- New "int8" entry in MOSSTTSLoader's dtype dropdown
- get_or_load branches on dtype_str == "int8" — passes BitsAndBytesConfig
+ device_map={"": resolved_device} to from_pretrained, skipping the
torch_dtype + .to(device) path (incompatible with bnb-loaded models)
- bitsandbytes is lazy-imported with a clear install message; only int8
users pay the dependency
- Validates that int8 is requested only on cuda (no cpu kernels)
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The dtype-mismatch error wording differs between CPU ("mat1 and mat2
must have the same dtype") and cuda ("expected scalar type Float but
found Half"). The pre-cast pytest.raises assertion was just a setup
sanity check and not necessary for the test's intent — verifying that
the hook makes the fp16 layer accept fp32 input.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
After casting the audio_tokenizer's weights to fp16/bf16 (saves ~3 GB
VRAM), incoming fp32 wav tensors hit Linear/Conv layers with fp16
weights and crash with "expected scalar type Float but found Half" deep
inside encode_audios_from_wav.
Register a forward-pre-hook on every leaf-ish module of the
audio_tokenizer that casts floating-point input tensors to the module's
parameter dtype on the fly. Covers both the encode (voice reference)
and decode paths without needing to enumerate the model's submodule
names.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
AutoProcessor.from_pretrained loads the audio_tokenizer model in fp32 by
default, ignoring the torch_dtype passed to the main model. .to(device)
without an explicit dtype only relocates, doesn't cast — leaving the
audio_tokenizer at fp32 on the GPU and roughly doubling its VRAM
footprint.
On a 20 GB card this is the difference between a comfortable headroom
and OOM during generation: observed peak reserved was 19.7 GB / 20.5 GB
total, despite expandable_segments already being on. Casting to the
model's resolved dtype on cuda recovers ~3 GB.
CPU path keeps default precision; fp16 ops are slow on most CPUs and
the VRAM concern doesn't apply off-GPU.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
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.
Previously, when model.generate() OOMed we re-raised the helpful error
without freeing anything. The post-success cleanup path (del intermediate
tensors, optional cleanup_all if keep_loaded=False) never ran. Result: the
13 GB model handle stayed resident on the GPU after the failed run,
guaranteeing the next workflow run also OOMs.
Now the OOM handler:
- Drops local refs to the input tensors so they are gc-eligible.
- Calls cleanup_all() if keep_loaded is False, mirroring the success path.
- Forces empty_cache + gc.collect before re-raising.
Tests still pass against the helper functions; this only changes the
error path in the node wrapper.
ComfyUI caches each node's output keyed by its inputs. With keep_loaded=False,
the previous run's Generate node freed the model and processor from the
cache (model=None, processor=None on the still-cached MOSS_TTS_MODEL handle).
Without IS_CHANGED, ComfyUI hands that empty handle to the next run's
Generate, which crashes with 'NoneType has no attribute build_user_message'.
Add IS_CHANGED that returns a fresh value (time.time()) when keep_loaded
is False, forcing ComfyUI to re-execute load() each run. get_or_load then
sees an empty cache and reloads the model from the local HuggingFace cache.
For keep_loaded=True the return value is stable, so ComfyUI keeps reusing
the same loaded model — preserves the original keep-resident behaviour.
After end-to-end testing it became clear MOSS-TTS' remote code is targeted
at a narrow transformers band: <4.50 misses MODALITY_TO_BASE_CLASS_MAPPING
and PreTrainedConfig, >=4.58 starts removing _get_initial_cache_position
and reworking ProcessorMixin internals. 4.50.x is the sweet spot.
- requirements.txt pins transformers >=4.50, <4.58.
- README explains the version constraint and the diagnostic for
noise-vs-speech output.
- README also notes that the model is Chinese-centric and English
output without a reference clip will sound Chinese-accented; this is
upstream model behaviour, not a node bug.
When generation fails (e.g. OOM), ComfyUI's execution.format_value calls
repr() on the workflow input dict. Default dict.__repr__ recursively reprs
the contained Processor, whose __repr__ calls to_dict() which deepcopies
every nested submodule — including the audio_tokenizer's full GPU weights.
That deepcopy triggers a second OOM during error formatting, swamps the
trace with copy.deepcopy frames, and masks the original error.
Wrap the model handle in a dict subclass with a tiny __repr__ that just
shows device / dtype / attn / model_id and never touches the heavy
contents. All existing entry['key'] / entry.get(...) access patterns are
preserved.
transformers 5.x removed GenerationMixin._get_initial_cache_position in
favour of new cache helpers. MOSS-TTS' custom _sample method still calls
self._get_initial_cache_position(cur_len, device, model_kwargs) so we have
to restore it. The reimplementation matches the old behaviour: populate
model_kwargs['cache_position'] with arange(past_length, cur_len).
The user upgraded transformers and the API got reworked — ProcessorMixin
no longer exposes 'attributes' / 'optional_attributes' as data classvars
(replaced by methods like get_attributes()). Our shim was unconditionally
patching ProcessorMixin.__init__ and crashing on AttributeError when its
inner code referenced cls.optional_attributes.
Skip the entire patch on transformers 5.x where the old classvars are
absent — the new release auto-derives attributes natively, no shim needed.
Same gate-by-feature-detection logic as the other shims in this function.
MossTTSDelayConfig only propagates hidden_size and vocab_size from its inner
language_config to the top-level config. Many decoder attrs that transformers'
generation/cache machinery requires (num_hidden_layers, num_attention_heads,
head_dim, rope_*, etc.) only exist on language_config and aren't surfaced.
Install a __getattr__ fallback on the config class so any otherwise-missing
attribute is fetched from language_config. Combined with the existing
token_id backfill (which has to overwrite explicit Nones, not missing
attrs), this lets transformers' internals find what they need without us
hand-listing every attribute they might touch.
MossTTSDelayConfig.__init__ sets self.pad_token_id = 151643 (default), then
calls super().__init__(**kwargs). PretrainedConfig's initializer overwrites
self.pad_token_id back to None when kwargs has no pad_token_id — which is
exactly the case here, since the repo's saved config.json only stores those
IDs under language_config, not top-level. The end result is
processor.model_config.pad_token_id == None, and MOSS-TTS' _pad crashes
with 'can't assign a NoneType to a torch.LongTensor' on the first
processor() call.
After loading the processor, copy pad/eos/bos token_id up from
language_config to the top-level model_config when missing.
Previous auto-derive logic counted any cls attribute ending in '_class' that
was non-None. ProcessorMixin has methods like check_argument_for_proper_class
and register_for_auto_class which match that suffix, so the derived attributes
list ended up as ['check_argument_for_proper', 'register_for_auto', 'tokenizer'].
Filter to data attributes whose value is a string — that's how transformers
declares processor classes (tokenizer_class = 'AutoTokenizer'). Methods'
function objects are not strings and are now excluded.
Newer transformers releases compute ProcessorMixin.attributes automatically
from each subclass' *_class members. 4.57.6 still hardcodes
['feature_extractor', 'tokenizer'] as the default, and MOSS-TTS' processor
inherits that default while only declaring tokenizer_class and
audio_tokenizer_class. ProcessorMixin.__init__ then asserts both attributes
were passed and crashes with 'This processor requires 2 arguments:
feature_extractor, tokenizer'.
Monkey-patch ProcessorMixin.__init__ to derive cls.attributes from *_class
class members for any subclass that didn't override attributes itself.
optional_attributes are excluded from the derived list to mirror the
new transformers behavior.
MOSS-TTS' processor declares audio_tokenizer_class = 'AutoModel'. transformers'
ProcessorMixin then runs check_argument_for_proper_class which translates
that string via AUTO_TO_BASE_CLASS_MAPPING. The mapping in 4.57.6 only knows
AutoTokenizer / AutoFeatureExtractor / AutoImageProcessor / AutoVideoProcessor;
AutoModel isn't there, so the lookup returns 'AutoModel' literally and the
isinstance check tries MossAudioTokenizerModel against the AutoModel factory
class (not its instances). That always fails.
Registering AutoModel -> PreTrainedModel makes the check resolve to the base
class that MossAudioTokenizerModel actually inherits from, so the processor
init succeeds.
MOSS-TTS' audio tokenizer config does:
from transformers.configuration_utils import PreTrainedConfig
In newer transformers (>= ~4.58) the class was renamed from PretrainedConfig
(lowercase t) to PreTrainedConfig (capital T). On 4.57.6 only the old name
exists, so the import fails. Aliasing the new name to the old one in the
compat shim lets MOSS-TTS' remote code import it regardless of version.
MOSS-TTS' remote processor code does:
processing_utils.MODALITY_TO_BASE_CLASS_MAPPING['audio_tokenizer'] = 'PreTrainedModel'
That attribute only exists in transformers releases newer than what most
ComfyUI envs ship (e.g. transformers 4.57.6 doesn't have it yet). Without it,
AutoProcessor.from_pretrained crashes with AttributeError before the model
can load. Add a compat helper that creates the dict if missing, so MOSS-TTS'
own line works regardless of the installed transformers version.
Setting --rootdir=tests in addopts keeps pytest's path resolution inside the
tests directory. Combined with --import-mode=importlib (which loads test files
via importlib.util without sys.path package walking), this fully avoids the
import of the package-level __init__.py during test discovery — that file uses
ComfyUI-style relative imports that only resolve under ComfyUI's loader.
--confcutdir=tests stops conftest.py search at the tests dir for the same
reason. The root conftest.py is left in place but is no longer needed for
sys.path setup; importlib mode handles that automatically.
collect_ignore in conftest.py was being evaluated after pytest's package-detection
already tried to import the root __init__.py. The hard fix:
1. Re-add --ignore=__init__.py (and the non-test top-level dirs) to addopts so
pytest skips them at the path-resolution layer.
2. Add a pytest_ignore_collect hook in the root conftest.py as a backup, so any
internal pytest pathway that still reaches __init__.py at rootdir is rejected.
Together they prevent pytest from ever exec'ing the ComfyUI-only relative
imports in __init__.py during test collection.
Two issues from final review verification:
1. Pytest's default importer walked into the root __init__.py (which uses
ComfyUI-relative imports) when collecting tests. Adding a root conftest.py
with collect_ignore for non-test top-level dirs and switching to
--import-mode=importlib stops it. The conftest also seeds sys.path so
tests can still 'from lib.x import' / 'from nodes.x import'.
2. lib/model_state.py used to do 'from transformers import AutoModel,
AutoProcessor' at module-level. transformers' lazy chain transitively
imports torchvision, which breaks at attribute access in environments
where torch and torchvision versions don't match (common in ComfyUI
installs). Lifting the import into a _get_auto_classes() factory keeps
the package importable; tests patch the factory instead of patching
transformers attributes.
Pytest's default 'prepend' import mode walks up from each test file looking
for parent packages. With a root __init__.py (required by ComfyUI for relative
imports) it tries to import the root package — which fails outside the ComfyUI
loader because of 'from .nodes import' having no parent.
Switching to --import-mode=importlib makes pytest load each test module via
importlib.util.spec_from_file_location without sys.path/package gymnastics.
Also explicitly ignore the non-test top-level dirs so collection never wanders
in there.
Tests still find lib/ and nodes/ via pythonpath = ['.'].
With both tests/__init__.py and a root __init__.py present, pytest treats
tests as a child of an implicit root package and tries to import the root
__init__.py — which uses 'from .nodes import' (relative). That fails outside
the ComfyUI loader because pytest invokes it with no parent package.
Removing tests/__init__.py switches pytest to rootdir-relative discovery,
which leaves the root __init__.py alone.
- 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.
