Investigate Higgs Audio transformers 5 compatibility

docs/Dev reports/TRANSFORMERS_5_HIGGS_AUDIO_REPORT.md

@@ -0,0 +1,123 @@
+TTS Audio Suite - Transformers 5.x Higgs Audio 2 Investigation Report
+
+Scope
+- Problem: Higgs Audio 2 regressed after the suite moved the main environment to Transformers 5.x.
+- Goal: determine whether Higgs Audio 2 is a good native Transformers 5 patch target, or whether it should move to runtime isolation.
+
+Environment
+- Windows / ComfyUI runtime.
+- Main env on Transformers 5.x.
+- Direct repro and validation performed in the real Windows Python environment outside the ComfyUI UI wrapper.
+
+Findings Summary
+- Higgs Audio 2 is not blocked at the shallow import/API level only.
+- Some real compatibility fixes were possible:
+  - bundled cache API adaptation
+  - cache-arg rename fixes
+  - major cold-load regression fix
+- However, the remaining failure is decode-path semantics, especially around StaticCache and CUDA graph replay.
+- Higgs can be made to load and "run" on Transformers 5, but the high-performance path is not trustworthy and produced incorrect audio.
+- This is not a good short-path patch target like IndexTTS.
+
+Reproduction
+1. Transformers 5 main environment.
+2. Higgs Audio 2 loaded through the normal TTS Text path with voice cloning.
+3. Observed one or more of:
+   - multi-minute cold-load stalls
+   - StaticCache API crashes
+   - CUDA graph capture/runtime failures
+   - runaway/gibberish generation
+   - near-empty or silent output
+
+Key Diagnostic Evidence
+- Initial cold-load timing showed:
+  - `⏱️ Higgs Audio: Core model load took 440.61s`
+- Tokenizer/HuBERT stages were not the bottleneck:
+  - text tokenizer around 1-2s
+  - audio tokenizer around 2-5s
+- The 440s stall was inside `HiggsAudioModel.from_pretrained(...)`, specifically Hugging Face missing-key initialization.
+- Direct checkpoint/model key comparison in runtime showed:
+  - `missing_count 0`
+  - `unexpected_count 0`
+- After bypassing unnecessary missing-key initialization for complete local checkpoints, direct timing dropped to:
+  - `⏱️ Higgs Audio: Core model load took 58.46s`
+  - later direct runs around `33-58s` depending on cache state
+
+- StaticCache/API drift evidence:
+  - `'StaticCache' object is not subscriptable`
+  - later mask/cache shape mismatches around static cache width `1024`
+
+- Cache-arg drift evidence:
+  - bundled Higgs/LLaMA path mixed `past_key_value` and `past_key_values`
+  - led to duplicated kwarg failures and broken decode semantics
+
+- CUDA graph replay evidence:
+  - direct runtime probe with graph replay active produced:
+    - immediate or near-immediate audio termination
+    - tiny garbage clips
+    - or invalid decode behavior
+  - direct runtime probe with StaticCache but without CUDA graph replay produced:
+    - normal-length audio token generation
+    - reasonable token spread
+    - real waveform output
+
+Attempts and Results (chronological)
+1) StaticCache compatibility shims
+   - Added bundled helpers to access cache tensors/device across old tuple caches and new Cache/StaticCache APIs.
+   - Result: fixed immediate StaticCache indexing crashes.
+
+2) CUDA graph failure fallback cleanup
+   - On graph capture failure, forced cache recreation and switched to DynamicCache instead of continuing with stale StaticCache assumptions.
+   - Result: removed some follow-on crashes after graph failure, but not the core correctness issue.
+
+3) Attention backend adjustment
+   - Stopped forcing eager attention on CUDA and preferred SDPA instead.
+   - Result: better alignment with the intended high-performance path, but did not fully restore correctness.
+
+4) TTS defaults forcing audio generation
+   - Set Higgs TTS defaults to force audio generation rather than letting the model drift in text mode.
+   - Result: prevented obvious free-running text-style generation, but did not solve bad output.
+
+5) Cold-load regression fix
+   - Added a local complete-checkpoint detector and skipped Hugging Face missing-key initialization when the local Higgs checkpoint was complete.
+   - Result: large real improvement to cold-load time.
+
+6) Cache-argument adaptation
+   - Patched bundled Higgs/LLaMA call sites for Transformers 5 cache-arg naming.
+   - Result: moved runtime further, but surfaced deeper semantic issues instead of fixing them outright.
+
+7) CUDA graph runner PyTorch compatibility
+   - Wrapped input-buffer updates in inference mode to satisfy current PyTorch restrictions.
+   - Result: removed one explicit runtime crash, but not the underlying decode corruption.
+
+8) StaticCache without CUDA graph replay direct comparison
+   - Direct Windows probes showed:
+     - `enable_cuda_graphs=True` with graph replay: corrupted/too-short audio output
+     - `enable_cuda_graphs=False`: long audio-token generation and real waveform output
+     - `StaticCache` with graph replay skipped: closer to sane behavior than replay path
+   - Result: isolated the remaining failure to the CUDA graph replay path on Transformers 5.
+
+Interpretation
+- Higgs Audio 2 on Transformers 5 is not blocked by a single missing symbol or a few renamed APIs.
+- The remaining problem is decode/runtime behavior, especially the interaction between:
+  - StaticCache
+  - bundled Higgs decode logic
+  - LLaMA attention/cache semantics on current Transformers 5
+  - CUDA graph replay
+- That is behavior-level drift, not just compatibility-shim work.
+
+Decision
+- Higgs Audio 2 on Transformers 5 is still theoretically patchable.
+- It is not a good use of time right now.
+- The engine reached the same bad zone as Qwen3-TTS, but with even more sensitivity around CUDA graphs and cache behavior.
+- Runtime isolation is the correct path for now.
+
+Operational Recommendation
+- Do not continue native Transformers 5 Higgs patching in the main branch.
+- Route Higgs Audio 2 into an isolated legacy runtime instead.
+- Prefer a dedicated Higgs runtime first rather than assuming shared legacy parity.
+- Revisit native Transformers 5 support only if Higgs becomes important enough to justify deeper decode/cuda-graph compatibility work.
+
+Notes
+- The cold-load fix was a real improvement and may still be worth revisiting later, even outside the final T5 decision.
+- The rest of this investigation should be treated as a branch-local investigation record, not as a stable main-branch solution.

