Fix Step Audio EditX transformers 4.54+ compatibility with simple config fix

Replace complex weight restoration workaround with a simple one-line config fix.

The root cause: Step-Audio-EditX model config was missing 'tie_word_embeddings'
setting, so transformers 4.54+ defaulted to True and tied lm_head to embed_tokens
even though the checkpoint has different weights for each.

Solution: Add 'tie_word_embeddings': false to config.json to prevent tying.

This is much simpler and cleaner than the previous safetensors weight restoration
approach. One line in config.json fixes the issue for all transformers versions.

docs/Dev reports/fails_to_avoid_TTS_Engine_Implementation.md

@@ -91,3 +91,9 @@
   - 2D `[channels, samples]` → 3D `[1, channels, samples]` (one unsqueeze)
 - **Wrong**: Squeezing 3D to 2D, or only adding one dimension to 1D
 - **Pattern**: `tensor.unsqueeze(0)` adds batch dim to 2D; check VibeVoice processor for reference
+
+### Step Audio EditX Sample Rate
+- **Symptom**: Pitch shift in edit/clone modes with external audio (44100 Hz, 48000 Hz)
+- **Cause**: Saving audio with wrong sample rate metadata → `torchaudio.load()` misinterprets data
+- **Fix**: Resample to 24000 Hz before saving temp files (edit node, TTS text node, TTS SRT node)
+- **Prevention**: Always resample to engine's native rate (24000 Hz for Step Audio EditX)

engines/chatterbox/audio_timing.py

@@ -476,10 +476,10 @@ class TimedAudioAssembler:
             else:
                 raise AudioTimingError(f"Invalid stretcher_type: {stretcher_type}. Use 'ffmpeg' or 'phase_vocoder'")
     
-    def assemble_timed_audio(self, audio_segments: List[torch.Tensor], 
+    def assemble_timed_audio(self, audio_segments: List[torch.Tensor],
                            target_timings: List[Tuple[float, float]],
                            total_duration: Optional[float] = None,
-                           fade_duration: float = 0.01) -> torch.Tensor:
+                           fade_duration: float = 0.01) -> Tuple[torch.Tensor, str]:
         """
         Assemble audio segments with precise timing
         
@@ -582,7 +582,15 @@ class TimedAudioAssembler:
             self._place_segment_with_fade(output, audio_segment, start_sample, end_sample, fade_duration)
         
         return output
-    
+
+    def get_stretch_method_used(self) -> str:
+        """Get which stretching method was used in the last assembly"""
+        return self.stretch_method_used or "none"
+
+    def get_time_stretcher(self):
+        """Get the time stretcher object used (for backward compatibility)"""
+        return self.time_stretcher
+
     def _place_segment_with_fade(self, output: torch.Tensor, segment: torch.Tensor,
                                start_sample: int, end_sample: int, fade_duration: float):
         """

engines/step_audio_editx/step_audio_editx_impl/model_loader.py

@@ -289,41 +289,6 @@ class UnifiedModelLoader:
             # Without this, the model generates garbage due to training-time randomness
             model.eval()
 
