660 lines
32 KiB
Python
660 lines
32 KiB
Python
"""
|
|
SRT Audio Assembly - Handles audio segment assembly and time-stretching
|
|
Extracted from the massive SRT TTS node for better maintainability
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
import tempfile
|
|
import os
|
|
from typing import Dict, Any, Optional, List, Tuple
|
|
from utils.audio.processing import AudioProcessingUtils
|
|
|
|
|
|
class AudioAssemblyEngine:
|
|
"""
|
|
Handles assembly of audio segments with different timing modes
|
|
Integrates with existing ChatterBox audio timing utilities
|
|
"""
|
|
|
|
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:
|
|
"""
|
|
Concatenate audio segments naturally with optional crossfading
|
|
|
|
Args:
|
|
audio_segments: List of audio tensors to concatenate
|
|
fade_duration: Duration for crossfading between segments (seconds)
|
|
|
|
Returns:
|
|
Concatenated audio tensor
|
|
"""
|
|
if not audio_segments:
|
|
return torch.empty((1, 0))
|
|
if len(audio_segments) == 1:
|
|
return audio_segments[0]
|
|
|
|
if fade_duration > 0.0:
|
|
return self._concatenate_with_crossfade(audio_segments, fade_duration)
|
|
else:
|
|
# Simple concatenation - handle tensor dimensions properly
|
|
if not audio_segments:
|
|
return torch.empty(0)
|
|
|
|
# Filter out empty segments first
|
|
non_empty_segments = [seg for seg in audio_segments if seg.numel() > 0]
|
|
if not non_empty_segments:
|
|
return torch.empty(0)
|
|
|
|
# Determine target dimension from first non-empty segment
|
|
first_segment = non_empty_segments[0]
|
|
target_dim = first_segment.dim()
|
|
|
|
# Normalize all segments to same dimension
|
|
normalized_segments = []
|
|
for segment in non_empty_segments:
|
|
if segment.dim() == target_dim:
|
|
normalized_segments.append(segment)
|
|
elif target_dim == 1 and segment.dim() == 2:
|
|
# Convert 2D to 1D (average channels if multi-channel)
|
|
if segment.shape[0] > 1:
|
|
normalized_segments.append(torch.mean(segment, dim=0))
|
|
else:
|
|
normalized_segments.append(segment.squeeze(0))
|
|
elif target_dim == 2 and segment.dim() == 1:
|
|
# Convert 1D to 2D
|
|
normalized_segment = segment.unsqueeze(0)
|
|
if first_segment.shape[0] > 1:
|
|
normalized_segment = normalized_segment.repeat(first_segment.shape[0], 1)
|
|
normalized_segments.append(normalized_segment)
|
|
|
|
# Concatenate on appropriate dimension
|
|
if target_dim == 1:
|
|
return torch.cat(normalized_segments, dim=0)
|
|
else:
|
|
return torch.cat(normalized_segments, dim=1)
|
|
|
|
def _concatenate_with_crossfade(self, audio_segments: List[torch.Tensor],
|
|
fade_duration: float) -> torch.Tensor:
|
|
"""Concatenate audio segments with crossfading between them"""
|
|
# Filter out empty segments first
|
|
non_empty_segments = [seg for seg in audio_segments if seg.numel() > 0]
|
|
if not non_empty_segments:
|
|
return torch.empty(0)
|
|
if len(non_empty_segments) == 1:
|
|
return non_empty_segments[0]
|
|
|
|
fade_samples = int(fade_duration * self.sample_rate)
|
|
result = non_empty_segments[0]
|
|
|
|
for i in range(1, len(non_empty_segments)):
|
|
current_segment = non_empty_segments[i]
|
|
|
|
# Handle tensor dimensions properly - check last dimension regardless of tensor shape
|
|
result_samples = result.size(-1)
|
|
current_samples = current_segment.size(-1)
|
|
|
|
if fade_samples > 0 and result_samples >= fade_samples and current_samples >= fade_samples:
|
|
# Create fade out for end of previous segment
|
|
fade_out = torch.linspace(1.0, 0.0, fade_samples, device=result.device)
|
|
fade_in = torch.linspace(0.0, 1.0, fade_samples, device=current_segment.device)
|
|
|
|
# Normalize both tensors to same dimension before applying fades
|
|
if result.dim() != current_segment.dim():
|
|
if result.dim() == 1 and current_segment.dim() == 2:
|
|
# Convert current_segment to 1D (average channels if multi-channel)
|
|
if current_segment.shape[0] > 1:
|
|
current_segment = torch.mean(current_segment, dim=0)
|
|
else:
|
|
current_segment = current_segment.squeeze(0)
|
|
elif result.dim() == 2 and current_segment.dim() == 1:
|
|
# Convert current_segment to 2D to match result
|
|
current_segment = current_segment.unsqueeze(0)
|
|
if result.shape[0] > 1:
|
|
current_segment = current_segment.repeat(result.shape[0], 1)
|
|
|
|
# Apply fades based on result tensor dimensions (both are now same dimension)
|
|
if result.dim() == 1:
|
|
# 1D tensors (mono)
|
|
result_end = result[-fade_samples:] * fade_out
|
|
current_start = current_segment[:fade_samples] * fade_in
|
|
|
|
# Crossfade by overlapping
|
|
crossfaded = result_end + current_start
|
|
|
|
# Combine: previous audio (minus fade region) + crossfaded region + rest of current segment
|
|
result = torch.cat([
|
|
result[:-fade_samples],
|
|
crossfaded,
|
|
current_segment[fade_samples:]
|
|
], dim=0)
|
|
else:
|
|
# 2D tensors (multi-channel)
|
|
result_end = result[:, -fade_samples:] * fade_out
|
|
current_start = current_segment[:, :fade_samples] * fade_in
|
|
|
|
# Crossfade by overlapping
|
|
crossfaded = result_end + current_start
|
|
|
|
# Combine: previous audio (minus fade region) + crossfaded region + rest of current segment
|
|
result = torch.cat([
|
|
result[:, :-fade_samples],
|
|
crossfaded,
|
|
current_segment[:, fade_samples:]
|
|
], dim=1)
|
|
else:
|
|
# No crossfading, just concatenate on the appropriate dimension
|
|
if result.dim() == 1 and current_segment.dim() == 1:
|
|
result = torch.cat([result, current_segment], dim=0)
|
|
elif result.dim() == 2 and current_segment.dim() == 2:
|
|
result = torch.cat([result, current_segment], dim=1)
|
|
else:
|
|
# Mixed dimensions - normalize to match result's dimensions
|
|
if result.dim() == 1 and current_segment.dim() == 2:
|
|
# Convert current_segment to 1D (average channels if multi-channel)
|
|
if current_segment.shape[0] > 1:
|
|
current_segment = torch.mean(current_segment, dim=0)
|
|
else:
|
|
current_segment = current_segment.squeeze(0)
|
|
result = torch.cat([result, current_segment], dim=0)
|
|
elif result.dim() == 2 and current_segment.dim() == 1:
|
|
# Convert current_segment to 2D to match result
|
|
current_segment = current_segment.unsqueeze(0)
|
|
if result.shape[0] > 1:
|
|
current_segment = current_segment.repeat(result.shape[0], 1)
|
|
result = torch.cat([result, current_segment], dim=1)
|
|
|
|
return result
|
|
|
|
def assemble_stretch_to_fit(self, audio_segments: List[torch.Tensor],
|
|
target_timings: List[Tuple[float, float]],
|
|
fade_duration: float = 0.01) -> torch.Tensor:
|
|
"""
|
|
Assemble audio segments using stretch-to-fit mode - ORIGINAL BEHAVIOR
|
|
Uses the original TimedAudioAssembler from chatterbox
|
|
"""
|
|
if len(audio_segments) != len(target_timings):
|
|
raise ValueError("Number of segments must match number of timings")
|
|
|
|
if not audio_segments:
|
|
return torch.empty(0)
|
|
|
|
# Use the original TimedAudioAssembler from chatterbox - EXACT ORIGINAL BEHAVIOR
|
|
try:
|
|
from engines.chatterbox.audio_timing import TimedAudioAssembler
|
|
|
|
assembler = TimedAudioAssembler(self.sample_rate)
|
|
final_audio, stretcher_method = assembler.assemble_timed_audio(
|
|
audio_segments, target_timings, fade_duration=fade_duration
|
|
)
|
|
# Track which stretcher method was used
|
|
self.last_stretch_method_used = stretcher_method
|
|
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],
|
|
subtitles: List, device) -> torch.Tensor:
|
|
"""
|
|
ORIGINAL: Assemble audio by placing segments at their SRT start times, allowing audible overlaps.
|
|
Silence is implicitly added in gaps. EXACT COPY FROM ORIGINAL LINES 1776-1875
|
|
"""
|
|
if not audio_segments:
|
|
return torch.empty(0) # Return empty tensor if no segments
|
|
|
|
# Use device directly for all operations within this function.
|
|
# device is the consistent target device (e.g., 'cuda' or 'cpu')
|
|
# established during model loading.
|
|
target_device = device
|
|
|
|
# Determine output buffer properties (num_channels, dtype) from the first segment,
|
|
# ensuring it's on the target_device first.
|
|
# audio_segments (which is normalized_segments) should already be on device.
|
|
_first_segment_for_props = audio_segments[0].to(target_device)
|
|
num_channels = _first_segment_for_props.shape[0] if _first_segment_for_props.dim() == 2 else 1
|
|
dtype = _first_segment_for_props.dtype
|
|
|
|
# Calculate total duration needed for the output buffer
|
|
max_end_time = 0.0
|
|
for i, (segment_from_list, subtitle) in enumerate(zip(audio_segments, subtitles)):
|
|
# Ensure segment is on target_device for size calculation
|
|
segment_on_target = segment_from_list.to(target_device)
|
|
segment_end_time = subtitle.start_time + (segment_on_target.size(-1) / self.sample_rate)
|
|
max_end_time = max(max_end_time, segment_end_time)
|
|
|
|
# Ensure the buffer is at least as long as the last subtitle's end time
|
|
if subtitles:
|
|
max_end_time = max(max_end_time, subtitles[-1].end_time)
|
|
|
|
total_samples = int(max_end_time * self.sample_rate)
|
|
|
|
# Initialize output buffer with zeros
|
|
if num_channels == 1:
|
|
output_audio = torch.zeros(total_samples, device=target_device, dtype=dtype)
|
|
else:
|
|
output_audio = torch.zeros(num_channels, total_samples, device=target_device, dtype=dtype)
|
|
|
|
for i, (original_segment_from_list, subtitle) in enumerate(zip(audio_segments, subtitles)):
|
|
# Process segment
|
|
|
|
# Explicitly move the current segment to target_device before any processing
|
|
current_processing_segment = original_segment_from_list.to(target_device)
|
|
|
|
# Ensure normalized_audio matches the channel dimension of output_audio
|
|
# current_processing_segment is now guaranteed to be on target_device.
|
|
normalized_audio_for_add = current_processing_segment
|
|
|
|
if num_channels == 1: # Output is mono (1D)
|
|
if normalized_audio_for_add.dim() == 2:
|
|
# If segment is 2D (e.g., [1, samples] or [channels, samples]), squeeze to 1D
|
|
if normalized_audio_for_add.shape[0] == 1: # If it's [1, samples], just squeeze
|
|
normalized_audio_for_add = normalized_audio_for_add.squeeze(0)
|
|
else: # If it's multi-channel, sum to mono
|
|
normalized_audio_for_add = torch.sum(normalized_audio_for_add, dim=0)
|
|
# If it's already 1D, no change needed
|
|
else: # Output is stereo/multi-channel (2D)
|
|
if normalized_audio_for_add.dim() == 1:
|
|
# If segment is mono (1D), expand to 2D and repeat channels
|
|
normalized_audio_for_add = normalized_audio_for_add.unsqueeze(0).repeat(num_channels, 1)
|
|
elif normalized_audio_for_add.dim() == 2 and normalized_audio_for_add.shape[0] != num_channels:
|
|
# If segment is 2D but has wrong channel count, raise error
|
|
raise RuntimeError(f"Channel mismatch: output buffer has {num_channels} channels, but segment {i} has {normalized_audio_for_add.shape[0]} channels.")
|
|
|
|
# Segment normalized
|
|
# At this point, normalized_audio_for_add is on target_device and has the correct dimensions for addition.
|
|
|
|
start_sample = int(subtitle.start_time * self.sample_rate)
|
|
end_sample_segment = start_sample + normalized_audio_for_add.size(-1)
|
|
|
|
# Resize output_audio if current segment extends beyond current buffer size
|
|
if end_sample_segment > output_audio.size(-1):
|
|
new_total_samples = end_sample_segment
|
|
if num_channels == 1:
|
|
new_output_audio = torch.zeros(new_total_samples, device=target_device, dtype=dtype)
|
|
else:
|
|
new_output_audio = torch.zeros(num_channels, new_total_samples, device=target_device, dtype=dtype)
|
|
|
|
# Copy existing audio to the new larger buffer
|
|
current_len = output_audio.size(-1)
|
|
if num_channels == 1:
|
|
new_output_audio[:current_len] = output_audio
|
|
else:
|
|
new_output_audio[:, :current_len] = output_audio
|
|
output_audio = new_output_audio
|
|
# Buffer resized
|
|
|
|
# Both output_audio and normalized_audio_for_add are on target_device.
|
|
if output_audio.dim() == 1:
|
|
output_audio[start_sample:end_sample_segment] += normalized_audio_for_add
|
|
else:
|
|
output_audio[:, start_sample:end_sample_segment] += normalized_audio_for_add
|
|
|
|
# Segment placed
|
|
|
|
# Assembly complete
|
|
return output_audio
|
|
|
|
def assemble_smart_natural(self, audio_segments: List[torch.Tensor],
|
|
processed_segments: List[torch.Tensor],
|
|
adjustments: List[Dict], subtitles: List, device: str) -> torch.Tensor:
|
|
"""
|
|
Assemble audio using smart natural timing with processed segments
|
|
ORIGINAL FINAL ASSEMBLY LOGIC FROM LINES 1703-1774
|
|
"""
|
|
if not processed_segments:
|
|
return torch.empty(0)
|
|
|
|
# Import required for assembly using import manager
|
|
try:
|
|
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")
|
|
SRTSubtitle = modules.get("SRTSubtitle")
|
|
|
|
if not all([AudioTimingUtils, SRTSubtitle]):
|
|
raise ImportError("Required SRT modules not found")
|
|
except ImportError:
|
|
# Fallback to simple concatenation
|
|
return AudioProcessingUtils.concatenate_audio_segments(processed_segments, "simple")
|
|
|
|
# ORIGINAL FINAL ASSEMBLY LOGIC - Extract mutable_subtitles from adjustments
|
|
# Since we don't have access to the mutable_subtitles from timing engine,
|
|
# we'll reconstruct the final timing from the adjustments
|
|
|
|
final_audio_parts = []
|
|
current_output_time = 0.0
|
|
|
|
for i, segment_audio in enumerate(processed_segments):
|
|
# Get timing info from adjustments
|
|
adj = adjustments[i] if i < len(adjustments) else {}
|
|
|
|
# Determine start time for this segment
|
|
segment_start_time = adj.get('final_srt_start', subtitles[i].start_time if i < len(subtitles) else current_output_time)
|
|
|
|
# Add silence if there's a gap between current_output_time and segment's start_time
|
|
if current_output_time < segment_start_time:
|
|
gap_duration = segment_start_time - current_output_time
|
|
# Add silence gap
|
|
silence = AudioTimingUtils.create_silence(gap_duration, self.sample_rate,
|
|
channels=segment_audio.shape[0] if segment_audio.dim() == 2 else 1,
|
|
device=segment_audio.device)
|
|
final_audio_parts.append(silence)
|
|
current_output_time += gap_duration
|
|
|
|
final_audio_parts.append(segment_audio)
|
|
current_output_time += AudioTimingUtils.get_audio_duration(segment_audio, self.sample_rate)
|
|
|
|
if not final_audio_parts:
|
|
# Handle case where no segments were processed (e.g., immediate interruption)
|
|
if processed_segments:
|
|
return torch.empty(0, device=processed_segments[0].device, dtype=processed_segments[0].dtype)
|
|
else:
|
|
return torch.empty(0)
|
|
|
|
# ORIGINAL DIMENSION NORMALIZATION LOGIC
|
|
# Ensure all parts have the same number of dimensions before concatenating
|
|
target_dim = processed_segments[0].dim() if processed_segments else 1
|
|
target_channels = processed_segments[0].shape[0] if target_dim == 2 else 1
|
|
|
|
normalized_final_audio_parts = []
|
|
target_device_for_concat = device # Use passed device parameter
|
|
|
|
for part in final_audio_parts:
|
|
# Ensure each part is on the target_device before normalization
|
|
part_on_device = part.to(target_device_for_concat)
|
|
|
|
processed_part_for_append = None
|
|
current_part_dim = part_on_device.dim()
|
|
|
|
if current_part_dim == target_dim:
|
|
if target_dim == 2 and part_on_device.shape[0] != target_channels:
|
|
if part_on_device.shape[0] == 1:
|
|
processed_part_for_append = part_on_device.repeat(target_channels, 1)
|
|
else:
|
|
# This branch raises an error, so no tensor is assigned here
|
|
raise RuntimeError(f"Channel mismatch in final assembly: Expected {target_channels} channels, got {part_on_device.shape[0]}")
|
|
else:
|
|
# Part is already on target_device_for_concat and correctly shaped
|
|
processed_part_for_append = part_on_device
|
|
elif current_part_dim == 1 and target_dim == 2: # Mono part, target is multi-channel
|
|
processed_part_for_append = part_on_device.unsqueeze(0).repeat(target_channels, 1)
|
|
elif current_part_dim == 2 and target_dim == 1: # Multi-channel part, target is mono
|
|
processed_part_for_append = torch.sum(part_on_device, dim=0)
|
|
else:
|
|
# This branch raises an error
|
|
raise RuntimeError(f"Dimension mismatch in final assembly: Expected {target_dim}D, got {current_part_dim}D for part with shape {part_on_device.shape}")
|
|
|
|
# Ensure the final tensor to be appended is on the target_device_for_concat.
|
|
if processed_part_for_append is not None:
|
|
normalized_final_audio_parts.append(processed_part_for_append.to(target_device_for_concat))
|
|
|
|
# All parts in normalized_final_audio_parts are on the target device
|
|
return torch.cat(normalized_final_audio_parts, dim=-1)
|
|
|
|
def apply_time_stretching(self, audio: torch.Tensor,
|
|
stretch_factor: float,
|
|
method: str = "auto") -> torch.Tensor:
|
|
"""
|
|
Apply time stretching to audio segment
|
|
|
|
Args:
|
|
audio: Input audio tensor
|
|
stretch_factor: Stretching factor (>1 = slower, <1 = faster)
|
|
method: Stretching method ("auto", "ffmpeg", "phase_vocoder")
|
|
|
|
Returns:
|
|
Time-stretched audio tensor
|
|
"""
|
|
if abs(stretch_factor - 1.0) < 0.01: # 1% tolerance
|
|
return audio
|
|
|
|
try:
|
|
# Try to use chatterbox audio timing utilities
|
|
if method == "auto" or method == "ffmpeg":
|
|
from engines.chatterbox.audio_timing import FFmpegTimeStretcher
|
|
|
|
cache_key = f"ffmpeg_{self.sample_rate}"
|
|
if cache_key not in self.stretcher_cache:
|
|
try:
|
|
self.stretcher_cache[cache_key] = FFmpegTimeStretcher()
|
|
except Exception:
|
|
# Fall back to phase vocoder if FFmpeg not available
|
|
cache_key = f"phase_vocoder_{self.sample_rate}"
|
|
if cache_key not in self.stretcher_cache:
|
|
from engines.chatterbox.audio_timing import PhaseVocoderTimeStretcher
|
|
self.stretcher_cache[cache_key] = PhaseVocoderTimeStretcher()
|
|
|
|
stretcher = self.stretcher_cache[cache_key]
|
|
return stretcher.time_stretch(audio, stretch_factor, self.sample_rate)
|
|
|
|
elif method == "phase_vocoder":
|
|
from engines.chatterbox.audio_timing import PhaseVocoderTimeStretcher
|
|
|
|
cache_key = f"phase_vocoder_{self.sample_rate}"
|
|
if cache_key not in self.stretcher_cache:
|
|
self.stretcher_cache[cache_key] = PhaseVocoderTimeStretcher()
|
|
|
|
stretcher = self.stretcher_cache[cache_key]
|
|
return stretcher.time_stretch(audio, stretch_factor, self.sample_rate)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown time stretching method: {method}")
|
|
|
|
except ImportError:
|
|
# Fallback to simple interpolation if chatterbox modules not available
|
|
return self._simple_time_stretch(audio, stretch_factor)
|
|
|
|
def create_silence(self, duration_seconds: float,
|
|
channels: int = 1,
|
|
device: Optional[torch.device] = None) -> torch.Tensor:
|
|
"""
|
|
Create silence tensor of specified duration
|
|
|
|
Args:
|
|
duration_seconds: Duration in seconds
|
|
channels: Number of audio channels
|
|
device: Target device for tensor
|
|
|
|
Returns:
|
|
Silence tensor
|
|
"""
|
|
try:
|
|
# Use chatterbox utilities if available
|
|
from engines.chatterbox.audio_timing import AudioTimingUtils
|
|
return AudioTimingUtils.create_silence(
|
|
duration_seconds, self.sample_rate, channels, device
|
|
)
|
|
except ImportError:
|
|
# Fallback implementation
|
|
num_samples = int(duration_seconds * self.sample_rate)
|
|
if channels == 1:
|
|
return torch.zeros(num_samples, device=device)
|
|
else:
|
|
return torch.zeros(channels, num_samples, device=device)
|
|
|
|
def _get_audio_duration(self, audio: torch.Tensor) -> float:
|
|
"""Get duration of audio tensor in seconds"""
|
|
try:
|
|
from engines.chatterbox.audio_timing import AudioTimingUtils
|
|
return AudioTimingUtils.get_audio_duration(audio, self.sample_rate)
|
|
except ImportError:
|
|
# Fallback implementation
|
|
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 _basic_stretch_assembly(self, audio_segments: List[torch.Tensor],
|
|
target_timings: List[Tuple[float, float]],
|
|
fade_duration: float) -> torch.Tensor:
|
|
"""
|
|
Basic stretch-to-fit assembly fallback implementation
|
|
"""
|
|
if not audio_segments:
|
|
return torch.empty(0)
|
|
|
|
total_duration = max(end_time for _, end_time in target_timings)
|
|
total_samples = int(total_duration * self.sample_rate)
|
|
|
|
# Create output buffer
|
|
first_segment = audio_segments[0]
|
|
if first_segment.dim() == 1:
|
|
output = torch.zeros(total_samples, device=first_segment.device, dtype=first_segment.dtype)
|
|
else:
|
|
output = torch.zeros(first_segment.size(0), total_samples,
|
|
device=first_segment.device, dtype=first_segment.dtype)
|
|
|
|
# Process each segment
|
|
for segment, (start_time, end_time) in zip(audio_segments, target_timings):
|
|
target_duration = end_time - start_time
|
|
current_duration = self._get_audio_duration(segment)
|
|
|
|
# Apply time stretching if needed
|
|
if abs(current_duration - target_duration) > 0.01:
|
|
stretch_factor = target_duration / current_duration if current_duration > 0 else 1.0
|
|
segment = self._simple_time_stretch(segment, stretch_factor)
|
|
|
|
# Place segment in output
|
|
start_sample = int(start_time * self.sample_rate)
|
|
end_sample = start_sample + segment.size(-1)
|
|
|
|
if output.dim() == 1:
|
|
output[start_sample:min(end_sample, output.size(0))] = segment[:min(segment.size(0), output.size(0) - start_sample)]
|
|
else:
|
|
output[:, start_sample:min(end_sample, output.size(-1))] = segment[:, :min(segment.size(-1), output.size(-1) - start_sample)]
|
|
|
|
return output
|
|
|
|
def _simple_time_stretch(self, audio: torch.Tensor, stretch_factor: float) -> torch.Tensor:
|
|
"""
|
|
Simple time stretching using interpolation (fallback method)
|
|
"""
|
|
if abs(stretch_factor - 1.0) < 0.01:
|
|
return audio
|
|
|
|
original_shape = audio.shape
|
|
if audio.dim() == 1:
|
|
audio = audio.unsqueeze(0)
|
|
|
|
stretched_channels = []
|
|
for channel_idx in range(audio.size(0)):
|
|
channel_audio = audio[channel_idx].cpu().numpy()
|
|
|
|
original_length = len(channel_audio)
|
|
new_length = int(original_length * stretch_factor)
|
|
|
|
# Create new time indices and interpolate
|
|
old_indices = np.linspace(0, original_length - 1, original_length)
|
|
new_indices = np.linspace(0, original_length - 1, new_length)
|
|
stretched = np.interp(new_indices, old_indices, channel_audio)
|
|
|
|
stretched_channels.append(torch.from_numpy(stretched))
|
|
|
|
# Combine channels and restore original shape
|
|
result = torch.stack(stretched_channels, dim=0).to(audio.device)
|
|
if len(original_shape) == 1:
|
|
result = result.squeeze(0)
|
|
|
|
return result
|
|
|
|
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(),
|
|
'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"""
|
|
methods = ['simple_interpolation'] # Always available fallback
|
|
|
|
try:
|
|
from engines.chatterbox.audio_timing import FFmpegTimeStretcher
|
|
FFmpegTimeStretcher() # Test if FFmpeg is available
|
|
methods.append('ffmpeg')
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
from engines.chatterbox.audio_timing import PhaseVocoderTimeStretcher
|
|
methods.append('phase_vocoder')
|
|
except Exception:
|
|
pass
|
|
|
|
return methods
|
|
|
|
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,
|
|
fade_duration: float = 0.01) -> torch.Tensor:
|
|
"""
|
|
Route to the correct assembly method based on timing mode
|
|
|
|
Args:
|
|
audio_segments: Raw audio segments
|
|
subtitles: SRT subtitle objects
|
|
timing_mode: The timing mode to use
|
|
device: Target device
|
|
adjustments: Timing adjustments (for smart_natural mode)
|
|
processed_segments: Processed segments (for smart_natural mode)
|
|
fade_duration: Fade duration for stretch modes
|
|
|
|
Returns:
|
|
Assembled audio tensor
|
|
"""
|
|
if timing_mode == "pad_with_silence":
|
|
# Use overlap assembly - places audio at SRT start times with overlaps allowed
|
|
print(f"🔧 Assembly: Using overlap assembly for {timing_mode} mode")
|
|
return self.assemble_with_overlaps(audio_segments, subtitles, device)
|
|
|
|
elif timing_mode == "smart_natural":
|
|
# Use smart natural assembly with processed segments and adjustments
|
|
if processed_segments is None or adjustments is None:
|
|
raise ValueError("smart_natural mode requires processed_segments and adjustments")
|
|
print(f"🔧 Assembly: Using smart natural assembly for {timing_mode} mode")
|
|
return self.assemble_smart_natural(audio_segments, processed_segments, adjustments, subtitles, device)
|
|
|
|
elif timing_mode == "stretch_to_fit":
|
|
# Use stretch-to-fit assembly
|
|
target_timings = [(sub.start_time, sub.end_time) for sub in subtitles]
|
|
print(f"🔧 Assembly: Using stretch-to-fit assembly for {timing_mode} mode")
|
|
return self.assemble_stretch_to_fit(audio_segments, target_timings, fade_duration)
|
|
|
|
elif timing_mode == "concatenate":
|
|
# Use simple concatenation
|
|
print(f"🔧 Assembly: Using concatenation assembly for {timing_mode} mode")
|
|
return self.assemble_concatenation(audio_segments, fade_duration)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown timing mode: {timing_mode}") |