engines/adapters/higgs_audio_adapter.py

@@ -138,7 +138,7 @@ class HiggsAudioEngineAdapter:
         top_p = params.get("top_p", 0.95)
         top_k = params.get("top_k", 50)
         max_new_tokens = params.get("max_new_tokens", 2048)
-        force_audio_gen = params.get("force_audio_gen", False)
+        force_audio_gen = params.get("force_audio_gen", True)
         ras_win_len = params.get("ras_win_len", 7)  # Default from boson_multimodal
         ras_max_num_repeat = params.get("ras_max_num_repeat", 2)  # Default from boson_multimodal
         enable_cuda_graphs = params.get("enable_cuda_graphs", True)  # CUDA Graph toggle
@@ -394,7 +394,7 @@ class HiggsAudioEngineAdapter:
         validated["max_new_tokens"] = max(1, min(4096, int(max_tokens)))
         
         # Force audio gen validation
-        force_audio_gen = params.get("force_audio_gen", False)
+        force_audio_gen = params.get("force_audio_gen", True)
         validated["force_audio_gen"] = bool(force_audio_gen)
         
         # RAS parameters validation
@@ -426,4 +426,4 @@ class HiggsAudioEngineAdapter:
         # Cache settings
         validated["enable_audio_cache"] = params.get("enable_audio_cache", True)
         
-        return validated
+        return validated

engines/higgs_audio/boson_multimodal/audio_processing/higgs_audio_tokenizer.py

@@ -2,9 +2,10 @@
 # Licensed under MIT License
 # Modifications by BosonAI
 
-import math
-import os
-import torch
+import math
+import os
+import time
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from typing import Optional, Union, Sequence
@@ -88,9 +89,11 @@ class HiggsAudioTokenizer(nn.Module):
             self.semantic_dim = 768
             self.encoder_semantic_dim = 768
 
