drphero
7d12f53ca6
- resample edited segments back to original sample rate - normalize 3D tensors after edit post-processing for SRT timing
704 lines
29 KiB
Python
704 lines
29 KiB
Python
"""
|
|
Step Audio EditX Inline Edit Post-Processor
|
|
|
|
Applies Step Audio EditX edits (emotion, style, speed, paralinguistic) to TTS-generated
|
|
audio segments that contain inline edit tags.
|
|
|
|
This post-processor works with ANY TTS engine - the edit tags are stripped before TTS
|
|
generation, then edits are applied as a post-processing step.
|
|
|
|
Usage:
|
|
# After TTS generation returns segments
|
|
segments = EditPostProcessor.process_segments(segments, engine_config)
|
|
|
|
Architecture:
|
|
1. Scans segments for edit_tags metadata
|
|
2. Uses the Step Audio EditX Audio Editor node for editing (reuses existing code)
|
|
3. Applies edits in sequence (emotion → style → speed → paralinguistic)
|
|
4. Returns modified segments for assembly
|
|
"""
|
|
|
|
import torch
|
|
from typing import List, Dict, Any, Optional
|
|
import comfy.model_management as model_management
|
|
|
|
# Cache for Audio Editor node instance
|
|
_audio_editor_node = None
|
|
|
|
# Global settings for inline edit tags (set from ComfyUI settings menu)
|
|
_inline_tag_settings = {
|
|
"precision": "auto",
|
|
"device": "auto",
|
|
"vc_engine": "chatterbox_23lang", # Default VC engine for <restore> tags
|
|
"cosyvoice_variant": "RL" # CosyVoice model variant: "RL" or "standard"
|
|
}
|
|
|
|
# Settings cache file to persist across module reloads
|
|
import os
|
|
import json
|
|
_settings_cache_file = os.path.join(os.path.dirname(__file__), ".inline_tag_settings_cache.json")
|
|
|
|
# Load settings from cache on module import
|
|
try:
|
|
if os.path.exists(_settings_cache_file):
|
|
with open(_settings_cache_file, 'r') as f:
|
|
cached_settings = json.load(f)
|
|
_inline_tag_settings.update(cached_settings)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def set_inline_tag_settings(precision: str = "auto", device: str = "auto", vc_engine: str = "chatterbox_23lang", cosyvoice_variant: str = "RL"):
|
|
"""
|
|
Set global settings for inline edit tag processing and voice restoration.
|
|
Called from the API endpoint when user changes settings in ComfyUI menu.
|
|
|
|
Args:
|
|
precision: Model precision (auto, fp32, fp16, bf16, int8, int4)
|
|
device: Device (auto, cuda, cpu, xpu)
|
|
vc_engine: Voice conversion engine for <restore> tags (chatterbox_23lang, chatterbox, cosyvoice)
|
|
cosyvoice_variant: CosyVoice model variant (RL, standard) - only used when vc_engine is cosyvoice
|
|
"""
|
|
global _inline_tag_settings
|
|
_inline_tag_settings["precision"] = precision
|
|
_inline_tag_settings["device"] = device
|
|
_inline_tag_settings["vc_engine"] = vc_engine
|
|
_inline_tag_settings["cosyvoice_variant"] = cosyvoice_variant
|
|
print(f"🎨 Step Audio EditX inline tags: precision={precision}, device={device}")
|
|
print(f"🔄 Voice restoration engine: {vc_engine}")
|
|
if vc_engine == "cosyvoice":
|
|
print(f" CosyVoice variant: {cosyvoice_variant}")
|
|
|
|
# Save to cache file to persist across module reloads
|
|
try:
|
|
with open(_settings_cache_file, 'w') as f:
|
|
json.dump(_inline_tag_settings, f)
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to save settings cache: {e}")
|
|
|
|
|
|
def get_inline_tag_settings() -> Dict[str, str]:
|
|
"""Get current inline tag settings"""
|
|
return _inline_tag_settings.copy()
|
|
|
|
|
|
def _get_audio_editor_node():
|
|
"""
|
|
Get Step Audio EditX Audio Editor node instance.
|
|
Lazy loads and caches the node for reuse.
|
|
"""
|
|
global _audio_editor_node
|
|
|
|
if _audio_editor_node is not None:
|
|
return _audio_editor_node
|
|
|
|
try:
|
|
import os
|
|
import sys
|
|
import importlib.util
|
|
|
|
# Get project root and construct path to audio editor node
|
|
current_dir = os.path.dirname(__file__) # utils/audio
|
|
utils_dir = os.path.dirname(current_dir) # utils
|
|
project_root = os.path.dirname(utils_dir) # project root
|
|
|
|
audio_editor_path = os.path.join(
|
|
project_root,
|
|
"nodes",
|
|
"step_audio_editx_special",
|
|
"step_audio_editx_audio_editor_node.py"
|
|
)
|
|
|
|
# Load module from file path
|
|
spec = importlib.util.spec_from_file_location("audio_editor_module", audio_editor_path)
|
|
audio_editor_module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(audio_editor_module)
|
|
|
|
StepAudioEditXAudioEditorNode = audio_editor_module.StepAudioEditXAudioEditorNode
|
|
_audio_editor_node = StepAudioEditXAudioEditorNode()
|
|
return _audio_editor_node
|
|
except Exception as e:
|
|
print(f"❌ EditPostProcessor: Failed to load Audio Editor node: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return None
|
|
|
|
|
|
def _resample_waveform(waveform: torch.Tensor, orig_sr: int, target_sr: int) -> torch.Tensor:
|
|
"""
|
|
Resample waveform to target sample rate, preserving shape conventions.
|
|
|
|
Args:
|
|
waveform: Audio tensor (1D, 2D [C,S], or 3D [B,C,S])
|
|
orig_sr: Original sample rate
|
|
target_sr: Target sample rate
|
|
|
|
Returns:
|
|
Resampled waveform tensor with same dimensionality as input (when possible)
|
|
"""
|
|
if orig_sr == target_sr:
|
|
return waveform
|
|
|
|
import torchaudio
|
|
|
|
# Normalize to 2D [C, S] for resampling
|
|
had_batch = False
|
|
if waveform.dim() == 3:
|
|
had_batch = True
|
|
waveform_2d = waveform.squeeze(0)
|
|
elif waveform.dim() == 1:
|
|
waveform_2d = waveform.unsqueeze(0)
|
|
else:
|
|
waveform_2d = waveform
|
|
|
|
resampler = torchaudio.transforms.Resample(orig_freq=orig_sr, new_freq=target_sr)
|
|
resampled = resampler(waveform_2d.cpu())
|
|
|
|
# Restore original dimensionality
|
|
if waveform.dim() == 1:
|
|
return resampled.squeeze(0)
|
|
if had_batch:
|
|
return resampled.unsqueeze(0)
|
|
return resampled
|
|
|
|
|
|
def _validate_audio_duration(audio_tensor: torch.Tensor, sample_rate: int) -> bool:
|
|
"""
|
|
Validate audio duration is within Step Audio EditX limits (0.5-30s).
|
|
|
|
Args:
|
|
audio_tensor: Audio waveform tensor
|
|
sample_rate: Sample rate in Hz
|
|
|
|
Returns:
|
|
True if valid, False otherwise
|
|
"""
|
|
# Get samples count from last dimension
|
|
if audio_tensor.dim() == 1:
|
|
samples = audio_tensor.shape[0]
|
|
elif audio_tensor.dim() == 2:
|
|
samples = audio_tensor.shape[-1]
|
|
elif audio_tensor.dim() == 3:
|
|
samples = audio_tensor.shape[-1]
|
|
else:
|
|
return False
|
|
|
|
duration = samples / sample_rate
|
|
|
|
if duration < 0.5:
|
|
print(f"⚠️ EditPostProcessor: Segment too short ({duration:.2f}s < 0.5s) - skipping edit")
|
|
return False
|
|
if duration > 30.0:
|
|
print(f"⚠️ EditPostProcessor: Segment too long ({duration:.2f}s > 30s) - skipping edit")
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def _apply_edit_via_node(
|
|
editor_node,
|
|
audio_dict: dict,
|
|
transcript: str,
|
|
edit_tag, # EditTag dataclass
|
|
tts_engine_data=None # Optional engine data to reuse pre-loaded engine
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply a single edit using the Audio Editor node.
|
|
|
|
Args:
|
|
editor_node: StepAudioEditXAudioEditorNode instance
|
|
audio_dict: ComfyUI audio dict with 'waveform' and 'sample_rate'
|
|
transcript: Transcript of the audio (required for editing)
|
|
edit_tag: EditTag object with edit_type, value, iterations
|
|
tts_engine_data: Optional engine data to reuse pre-loaded engine (avoids duplicate loading)
|
|
|
|
Returns:
|
|
Edited audio dict (ComfyUI format)
|
|
"""
|
|
edit_type = edit_tag.edit_type
|
|
value = edit_tag.value
|
|
iterations = edit_tag.iterations
|
|
|
|
# Get inline tag settings from global config
|
|
settings = get_inline_tag_settings()
|
|
precision = settings.get("precision", "auto")
|
|
device = settings.get("device", "auto")
|
|
|
|
# Debug: ALWAYS print settings to verify they're being read
|
|
print(f" 🎨 DEBUG: Inline tag settings: precision={precision}, device={device}")
|
|
|
|
# Prepare audio_text for the editor node
|
|
# For paralinguistic, insert tag at position
|
|
if edit_type == "paralinguistic":
|
|
position = edit_tag.position if edit_tag.position is not None else len(transcript)
|
|
audio_text = transcript[:position] + f"<{value}>" + transcript[position:]
|
|
else:
|
|
audio_text = transcript
|
|
|
|
# Map edit parameters to node inputs
|
|
emotion = value if edit_type == "emotion" else "none"
|
|
style = value if edit_type == "style" else "none"
|
|
speed = value if edit_type == "speed" else "none"
|
|
|
|
# Call the Audio Editor node (has progress bar built-in)
|
|
# Pass inline tag settings and engine data for model loading
|
|
edited_audio, _ = editor_node.edit_audio(
|
|
input_audio=audio_dict,
|
|
audio_text=audio_text,
|
|
edit_type=edit_type,
|
|
emotion=emotion,
|
|
style=style,
|
|
speed=speed,
|
|
n_edit_iterations=iterations,
|
|
tts_engine=tts_engine_data, # Reuse pre-loaded engine to avoid duplicate loading
|
|
suppress_progress=True, # We show our own iteration progress
|
|
inline_tag_precision=precision,
|
|
inline_tag_device=device
|
|
)
|
|
|
|
return edited_audio
|
|
|
|
|
|
# Global cache for Voice Changer node (lazy-loaded)
|
|
_cached_vc_node = None
|
|
|
|
def _restore_voice_via_vc(edited_audio_dict, original_voice_dict, iterations=1, language="English", vc_engine=None):
|
|
"""
|
|
Apply voice conversion to restore original voice.
|
|
Reuses existing UnifiedVoiceChangerNode logic with its cache system.
|
|
|
|
Supports multiple VC engines:
|
|
- chatterbox_23lang: ChatterBox Official 23-Lang (default, 23 languages)
|
|
- chatterbox: ChatterBox Classic (English, German, Norwegian)
|
|
- cosyvoice: CosyVoice3 native VC (24kHz output, multilingual)
|
|
|
|
Args:
|
|
edited_audio_dict: Audio after edits {waveform, sample_rate}
|
|
original_voice_dict: Original voice reference {waveform, sample_rate}
|
|
iterations: Number of VC refinement passes (1-5)
|
|
language: Language for VC model (e.g., "English", "Polish", etc.)
|
|
vc_engine: VC engine to use (defaults to global setting)
|
|
|
|
Returns:
|
|
dict: Restored audio {waveform, sample_rate}
|
|
"""
|
|
global _cached_vc_node
|
|
|
|
# Get VC engine from settings if not specified
|
|
if vc_engine is None:
|
|
settings = get_inline_tag_settings()
|
|
vc_engine = settings.get("vc_engine", "chatterbox_23lang")
|
|
|
|
# Lazy-load VC node instance (cache globally like Audio Editor)
|
|
if _cached_vc_node is None:
|
|
import os
|
|
import sys
|
|
import importlib.util
|
|
|
|
# Get project root and construct path to VC node
|
|
current_dir = os.path.dirname(__file__) # utils/audio
|
|
utils_dir = os.path.dirname(current_dir) # utils
|
|
project_root = os.path.dirname(utils_dir) # project root
|
|
|
|
vc_node_path = os.path.join(
|
|
project_root, 'nodes', 'unified', 'voice_changer_node.py'
|
|
)
|
|
|
|
# Load module using importlib
|
|
spec = importlib.util.spec_from_file_location("voice_changer_node", vc_node_path)
|
|
vc_module = importlib.util.module_from_spec(spec)
|
|
sys.modules["voice_changer_node"] = vc_module
|
|
spec.loader.exec_module(vc_module)
|
|
|
|
# Instantiate the VC node
|
|
_cached_vc_node = vc_module.UnifiedVoiceChangerNode()
|
|
|
|
# Create engine config based on selected VC engine
|
|
if vc_engine == "cosyvoice":
|
|
# Get CosyVoice variant from settings
|
|
settings = get_inline_tag_settings()
|
|
variant = settings.get("cosyvoice_variant", "RL")
|
|
|
|
# Map variant to model path
|
|
if variant == "RL":
|
|
model_path = "local:Fun-CosyVoice3-0.5B-RL"
|
|
else:
|
|
model_path = "local:Fun-CosyVoice3-0.5B"
|
|
|
|
engine_config = {
|
|
'engine_type': 'cosyvoice',
|
|
'config': {
|
|
'device': 'auto',
|
|
'mode': 'zero_shot', # Use zero_shot for VC
|
|
'speed': 1.0,
|
|
'model_path': model_path
|
|
}
|
|
}
|
|
elif vc_engine == "chatterbox":
|
|
engine_config = {
|
|
'engine_type': 'chatterbox',
|
|
'config': {
|
|
'device': 'auto',
|
|
'language': language
|
|
}
|
|
}
|
|
else: # chatterbox_23lang (default)
|
|
engine_config = {
|
|
'engine_type': 'chatterbox_official_23lang',
|
|
'config': {
|
|
'device': 'auto',
|
|
'language': language
|
|
}
|
|
}
|
|
|
|
# Call existing VC node logic with refinement passes
|
|
result = _cached_vc_node.convert_voice(
|
|
TTS_engine=engine_config,
|
|
source_audio=edited_audio_dict,
|
|
narrator_target=original_voice_dict,
|
|
refinement_passes=iterations,
|
|
max_chunk_duration=30,
|
|
chunk_method="smart"
|
|
)
|
|
|
|
# Return first element (audio dict) from tuple result
|
|
return result[0] if isinstance(result, tuple) else result
|
|
|
|
|
|
def process_segments(
|
|
segments: List[Dict],
|
|
engine_config: Optional[Dict] = None,
|
|
pre_loaded_engine = None # Pre-loaded Step Audio EditX engine to reuse (avoids loading duplicate)
|
|
) -> List[Dict]:
|
|
"""
|
|
Process all segments, applying Step Audio EditX edits where edit_tags exist.
|
|
|
|
Called ONCE after all TTS generation completes.
|
|
|
|
Args:
|
|
segments: List of segment dicts with 'waveform', 'sample_rate', 'text', 'edit_tags' keys
|
|
engine_config: Optional engine configuration
|
|
pre_loaded_engine: Pre-loaded Step Audio EditX engine from TTS processor (for reuse)
|
|
|
|
Returns:
|
|
List of processed segments (modified in place for efficiency)
|
|
"""
|
|
from utils.text.step_audio_editx_special_tags import (
|
|
parse_edit_tags_with_iterations,
|
|
sort_edit_tags_for_processing
|
|
)
|
|
|
|
# Get inline tag settings once at the start
|
|
settings = get_inline_tag_settings()
|
|
inline_precision = settings.get("precision", "auto")
|
|
inline_device = settings.get("device", "auto")
|
|
print(f"🔧 EditPostProcessor: Using inline tag settings - precision={inline_precision}, device={inline_device}")
|
|
|
|
# Find segments that need editing
|
|
segments_to_edit = []
|
|
for i, segment in enumerate(segments):
|
|
# Check if segment already has edit_tags (pre-parsed)
|
|
if 'edit_tags' in segment and segment['edit_tags']:
|
|
segments_to_edit.append((i, segment, segment['edit_tags']))
|
|
else:
|
|
# Try to extract edit tags from text (fallback)
|
|
text = segment.get('text', '')
|
|
if text:
|
|
clean_text, edit_tags = parse_edit_tags_with_iterations(text)
|
|
if edit_tags:
|
|
# Update segment with clean text and tags
|
|
segment['original_text'] = text
|
|
segment['text'] = clean_text
|
|
segment['edit_tags'] = edit_tags
|
|
segments_to_edit.append((i, segment, edit_tags))
|
|
|
|
if not segments_to_edit:
|
|
return segments
|
|
|
|
print(f"\n🎨 EditPostProcessor: Found {len(segments_to_edit)} segment(s) with edit tags")
|
|
|
|
# Get Audio Editor node
|
|
editor_node = _get_audio_editor_node()
|
|
if editor_node is None:
|
|
print("⚠️ EditPostProcessor: Audio Editor node not available - skipping edits")
|
|
return segments
|
|
|
|
# Prepare engine data for Audio Editor (reuses pre-loaded engine if available)
|
|
tts_engine_data = None
|
|
if pre_loaded_engine is not None:
|
|
tts_engine_data = {
|
|
"engine_type": "step_audio_editx",
|
|
"config": engine_config if engine_config else {},
|
|
"pre_loaded_engine": pre_loaded_engine # Pass engine directly to avoid re-loading
|
|
}
|
|
|
|
# Track segments that need voice restoration (collect for batch processing)
|
|
segments_needing_restore = []
|
|
|
|
# Process each segment with edits (NON-RESTORE edits only)
|
|
for idx, segment, edit_tags in segments_to_edit:
|
|
if model_management.interrupt_processing:
|
|
print("⚠️ EditPostProcessor: Processing interrupted")
|
|
break
|
|
|
|
waveform = segment.get('waveform')
|
|
sample_rate = segment.get('sample_rate', 24000)
|
|
original_sample_rate = sample_rate
|
|
transcript = segment.get('text', '')
|
|
|
|
if waveform is None:
|
|
print(f"⚠️ EditPostProcessor: Segment {idx} has no waveform - skipping")
|
|
continue
|
|
|
|
# Validate duration
|
|
if not _validate_audio_duration(waveform, sample_rate):
|
|
continue
|
|
|
|
# Store original PRE-EDIT audio as voice reference for restore tag
|
|
import torch
|
|
original_audio_dict = {
|
|
'waveform': waveform.clone() if hasattr(waveform, 'clone') else waveform,
|
|
'sample_rate': sample_rate
|
|
}
|
|
segment['_inline_edit_orig_sr'] = original_sample_rate
|
|
|
|
# Show segment being edited with clean formatting
|
|
print(f"\n📝 Segment {idx + 1} - Applying edit tags:")
|
|
print("="*60)
|
|
print(transcript)
|
|
print("="*60)
|
|
|
|
# Sort tags for optimal processing order
|
|
sorted_tags = sort_edit_tags_for_processing(edit_tags)
|
|
|
|
# Group tags by type: non-paralinguistic → paralinguistic → restore
|
|
restore_tags = [t for t in sorted_tags if t.edit_type == 'restore']
|
|
paralinguistic_tags = [t for t in sorted_tags if t.edit_type == 'paralinguistic']
|
|
non_paralinguistic_tags = [t for t in sorted_tags if t.edit_type not in ['paralinguistic', 'restore']]
|
|
|
|
# Create ComfyUI audio dict for the editor node
|
|
current_audio_dict = {
|
|
'waveform': waveform,
|
|
'sample_rate': sample_rate
|
|
}
|
|
|
|
# Track intermediate audio after each edit for restore reference
|
|
iteration_audio_snapshots = {} # iteration_number -> audio_dict
|
|
current_iteration = 0
|
|
|
|
# Apply non-paralinguistic edits first (emotion, style, speed)
|
|
for tag in non_paralinguistic_tags:
|
|
try:
|
|
print(f" 🎨 Applying {tag.edit_type}:{tag.value} ({tag.iterations} iteration{'s' if tag.iterations > 1 else ''})")
|
|
edited_audio_dict = _apply_edit_via_node(
|
|
editor_node=editor_node,
|
|
audio_dict=current_audio_dict,
|
|
transcript=transcript,
|
|
edit_tag=tag,
|
|
tts_engine_data=tts_engine_data
|
|
)
|
|
# Use edited audio as input for next edit
|
|
current_audio_dict = edited_audio_dict
|
|
|
|
# Store snapshot for each iteration
|
|
for i in range(tag.iterations):
|
|
current_iteration += 1
|
|
iteration_audio_snapshots[current_iteration] = {
|
|
'waveform': current_audio_dict['waveform'].clone() if hasattr(current_audio_dict['waveform'], 'clone') else current_audio_dict['waveform'],
|
|
'sample_rate': current_audio_dict['sample_rate']
|
|
}
|
|
except Exception as e:
|
|
print(f" ❌ Failed: {e}")
|
|
|
|
# Apply all paralinguistic tags with proper iteration handling
|
|
if paralinguistic_tags:
|
|
try:
|
|
# Sort by position (descending) for insertion
|
|
sorted_para = sorted(paralinguistic_tags, key=lambda t: t.position or 0, reverse=True)
|
|
max_iterations = max(tag.iterations for tag in paralinguistic_tags)
|
|
|
|
# Print what we're applying
|
|
tag_summary = ", ".join([f"{t.value}:{t.iterations}" for t in paralinguistic_tags])
|
|
print(f" 🎨 Applying paralinguistic: {tag_summary} ({max_iterations} total iteration{'s' if max_iterations > 1 else ''})")
|
|
|
|
# Apply iterations - each iteration includes tags that still need processing
|
|
for iteration in range(1, max_iterations + 1):
|
|
# Determine which tags are active in this iteration
|
|
active_tags = [t for t in sorted_para if iteration <= t.iterations]
|
|
|
|
if not active_tags:
|
|
break # All tags exhausted
|
|
|
|
# Build audio_text with active tags (use angle brackets for Audio Editor)
|
|
audio_text = transcript
|
|
|
|
# Since tags are sorted descending (high to low position), we insert from right to left
|
|
# This means each insertion doesn't affect positions of tags that come before it
|
|
# So we don't need position offset tracking
|
|
|
|
for tag in active_tags:
|
|
position = tag.position if tag.position is not None else len(audio_text)
|
|
position = min(position, len(audio_text))
|
|
|
|
# Check if we need space before tag (if previous char exists and isn't already a space)
|
|
needs_space_before = (position > 0 and audio_text[position - 1] != ' ')
|
|
|
|
# Check if we need space after tag (if next char exists and isn't already a space)
|
|
needs_space_after = (position < len(audio_text) and audio_text[position] != ' ')
|
|
|
|
# Build tag text with appropriate spacing
|
|
space_before = " " if needs_space_before else ""
|
|
space_after = " " if needs_space_after else ""
|
|
tag_text = f"{space_before}<{tag.value}>{space_after}"
|
|
|
|
audio_text = audio_text[:position] + tag_text + audio_text[position:]
|
|
|
|
active_summary = ", ".join([t.value for t in reversed(active_tags)])
|
|
print(f" → Iteration {iteration}/{max_iterations}: [{active_summary}]")
|
|
|
|
# Apply this iteration via Audio Editor node
|
|
edited_audio_dict, _ = editor_node.edit_audio(
|
|
input_audio=current_audio_dict,
|
|
audio_text=audio_text,
|
|
edit_type="paralinguistic",
|
|
emotion="none",
|
|
style="none",
|
|
speed="none",
|
|
n_edit_iterations=1, # Always 1 iteration, we loop ourselves
|
|
tts_engine=tts_engine_data, # Reuse pre-loaded engine to avoid duplicate loading
|
|
suppress_progress=True, # We show our own iteration progress
|
|
inline_tag_precision=inline_precision,
|
|
inline_tag_device=inline_device
|
|
)
|
|
current_audio_dict = edited_audio_dict
|
|
|
|
# Store snapshot after this iteration
|
|
current_iteration += 1
|
|
iteration_audio_snapshots[current_iteration] = {
|
|
'waveform': current_audio_dict['waveform'].clone() if hasattr(current_audio_dict['waveform'], 'clone') else current_audio_dict['waveform'],
|
|
'sample_rate': current_audio_dict['sample_rate']
|
|
}
|
|
|
|
except Exception as e:
|
|
print(f" ❌ Failed: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
# Continue with unmodified audio for this tag
|
|
continue
|
|
|
|
# If segment has restore tags, collect it for batch restoration later
|
|
if restore_tags:
|
|
segments_needing_restore.append({
|
|
'idx': idx,
|
|
'segment': segment,
|
|
'restore_tags': restore_tags,
|
|
'edited_audio': current_audio_dict,
|
|
'original_audio': original_audio_dict,
|
|
'iteration_snapshots': iteration_audio_snapshots # Pass iteration snapshots
|
|
})
|
|
|
|
# Update segment with edited audio (pre-restore)
|
|
if not restore_tags:
|
|
edited_sr = int(current_audio_dict.get('sample_rate', original_sample_rate))
|
|
if edited_sr != original_sample_rate:
|
|
segment['waveform'] = _resample_waveform(
|
|
current_audio_dict['waveform'], edited_sr, original_sample_rate
|
|
)
|
|
segment['sample_rate'] = original_sample_rate
|
|
else:
|
|
segment['waveform'] = current_audio_dict['waveform']
|
|
segment['sample_rate'] = original_sample_rate
|
|
segment.pop('_inline_edit_orig_sr', None)
|
|
else:
|
|
# Store pre-restore audio; final resample happens after restore
|
|
segment['waveform'] = current_audio_dict['waveform']
|
|
|
|
# BATCH RESTORE: Process all voice restorations at once (loads VC model only once)
|
|
if segments_needing_restore:
|
|
print(f"\n🎨🔄 Batch voice restoration for {len(segments_needing_restore)} segment(s)...")
|
|
|
|
for restore_info in segments_needing_restore:
|
|
if model_management.interrupt_processing:
|
|
print("⚠️ Voice restoration interrupted")
|
|
break
|
|
|
|
idx = restore_info['idx']
|
|
segment = restore_info['segment']
|
|
restore_tags = restore_info['restore_tags']
|
|
edited_audio = restore_info['edited_audio']
|
|
original_audio = restore_info['original_audio']
|
|
iteration_snapshots = restore_info['iteration_snapshots']
|
|
|
|
for tag in restore_tags:
|
|
try:
|
|
# Determine which audio to use as reference
|
|
reference_iteration = tag.position # We stored reference iteration in position field
|
|
|
|
if reference_iteration is not None:
|
|
# Use specific iteration snapshot as reference
|
|
if reference_iteration in iteration_snapshots:
|
|
reference_audio = iteration_snapshots[reference_iteration]
|
|
print(f" 🔄 Segment {idx + 1}: Restoring voice using iteration {reference_iteration} as reference ({tag.iterations} VC pass{'es' if tag.iterations > 1 else ''})")
|
|
else:
|
|
print(f" ⚠️ Segment {idx + 1}: Iteration {reference_iteration} not found, using original audio")
|
|
reference_audio = original_audio
|
|
else:
|
|
# Use PRE-EDIT audio as reference (default behavior)
|
|
reference_audio = original_audio
|
|
print(f" 🔄 Segment {idx + 1}: Restoring voice ({tag.iterations} VC pass{'es' if tag.iterations > 1 else ''})")
|
|
|
|
language = engine_config.get('language', 'English') if engine_config else 'English'
|
|
|
|
# Ensure reference audio waveform is 2D for VC compatibility
|
|
ref_wf = reference_audio['waveform']
|
|
if ref_wf.dim() == 1:
|
|
ref_wf = ref_wf.unsqueeze(0)
|
|
elif ref_wf.dim() == 3:
|
|
ref_wf = ref_wf.squeeze(0)
|
|
|
|
reference_audio_fixed = {
|
|
'waveform': ref_wf,
|
|
'sample_rate': reference_audio['sample_rate']
|
|
}
|
|
|
|
# Restore voice using selected reference
|
|
restored_audio = _restore_voice_via_vc(
|
|
edited_audio_dict=edited_audio,
|
|
original_voice_dict=reference_audio_fixed,
|
|
iterations=tag.iterations,
|
|
language=language
|
|
)
|
|
edited_audio = restored_audio # Chain for multiple restore tags
|
|
|
|
except Exception as e:
|
|
print(f" ❌ Voice restoration failed: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
# Continue with current audio (skip restoration)
|
|
|
|
# Update segment with restored audio, resampling if needed
|
|
orig_sr = segment.pop('_inline_edit_orig_sr', segment.get('sample_rate', 24000))
|
|
edited_sr = int(edited_audio.get('sample_rate', orig_sr))
|
|
if edited_sr != orig_sr:
|
|
segment['waveform'] = _resample_waveform(edited_audio['waveform'], edited_sr, orig_sr)
|
|
segment['sample_rate'] = orig_sr
|
|
else:
|
|
segment['waveform'] = edited_audio['waveform']
|
|
segment['sample_rate'] = orig_sr
|
|
|
|
print(f"\n✅ EditPostProcessor: Completed edit post-processing")
|
|
return segments
|
|
|
|
|
|
def cleanup():
|
|
"""Clean up cached engine."""
|
|
global _step_audio_engine, _step_audio_loaded
|
|
|
|
if _step_audio_engine is not None:
|
|
try:
|
|
_step_audio_engine.to("cpu")
|
|
except Exception:
|
|
pass
|
|
_step_audio_engine = None
|
|
_step_audio_loaded = False
|