diodiogod-TTS-Audio-Suite/docs/Dev reports/HIGGS_AUDIO_STATELESS_WRAPPER_IMPLEMENTATION.md

Higgs Audio Stateless Wrapper Implementation - Complete Journey

Problem Statement

Higgs Audio 2 could not integrate with ComfyUI's native model management system due to CUDA Graph memory corruption during model unloading. The goal was to apply the same stateless wrapper pattern used for ChatterBox streaming to enable safe model unloading while preserving performance.

Original Performance Baseline

• 50 it/s (tokens/second) generation speed
• CUDA Graph optimization enabled
• 11GB VRAM usage
• Voice cloning functional
• No safe model unloading - crashes on "Unload Models"

---

Implementation Attempts & Results

Attempt 1: Basic Stateless Wrapper Creation

Approach: Created StatelessHiggsAudioWrapper following ChatterBox pattern

• Created engines/higgs_audio/stateless_wrapper.py
• Implemented tensor isolation with .detach().clone()
• Wrapped HiggsAudioEngine with stateless methods

Result: ❌ Circular dependency causing infinite loop Fix: Modified unified interface to create HiggsAudioServeEngine directly

Attempt 2: Direct Serve Engine Integration

Approach: Bypass HiggsAudioEngine.initialize_engine() to avoid circular calls

• Modified utils/models/unified_model_interface.py
• Direct HiggsAudioServeEngine creation
• Manual component setup

Result: ❌ TypeError: 'generator' object is not callable Fix: Changed @property decorator to regular method for parameters()

Attempt 3: Device Mismatch Fixes - Round 1

Approach: Fix "Input type torch.FloatTensor and weight type torch.cuda.FloatTensor" errors

• Added device enforcement in serve engine input preparation
• Modified audio tokenizer to ensure tensors on correct device

Result: ❌ Still device mismatches in semantic model Progress: Error moved from one component to another (good sign)

Attempt 4: Deep Device Consistency

Approach: Comprehensive device fixing throughout audio pipeline

• Fixed get_regress_target() method in audio tokenizer
• Added next(model.parameters()).device for accurate device detection
• Fixed encoder device handling in _xcodec_encode()

Result: ❌ Device mismatch moved to KV cache operations Lesson: Device issues were systemic, requiring comprehensive fixes

Attempt 5: KV Cache Device Synchronization

Approach: Fix cache_position tensor CPU/CUDA mismatch in transformers library

• Added comprehensive KV cache device enforcement
• Implemented cache tensor movement to model device
• Enhanced input tensor device verification

Result: ❌ Persistent device mismatch in cache_utils.py:1220 Issue: Deep transformers library issue with cache_position tensor creation

Attempt 6: Disable KV Cache (Temporary)

Approach: Remove StaticCache to bypass device issues

• Set use_cache=False
• Removed past_key_values_buckets parameter

Result: ✅ No crashes BUT ❌ Performance degraded to ~11 it/s Lesson: StaticCache is critical for performance

Attempt 7: Smart KV Cache Recreation

Approach: Recreate caches when device mismatch detected

• Device mismatch detection in _prepare_kv_caches()
• Dynamic cache recreation with correct device
• Restore StaticCache functionality

Result: ❌ Still ~12 it/s performance with occasional device issues

Attempt 8: Performance Optimization - Remove Tensor Cloning

Approach: Optimize stateless wrapper by removing expensive operations

• Changed .detach().clone() to .detach() only
• Removed redundant device operations
• Added performance caching

Result: ❌ Still ~12 it/s performance Discovery: Cloning wasn't the main bottleneck

Attempt 9: Bypass Stateless Wrapper Test

Approach: Temporarily disable stateless wrapper to isolate performance issue

• Modified adapter to use original generate() method
• Test if wrapper was causing slowdown

Result: ❌ Still ~12 it/s performance even without wrapper Conclusion: Performance issue was deeper in the loading approach

Attempt 10: Root Cause Discovery - CPU→GPU Loading

Approach: Identified that CPU-first loading was breaking CUDA Graph optimization

• Original approach: Load directly on GPU
• Our approach: Load on CPU → Move to GPU
• CUDA Graphs are compiled for specific device configurations

