""" 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 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 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