-        elif semantic_techer == "hubert_base_general":
-            # Use unified downloader to follow TTS folder > cache > download policy
-            from utils.downloads.unified_downloader import unified_downloader
+        elif semantic_techer == "hubert_base_general":
+            print("🔄 Higgs Audio tokenizer: Loading HuBERT semantic model...")
+            stage_start = time.perf_counter()
+            # Use unified downloader to follow TTS folder > cache > download policy
+            from utils.downloads.unified_downloader import unified_downloader
 
             # Define HuBERT model for unified downloader - actual files that exist
             hubert_files = [
@@ -111,11 +114,12 @@ class HiggsAudioTokenizer(nn.Module):
             if not hubert_path:
                 raise RuntimeError("Failed to download HuBERT model for Higgs Audio")
 
-            self.semantic_model = AutoModel.from_pretrained(
-                hubert_path,
-                trust_remote_code=True,
-                local_files_only=True
-            )
+            self.semantic_model = AutoModel.from_pretrained(
+                hubert_path,
+                trust_remote_code=True,
+                local_files_only=True
+            )
+            print(f"✅ Higgs Audio tokenizer: HuBERT semantic model ready ({time.perf_counter() - stage_start:.2f}s)")
             
             self.semantic_sample_rate = 16000
             self.semantic_dim = 768
@@ -344,21 +348,29 @@ class HiggsAudioTokenizer(nn.Module):
         return o.cpu().numpy()
 
 
-def load_higgs_audio_tokenizer(tokenizer_name_or_path, device="cuda"):
-    is_local = os.path.exists(tokenizer_name_or_path)
-    if not is_local:
-        tokenizer_path = snapshot_download(tokenizer_name_or_path)
-    else:
-        tokenizer_path = tokenizer_name_or_path
+def load_higgs_audio_tokenizer(tokenizer_name_or_path, device="cuda"):
+    print("🔄 Higgs Audio tokenizer: Resolving local assets...")
+    overall_start = time.perf_counter()
+    is_local = os.path.exists(tokenizer_name_or_path)
+    if not is_local:
+        tokenizer_path = snapshot_download(tokenizer_name_or_path)
+    else:
+        tokenizer_path = tokenizer_name_or_path
     config_path = os.path.join(tokenizer_path, "config.json")
     model_path = os.path.join(tokenizer_path, "model.pth")
-    config = json.load(open(config_path))
-    model = HiggsAudioTokenizer(
-        **config,
-        device=device,
-    )
-    parameter_dict = torch.load(model_path, map_location=device)
-    model.load_state_dict(parameter_dict, strict=False)
-    model.to(device)
-    model.eval()
-    return model
+    config = json.load(open(config_path))
+    print("🔄 Higgs Audio tokenizer: Building tokenizer modules...")
+    stage_start = time.perf_counter()
+    model = HiggsAudioTokenizer(
+        **config,
+        device=device,
+    )
+    print(f"⏱️ Higgs Audio tokenizer: Module construction took {time.perf_counter() - stage_start:.2f}s")
+    print("🔄 Higgs Audio tokenizer: Loading tokenizer weights...")
+    stage_start = time.perf_counter()
+    parameter_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(parameter_dict, strict=False)
+    model.to(device)
+    model.eval()
+    print(f"✅ Higgs Audio tokenizer: Ready ({time.perf_counter() - stage_start:.2f}s for weights, {time.perf_counter() - overall_start:.2f}s total)")
+    return model

engines/higgs_audio/boson_multimodal/model/higgs_audio/common.py

@@ -1,11 +1,20 @@
-from torch import nn
-
-from transformers.modeling_utils import PreTrainedModel
-
-from .configuration_higgs_audio import HiggsAudioConfig
-
-
-class HiggsAudioPreTrainedModel(PreTrainedModel):
+import glob
+import json
+import os
+from contextlib import contextmanager
+
+from torch import nn
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers import AutoConfig
+
+from .configuration_higgs_audio import HiggsAudioConfig
+
+
+_HIGGS_CHECKPOINT_COMPLETENESS_CACHE = {}
+
+
+class HiggsAudioPreTrainedModel(PreTrainedModel):
     config_class = HiggsAudioConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
@@ -14,14 +23,79 @@ class HiggsAudioPreTrainedModel(PreTrainedModel):
     _supports_flash_attn_2 = True
     _supports_sdpa = True
 
-    def _init_weights(self, module):
-        std = self.config.init_std if hasattr(self.config, "init_std") else self.config.audio_encoder_config.init_std
+    def _init_weights(self, module):
+        std = self.config.init_std if hasattr(self.config, "init_std") else self.config.audio_encoder_config.init_std
 
         if isinstance(module, (nn.Linear, nn.Conv1d)):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
                 module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+def _get_local_checkpoint_keys(checkpoint_dir: str) -> set[str]:
+    index_path = os.path.join(checkpoint_dir, "model.safetensors.index.json")
+    if os.path.exists(index_path):
+        with open(index_path, "r", encoding="utf-8") as f:
+            index_data = json.load(f)
+        return set(index_data.get("weight_map", {}).keys())
+
+    key_set: set[str] = set()
+    try:
+        from safetensors import safe_open
+    except ImportError:
+        return key_set
+
+    for shard_path in sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))):
+        with safe_open(shard_path, framework="pt", device="cpu") as shard:
+            key_set.update(shard.keys())
+    return key_set
+
+
+def _is_complete_local_higgs_checkpoint(checkpoint_dir: str) -> bool:
+    cached = _HIGGS_CHECKPOINT_COMPLETENESS_CACHE.get(checkpoint_dir)
+    if cached is not None:
+        return cached
+
+    checkpoint_keys = _get_local_checkpoint_keys(checkpoint_dir)
+    if not checkpoint_keys:
+        _HIGGS_CHECKPOINT_COMPLETENESS_CACHE[checkpoint_dir] = False
+        return False
+
+    original_init_weights = HiggsAudioPreTrainedModel._init_weights
+    try:
+        # TTS Audio Suite patch: build the key map without paying the 3B random-init cost.
+        HiggsAudioPreTrainedModel._init_weights = lambda self, module: None
+        config = AutoConfig.from_pretrained(checkpoint_dir)
+        from .modeling_higgs_audio import HiggsAudioModel
+
+        model = HiggsAudioModel(config)
+        model_keys = set(model.state_dict().keys())
+        is_complete = model_keys == checkpoint_keys
+        _HIGGS_CHECKPOINT_COMPLETENESS_CACHE[checkpoint_dir] = is_complete
+        return is_complete
+    finally:
+        HiggsAudioPreTrainedModel._init_weights = original_init_weights
+
+
+@contextmanager
+def skip_hf_missing_key_initialization_if_checkpoint_complete(checkpoint_dir: str):
+    """
+    TTS Audio Suite patch:
+    Skip HF's expensive `_initialize_missing_keys()` pass when the local Higgs checkpoint
+    already contains the full model key set. On transformers 5 this unnecessary pass can
+    spend several minutes random-initializing a 3B model despite a complete checkpoint.
+    """
+    if not _is_complete_local_higgs_checkpoint(checkpoint_dir):
+        yield
+        return
+
+    original_initialize_missing_keys = HiggsAudioPreTrainedModel._initialize_missing_keys
+    try:
+        HiggsAudioPreTrainedModel._initialize_missing_keys = lambda self, is_quantized: None
+        yield
+    finally:
+        HiggsAudioPreTrainedModel._initialize_missing_keys = original_initialize_missing_keys

engines/higgs_audio/boson_multimodal/model/higgs_audio/cuda_graph_runner.py

@@ -103,27 +103,29 @@ class CUDAGraphRunner(nn.Module):
             "all_self_attns": all_self_attns,
         }
 
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        causal_mask: torch.Tensor,
-        position_ids: torch.Tensor,
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        causal_mask: torch.Tensor,
+        position_ids: torch.Tensor,
         audio_discrete_codes_mask: torch.Tensor,
         cache_position: torch.Tensor,
         audio_attention_mask: torch.Tensor,