-            # CRITICAL FIX for transformers 4.54+: Restore correct lm_head weights
-            # Bug: transformers 4.54+ incorrectly ties lm_head to embed_tokens even when they should be separate
-            # The Step-Audio-EditX model has different weights for lm_head (norm=227) vs embed_tokens (norm=255)
-            # but transformers 4.54+ ties them, causing wrong predictions (text tokens instead of audio tokens)
-            if hasattr(model, 'lm_head') and hasattr(model, 'model') and hasattr(model.model, 'embed_tokens'):
-                if model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr():
-                    # Weights are tied - need to restore correct lm_head from checkpoint
-                    self.logger.info("Detected incorrect weight tying in transformers 4.54+, restoring lm_head weights...")
-
-                    # Load correct lm_head weights from safetensors
-                    from safetensors import safe_open
-                    import glob
-
-                    safetensors_files = glob.glob(os.path.join(model_path, "model*.safetensors"))
-                    if safetensors_files:
-                        # Try to find lm_head.weight in safetensors files
-                        lm_head_weight = None
-                        for st_file in sorted(safetensors_files):
-                            try:
-                                with safe_open(st_file, framework="pt", device="cpu") as f:
-                                    if "lm_head.weight" in f.keys():
-                                        lm_head_weight = f.get_tensor("lm_head.weight")
-                                        break
-                            except:
-                                continue
-
-                        if lm_head_weight is not None:
-                            # Untie the weights and restore correct lm_head
-                            model.lm_head.weight = torch.nn.Parameter(
-                                lm_head_weight.to(model.lm_head.weight.dtype).to(model.lm_head.weight.device).clone()
-                            )
-                            self.logger.info(f"Restored lm_head weights (norm: {model.lm_head.weight.norm().item():.2f})")
-                        else:
-                            self.logger.warning("Could not find lm_head.weight in safetensors, skipping fix")
-
             self.logger.debug(f"Successfully loaded model from {source}")
             return model, tokenizer, model_path
 

nodes/chatterbox/chatterbox_srt_node.py

@@ -915,8 +915,22 @@ The audio will match these exact timings.""",
                 )
                 adjustments = smart_adjustments
             
+            # Get stretcher method info
+            stretch_method = None
+            if current_timing_mode == "stretch_to_fit" and 'assembler' in locals():
+                stretch_method = assembler.get_stretch_method_used()
+            elif current_timing_mode == "smart_natural" and hasattr(self, '_smart_natural_stretcher'):
+                stretch_method = self._smart_natural_stretcher
+            elif current_timing_mode == "concatenate" and 'assembler' in locals():
+                stretch_method = assembler.get_stretch_method_used()
+
             # Generate reports
-            timing_report = self._generate_timing_report(subtitles, adjustments, current_timing_mode, has_overlaps, mode_switched, timing_mode if mode_switched else None)
+            timing_report = self._generate_timing_report(
+                subtitles, adjustments, current_timing_mode,
+                has_overlaps, mode_switched,
+                timing_mode if mode_switched else None,
+                stretch_method
+            )
             adjusted_srt_string = self._generate_adjusted_srt_string(subtitles, adjustments, current_timing_mode)
             
             # Generate info with cache status and stretching method - ORIGINAL LOGIC FROM LINES 1141-1168

nodes/f5tts/f5tts_srt_node.py

@@ -769,12 +769,12 @@ Hello! This is F5-TTS SRT with character switching.
         
         return final_audio, adjustments
     
-    def _generate_timing_report(self, subtitles: List, adjustments: List[Dict], timing_mode: str, has_original_overlaps: bool = False, mode_switched: bool = False, original_mode: str = None) -> str:
+    def _generate_timing_report(self, subtitles: List, adjustments: List[Dict], timing_mode: str, has_original_overlaps: bool = False, mode_switched: bool = False, original_mode: str = None, stretch_method: str = None) -> str:
         """Generate detailed timing report."""
         # Delegate to reporting module
         from utils.timing.reporting import SRTReportGenerator
         reporter = SRTReportGenerator()
-        return reporter.generate_timing_report(subtitles, adjustments, timing_mode, has_original_overlaps, mode_switched, original_mode)
+        return reporter.generate_timing_report(subtitles, adjustments, timing_mode, has_original_overlaps, mode_switched, original_mode, stretch_method)
     
     def _generate_adjusted_srt_string(self, subtitles: List, adjustments: List[Dict], timing_mode: str) -> str:
         """Generate adjusted SRT string from final timings."""

utils/timing/assembly.py

@@ -20,12 +20,13 @@ class AudioAssemblyEngine:
     def __init__(self, sample_rate: int):
         """
         Initialize audio assembly engine
