Create inference_cli.py

inference_cli.py

@@ -0,0 +1,895 @@
+import argparse
+import codecs
+import re
+import tempfile
+from pathlib import Path
+import logging
+import numpy as np
+import soundfile as sf
+import tomli
+import torch
+import torchaudio
+from tqdm import tqdm
+from einops import rearrange
+from pydub import AudioSegment, silence
+from transformers import pipeline
+from huggingface_hub import login
+from cached_path import cached_path
+import matplotlib.pyplot as plt # Needed for save_spectrogram
+
+# --- Import Model Architectures ---
+# !! Ensure these models are defined in your project's 'model' module !!
+try:
+    from model import UNetT, DiT
+except ImportError:
+    print("Warning: Could not import UNetT, DiT from 'model'. Using placeholders.")
+    # Placeholder classes if import fails (script might not work correctly)
+    class MockModel:
+        def __init__(self, *args, **kwargs): pass
+        def to(self, device): return self
+        def eval(self): pass
+        def sample(self, *args, **kwargs):
+            duration = kwargs.get('duration', 500); mel_dim = 100
+            return torch.randn(1, duration, mel_dim), None
+        @property
+        def device(self): return torch.device("cpu")
+    DiT = MockModel
+    UNetT = MockModel
+
+# --- Import/Define Utility Functions ---
+
+from tokenizers import Tokenizer
+from phonemizer import phonemize
+
+# --- Functions copied/adapted from app.py ---
+
+# Function to load vocoder (from app.py context, may need adjustment)
+def load_vocoder(device='cpu'):
+    """Loads the Vocos vocoder."""
+    print("Loading Vocos vocoder (charactr/vocos-mel-24khz)...")
+    try:
+        # Ensure vocos library is installed: pip install vocos
+        from vocos import Vocos
+        # Determine torch dtype based on device for potential efficiency
+        # Note: Vocos might internally cast, but being explicit can help.
+        # Using float32 as a safe default unless on CUDA where float16 might work.
+        vocos_dtype = torch.float16 if str(device) == 'cuda' else torch.float32
+
+        vocos_model = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device)
+        # Cast to appropriate dtype if needed, although Vocos might handle this.
+        # vocos_model = vocos_model.to(dtype=vocos_dtype) # Optional casting
+        vocos_model.eval()
+        print("Vocos vocoder loaded successfully.")
+        return vocos_model
+    except ImportError:
+        print("Error: 'vocos' library not found. Please install it: pip install vocos")
+        raise
+    except Exception as e:
+        print(f"Error loading Vocos model: {e}")
+        raise
+
+# Function to remove silence from edges (from app.py)
+def remove_silence_edges(aseg):
+    """Removes silence from the beginning and end of an AudioSegment."""
+    print("Removing silence from audio edges...")
+    start_trim = silence.detect_leading_silence(aseg)
+    end_trim = silence.detect_leading_silence(aseg.reverse())
+    duration = len(aseg)
+    trimmed_aseg = aseg[start_trim:duration-end_trim]
+    print(f"Removed {start_trim}ms from start, {end_trim}ms from end.")
+    return trimmed_aseg
+
+# Function to save spectrogram (from app.py)
+def save_spectrogram(spectrogram, file_path):
+    """Saves a spectrogram visualization to a file."""
+    if spectrogram is None:
+        print("Spectrogram data is None, cannot save.")
+        return
+    try:
+        print(f"Saving spectrogram to {file_path}...")
+        plt.figure(figsize=(10, 4))
+        plt.imshow(spectrogram, aspect='auto', origin='lower', cmap='viridis')
+        plt.colorbar(label='Mel power')
+        plt.xlabel('Frames')
+        plt.ylabel('Mel bins')
+        plt.title('Generated Mel Spectrogram')
+        plt.tight_layout()
+        plt.savefig(file_path)
+        plt.close() # Close the figure to free memory
+        print("Spectrogram saved.")
+    except Exception as e:
+        print(f"Error saving spectrogram: {e}")
+
+# Helper function to load checkpoint (from app.py, slightly modified for CLI)
+def load_checkpoint(model, ckpt_path, device, use_ema=False):
+    """Loads model weights from a checkpoint file (.pt or .safetensors)."""
+    print(f"Loading checkpoint from {ckpt_path}...")
+    try:
+        if ckpt_path.endswith(".safetensors"):
+            # Ensure safetensors is installed: pip install safetensors
+            from safetensors.torch import load_file
+            state_dict = load_file(ckpt_path, device="cpu")
+        elif ckpt_path.endswith(".pt"):
+            state_dict = torch.load(ckpt_path, map_location="cpu")
+        else:
+             raise ValueError(f"Unsupported checkpoint format: {ckpt_path}. Must be .pt or .safetensors")
+
+        # Standardize state_dict format (e.g., remove 'state_dict' key if present)
+        if "state_dict" in state_dict:
+            state_dict = state_dict["state_dict"]
+
+        # Handle EMA weights
+        ema_key_prefix = "ema_model." # Adjust if your EMA keys have a different prefix
+        final_state_dict = {}
+        has_ema = any(k.startswith(ema_key_prefix) for k in state_dict.keys())
+
+        if use_ema:
+            if has_ema:
+                print("Attempting to load EMA weights.")
+                ema_state_dict = {k[len(ema_key_prefix):]: v for k, v in state_dict.items() if k.startswith(ema_key_prefix)}
+                if ema_state_dict:
+                    final_state_dict = ema_state_dict
+                    print("Using EMA weights.")
+                else:
+                    # This case shouldn't happen if has_ema is true, but as a safeguard:
+                    print("Warning: EMA weights requested but none found starting with prefix. Using regular weights.")
+                    final_state_dict = {k: v for k, v in state_dict.items() if not k.startswith(ema_key_prefix)}
+            else:
+                print("Warning: EMA weights requested but no keys found with EMA prefix. Using regular weights.")
+                final_state_dict = state_dict # Use the original dict if no EMA keys exist
+        else:
+            print("Loading non-EMA weights.")
+            # Filter out EMA weights if they exist and we explicitly don't want them
+            final_state_dict = {k: v for k, v in state_dict.items() if not k.startswith(ema_key_prefix)}
+
+
+        # Load into model, handling potential 'module.' prefix from DDP
+        model_state_dict = model.state_dict()
+        processed_state_dict = {}
+        for k, v in final_state_dict.items():
+            if k.startswith("module."):
+                k_proc = k[len("module."):]
+            else:
+                k_proc = k
+
+            if k_proc in model_state_dict:
+                if model_state_dict[k_proc].shape == v.shape:
+                    processed_state_dict[k_proc] = v
+                else:
+                    print(f"Warning: Shape mismatch for key {k_proc}. Checkpoint: {v.shape}, Model: {model_state_dict[k_proc].shape}. Skipping.")
+            # else: # Optional: Log unexpected keys
+            #     print(f"Warning: Key {k_proc} from checkpoint not found in model. Skipping.")
+
+        missing_keys, unexpected_keys = model.load_state_dict(processed_state_dict, strict=False)
+
+        if missing_keys:
+            print(f"Warning: Missing keys in model not found in checkpoint: {missing_keys}")
+        if unexpected_keys:
+            # This should ideally be empty if we filter correctly, but good to check.
+            print(f"Warning: Unexpected keys (should not happen with filtering): {unexpected_keys}")
+
+        print(f"Checkpoint loaded successfully from {ckpt_path}")
+
+    except FileNotFoundError:
+        print(f"Error: Checkpoint file not found at {ckpt_path}")
+        raise
+    except Exception as e:
+        print(f"Error loading checkpoint from {ckpt_path}: {e}")
+        raise # Re-raise the exception
+
+    model.eval()
+    return model.to(device)
+
+# Primary model loading function (from app.py)
+def load_custom(model_cls, model_cfg, ckpt_path: str, vocab_size: int, device='cpu', use_ema=True):
+    """Loads a custom TTS model (DiT or UNetT) with specified config and checkpoint."""
+    ckpt_path = ckpt_path.strip()
+
+    if ckpt_path.startswith("hf://"):
+        print(f"Downloading checkpoint from Hugging Face Hub: {ckpt_path}")
+        try:
+            ckpt_path = str(cached_path(ckpt_path))
+            print(f"Checkpoint downloaded to: {ckpt_path}")
+        except Exception as e:
+            print(f"Error downloading checkpoint {ckpt_path}: {e}")
+            raise
+
+    if not Path(ckpt_path).exists():
+         raise FileNotFoundError(f"Checkpoint file not found: {ckpt_path}")
+
+    # Ensure necessary config keys are present (add defaults if missing)
+    if 'mel_dim' not in model_cfg:
+        model_cfg['mel_dim'] = 100 # Default mel channels
+        print(f"Warning: 'mel_dim' not in model_cfg, defaulting to {model_cfg['mel_dim']}")
+    if 'text_num_embeds' not in model_cfg:
+         model_cfg['text_num_embeds'] = vocab_size
+         print(f"Setting 'text_num_embeds' in model_cfg to vocab size: {vocab_size}")
+
+    print(f"Instantiating model: {model_cls.__name__} with config: {model_cfg}")
+    try:
+        model = model_cls(**model_cfg).to(device) # Instantiate the model
+    except Exception as e:
+        print(f"Error instantiating model {model_cls.__name__} with config {model_cfg}: {e}")
+        raise
+
+    # Load weights using the helper function
+    model = load_checkpoint(model, ckpt_path, device, use_ema=use_ema)
+    model.eval() # Ensure model is in evaluation mode
+    return model
+
+
+# Text chunking function (from app.py)
+def chunk_text(text, max_chars):
+    """
+    Splits the input text into chunks based on punctuation and length limits.
+    (Copied from previous answer, assumed correct)
+    """
+    if not isinstance(text, str):
+         print("Warning: Input to chunk_text is not a string. Returning empty list.")
+         return []
+
+    if max_chars > 135:
+        print(f"Warning: Calculated max_chars ({max_chars}) > 135. Capping at 135.")
+        max_chars = 135
+    if max_chars < 50:
+        print(f"Warning: Calculated max_chars ({max_chars}) < 50. Setting to 50.")
+        max_chars = 50
+
+    split_after_space_chars = max_chars + int(max_chars * 0.33)
+    chunks = []
+    current_chunk = ""
+
+    # Split the text into sentences based on punctuation followed by whitespace
+    sentences = re.split(r"(?<=[;:,.!?])\s+|(?<=[；：，。！？])\s*", text) # Added \s* after CJK punc
+
+    for sentence in sentences:
+        sentence = sentence.strip()
+        if not sentence:
+            continue
+
+        # Estimate potential length increase due to space
+        estimated_len = len(current_chunk) + len(sentence) + (1 if current_chunk else 0)
+
+        if estimated_len <= max_chars:
+            current_chunk += (" " + sentence) if current_chunk else sentence
+        else:
+            # Process the current_chunk if adding the new sentence exceeds max_chars
+            while len(current_chunk) > split_after_space_chars:
+                split_index = current_chunk.rfind(" ", 0, split_after_space_chars)
+                if split_index == -1: split_index = split_after_space_chars
+                chunks.append(current_chunk[:split_index].strip())
+                current_chunk = current_chunk[split_index:].strip()
+
+            if current_chunk:
+                chunks.append(current_chunk)
+
+            # Start new chunk, handle if sentence itself is too long
+            while len(sentence) > split_after_space_chars:
+                 split_index = sentence.rfind(" ", 0, split_after_space_chars)
+                 if split_index == -1: split_index = split_after_space_chars
+                 chunks.append(sentence[:split_index].strip())
+                 sentence = sentence[split_index:].strip()
+            current_chunk = sentence
+
+    # Handle the last chunk
+    while len(current_chunk) > split_after_space_chars:
+        split_index = current_chunk.rfind(" ", 0, split_after_space_chars)
+        if split_index == -1: split_index = split_after_space_chars
+        chunks.append(current_chunk[:split_index].strip())
+        current_chunk = current_chunk[split_index:].strip()
+
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
+    return [c for c in chunks if c] # Filter empty chunks
+
+
+# Text to IPA function (from app.py)
+def text_to_ipa(text, language):
+    """Converts text to IPA using phonemizer with espeak backend."""
+    if not isinstance(text, str) or not text.strip():
+         print(f"Warning: Invalid input text for IPA conversion: {text}")
+         return "" # Return empty string for invalid input
+    try:
+        # Ensure phonemizer is installed: pip install phonemizer
+        # Ensure espeak-ng is installed: sudo apt-get install espeak-ng (or equivalent)
+        ipa_text = phonemize(
+            text,
+            language=language,
+            backend='espeak',
+            strip=False, # Keep punctuation
+            preserve_punctuation=True,
+            with_stress=True,
+            language_switch='remove-flags', # Use this instead of regex removal
+            njobs=1 # Set njobs=1 for potentially better stability/simpler debugging
+        )
+        # Specific removals (might be redundant with remove-flags, but kept for consistency)
+        ipa_text = re.sub(r'tʃˈaɪniːzlˈe̞tə', '', ipa_text)
+        ipa_text = re.sub(r'tʃˈaɪniːzɭˈetə', '', ipa_text)
+        ipa_text = re.sub(r'dʒˈapəniːzlˈe̞tə', '', ipa_text)
+        ipa_text = re.sub(r'dʒˈapəniːzɭˈetə', '', ipa_text)
+
+        ipa_text = ipa_text.strip()
+        # Replace multiple spaces with single space
+        ipa_text = re.sub(r'\s+', ' ', ipa_text)
+
+        print(f"Text: '{text}' | Lang: {language} | IPA: '{ipa_text}'")
+        return ipa_text
+    except ImportError:
+         print("Error: 'phonemizer' library not found. Please install it: pip install phonemizer")
+         raise
+    except Exception as e:
+        # Check if it's an espeak error (often happens if language is unsupported)
+        if "espeak" in str(e).lower():
+             print(f"Error: Espeak backend failed for language '{language}'. Is the language code valid and espeak-ng installed/supporting it?")
+             print(f"  Original error: {e}")
+        else:
+             print(f"Error phonemizing text: '{text}' with language '{language}'. Error: {e}")
+        # Decide how to handle error
+        raise ValueError(f"Phonemization failed for '{text}' ({language})") from e
+
+
+# --- End of functions from app.py ---
+
+# --- Argument Parser Setup ---
+# (Parser definition remains the same as previous refactored version)
+parser = argparse.ArgumentParser(
+    prog="python3 inference-cli.py",
+    description="Commandline interface for F5/E2 TTS.",
+)
+parser.add_argument(
+    "-c", "--config", type=str, default="inference-cli.toml",
+    help="Path to configuration file (TOML format). Default: inference-cli.toml"
+)
+# --- Arguments overriding config or providing inputs ---
+parser.add_argument( "--ckpt_path", type=str, default=None, help="Path or Hub ID (hf://...) to the TTS model checkpoint (.pt/.safetensors). Overrides config.")
+parser.add_argument( "--ref_audio", type=str, default=None, help="Path to the reference audio file (<10s recommended). Overrides config.")
+parser.add_argument( "--ref_text", type=str, default=None, help="Reference text. If omitted, Whisper transcription is used. Overrides config.")
+parser.add_argument( "--gen_text", type=str, default=None, help="Text to synthesize. Overrides config.")
+parser.add_argument( "--gen_file", type=str, default=None, help="File containing text to synthesize (overrides --gen_text and config).")
+parser.add_argument( "--output_dir", type=str, default=None, help="Directory to save output audio and spectrogram. Overrides config.")
+parser.add_argument( "--output_name", type=str, default="out", help="Base name for output files (e.g., 'my_speech' -> my_speech.wav, my_speech.png). Default: out.")
+# --- Parameter Arguments ---
+parser.add_argument( "--ref_language", type=str, default=None, help="Language code for reference text phonemization (e.g., 'en-us', 'pl', 'de'). Overrides config.")
+parser.add_argument( "--language", type=str, default=None, help="Language code for generated text phonemization (e.g., 'en-us', 'pl', 'de'). Overrides config.")
+parser.add_argument( "--speed", type=float, default=None, help="Speech speed multiplier. Overrides config.")
+parser.add_argument( "--nfe", type=int, default=None, help="Number of function evaluations (sampling steps). Overrides config.")
+parser.add_argument( "--cfg", type=float, default=None, help="Classifier-Free Guidance strength. Overrides config.")
+parser.add_argument( "--sway", type=float, default=None, help="Sway sampling coefficient. Overrides config.")
+parser.add_argument( "--cross_fade", type=float, default=None, help="Cross-fade duration between batches (seconds). Overrides config.")
+parser.add_argument( "--remove_silence", action=argparse.BooleanOptionalAction, default=None, help="Enable/disable final silence removal. Overrides config.")
+parser.add_argument( "--hf_token", type=str, default=None, help="Hugging Face API token (for downloading private models/checkpoints).")
+parser.add_argument( "--tokenizer_path", type=str, default=None, help="Path to the tokenizer.json file. Overrides config.")
+parser.add_argument( "--device", type=str, default=None, help="Device to use ('cuda', 'cpu', 'mps'). Auto-detects if not set.")
+parser.add_argument( "--dtype", type=str, default=None, help="Data type ('float16', 'bfloat16', 'float32'). Auto-selects if not set.")
+
+args = parser.parse_args()
+
+# --- Load Configuration ---
+config = {}
+if Path(args.config).exists():
+    try:
+        with open(args.config, "rb") as f:
+            config = tomli.load(f)
+        print(f"Loaded configuration from {args.config}")
+    except Exception as e:
+        print(f"Warning: Could not load config file {args.config}. Error: {e}")
+else:
+    print(f"Warning: Config file {args.config} not found. Using defaults and CLI args.")
+
+# --- Determine Parameters (CLI > Config > Defaults) ---
+# (Parameter determination remains the same)
+ckpt_path = args.ckpt_path or config.get("ckpt_path", "hf://Gregniuki/F5-tts_English_German_Polish/multi3/model_900000.pt")
+ref_audio_path = args.ref_audio or config.get("ref_audio")
+ref_text = args.ref_text if args.ref_text is not None else config.get("ref_text", "")
+gen_text = args.gen_text or config.get("gen_text")
+gen_file = args.gen_file or config.get("gen_file")
+output_dir = Path(args.output_dir or config.get("output_dir", "."))
+output_name = args.output_name or config.get("output_name", "out")
+
+ref_language = args.ref_language or config.get("ref_language", "en-us")
+language = args.language or config.get("language", "en-us")
+speed = args.speed if args.speed is not None else config.get("speed", 1.0)
+nfe_step = args.nfe if args.nfe is not None else config.get("nfe", 32)
+cfg_strength = args.cfg if args.cfg is not None else config.get("cfg", 2.0)
+sway_sampling_coef = args.sway if args.sway is not None else config.get("sway", -1.0)
+cross_fade_duration = args.cross_fade if args.cross_fade is not None else config.get("cross_fade", 0.15)
+remove_silence_flag = args.remove_silence if args.remove_silence is not None else config.get("remove_silence", False)
+hf_token = args.hf_token or config.get("hf_token")
+tokenizer_path = args.tokenizer_path or config.get("tokenizer_path", "data/Emilia_ZH_EN_pinyin/tokenizer.json")
+
+
+# --- Validate Required Arguments ---
+if not ckpt_path: raise ValueError("Missing required argument/config: --ckpt_path")
+if not ref_audio_path: raise ValueError("Missing required argument/config: --ref_audio")
+if not gen_text and not gen_file: raise ValueError("Missing required argument/config: --gen_text or --gen_file")
+
+# --- Read gen_text from file if provided ---
+if gen_file:
+    try:
+        with codecs.open(gen_file, "r", "utf-8") as f: gen_text = f.read()
+        print(f"Loaded generation text from {gen_file}")
+    except Exception as e: raise ValueError(f"Error reading generation text file {gen_file}: {e}")
+
+# --- Setup Device and Dtype ---
+# (Device/Dtype setup remains the same)
+cli_device = args.device or config.get("device")
+if cli_device:
+    device = torch.device(cli_device)
+else:
+    device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+
+cli_dtype = args.dtype or config.get("dtype")
+if cli_dtype:
+    dtype_map = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
+    if cli_dtype in dtype_map: dtype = dtype_map[cli_dtype]
+    else: raise ValueError(f"Unsupported dtype: {cli_dtype}")
+else:
+    if device.type == "cuda": dtype = torch.float16
+    elif device.type == "cpu" and hasattr(torch.backends, 'cpu') and torch.backends.cpu.supports_bfloat16: dtype = torch.bfloat16
+    else: dtype = torch.float32
+
+print(f"Using device: {device}, dtype: {dtype}")
+
+# --- Hugging Face Login ---
+if hf_token:
+    print("Logging in to Hugging Face Hub...")
+    try:
+        login(token=hf_token)
+        print("Logged in successfully.")
+    except Exception as e:
+        print(f"Warning: Hugging Face login failed: {e}")
+
+
+# --- Create Output Directory ---
+output_dir.mkdir(parents=True, exist_ok=True)
+wave_path = output_dir / f"{output_name}.wav"
+spectrogram_path = output_dir / f"{output_name}.png"
+
+# --- Load Models and Tokenizer ---
+print("Loading Tokenizer...")
+try:
+    if not Path(tokenizer_path).exists():
+        raise FileNotFoundError(f"Tokenizer file not found: {tokenizer_path}")
+    tokenizer = Tokenizer.from_file(tokenizer_path)
+    vocab_size = tokenizer.get_vocab_size()
+    print(f"Tokenizer loaded successfully. Vocab size: {vocab_size}")
+except Exception as e:
+    raise ValueError(f"Error loading tokenizer from {tokenizer_path}: {e}")
+
+print("Loading Vocoder...")
+# Pass device to load_vocoder
+vocos = load_vocoder(device=device) # Already includes .to(device).eval()
+
+print("Loading ASR Model (Whisper)...")
+try:
+    whisper_dtype = torch.float16 if device.type == 'cuda' else torch.float32
+    # Reduce default batch_size for Whisper CLI use
+    pipe = pipeline(
+        "automatic-speech-recognition",
+        model="openai/whisper-large-v3-turbo",
+        torch_dtype=whisper_dtype,
+        device=device,
+        model_kwargs={"attn_implementation": "sdpa"} # Use SDPA if available
+    )
+    print("Whisper model loaded.")
+except Exception as e:
+    print(f"Warning: Could not load Whisper ASR model: {e}. Transcription will not be available.")
+    pipe = None
+
+print("Loading TTS Model...")
+# --- Determine Model Class and Config ---
+# Example configs (ensure they match your actual model requirements)
+F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
+E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4) # Add mel_dim/text_num_embeds if needed by class
+
+# Heuristic to determine model class (improve if needed)
+if "E2TTS" in ckpt_path or "UNetT" in ckpt_path:
+     model_cls = UNetT
+     model_cfg = E2TTS_model_cfg
+     print(f"Assuming E2-TTS (UNetT) architecture for {ckpt_path}.")
+elif "F5TTS" in ckpt_path or "DiT" in ckpt_path:
+     model_cls = DiT
+     model_cfg = F5TTS_model_cfg
+     print(f"Assuming F5-TTS (DiT) architecture for {ckpt_path}.")
+else:
+     # Default or raise error if model type cannot be inferred
+     print(f"Warning: Cannot infer model type from '{ckpt_path}'. Defaulting to DiT/F5TTS.")
+     model_cls = DiT
+     model_cfg = F5TTS_model_cfg
+
+
+try:
+    # Pass vocab_size needed by load_custom
+    ema_model = load_custom(model_cls, model_cfg, ckpt_path, vocab_size=vocab_size, device=device, use_ema=True)
+    # Ensure model is using the target runtime dtype
+    ema_model = ema_model.to(dtype=dtype)
+    print(f"TTS Model loaded successfully ({model_cls.__name__}).")
+except Exception as e:
+    print(f"Critical Error: Failed to load TTS model from {ckpt_path}: {e}")
+    raise
+
+# --- Settings from app.py ---
+target_sample_rate = 24000
+n_mel_channels = model_cfg.get('mel_dim', 100) # Use mel_dim from config if available
+hop_length = 256
+target_rms = 0.1
+
+# --- Main Inference Logic ---
+
+def infer_batch(ref_audio_tuple, ref_text_ipa, gen_text_ipa_batches,
+                ema_model, vocos, tokenizer,
+                remove_silence_post, cross_fade_duration,
+                nfe_step, cfg_strength, sway_sampling_coef, speed,
+                target_sample_rate, hop_length, target_rms, device, dtype):
+    """
+    Generates audio batches based on reference and text inputs.
+    (Function body remains the same as previous refactored version)
+    """
+    audio, sr = ref_audio_tuple
+    audio = audio.to(device, dtype=dtype)
+
+    # Preprocess reference audio (resample, RMS norm)
+    if audio.shape[0] > 1: audio = torch.mean(audio, dim=0, keepdim=True)
+    current_rms = torch.sqrt(torch.mean(torch.square(audio)))
+    rms_applied_factor = 1.0 # Track scaling factor applied to ref
+    if current_rms < target_rms and current_rms > 1e-5: # Add safety check for near-silent audio
+        print(f"Reference audio RMS ({current_rms:.3f}) below target ({target_rms}). Normalizing.")
+        rms_applied_factor = target_rms / current_rms
+        audio = audio * rms_applied_factor
+    elif current_rms <= 1e-5:
+         print("Warning: Reference audio is near silent. Skipping RMS normalization.")
+    else:
+        print(f"Reference audio RMS ({current_rms:.3f}) >= target ({target_rms}). No normalization.")
+
+    if sr != target_sample_rate:
+        print(f"Resampling reference audio from {sr} Hz to {target_sample_rate} Hz.")
+        resampler = torchaudio.transforms.Resample(sr, target_sample_rate).to(device)
+        audio = resampler(audio)
+
+    ref_audio_len_frames = audio.shape[-1] // hop_length
+    print(f"Reference audio length: {audio.shape[-1]/target_sample_rate:.2f}s ({ref_audio_len_frames} frames)")
+
+    generated_waves = []
+    spectrograms = []
+
+    progress_bar = tqdm(gen_text_ipa_batches, desc="Generating Batches")
+    for i, gen_text_ipa in enumerate(progress_bar):
+        progress_bar.set_postfix({"Batch": f"{i+1}/{len(gen_text_ipa_batches)}"})
+
+        # Combine reference and generated IPA text
+        combined_ipa_text = ref_text_ipa + " " + gen_text_ipa
+        # print(f"Batch {i+1} Combined IPA: {combined_ipa_text}") # Debug
+
+        # Tokenize
+        try:
+            # Tokenizer expects single string or list of strings
+            encoding = tokenizer.encode(combined_ipa_text)
+            tokens = encoding.ids
+            token_str = encoding.tokens # For logging/debug
+
+            # --- Model Input Formatting ---
+            # Check how your specific model's `sample` method expects the 'text' input.
+            # Option 1 (like app.py): String of space-separated tokens
+            # token_input_string = ' '.join(map(str, token_str))
+            # final_text_list = [token_input_string]
+
+            # Option 2: List of token IDs (might be more common)
+            # final_text_list = [tokens] # List containing the list/tensor of IDs
+
+            # Option 3: Tensor of token IDs (check model docs)
+            # Assuming model expects Option 1 based on app.py:
+            token_input_string = ' '.join(map(str, token_str))
+            final_text_list = [token_input_string]
+            # print(f"Batch {i+1} Input Text List for Model: {final_text_list}")
+
+        except Exception as e:
+            print(f"Error tokenizing batch {i+1}: '{combined_ipa_text}'. Error: {e}")
+            continue
+
+        # Calculate duration
+        ref_ipa_len = len(ref_text_ipa)
+        gen_ipa_len = len(gen_text_ipa)
+        if ref_ipa_len == 0: ref_ipa_len = 1 # Avoid division by zero
+
+        duration_frames = ref_audio_len_frames + int(((ref_audio_len_frames / ref_ipa_len) * gen_ipa_len) / speed)
+        min_duration_frames = max(10, target_sample_rate // hop_length // 4) # Shorter min duration (e.g. 0.25s)
+        duration_frames = max(min_duration_frames, duration_frames)
+        max_duration_frames = 40 * target_sample_rate // hop_length # Increase max duration slightly?
+        if duration_frames > max_duration_frames:
+            print(f"Warning: Calculated duration {duration_frames} frames exceeds max {max_duration_frames}. Capping.")
+            duration_frames = max_duration_frames
+
+        # print(f"Batch {i+1}: Duration={duration_frames} frames")
+
+        # Inference
+        try:
+            with torch.inference_mode():
+                cond_audio = audio.to(ema_model.device, dtype=dtype) # Match model device/dtype
+                # print(f"Model device: {ema_model.device}, Cond audio device: {cond_audio.device}, dtype: {cond_audio.dtype}")
+
+                generated_mel, _ = ema_model.sample(
+                    cond=cond_audio,
+                    text=final_text_list, # Pass formatted text input
+                    duration=duration_frames,
+                    steps=nfe_step,
+                    cfg_strength=cfg_strength,
+                    sway_sampling_coef=sway_sampling_coef,
+                )
+
+            # Process generated mel
+            generated_mel = generated_mel.to(device, dtype=dtype) # Back to main device/dtype
+            generated_mel = generated_mel[:, ref_audio_len_frames:, :]
+            generated_mel_spec = rearrange(generated_mel, "1 n d -> 1 d n")
+
+            # Vocoding
+            # Vocos usually expects float32
+            vocos_input_mel = generated_mel_spec.to(vocos.device, dtype=torch.float32)
+            generated_wave = vocos.decode(vocos_input_mel)
+            generated_wave = generated_wave.to(device, dtype=torch.float32)
+
+            # Adjust RMS (Scale generated audio by the same factor applied to reference)
+            generated_wave = generated_wave * rms_applied_factor
+
+            # Convert to numpy
+            generated_wave_np = generated_wave.squeeze().cpu().numpy()
+            generated_waves.append(generated_wave_np)
+            spectrograms.append(generated_mel_spec[0].cpu().to(torch.float32).numpy())
+
+        except Exception as e:
+            logging.exception(f"Error during inference/processing for batch {i+1}:") # Log traceback
+            print(f"Error details: {e}")
+            continue
+
+    if not generated_waves:
+        print("No audio waves were generated.")
+        return None, None
+
+    # Combine batches
+    print(f"Combining {len(generated_waves)} generated batches...")
+    if cross_fade_duration <= 0 or len(generated_waves) == 1:
+        final_wave = np.concatenate(generated_waves)
+    else:
+        # (Cross-fading logic remains the same)
+        final_wave = generated_waves[0]
+        for i in range(1, len(generated_waves)):
+            prev_wave = final_wave; next_wave = generated_waves[i]
+            cf_samples = min(int(cross_fade_duration * target_sample_rate), len(prev_wave), len(next_wave))
+            if cf_samples <= 0: final_wave = np.concatenate([prev_wave, next_wave]); continue
+            p_olap = prev_wave[-cf_samples:]; n_olap = next_wave[:cf_samples]
+            f_out = np.linspace(1, 0, cf_samples, dtype=p_olap.dtype); f_in = np.linspace(0, 1, cf_samples, dtype=n_olap.dtype)
+            cf_olap = p_olap * f_out + n_olap * f_in
+            final_wave = np.concatenate([prev_wave[:-cf_samples], cf_olap, next_wave[cf_samples:]])
+        print(f"Applied cross-fade of {cross_fade_duration:.2f}s between batches.")
+
+    # Optional: Remove silence post-combination
+    if remove_silence_post:
+        print("Removing silence from final output...")
+        try:
+            final_wave_float32 = final_wave.astype(np.float32)
+            with tempfile.NamedTemporaryFile(delete=True, suffix=".wav") as tmp_wav:
+                sf.write(tmp_wav.name, final_wave_float32, target_sample_rate)
+                aseg = AudioSegment.from_file(tmp_wav.name)
+                non_silent_segs = silence.split_on_silence(
+                    aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500
+                )
+                if not non_silent_segs:
+                    print("Warning: Silence removal resulted in empty audio. Keeping original.")
+                else:
+                    non_silent_wave = sum(non_silent_segs, AudioSegment.silent(duration=0))
+                    non_silent_wave.export(tmp_wav.name, format="wav")
+                    final_wave_tensor, _ = torchaudio.load(tmp_wav.name)
+                    final_wave = final_wave_tensor.squeeze().cpu().numpy()
+                    print("Silence removal applied.")
+        except Exception as e:
+            print(f"Warning: Failed to remove silence: {e}. Using original.")
+
+    # Combine spectrograms
+    print("Combining spectrograms...")
+    try:
+        if spectrograms:
+             combined_spectrogram = np.concatenate(spectrograms, axis=1)
+        else:
+             combined_spectrogram = None
+    except ValueError as e:
+        print(f"Warning: Could not concatenate spectrograms: {e}. Skipping.")
+        combined_spectrogram = None
+
+    return final_wave, combined_spectrogram
+
+
+def main_infer(ref_audio_orig_path, ref_text_input, gen_text_full,
+               ema_model, vocos, tokenizer, pipe_asr, # Loaded models/utils
+               ref_language, language, # Languages
+               speed, nfe_step, cfg_strength, sway_sampling_coef, # Sampling params
+               remove_silence_flag, cross_fade_duration, # Postprocessing
+               target_sample_rate, hop_length, target_rms, # Audio params
+               device, dtype): # System params
+    """
+    Main inference function coordinating preprocessing, batching, and generation.
+    (Function body remains the same as previous refactored version)
+    """
+    print(f"Starting inference for text: '{gen_text_full[:100]}...'")
+
+    # --- Reference Audio Preprocessing ---
+    print("Processing reference audio...")
+    processed_ref_path = None
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_ref_wav:
+            processed_ref_path = temp_ref_wav.name # Store path for potential use
+            aseg = AudioSegment.from_file(ref_audio_orig_path)
+            print(f"Original ref duration: {len(aseg)/1000:.2f}s")
+
+            # Edge silence removal + padding
+            aseg = remove_silence_edges(aseg)
+            aseg += AudioSegment.silent(duration=150)
+
+            # Split/recombine on silence
+            non_silent_segs = silence.split_on_silence(
+                aseg, min_silence_len=700, silence_thresh=-50, keep_silence=700
+            )
+            if non_silent_segs:
+                 aseg = sum(non_silent_segs, AudioSegment.silent(duration=0)) # Use sum for conciseness
+            else:
+                 print("Warning: Silence splitting/recombining resulted in empty audio. Using edge-trimmed.")
+
+            # Clip to 10s
+            max_ref_duration_ms = 10000
+            if len(aseg) > max_ref_duration_ms:
+                print(f"Reference audio exceeds {max_ref_duration_ms/1000}s. Clipping...")
+                aseg = aseg[:max_ref_duration_ms]
+
+            aseg.export(processed_ref_path, format="wav")
+            print(f"Processed ref duration: {len(aseg)/1000:.2f}s. Saved to temp file: {processed_ref_path}")
+
+            # Load processed audio tensor
+            ref_audio_tensor, sr_ref = torchaudio.load(processed_ref_path)
+
+    except Exception as e:
+        print(f"Error processing reference audio {ref_audio_orig_path}: {e}")
+        if processed_ref_path and Path(processed_ref_path).exists():
+            Path(processed_ref_path).unlink() # Clean up temp file on error
+        raise
+
+    # --- Reference Text Handling ---
+    ref_text_processed = ""
+    if not ref_text_input or ref_text_input.strip() == "":
+        print("No reference text provided. Transcribing reference audio...")
+        if pipe_asr is None:
+             raise ValueError("Whisper ASR model not loaded. Cannot transcribe. Please provide --ref_text.")
+        if not processed_ref_path:
+             raise ValueError("Processed reference audio path is missing for transcription.")
+        try:
+            # Ensure Whisper input dtype matches its loaded dtype
+            whisper_input_dtype = pipe_asr.model.dtype
+
+            # Load audio specifically for Whisper if dtypes differ significantly
+            # Or rely on pipeline handling. Assuming pipeline handles it for now.
+            print(f"Transcribing: {processed_ref_path}")
+            transcription_result = pipe_asr(
+                processed_ref_path,
+                chunk_length_s=15,
+                batch_size=8, # Smaller batch size for CLI
+                generate_kwargs={"task": "transcribe", "language": None}, # Whisper language detection
+                return_timestamps=False,
+            )
+            ref_text_processed = transcription_result["text"].strip()
+            print(f"Transcription finished: '{ref_text_processed}'")
+            if not ref_text_processed:
+                 print("Warning: Transcription resulted in empty text. Using placeholder.")
+                 ref_text_processed = "Reference audio"
+        except Exception as e:
+             logging.exception("Error during transcription:")
+             raise ValueError("Transcription failed. Please provide --ref_text.")
+    else:
+        print("Using provided reference text.")
+        ref_text_processed = ref_text_input
+
+    # Clean up the temporary processed reference audio file
+    if processed_ref_path and Path(processed_ref_path).exists():
+        try:
+            Path(processed_ref_path).unlink()
+            # print(f"Cleaned up temp ref file: {processed_ref_path}") # Debug
+        except OSError as e:
+            print(f"Warning: Could not delete temp ref file {processed_ref_path}: {e}")
+
+
+    # Ensure reference text ends with ". "
+    if not ref_text_processed.endswith(". "):
+        ref_text_processed = ref_text_processed.rstrip('. ') + ". " # More robust way
+    print(f"Final Reference Text: '{ref_text_processed}'")
+
+    # --- Phonemize Reference Text ---
+    print(f"Phonemizing reference text with language: {ref_language}")
+    ref_text_ipa = text_to_ipa(ref_text_processed, language=ref_language)
+    if not ref_text_ipa: raise ValueError("Reference text phonemization failed.")
+
+    # --- Chunk and Phonemize Generation Text ---
+    ref_audio_duration_sec = ref_audio_tensor.shape[-1] / sr_ref if sr_ref > 0 else 1.0
+    if ref_audio_duration_sec <= 0: ref_audio_duration_sec = 1.0
+    chars_per_sec = len(ref_text_processed.encode('utf-8')) / ref_audio_duration_sec if ref_audio_duration_sec > 0 else 10.0
+    if chars_per_sec <= 0: chars_per_sec = 10.0
+    target_chunk_duration_sec = max(5.0, 20.0 - ref_audio_duration_sec)
+    max_chars = int(chars_per_sec * target_chunk_duration_sec)
+
+    print(f"Ref duration: {ref_audio_duration_sec:.2f}s => Calculated max_chars/batch: {max_chars}")
+    gen_text_batches_plain = chunk_text(gen_text_full, max_chars=max_chars)
+    if not gen_text_batches_plain: raise ValueError("Text chunking resulted in zero batches.")
+    print(f"Split generation text into {len(gen_text_batches_plain)} batches.")
+
+    print(f"Phonemizing generation text batches with language: {language}")
+    gen_text_ipa_batches = []
+    for i, batch_text in enumerate(gen_text_batches_plain):
+        # print(f" Phonemizing batch {i+1}/{len(gen_text_batches_plain)}...") # Verbose
+        batch_ipa = text_to_ipa(batch_text, language=language)
+        if batch_ipa: gen_text_ipa_batches.append(batch_ipa)
+        else: print(f"Warning: Skipping batch {i+1} due to phonemization failure.")
+
+    if not gen_text_ipa_batches: raise ValueError("Phonemization failed for all generation text batches.")
+
+    # --- Run Batched Inference ---
+    print(f"Starting batch inference process ({len(gen_text_ipa_batches)} batches)...")
+    final_wave, combined_spectrogram = infer_batch(
+        (ref_audio_tensor, sr_ref), ref_text_ipa, gen_text_ipa_batches,
+        ema_model, vocos, tokenizer,
+        remove_silence_flag, cross_fade_duration,
+        nfe_step, cfg_strength, sway_sampling_coef, speed,
+        target_sample_rate, hop_length, target_rms,
+        device, dtype
+    )
+
+    return final_wave, combined_spectrogram
+
+
+# --- Execution ---
+if __name__ == "__main__":
+    # Setup logging
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+    try:
+        final_wave_np, combined_spectrogram_np = main_infer(
+            ref_audio_path, ref_text, gen_text,
+            ema_model, vocos, tokenizer, pipe,
+            ref_language, language,
+            speed, nfe_step, cfg_strength, sway_sampling_coef,
+            remove_silence_flag, cross_fade_duration,
+            target_sample_rate, hop_length, target_rms,
+            device, dtype
+        )
+
+        # --- Save Outputs ---
+        output_saved = False
+        if final_wave_np is not None and len(final_wave_np) > 0:
+            print(f"Saving final audio ({len(final_wave_np)/target_sample_rate:.2f}s) to {wave_path}...")
+            final_wave_float32 = final_wave_np.astype(np.float32) # Ensure float32 for sf
+            sf.write(str(wave_path), final_wave_float32, target_sample_rate)
+            print("Audio saved successfully.")
+            output_saved = True
+        else:
+            print("Inference did not produce a valid audio wave.")
+
+        if combined_spectrogram_np is not None:
+            print(f"Saving combined spectrogram to {spectrogram_path}...")
+            save_spectrogram(combined_spectrogram_np, str(spectrogram_path))
+            print("Spectrogram saved successfully.")
+            output_saved = True
+        # else: # No need to print if spectrogram was None
+        #     print("Spectrogram generation failed or was skipped.")
+
+        if not output_saved:
+             print("No output files were generated.")
+
+    except FileNotFoundError as e:
+         logging.error(f"File not found: {e}")
+         print(f"\nError: A required file was not found. Please check paths. Details: {e}")
+         exit(1)
+    except ValueError as e:
+         logging.error(f"Value error: {e}")
+         print(f"\nError: An invalid value or configuration was encountered. Details: {e}")
+         exit(1)
+    except Exception as e:
+        logging.exception("An unexpected error occurred during inference:") # Log traceback
+        print(f"\nAn unexpected error occurred: {e}")
+        exit(1)
+
+    print("\nInference completed.")