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>
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.
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.
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.
- 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.