Root Cause Found: CPU→GPU loading breaks CUDA Graph optimization

Attempt 11: Direct GPU Loading

Approach: Restore original loading pattern

• Modified unified_model_interface.py to load directly on target device
• Removed CPU→GPU transfer code
• Fixed data collator device mismatch (torch.full() device parameter)

Result: ✅ 54.92 tok/s performance restored! BUT ❌ Crashes on model unload Issue: CUDA Graph warnings → Process termination

Attempt 12: Comprehensive CUDA Graph Cleanup

Approach: Multi-layer CUDA Graph cleanup with official PyTorch APIs

• Disabled CUDA Graph compilation during cleanup
• Used torch.cuda.reset_accumulated_memory_stats()
• Multiple CUDA synchronization passes
• Complete KV cache tensor replacement

Result: ❌ Still crashes with CUDA warnings Issue: CUDA Graph allocations persist at PyTorch level

Attempt 13: Complete Engine Destruction

Approach: Nuclear option - destroy and recreate entire serve engine

• Store model paths for recreation
• Destroy all engine components (model, tokenizer, caches)
• Set engine to None to force complete recreation
• Mark for recreation with stored paths

Result: ❌ Still crashes with identical CUDA warnings Issue: CUDA Graph captures happen at memory allocator level, persist even after object destruction

Attempt 14: Surgical CUDA Graph Preservation (FAILED)

Approach: Keep CUDA Graphs but clean them surgically

• Exit inference mode for tensor modifications
• Replace cache tensors with CPU versions
• Use official PyTorch CUDA Graph management APIs

Result: ❌ Cache structure errors and persistent CUDA warnings Issue: Cache structure is complex (lists), CUDA Graph captures are too deep

Attempt 15: PyTorch 2.2+ Modern CUDA Graph APIs (FAILED)

Approach: Research and implement modern PyTorch CUDA Graph cleanup APIs

• Researched PyTorch 2.2, 2.3, 2.4 improvements
• Implemented torch.cuda.graph() context manager approach
• Added memory pool context managers and proper cleanup sequences

Result: ❌ Context manager API errors and persistent CUDA warnings Issue: Modern APIs require different parameters, CUDA Graph captures persist at allocator level

Attempt 16: CUDA Graph Toggle Implementation (PARTIAL SUCCESS)

Approach: Accept fundamental incompatibility and implement user choice toggle

• Added enable_cuda_graphs boolean parameter to Higgs Audio Engine node
• Implemented conditional loading: StaticCache vs DynamicCache
• Added parameter passing chain: Node → Config → Adapter → Engine → Factory
• Conditional memory unloading protection based on toggle setting

Result: ✅ Toggle UI working BUT ❌ CUDA Graphs still created in Memory Safe mode Progress:

• ✅ Parameter passing chain complete
• ✅ "Memory Safe (CUDA Graphs OFF)" vs "High Performance (CUDA Graphs ON)" modes working
• ✅ Conditional cleanup logic implemented
• ✅ Memory unloading protection working
• ❌ Critical Issue: StaticCache still created even when CUDA Graphs disabled
• ❌ Still crashes: CUDA allocation warnings persist in "safe" mode

Technical Discovery: Environment variables set too late - serve engine constructor creates StaticCache before disable flags take effect

---

CONCLUSION: Fundamental Architectural Challenge

After 16 different implementation attempts, the core issue remains unsolved:

Root Cause Analysis:

1. CUDA Graphs are compiled at the PyTorch C++ level - cannot be disabled via Python APIs alone
2. StaticCache constructor immediately allocates CUDA Graph memory - happens before environment variables take effect
3. Memory allocator captures references during first generation - these become "captured allocations"
4. ComfyUI's model unloading tries to free captured allocations - triggers fatal CUDA warnings
5. Even DynamicCache approach fails - underlying model still creates CUDA Graph optimizations

Current Status: Toggle Implementation (Partial Success)

✅ Successfully Implemented:

• User toggle UI with clear performance/safety tradeoffs
• Complete parameter passing chain from node to factory
• Conditional cleanup logic that prevents crashes when CUDA Graphs enabled
• Memory unloading protection with clear user messaging
• Mode detection and appropriate warnings