-        fast_forward_attention_mask: torch.Tensor,
-        **kwargs,
-    ) -> torch.Tensor:
-        # Copy input tensors to buffers
-        self.input_buffers["hidden_states"].copy_(hidden_states, non_blocking=True)
-        self.input_buffers["causal_mask"].copy_(causal_mask, non_blocking=True)
-        self.input_buffers["position_ids"].copy_(position_ids, non_blocking=True)
-        self.input_buffers["audio_discrete_codes_mask"].copy_(audio_discrete_codes_mask, non_blocking=True)
-        self.input_buffers["cache_position"].copy_(cache_position, non_blocking=True)
-        self.input_buffers["audio_attention_mask"].copy_(audio_attention_mask, non_blocking=True)
-        self.input_buffers["fast_forward_attention_mask"].copy_(fast_forward_attention_mask, non_blocking=True)
-
-        # Run the captured graph
-        self.graph.replay()
+        fast_forward_attention_mask: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        # TTS Audio Suite patch: PyTorch 2.10 rejects inplace updates to inference tensors
+        # unless the updates also happen under inference_mode.
+        with torch.inference_mode():
+            self.input_buffers["hidden_states"].copy_(hidden_states, non_blocking=True)
+            self.input_buffers["causal_mask"].copy_(causal_mask, non_blocking=True)
+            self.input_buffers["position_ids"].copy_(position_ids, non_blocking=True)
+            self.input_buffers["audio_discrete_codes_mask"].copy_(audio_discrete_codes_mask, non_blocking=True)
+            self.input_buffers["cache_position"].copy_(cache_position, non_blocking=True)
+            self.input_buffers["audio_attention_mask"].copy_(audio_attention_mask, non_blocking=True)
+            self.input_buffers["fast_forward_attention_mask"].copy_(fast_forward_attention_mask, non_blocking=True)
+
+        # Run the captured graph
+        self.graph.replay()
 
         return self.output_buffers["hidden_states"], None, None

engines/higgs_audio/boson_multimodal/model/higgs_audio/modeling_higgs_audio.py

@@ -78,7 +78,7 @@ def get_cache_seq_length(cache):
         seq_len = cache.get_seq_length()
         return seq_len.item() if hasattr(seq_len, 'item') else seq_len
 
-def get_cache_max_length(cache):
+def get_cache_max_length(cache):
     """Get cache max length with compatibility for different transformers versions"""
     # Handle DynamicCache specifically - it doesn't have a fixed max length
     from transformers.cache_utils import DynamicCache
@@ -105,9 +105,46 @@ def get_cache_max_length(cache):
         return cache.max_length
     else:
         # Debug fallback - return a reasonable default instead of 0