-        
+
         Args:
             sample_rate: Audio sample rate
         """
         self.sample_rate = sample_rate
         self.stretcher_cache = {}  # Cache for time stretching instances
+        self.last_stretch_method_used = None  # Track last stretcher method used
         
     def assemble_concatenation(self, audio_segments: List[torch.Tensor],
                              fade_duration: float = 0.0) -> torch.Tensor:
@@ -192,15 +193,18 @@ class AudioAssemblyEngine:
         # Use the original TimedAudioAssembler from chatterbox - EXACT ORIGINAL BEHAVIOR
         try:
             from engines.chatterbox.audio_timing import TimedAudioAssembler
-            
+
             assembler = TimedAudioAssembler(self.sample_rate)
             final_audio = assembler.assemble_timed_audio(
                 audio_segments, target_timings, fade_duration=fade_duration
             )
+            # Track which stretcher method was used
+            self.last_stretch_method_used = assembler.get_stretch_method_used()
             return final_audio
-            
+
         except ImportError as e:
             # Fallback to basic assembly if chatterbox modules not available
+            self.last_stretch_method_used = "basic_fallback"
             return self._basic_stretch_assembly(audio_segments, target_timings, fade_duration)
     
     def assemble_with_overlaps(self, audio_segments: List[torch.Tensor],
@@ -575,15 +579,20 @@ class AudioAssemblyEngine:
     def get_assembly_info(self) -> Dict[str, Any]:
         """
         Get information about the assembly engine configuration
-        
+
         Returns:
             Dictionary with assembly engine info
         """
         return {
             'sample_rate': self.sample_rate,
             'cached_stretchers': list(self.stretcher_cache.keys()),
-            'available_methods': self._get_available_methods()
+            'available_methods': self._get_available_methods(),
+            'last_stretch_method_used': self.last_stretch_method_used
         }
+
+    def get_stretch_method_used(self) -> str:
+        """Get which stretching method was used in the last assembly"""
+        return self.last_stretch_method_used or "none"
     
     def _get_available_methods(self) -> List[str]:
         """Get list of available time stretching methods"""
@@ -604,7 +613,7 @@ class AudioAssemblyEngine:
         
         return methods
     
-    def assemble_by_timing_mode(self, audio_segments: List[torch.Tensor], 
+    def assemble_by_timing_mode(self, audio_segments: List[torch.Tensor],
                                subtitles: List, timing_mode: str, device,
                                adjustments: Optional[List[Dict]] = None,
                                processed_segments: Optional[List[torch.Tensor]] = None,

utils/timing/reporting.py

@@ -26,7 +26,7 @@ class SRTReportGenerator:
                 overlapping_segments.add(i)
         return overlapping_segments
     
-    def generate_timing_report(self, subtitles: List, adjustments: List[Dict], timing_mode: str, has_original_overlaps: bool = False, mode_switched: bool = False, original_mode: str = None) -> str:
+    def generate_timing_report(self, subtitles: List, adjustments: List[Dict], timing_mode: str, has_original_overlaps: bool = False, mode_switched: bool = False, original_mode: str = None, stretch_method: str = None) -> str:
         """Generate detailed timing report with original vs generated overlap distinction"""
         import numpy as np
         
@@ -55,9 +55,13 @@ No segments were processed due to immediate interruption.
             "=" * 50,
             f"Total subtitles: {len(subtitles)}",
             f"Total duration: {subtitles[-1].end_time:.3f}s",
-            "",
-            "Per-subtitle analysis:"
         ]
+
+        # Add stretcher method info if available
+        if stretch_method and timing_mode == "stretch_to_fit":
+            report_lines.append(f"Time stretcher: {stretch_method}")
+
+        report_lines.extend(["", "Per-subtitle analysis:"])
         
         if timing_mode == "smart_natural":
             # For smart_natural mode, iterate directly over the detailed adjustments report