❌ Still Failing:

• StaticCache still created even in "Memory Safe" mode
• CUDA Graph allocations still captured during generation
• Crashes still occur on model unload in both modes
• Environment variable timing - set after engine constructor runs

Technical Findings:

1. Performance Impact: Confirmed ~40% speed reduction when CUDA Graphs disabled (55+ tok/s → ~35 tok/s)
2. Memory Management: 11GB VRAM freed successfully, but CUDA state corruption persists
3. Architectural Conflict: ComfyUI's dynamic model management fundamentally incompatible with CUDA Graph memory capture
4. PyTorch Limitation: No clean API to disable CUDA Graphs after model initialization

Final Assessment:

The toggle provides user choice but both modes have limitations:

• High Performance Mode: Full speed, but model cannot be unloaded (restart required)
• Memory Safe Mode: Still crashes due to architectural conflicts (implementation incomplete)

Recommendation: Document current state as "partial solution" - users can choose performance vs theoretical memory safety, understanding both have constraints.

---

Key Technical Insights Discovered

1. CUDA Graph Optimization Requirements

• Must load directly on target GPU device
• CPU→GPU transfer breaks CUDA Graph compilation
• StaticCache is essential for performance (50+ it/s vs 12 it/s)

2. Device Consistency Challenges

• Device mismatches occur at multiple pipeline stages:
  • Audio tokenizer input processing
  • Semantic model tensor operations
  • Encoder operations
  • KV cache tensor operations
  • Data collator tensor concatenation
• Requires systematic fixes throughout entire pipeline

3. Memory Management Conflicts

• CUDA Graphs capture memory allocations during generation
• Model unloading tries to free "captured" allocations
• Results in warnings → crashes
• Requires explicit CUDA Graph cleanup before unloading

4. Performance Bottlenecks Identified

• Primary: CPU→GPU loading approach (54 it/s → 12 it/s impact)
• Secondary: Tensor cloning operations (moderate impact)
• Tertiary: Excessive device synchronization (minor impact)

---

Files Modified

Core Implementation

• engines/higgs_audio/stateless_wrapper.py - Main stateless wrapper
• utils/models/unified_model_interface.py - Factory integration
• engines/adapters/higgs_audio_adapter.py - Adapter integration
• utils/models/comfyui_model_wrapper.py - ComfyUI integration

Device Consistency Fixes

• engines/higgs_audio/boson_multimodal/serve/serve_engine.py - Input/cache device handling
• engines/higgs_audio/boson_multimodal/audio_processing/higgs_audio_tokenizer.py - Tokenizer device handling
• engines/higgs_audio/boson_multimodal/data_collator/higgs_audio_collator.py - Collator device handling

---

Current Status (As of Final Context)

✅ Confirmed Working

• Voice cloning functionality - Full feature preservation
• Memory management - 11GB VRAM freed successfully
• Device consistency - No device mismatch crashes
• High performance potential - 54.92 tok/s achieved

❌ Outstanding Issues

• CUDA Graph cleanup crashes - Hybrid solution untested
• Performance vs Safety tradeoff - Need both simultaneously

🔬 Untested Solution

• Hybrid approach with comprehensive CUDA Graph cleanup
• Expected outcome: 54+ it/s performance + no crashes + safe unloading
• Status: Implementation complete, testing required

---

Recommended Next Steps

1. Test hybrid solution - Verify 54 it/s + no crashes
2. If crashes persist: Investigate alternative CUDA Graph cleanup methods
3. Performance monitoring: Ensure no regression from cleanup operations
4. Stress testing: Multiple load/unload cycles
5. Documentation update: Record final working solution

---

Lessons Learned

1. Device loading approach is critical for CUDA Graph optimization
2. Systematic device consistency required throughout pipeline
3. CUDA Graph cleanup must be explicit and comprehensive
4. Performance vs Safety requires sophisticated balancing, not simple tradeoffs
5. Deep integration challenges require understanding of underlying frameworks (transformers, CUDA)

This implementation represents a complex balance between PyTorch CUDA Graph optimization, device management, and ComfyUI model lifecycle integration.