-        print(f"⚠️ Cache type {type(cache)} has unknown max length attributes: {[attr for attr in dir(cache) if 'max' in attr.lower() or 'len' in attr.lower()]}")
-        return 2048  # Reasonable default instead of 0
-from transformers.utils import logging, ModelOutput
+        print(f"⚠️ Cache type {type(cache)} has unknown max length attributes: {[attr for attr in dir(cache) if 'max' in attr.lower() or 'len' in attr.lower()]}")
+        return 2048  # Reasonable default instead of 0
+
+
+def get_cache_layer_pair(cache, layer_idx: int):
+    """TTS Audio Suite patch: access cache tensors across old tuple and new Cache APIs."""
+    if cache is None:
+        return None, None
+
+    key_layers = getattr(cache, "key_cache", None)
+    value_layers = getattr(cache, "value_cache", None)
+    if key_layers is not None and value_layers is not None:
+        if layer_idx < len(key_layers) and layer_idx < len(value_layers):
+            return key_layers[layer_idx], value_layers[layer_idx]
+        return None, None
+
+    if hasattr(cache, "__len__") and layer_idx < len(cache):
+        cache_tuple = cache[layer_idx]
+        if cache_tuple is not None and len(cache_tuple) >= 2:
+            return cache_tuple[0], cache_tuple[1]
+
+    return None, None
+
+
+def get_cache_device(cache):
+    """TTS Audio Suite patch: resolve cache device across old tuple and new StaticCache APIs."""
+    if cache is None:
+        return None
+
+    key_layers = getattr(cache, "key_cache", None)
+    if key_layers:
+        first_key = key_layers[0]
+        return getattr(first_key, "device", None)
+
+    key_cache, _ = get_cache_layer_pair(cache, 0)
+    if key_cache is not None:
+        return getattr(key_cache, "device", None)
+
+    return None
+from transformers.utils import logging, ModelOutput
 
 from .common import HiggsAudioPreTrainedModel
 from .utils import (
@@ -711,7 +748,7 @@ class HiggsAudioDualFFNDecoderLayer(nn.Module):
                 hidden_states=hidden_states,
                 attention_mask=audio_attention_mask,
                 position_ids=position_ids,
-                past_key_value=past_key_value,
+                past_key_values=past_key_value,  # TTS Audio Suite patch: transformers 5 LlamaAttention renamed this arg
                 output_attentions=output_attentions,
                 use_cache=use_cache,
                 cache_position=cache_position,
@@ -738,7 +775,7 @@ class HiggsAudioDualFFNDecoderLayer(nn.Module):
             hidden_states=hidden_states,
             attention_mask=attention_mask,
             position_ids=position_ids,
-            past_key_value=past_key_value,
+            past_key_values=past_key_value,  # TTS Audio Suite patch: transformers 5 LlamaAttention renamed this arg
             output_attentions=output_attentions,
             use_cache=use_cache,
             cache_position=cache_position,
@@ -1195,27 +1232,27 @@ class HiggsAudioModel(HiggsAudioPreTrainedModel, GenerationMixin):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
             if isinstance(decoder_layer, HiggsAudioDualFFNDecoderLayer):
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=causal_mask,
-                    audio_attention_mask=audio_attention_mask,
-                    fast_forward_attention_mask=fast_forward_attention_mask,
-                    position_ids=position_ids,
-                    audio_out_mask=audio_discrete_codes_mask,
-                    is_decoding_audio_token=is_decoding_audio_token,
-                    past_key_value=past_key_values,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                    cache_position=cache_position,
-                    position_embeddings=position_embeddings,
-                    is_using_cuda_graph=is_using_cuda_graph,
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    audio_attention_mask=audio_attention_mask,
+                    fast_forward_attention_mask=fast_forward_attention_mask,
+                    position_ids=position_ids,
+                    audio_out_mask=audio_discrete_codes_mask,
+                    is_decoding_audio_token=is_decoding_audio_token,
+                    past_key_value=past_key_values,  # TTS Audio Suite patch: bundled dual layer still uses singular arg name
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    position_embeddings=position_embeddings,
+                    is_using_cuda_graph=is_using_cuda_graph,
                 )
             else:
                 layer_outputs = decoder_layer(
                     hidden_states,
                     attention_mask=causal_mask,
                     position_ids=position_ids,
-                    past_key_value=past_key_values,
+                    past_key_values=past_key_values,  # TTS Audio Suite patch: transformers 5 LlamaDecoderLayer renamed this arg
                     output_attentions=output_attentions,
                     use_cache=use_cache,
                     cache_position=cache_position,
@@ -1436,16 +1473,18 @@ class HiggsAudioModel(HiggsAudioPreTrainedModel, GenerationMixin):
 
         # Use the captured cuda graph runner for decoding
         # if it exists, otherwise use the normal forward pass
-        if (
-            past_key_values is not None
-            and get_cache_max_length(past_key_values) in self.decode_graph_runners
-            and (input_ids.shape[-1] == 1)
-        ):
-            _forward_core = self.decode_graph_runners[get_cache_max_length(past_key_values)][is_decoding_audio_token]
-            is_using_cuda_graph = True
-        else:
-            _forward_core = self._forward_core
-            is_using_cuda_graph = False
+        if past_key_values is not None and (input_ids.shape[-1] == 1):
+            graph_runners = self.decode_graph_runners.get(get_cache_max_length(past_key_values), {})
+            if is_decoding_audio_token in graph_runners:
+                _forward_core = graph_runners[is_decoding_audio_token]
+                is_using_cuda_graph = True
+            else:
+                # TTS Audio Suite patch: fall back cleanly when a decode graph for this mode is unavailable.
+                _forward_core = self._forward_core
+                is_using_cuda_graph = False
+        else:
+            _forward_core = self._forward_core
+            is_using_cuda_graph = False
 
         hidden_states, all_hidden_states, all_self_attns = _forward_core(
             hidden_states=hidden_states,
@@ -1569,11 +1608,11 @@ class HiggsAudioModel(HiggsAudioPreTrainedModel, GenerationMixin):
             num_layers = len(from_cache) if hasattr(from_cache, '__len__') else 0
             to_cache.crop(0)  # Clear existing cache
             for layer_idx in range(num_layers):
-                if hasattr(from_cache, '__len__') and layer_idx < len(from_cache):
-                    key_cache, value_cache = from_cache[layer_idx]
-                    if key_cache is not None and value_cache is not None:
-                        to_cache.update(key_cache, value_cache, layer_idx)
-            return
+                if hasattr(from_cache, '__len__') and layer_idx < len(from_cache):
+                    key_cache, value_cache = get_cache_layer_pair(from_cache, layer_idx)
+                    if key_cache is not None and value_cache is not None:
+                        to_cache.update(key_cache, value_cache, layer_idx)
+            return
         
         # Handle DynamicCache to StaticCache or vice versa - skip copying for now
         if isinstance(from_cache, DynamicCache) or isinstance(to_cache, DynamicCache):
@@ -1587,16 +1626,12 @@ class HiggsAudioModel(HiggsAudioPreTrainedModel, GenerationMixin):
                 f"The target cache size {get_cache_max_length(to_cache)} is smaller than the source cache size {from_cache_size}."
             )
             # Use new cache API - access by layer index with safety checks
-            from_cache_tuple = from_cache[layer_idx] if hasattr(from_cache, '__len__') and layer_idx < len(from_cache) else (None, None)
-            to_cache_tuple = to_cache[layer_idx] if hasattr(to_cache, '__len__') and layer_idx < len(to_cache) else (None, None)
-            
-            if from_cache_tuple is not None and to_cache_tuple is not None:
-                from_key, from_value = from_cache_tuple
-                to_key, to_value = to_cache_tuple
-                
-                if from_key is not None and to_key is not None and from_value is not None and to_value is not None:
-                    to_key[:, :, :from_cache_size, :] = from_key
-                    to_value[:, :, :from_cache_size, :] = from_value
+            from_key, from_value = get_cache_layer_pair(from_cache, layer_idx)
+            to_key, to_value = get_cache_layer_pair(to_cache, layer_idx)
+
+            if from_key is not None and to_key is not None and from_value is not None and to_value is not None:
+                to_key[:, :, :from_cache_size, :] = from_key
+                to_value[:, :, :from_cache_size, :] = from_value
 
     def _prepare_kv_cache(
         self,

engines/higgs_audio/boson_multimodal/serve/serve_engine.py

@@ -2,13 +2,14 @@ import asyncio
 import base64
 import json
 import os
+import time
 import torch
 import numpy as np
 from io import BytesIO
 from dataclasses import dataclass
 from typing import List, Optional, Union
 from copy import deepcopy
-from transformers import AutoTokenizer, AutoProcessor, PreTrainedTokenizerFast
+from transformers import AutoTokenizer, AutoProcessor, PreTrainedTokenizerFast, __version__ as transformers_version
 from transformers.cache_utils import StaticCache, DynamicCache
 from transformers.generation.streamers import BaseStreamer
 from transformers.generation.stopping_criteria import StoppingCriteria
@@ -20,9 +21,11 @@ import torchaudio
 from tqdm import tqdm
 
 
-from ..dataset.chatml_dataset import ChatMLSample, ChatMLDatasetSample, prepare_chatml_sample
-from ..model.higgs_audio import HiggsAudioModel
-from ..model.higgs_audio.utils import revert_delay_pattern
+from ..dataset.chatml_dataset import ChatMLSample, ChatMLDatasetSample, prepare_chatml_sample
+from ..model.higgs_audio import HiggsAudioModel
+from ..model.higgs_audio.modeling_higgs_audio import get_cache_device  # TTS Audio Suite patch
+from ..model.higgs_audio.common import skip_hf_missing_key_initialization_if_checkpoint_complete  # TTS Audio Suite patch
+from ..model.higgs_audio.utils import revert_delay_pattern
 from ..data_collator.higgs_audio_collator import HiggsAudioSampleCollator
 from ..audio_processing.higgs_audio_tokenizer import load_higgs_audio_tokenizer
 
@@ -329,20 +332,30 @@ class HiggsAudioServeEngine:
         self.audio_tokenizer_name_or_path = audio_tokenizer_name_or_path
         self.torch_dtype = torch_dtype
         
-        # Store CUDA Graph setting IMMEDIATELY for cache creation
-        self._cuda_graphs_enabled = enable_cuda_graphs
-        self._cache_type = "StaticCache" if enable_cuda_graphs else "DynamicCache"
-        
-        print(f"🔧 HiggsAudioServeEngine: CUDA Graphs {'ENABLED' if enable_cuda_graphs else 'DISABLED'}")
-        print(f"🔧 Cache Type: {self._cache_type}")
-
-        # Initialize model and tokenizer
-        # Load with attention implementation fix for newer transformers
-        self.model = HiggsAudioModel.from_pretrained(
-            model_name_or_path,
-            dtype=torch_dtype,  # Changed from torch_dtype (deprecated in transformers 4.57.1+)
-            attn_implementation="eager"  # Force eager attention for compatibility
-        ).to(device)
+        # Store CUDA Graph setting IMMEDIATELY for cache creation
+        self._cuda_graphs_enabled = enable_cuda_graphs
+        self._cuda_graph_replay_enabled = enable_cuda_graphs
+        self._cache_type = "StaticCache" if enable_cuda_graphs else "DynamicCache"
+        
+        print(f"🔧 HiggsAudioServeEngine: CUDA Graphs {'ENABLED' if enable_cuda_graphs else 'DISABLED'}")
+        print(f"🔧 Cache Type: {self._cache_type}")
+        if enable_cuda_graphs and int(transformers_version.split(".", 1)[0]) >= 5:
+            # TTS Audio Suite patch: current Higgs CUDA-graph replay corrupts decode on transformers 5.
+            # Keep StaticCache, but skip graph replay until the replay path itself is fixed.
+            self._cuda_graph_replay_enabled = False
+            print("⚠️ Higgs Audio: CUDA graph replay disabled on Transformers 5; using StaticCache decode path")
+        self._preferred_attn_implementation = "sdpa" if device == "cuda" else "eager"
+
+        # Initialize model and tokenizer
+        # TTS Audio Suite patch: prefer SDPA on CUDA to preserve Higgs static-cache/CUDA-graph behavior.
+        stage_start = time.perf_counter()
+        with skip_hf_missing_key_initialization_if_checkpoint_complete(model_name_or_path):
+            self.model = HiggsAudioModel.from_pretrained(
+                model_name_or_path,
+                dtype=torch_dtype,  # Changed from torch_dtype (deprecated in transformers 4.57.1+)
+                attn_implementation=self._preferred_attn_implementation,
+            ).to(device)
+        print(f"⏱️ Higgs Audio: Core model load took {time.perf_counter() - stage_start:.2f}s")
         
         # Fix attention implementation for transformers compatibility
         self._fix_attention_implementation()
@@ -352,6 +365,8 @@ class HiggsAudioServeEngine:
         if tokenizer_name_or_path is None:
             tokenizer_name_or_path = model_name_or_path
         # logger.info(f"Loading tokenizer from {tokenizer_name_or_path}")
+        print("🔄 Higgs Audio: Loading text tokenizer...")
+        stage_start = time.perf_counter()
         try:
             # First try auto-detection
             self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
@@ -368,9 +383,12 @@ class HiggsAudioServeEngine:
                 f"falling back to local PreTrainedTokenizerFast"
             )
             self.tokenizer = _load_higgs_tokenizer_fast_fallback(tokenizer_name_or_path)
+        print(f"✅ Higgs Audio: Text tokenizer ready ({time.perf_counter() - stage_start:.2f}s)")
 
-        # logger.info(f"Initializing Higgs Audio Tokenizer")
+        print("🔄 Higgs Audio: Loading audio tokenizer and semantic model...")
+        stage_start = time.perf_counter()
         self.audio_tokenizer = load_higgs_audio_tokenizer(audio_tokenizer_name_or_path, device=device)
+        print(f"✅ Higgs Audio: Audio tokenizer ready ({time.perf_counter() - stage_start:.2f}s)")
 
         self.audio_num_codebooks = self.model.config.audio_num_codebooks
         self.audio_codebook_size = self.model.config.audio_codebook_size
@@ -429,15 +447,17 @@ class HiggsAudioServeEngine:
             }
             print(f"🛡️ Created DynamicCache buckets for memory safety")
 
