diodiogod-TTS-Audio-Suite/utils/timing/engine.py

"""
SRT Timing Engine - Handles complex timing calculations and adjustments
Extracted from the massive SRT TTS node for better maintainability
"""

import torch
import numpy as np
from typing import Dict, Any, Optional, List, Tuple
from utils.audio.processing import AudioProcessingUtils


class TimingEngine:
    """
    Advanced timing engine for SRT audio synchronization
    Handles smart timing calculations, adjustments, and optimizations
    """
    
    def __init__(self, sample_rate: int):
        """
        Initialize timing engine
        
        Args:
            sample_rate: Audio sample rate for calculations
        """
        self.sample_rate = sample_rate
        
    def calculate_concatenation_adjustments(self, audio_segments: List[torch.Tensor],
                                          subtitles: List) -> List[Dict]:
        """
        Calculate timing adjustments for concatenation mode
        Ignores original SRT timings and creates new sequential timings
        """
        adjustments = []
        current_time = 0.0
        
        for i, (audio_segment, subtitle) in enumerate(zip(audio_segments, subtitles)):
            # Get natural duration of this segment
            natural_duration = self._get_audio_duration(audio_segment)
            
            # Calculate new timing for this segment
            start_time = current_time
            end_time = current_time + natural_duration
            
            # Create adjustment record
            adjustment = {
                'sequence': subtitle.sequence,
                'start_time': start_time,
                'end_time': end_time,
                'natural_duration': natural_duration,
                'original_text': subtitle.text,
                'original_srt_start': subtitle.start_time,
                'original_srt_end': subtitle.end_time,
                'original_srt_duration': subtitle.duration,
                'timing_change': end_time - subtitle.end_time,  # How much timing changed
                'needs_stretching': False,  # No stretching in concatenate mode
                'mode': 'concatenate'
            }
            adjustments.append(adjustment)
            
            # Update current time for next segment
            current_time = end_time
        
        return adjustments
        
    def calculate_smart_timing_adjustments(self, audio_segments: List[torch.Tensor],
                                         subtitles: List, tolerance: float,
                                         max_stretch_ratio: float,
                                         min_stretch_ratio: float, device: str) -> Tuple[List[Dict], List[torch.Tensor]]:
        """
        ORIGINAL Smart balanced timing: Adjusts SRT segment timings based on actual spoken duration
        and a user-defined timing_tolerance - EXACT ORIGINAL IMPLEMENTATION FROM LINES 1497-1774
        """
        # Import required modules using the same pattern as SRT node
        try:
            # Use the import manager to get the modules
            from utils.system.import_manager import import_manager
            success, modules, source = import_manager.import_srt_modules()
            if not success:
                raise ImportError("SRT modules not available")
            
            AudioTimingUtils = modules.get("AudioTimingUtils")
            FFmpegTimeStretcher = modules.get("FFmpegTimeStretcher")
            PhaseVocoderTimeStretcher = modules.get("PhaseVocoderTimeStretcher")
            AudioTimingError = modules.get("AudioTimingError")
            SRTSubtitle = modules.get("SRTSubtitle")
            
            if not all([AudioTimingUtils, FFmpegTimeStretcher, PhaseVocoderTimeStretcher, AudioTimingError, SRTSubtitle]):
                raise ImportError("Required SRT modules not found")
                
            import comfy.model_management as model_management
        except ImportError:
            # Fallback if modules not available
            return self._fallback_smart_timing(audio_segments, subtitles)
        
        processed_segments = []
        smart_adjustments_report = []
        
        # Create a mutable copy of subtitles to adjust start/end times
        mutable_subtitles = [SRTSubtitle(s.sequence, s.start_time, s.end_time, s.text) for s in subtitles]
        
        # Initialize stretcher for smart_natural mode
        try:
            # Try FFmpeg first
            print("Smart natural mode: Trying FFmpeg stretcher...")
            time_stretcher = FFmpegTimeStretcher()
            self._smart_natural_stretcher = "ffmpeg"
            print("Smart natural mode: Using FFmpeg stretcher")
        except AudioTimingError as e:
            # Fall back to Phase Vocoder
            print(f"Smart natural mode: FFmpeg initialization failed ({str(e)}), falling back to Phase Vocoder")
            time_stretcher = PhaseVocoderTimeStretcher()
            self._smart_natural_stretcher = "phase_vocoder"
            print("Smart natural mode: Using Phase Vocoder stretcher")
        
        # Process audio with smart natural timing
        for i, audio in enumerate(audio_segments):
            # Check for interruption during smart natural processing
            if model_management.interrupt_processing:
                print(f"⚠️ Smart natural processing interrupted at segment {i+1}/{len(audio_segments)}")
                raise InterruptedError(f"Smart natural processing interrupted at segment {i+1}/{len(audio_segments)}")
                
            current_subtitle = mutable_subtitles[i]
            natural_duration = AudioTimingUtils.get_audio_duration(audio, self.sample_rate)
            
            original_srt_start = subtitles[i].start_time # Use original for reference
            original_srt_end = subtitles[i].end_time
            
            # Step 1: Check if natural duration fits within current SRT slot
            # This is the duration the SRT *originally* allocated for this segment
            initial_target_duration = original_srt_end - original_srt_start
            
            segment_report = {
                'segment_index': i,
                'sequence': current_subtitle.sequence,
                'original_srt_start': original_srt_start,
                'original_srt_end': original_srt_end,
                'original_srt_duration': initial_target_duration,
                'natural_audio_duration': natural_duration,
                'next_segment_shifted_by': 0.0,
                'stretch_factor_applied': 1.0,
                'padding_added': 0.0,
                'truncated_by': 0.0,
                'final_segment_duration': natural_duration, # Will be updated
                'final_srt_start': original_srt_start, # Will be updated
                'final_srt_end': original_srt_end, # Will be updated
                'original_text': subtitles[i].text, # Add original subtitle text
                'actions': []
            }
            
            # Process segment timing
            
            # Calculate how much extra time is needed for the natural audio
            time_needed_beyond_srt = natural_duration - initial_target_duration
            
            adjusted_current_segment_end = original_srt_end # This will be updated
            
            # Step 2 & 3: Adjust Next Segment Start (if needed)
            if time_needed_beyond_srt > 0: # Natural audio is longer than original SRT slot
                # Audio is longer than slot
                segment_report['actions'].append(f"Natural audio ({natural_duration:.3f}s) is longer than original SRT slot ({initial_target_duration:.3f}s) by {time_needed_beyond_srt:.3f}s.")
                
                if i + 1 < len(mutable_subtitles):
                    next_subtitle = mutable_subtitles[i+1]
                    original_next_srt_start = subtitles[i+1].start_time # Use original for reference
                    
                    # First, try to consume any existing gap to the next subtitle
                    existing_gap = original_next_srt_start - original_srt_end
                    if existing_gap > 0:
                        time_to_consume_from_gap = min(time_needed_beyond_srt, existing_gap)
                        time_needed_beyond_srt -= time_to_consume_from_gap
                        adjusted_current_segment_end += time_to_consume_from_gap
                        segment_report['actions'].append(f"Consumed {time_to_consume_from_gap:.3f}s from existing gap. Remaining excess: {time_needed_beyond_srt:.3f}s.")
                        # Gap consumed
                    
                    if time_needed_beyond_srt > 0: # Still need more time after consuming gap
                        next_natural_audio_duration = AudioTimingUtils.get_audio_duration(audio_segments[i+1], self.sample_rate)
                        
                        # Calculate "room" in the next segment: how much shorter its natural audio is than its SRT slot
                        next_segment_room = max(0.0, next_subtitle.duration - next_natural_audio_duration)
                        # Calculate available room
                        segment_report['actions'].append(f"Next segment (Seq {next_subtitle.sequence}) has {next_segment_room:.3f}s room.")

                        # How much can we shift the next subtitle without exceeding tolerance?
                        # This is the amount of time we can "borrow" from the next segment's start.
                        max_shift_allowed = min(tolerance, next_segment_room) # Only shift into its room, within tolerance
                        
                        # How much do we *want* to shift the next subtitle?
                        desired_shift = time_needed_beyond_srt
                        
                        actual_shift = min(desired_shift, max_shift_allowed)
                        
                        if actual_shift > 0:
                            # Shift the next subtitle's start and end times
                            next_subtitle.start_time += actual_shift
                            next_subtitle.end_time += actual_shift
                            adjusted_current_segment_end += actual_shift # Add to the already adjusted end
                            segment_report['next_segment_shifted_by'] = actual_shift
                            segment_report['actions'].append(f"Shifted next subtitle (Seq {next_subtitle.sequence}) by {actual_shift:.3f}s. New next SRT start: {next_subtitle.start_time:.3f}s.")
                            # Subtitle shifted
                        else:
                            segment_report['actions'].append("Cannot shift next subtitle within tolerance/available room.")
                            # Cannot shift subtitle
                    else: # No next subtitle or no excess after consuming gap
                        segment_report['actions'].append("No next subtitle to shift or excess consumed by gap.")
                        print("   No next subtitle to shift or excess consumed by gap.")
                else:
                    segment_report['actions'].append("No next subtitle to shift.")
                    print("   No next subtitle to shift.")
            
            # Step 4: Stretch/Shrink (if still needed)
            # The new target duration for the current segment is from its original start to its (potentially) adjusted end
            new_target_duration = adjusted_current_segment_end - original_srt_start
            
            # Calculate stretch factor needed to fit natural audio into the new target duration
            stretch_factor = new_target_duration / natural_duration if natural_duration > 0 else 1.0
            
            # Apply stretch factor limits based on max_stretch_ratio and min_stretch_ratio
            clamped_stretch_factor = max(min_stretch_ratio, min(max_stretch_ratio, stretch_factor))
            
            # Check if stretching is actually needed and if it's within acceptable limits
            if abs(clamped_stretch_factor - 1.0) > 0.01: # Apply stretch if deviation is more than 1%
                # Apply audio stretching
                segment_report['actions'].append(f"⏱️ Applying stretch/shrink: natural {natural_duration:.3f}s -> target {new_target_duration:.3f}s (factor: {clamped_stretch_factor:.3f}x).")
                segment_report['stretch_factor_applied'] = clamped_stretch_factor
                try:
                    stretched_audio = time_stretcher.time_stretch(audio, clamped_stretch_factor, self.sample_rate)
                    processed_audio = stretched_audio
                except Exception as e:
                    segment_report['actions'].append("Time stretching failed, using padding/truncation")
                    # Time stretching failed, use fallback
                    processed_audio = audio # Use original audio if stretching fails
            else:
                processed_audio = audio
                segment_report['actions'].append("No significant stretch/shrink needed.")
                # No stretching needed
            
            # Step 5: Pad with Silence (last resort) or Truncate
            final_processed_duration = AudioTimingUtils.get_audio_duration(processed_audio, self.sample_rate)
            
            if final_processed_duration < new_target_duration:
                padding_needed = new_target_duration - final_processed_duration
                if padding_needed > 0:
                    segment_report['padding_added'] = padding_needed
                    segment_report['actions'].append(f"Padding with {padding_needed:.3f}s silence to reach target duration.")
                    # Add silence padding
                    processed_audio = AudioTimingUtils.pad_audio_to_duration(processed_audio, new_target_duration, self.sample_rate, "end")
            elif final_processed_duration > new_target_duration:
                # Truncate if still too long
                truncated_by = final_processed_duration - new_target_duration
                segment_report['truncated_by'] = truncated_by
                
                # Define threshold for insignificant truncations (50ms)
                INSIGNIFICANT_TRUNCATION_THRESHOLD = 0.05
                
                if truncated_by > INSIGNIFICANT_TRUNCATION_THRESHOLD:
                    # Significant truncation - show warning emoji
                    segment_report['actions'].append(f"🚧 Truncating audio by {truncated_by:.3f}s.")
                else:
                    # Insignificant truncation - show without alarming emoji
                    segment_report['actions'].append(f"Fine-tuning audio duration (-{truncated_by:.3f}s for precision).")
                
                # Truncate audio
                target_samples = AudioTimingUtils.seconds_to_samples(new_target_duration, self.sample_rate)
                processed_audio = processed_audio[..., :target_samples]
            
            processed_segments.append(processed_audio)
            
            # Update the current subtitle's end time to reflect the final processed duration
            # This is important for calculating the gap to the *next* segment in the final assembly
            current_subtitle.end_time = original_srt_start + AudioTimingUtils.get_audio_duration(processed_audio, self.sample_rate)
            
            segment_report['final_segment_duration'] = AudioTimingUtils.get_audio_duration(processed_audio, self.sample_rate)
            segment_report['final_srt_start'] = current_subtitle.start_time
            segment_report['final_srt_end'] = current_subtitle.end_time
            
            smart_adjustments_report.append(segment_report)
        
        return smart_adjustments_report, processed_segments
    
    def _fallback_smart_timing(self, audio_segments: List[torch.Tensor], subtitles: List) -> Tuple[List[Dict], List[torch.Tensor]]:
        """Fallback when chatterbox modules not available - SHOULD NOT HAPPEN IN NORMAL OPERATION"""
        # This should not show in reports as it means the import failed completely
        raise ImportError("Smart timing not available - missing chatterbox modules")
    
    def calculate_overlap_timing(self, audio_segments: List[torch.Tensor], 
                               subtitles: List) -> Tuple[torch.Tensor, List[Dict]]:
        """
        Calculate timing for pad_with_silence mode allowing overlaps
        
        Args:
            audio_segments: List of audio tensors
            subtitles: List of SRTSubtitle objects
            
        Returns:
            Tuple of (final assembled audio, timing adjustments)
        """
        if not audio_segments:
            return torch.empty(0), []
        
        # Calculate total duration needed
        max_end_time = 0.0
        adjustments = []
        
        for i, (segment, subtitle) in enumerate(zip(audio_segments, subtitles)):
            natural_duration = self._get_audio_duration(segment)
            segment_end_time = subtitle.start_time + natural_duration
            max_end_time = max(max_end_time, segment_end_time)
            
            # Create timing adjustment record
            adjustment = {
                'segment_index': i,
                'sequence': subtitle.sequence,
                'start_time': subtitle.start_time,
                'end_time': segment_end_time,
                'natural_duration': natural_duration,
                'target_duration': subtitle.duration,
                'stretch_factor': 1.0,  # No stretching in overlap mode
                'needs_stretching': False,
                'stretch_type': 'none',
                'placement': 'natural_at_start_time'
            }
            adjustments.append(adjustment)
        
        if subtitles:
            max_end_time = max(max_end_time, subtitles[-1].end_time)
        
        # Assemble audio with overlaps
        final_audio = self._assemble_overlapping_audio(audio_segments, subtitles, max_end_time)
        
        return final_audio, adjustments
    
    def _get_audio_duration(self, audio: torch.Tensor) -> float:
        """Get duration of audio tensor in seconds"""
        if audio.dim() == 1:
            return audio.size(0) / self.sample_rate
        elif audio.dim() == 2:
            return audio.size(-1) / self.sample_rate
        else:
            raise ValueError(f"Unsupported audio tensor dimensions: {audio.dim()}")
    
    def _apply_time_stretch(self, audio: torch.Tensor, stretch_factor: float) -> torch.Tensor:
        """Apply time stretching to audio (placeholder - would use actual stretching)"""
        # This would integrate with the actual time stretching utilities
        # For now, return original audio (the actual implementation would be in audio_assembly.py)
        return audio
    
    def _add_silence_padding(self, audio: torch.Tensor, padding_duration: float) -> torch.Tensor:
        """Add silence padding to audio"""
        padding_samples = int(padding_duration * self.sample_rate)
        
        if audio.dim() == 1:
            padding = torch.zeros(padding_samples, device=audio.device, dtype=audio.dtype)
            return torch.cat([audio, padding], dim=0)
        else:
            padding = torch.zeros(audio.size(0), padding_samples, device=audio.device, dtype=audio.dtype)
            return torch.cat([audio, padding], dim=-1)
    
    def _assemble_overlapping_audio(self, audio_segments: List[torch.Tensor], 
                                  subtitles: List, max_end_time: float) -> torch.Tensor:
        """Assemble audio segments allowing overlaps at their SRT start times"""
        total_samples = int(max_end_time * self.sample_rate)
        
        # Initialize output buffer
        device = audio_segments[0].device
        dtype = audio_segments[0].dtype
        
        if audio_segments[0].dim() == 1:
            output_audio = torch.zeros(total_samples, device=device, dtype=dtype)
        else:
            channels = audio_segments[0].shape[0]
            output_audio = torch.zeros(channels, total_samples, device=device, dtype=dtype)
        
        # Place each segment at its start time
        for segment, subtitle in zip(audio_segments, subtitles):
            start_sample = int(subtitle.start_time * self.sample_rate)
            end_sample = start_sample + segment.size(-1)
            
            # Resize output if needed
            if end_sample > output_audio.size(-1):
                new_size = end_sample
                if output_audio.dim() == 1:
                    new_output = torch.zeros(new_size, device=device, dtype=dtype)
                    new_output[:output_audio.size(-1)] = output_audio
                else:
                    new_output = torch.zeros(output_audio.shape[0], new_size, device=device, dtype=dtype)
                    new_output[:, :output_audio.size(-1)] = output_audio
                output_audio = new_output
            
            # Add segment to output (allowing overlaps)
            if output_audio.dim() == 1:
                output_audio[start_sample:end_sample] += segment
            else:
                output_audio[:, start_sample:end_sample] += segment
        
        return output_audio
    
    def validate_timing_feasibility(self, subtitles: List, 
                                  max_stretch_ratio: float, 
                                  min_stretch_ratio: float) -> List[str]:
        """
        Validate if SRT timing is feasible for TTS generation
        
        Args:
            subtitles: List of SRTSubtitle objects
            max_stretch_ratio: Maximum allowed stretch ratio
            min_stretch_ratio: Minimum allowed stretch ratio
            
        Returns:
            List of warning messages
        """
        warnings = []
        
        for subtitle in subtitles:
            # Estimate natural speech duration (rough: 150 words per minute)
            word_count = len(subtitle.text.split())
            estimated_duration = word_count / 2.5  # 150 words/min = 2.5 words/sec
            
            if estimated_duration > 0:
                required_stretch = subtitle.duration / estimated_duration
                
                if required_stretch > max_stretch_ratio:
                    warnings.append(
                        f"Subtitle {subtitle.sequence}: May need extreme compression "
                        f"({required_stretch:.2f}x, text: '{subtitle.text[:50]}...')"
                    )
                elif required_stretch < min_stretch_ratio:
                    warnings.append(
                        f"Subtitle {subtitle.sequence}: May need extreme stretching "
                        f"({required_stretch:.2f}x, text: '{subtitle.text[:50]}...')"
                    )
        
        return warnings