-        if self.model.config.encode_whisper_embed:
-            logger.info(f"Loading whisper processor")
-            whisper_processor = AutoProcessor.from_pretrained(
-                "openai/whisper-large-v3-turbo",
-                trust_remote=True,
-                device=self.device,
-            )
-        else:
-            whisper_processor = None
+        if self.model.config.encode_whisper_embed:
+            logger.info(f"Loading whisper processor")
+            stage_start = time.perf_counter()
+            whisper_processor = AutoProcessor.from_pretrained(
+                "openai/whisper-large-v3-turbo",
+                trust_remote=True,
+                device=self.device,
+            )
+            print(f"⏱️ Higgs Audio: Whisper processor load took {time.perf_counter() - stage_start:.2f}s")
+        else:
+            whisper_processor = None
 
         # Reuse collator to prepare inference samples
         self.collator = HiggsAudioSampleCollator(
@@ -466,28 +486,38 @@ class HiggsAudioServeEngine:
             # logger.info("CUDA graph capture skipped (not using CUDA device)")
             pass
     
-    def _ensure_cuda_graphs(self):
-        """Initialize CUDA graphs on first use if needed"""
-        if self.device == "cuda" and self.enable_cuda_graphs and not self.cuda_graphs_initialized:
-            try:
-                # logger.info(f"🚀 Initializing CUDA graphs on first inference")
-                self.model.capture_model(self.kv_caches.values())
-                self.cuda_graphs_initialized = True
-                # logger.info("✅ CUDA graph capture successful")
-            except Exception as e:
-                logger.warning(f"⚠️ CUDA graph capture failed: {e}")
-                logger.info("Continuing without CUDA graphs (performance may be slower)")
-                self.enable_cuda_graphs = False
+    def _ensure_cuda_graphs(self):
+        """Initialize CUDA graphs on first use if needed"""
+        if self.device == "cuda" and self.enable_cuda_graphs and not self.cuda_graphs_initialized:
+            if not self._cuda_graph_replay_enabled:
+                self.cuda_graphs_initialized = True
+                return
+            try:
+                # logger.info(f"🚀 Initializing CUDA graphs on first inference")
+                self.model.capture_model(self.kv_caches.values())
+                self.cuda_graphs_initialized = True
+                # logger.info("✅ CUDA graph capture successful")
+            except Exception as e:
+                logger.warning(f"⚠️ CUDA graph capture failed: {e}")
+                logger.info("Continuing without CUDA graphs (performance may be slower)")
+                self.enable_cuda_graphs = False
+                self._cuda_graph_replay_enabled = False
+                # TTS Audio Suite patch: graph capture failure must also stop StaticCache reuse.
+                self._cuda_graphs_enabled = False
+                self._cache_type = "DynamicCache"
+                self._force_cache_recreation = True
+                self.current_past_key_values_bucket = None
 
     def _fix_attention_implementation(self):
         """Fix attention implementation for all LLaMA layers to avoid None errors"""
-        # logger.info("Fixing attention implementation for transformers compatibility")
+        # logger.info("Fixing attention implementation for transformers compatibility")
+        preferred_attn = getattr(self, "_preferred_attn_implementation", "eager")
         
         # Fix the main config
         if hasattr(self.model.config, 'text_config'):
             if hasattr(self.model.config.text_config, '_attn_implementation'):
-                if self.model.config.text_config._attn_implementation is None:
-                    self.model.config.text_config._attn_implementation = "eager"
+                if self.model.config.text_config._attn_implementation is None:
+                    self.model.config.text_config._attn_implementation = preferred_attn
         
         # Fix all attention layers directly
         if hasattr(self.model, 'language_model') and hasattr(self.model.language_model, 'model'):
@@ -496,8 +526,8 @@ class HiggsAudioServeEngine:
                 for layer in llm_model.layers:
                     if hasattr(layer, 'self_attn') and hasattr(layer.self_attn, 'config'):
                         if hasattr(layer.self_attn.config, '_attn_implementation'):
-                            if layer.self_attn.config._attn_implementation is None:
-                                layer.self_attn.config._attn_implementation = "eager"
+                            if layer.self_attn.config._attn_implementation is None:
+                                layer.self_attn.config._attn_implementation = preferred_attn
 
     def disable_cuda_graphs_permanently(self):
         """
@@ -507,9 +537,10 @@ class HiggsAudioServeEngine:
         try:
             logger.info("🚫 Permanently disabling CUDA graphs for resurrected model")
             
-            # Disable CUDA graph functionality completely
-            self.enable_cuda_graphs = False
-            self.cuda_graphs_initialized = False
+            # Disable CUDA graph functionality completely
+            self.enable_cuda_graphs = False
+            self.cuda_graphs_initialized = False
+            self._cuda_graph_replay_enabled = False
             
             # Clear any existing CUDA graphs that might be corrupted
             if hasattr(self.model, 'decode_graph_runners'):
@@ -556,9 +587,10 @@ class HiggsAudioServeEngine:
         Disable CUDA graphs specifically for memory management scenarios.
         Called when cache invalidation indicates models will be unloaded/reloaded.
         """
-        logger.info("🚫 Disabling CUDA graphs for memory management cycle")
-        self.enable_cuda_graphs = False
-        self.cuda_graphs_initialized = False
+        logger.info("🚫 Disabling CUDA graphs for memory management cycle")
+        self.enable_cuda_graphs = False
+        self.cuda_graphs_initialized = False
+        self._cuda_graph_replay_enabled = False
 
     def _prepare_inputs(self, chat_ml_sample: ChatMLSample, force_audio_gen: bool = False):
         input_tokens, _, audio_contents, _ = prepare_chatml_sample(
@@ -661,23 +693,17 @@ class HiggsAudioServeEngine:
         cache_device_mismatch = False
         force_cache_recreation = getattr(self, '_force_cache_recreation', False)
         
-        if self.kv_caches and not force_cache_recreation:
-            # Check if any cache is on wrong device
-            for kv_cache in self.kv_caches.values():
-                if hasattr(kv_cache, '__len__') and len(kv_cache) > 0:
-                    try:
-                        cache_tuple = kv_cache[0]  # Get first layer's cache
-                        if cache_tuple is not None and len(cache_tuple) >= 2:
-                            key_cache, _ = cache_tuple
-                            cache_device = key_cache.device if key_cache is not None else model_device
-                        else:
-                            cache_device = model_device
-                        if cache_device != model_device:
-                            cache_device_mismatch = True
-                            break
-                    except (AttributeError, IndexError):
-                        cache_device_mismatch = True
-                        break
+        if self.kv_caches and not force_cache_recreation:
+            # Check if any cache is on wrong device
+            for kv_cache in self.kv_caches.values():
+                try:
+                    cache_device = get_cache_device(kv_cache) or model_device
+                    if cache_device != model_device:
+                        cache_device_mismatch = True
+                        break
+                except (AttributeError, IndexError, TypeError):
+                    cache_device_mismatch = True
+                    break
         
         # Recreate caches if device mismatch detected or forced recreation requested
         if cache_device_mismatch or force_cache_recreation:
@@ -847,23 +873,26 @@ class HiggsAudioServeEngine:
             # Clean up progress bar
             progress_streamer.end()
 
-            if len(outputs[1]) > 0:
-                wv_list = []
-                for output_audio in outputs[1]:
-                    vq_code = revert_delay_pattern(output_audio).clip(0, self.audio_codebook_size - 1)[:, 1:-1]
-                    wv_numpy = self.audio_tokenizer.decode(vq_code.unsqueeze(0))[0, 0]
-                    wv_list.append(wv_numpy)
-                wv_numpy = np.concatenate(wv_list)
-            else:
-                wv_numpy = None
-
-            # We only support one request at a time now
-            generated_text_tokens = outputs[0][0].cpu().numpy()[len(prompt_token_ids) :]
-            generated_text = self.tokenizer.decode(generated_text_tokens)
-            generated_audio_tokens = outputs[1][0].cpu().numpy()
-            return HiggsAudioResponse(
-                audio=wv_numpy,
-                generated_audio_tokens=generated_audio_tokens,
+            generated_text_tokens = outputs[0][0].cpu().numpy()[len(prompt_token_ids) :]
+            generated_text = self.tokenizer.decode(generated_text_tokens)
+
+            if len(outputs[1]) > 0:
+                wv_list = []
+                for output_audio in outputs[1]:
+                    vq_code = revert_delay_pattern(output_audio).clip(0, self.audio_codebook_size - 1)[:, 1:-1]
+                    wv_numpy = self.audio_tokenizer.decode(vq_code.unsqueeze(0))[0, 0]
+                    wv_list.append(wv_numpy)
+                wv_numpy = np.concatenate(wv_list)
+                generated_audio_tokens = outputs[1][0].cpu().numpy()
+            else:
+                # TTS Audio Suite patch: surface missing audio generation instead of crashing on outputs[1][0].
+                raise RuntimeError(
+                    f"Higgs generated no audio tokens. Generated text tail: {generated_text[-200:]!r}"
+                )
+
+            return HiggsAudioResponse(
+                audio=wv_numpy,
+                generated_audio_tokens=generated_audio_tokens,
                 sampling_rate=self.audio_tokenizer.sampling_rate,
                 generated_text=generated_text,
                 generated_text_tokens=generated_text_tokens,

nodes/engines/higgs_audio_engine_node.py

@@ -173,8 +173,8 @@ class HiggsAudioEngineNode(BaseTTSNode):
                     "tooltip": "🔤 Maximum token limit - safety cap on generation length:\n\n⚠️ This is a LIMIT, not a target. Model stops when audio is complete OR limit is reached.\n\n• <10 tokens: ⚠️ May cause errors or cut off mid-word\n• 200-500: Safe for short sentences, faster processing\n• 1000-2048: 🎯 RECOMMENDED - Handles most content safely\n• 3000-4096: For very long paragraphs only\n\nFor normal text like 'Hello Bob', 200 vs 2048 makes no difference - same quality and length. Only matters for very short limits (causes truncation) or very long text (needs higher limits)."
                 }),
                 "force_audio_gen": ("BOOLEAN", {
-                    "default": False,
-                    "tooltip": "🎵 Force Audio Generation:\n\n• False: 🎯 RECOMMENDED - Model naturally chooses to generate audio tokens\n• True: Force model to generate audio tokens rather than text tokens\n\n⚠️ Only enable if model is generating text instead of audio. Usually not needed as the model should naturally generate audio for TTS requests."
+                    "default": True,
+                    "tooltip": "🎵 Force Audio Generation:\n\n• True: 🎯 RECOMMENDED for TTS - keeps Higgs in audio generation mode instead of free-running in text mode\n• False: Allow mixed/text generation behavior\n\nIf disabled, Higgs may burn the full token budget instead of terminating cleanly in TTS workflows."
                 }),
                 "ras_win_len": ("INT", {
                     "default": 7,
@@ -244,4 +244,4 @@ NODE_CLASS_MAPPINGS = {
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "HiggsAudioEngineNode": "⚙️ Higgs Audio 2 Engine"
-}
+}

nodes/higgs_audio/higgs_audio_tts_processor.py

@@ -67,7 +67,7 @@ class HiggsAudioTTSProcessor:
                 "top_p": self.engine_wrapper.config.get("top_p", 0.6),
                 "top_k": self.engine_wrapper.config.get("top_k", 80),
                 "max_new_tokens": self.engine_wrapper.config.get("max_new_tokens", 2048),
-                "force_audio_gen": self.engine_wrapper.config.get("force_audio_gen", False),
+                "force_audio_gen": self.engine_wrapper.config.get("force_audio_gen", True),
                 "ras_win_len": self.engine_wrapper.config.get("ras_win_len", 7),
                 "ras_max_num_repeat": self.engine_wrapper.config.get("ras_max_num_repeat", 2),
                 "system_prompt": self.engine_wrapper.config.get("system_prompt", "Generate audio following instruction.")
@@ -406,4 +406,4 @@ class HiggsAudioTTSProcessor:
             empty_audio = AudioProcessingUtils.create_silence(1.0, self.sample_rate)
             empty_comfy = AudioProcessingUtils.format_for_comfyui(empty_audio, self.sample_rate)
             
-            return (empty_comfy, f"Higgs Audio TTS error: {e}")
+            return (empty_comfy, f"Higgs Audio TTS error: {e}")