Upload model/f5tts_wrapper.py with huggingface_hub

model/f5tts_wrapper.py

@@ -0,0 +1,549 @@
+import os
+import torch
+import torchaudio
+import numpy as np
+from pathlib import Path
+from typing import Optional, Union, List, Tuple, Dict
+
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+from importlib.resources import files
+from pydub import AudioSegment, silence
+
+from f5_tts.model import CFM
+from f5_tts.model.utils import (
+    get_tokenizer,
+    convert_char_to_pinyin,
+)
+from f5_tts.infer.utils_infer import (
+    chunk_text,
+    load_vocoder,
+    transcribe,
+    initialize_asr_pipeline,
+)
+
+
+class F5TTSWrapper:
+    """
+    A wrapper class for F5-TTS that preprocesses reference audio once 
+    and allows for repeated TTS generation.
+    """
+    
+    def __init__(
+        self, 
+        model_name: str = "F5TTS_v1_Base", 
+        ckpt_path: Optional[str] = None,
+        vocab_file: Optional[str] = None,
+        vocoder_name: str = "vocos",
+        use_local_vocoder: bool = False,
+        vocoder_path: Optional[str] = None,
+        device: Optional[str] = None,
+        hf_cache_dir: Optional[str] = None,
+        target_sample_rate: int = 24000,
+        n_mel_channels: int = 100,
+        hop_length: int = 256,
+        win_length: int = 1024,
+        n_fft: int = 1024,
+        ode_method: str = "euler",
+        use_ema: bool = True,
+    ):
+        """
+        Initialize the F5-TTS wrapper with model configuration.
+        
+        Args:
+            model_name: Name of the F5-TTS model variant (e.g., "F5TTS_v1_Base")
+            ckpt_path: Path to the model checkpoint file. If None, will use default path.
+            vocab_file: Path to the vocab file. If None, will use default.
+            vocoder_name: Name of the vocoder to use ("vocos" or "bigvgan")
+            use_local_vocoder: Whether to use a local vocoder or download from HF
+            vocoder_path: Path to the local vocoder. Only used if use_local_vocoder is True.
+            device: Device to run the model on. If None, will automatically determine.
+            hf_cache_dir: Directory to cache HuggingFace models
+            target_sample_rate: Target sample rate for audio
+            n_mel_channels: Number of mel channels
+            hop_length: Hop length for the mel spectrogram
+            win_length: Window length for the mel spectrogram
+            n_fft: FFT size for the mel spectrogram
+            ode_method: ODE method for sampling ("euler" or "midpoint")
+            use_ema: Whether to use EMA weights from the checkpoint
+        """
+        # Set device
+        if device is None:
+            self.device = (
+                "cuda" if torch.cuda.is_available()
+                else "xpu" if torch.xpu.is_available()
+                else "mps" if torch.backends.mps.is_available()
+                else "cpu"
+            )
+        else:
+            self.device = device
+            
+        # Audio processing parameters
+        self.target_sample_rate = target_sample_rate
+        self.n_mel_channels = n_mel_channels
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_fft = n_fft
+        self.mel_spec_type = vocoder_name
+        
+        # Sampling parameters
+        self.ode_method = ode_method
+        
+        # Initialize ASR for transcription if needed
+        initialize_asr_pipeline(device=self.device)
+        
+        # Load model configuration
+        if ckpt_path is None:
+            repo_name = "F5-TTS"
+            ckpt_step = 1250000
+            ckpt_type = "safetensors"
+            
+            # Adjust for previous models
+            if model_name == "F5TTS_Base":
+                if vocoder_name == "vocos":
+                    ckpt_step = 1200000
+                elif vocoder_name == "bigvgan":
+                    model_name = "F5TTS_Base_bigvgan"
+                    ckpt_type = "pt"
+            elif model_name == "E2TTS_Base":
+                repo_name = "E2-TTS"
+                ckpt_step = 1200000
+                
+            ckpt_path = str(cached_path(f"hf://SWivid/{repo_name}/{model_name}/model_{ckpt_step}.{ckpt_type}"))
+        
+        # Load model configuration
+        config_path = str(files("f5_tts").joinpath(f"configs/{model_name}.yaml"))
+        model_cfg = OmegaConf.load(config_path)
+        model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+        model_arc = model_cfg.model.arch
+        
+        # Load tokenizer
+        if vocab_file is None:
+            vocab_file = str(files("f5_tts").joinpath("infer/examples/vocab.txt"))
+        tokenizer_type = "custom"
+        self.vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer_type)
+        
+        # Create model
+        self.model = CFM(
+            transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
+            mel_spec_kwargs=dict(
+                n_fft=n_fft,
+                hop_length=hop_length,
+                win_length=win_length,
+                n_mel_channels=n_mel_channels,
+                target_sample_rate=target_sample_rate,
+                mel_spec_type=vocoder_name,
+            ),
+            odeint_kwargs=dict(
+                method=ode_method,
+            ),
+            vocab_char_map=self.vocab_char_map,
+        ).to(self.device)
+        
+        # Load checkpoint
+        dtype = torch.float32 if vocoder_name == "bigvgan" else None
+        self._load_checkpoint(self.model, ckpt_path, dtype=dtype, use_ema=use_ema)
+        
+        # Load vocoder
+        if vocoder_path is None:
+            if vocoder_name == "vocos":
+                vocoder_path = "../checkpoints/vocos-mel-24khz"
+            elif vocoder_name == "bigvgan":
+                vocoder_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"
+        
+        self.vocoder = load_vocoder(
+            vocoder_name=vocoder_name,
+            is_local=use_local_vocoder,
+            local_path=vocoder_path,
+            device=self.device,
+            hf_cache_dir=hf_cache_dir
+        )
+        
+        # Storage for reference data
+        self.ref_audio_processed = None
+        self.ref_text = None
+        self.ref_audio_len = None
+        
+        # Default inference parameters
+        self.target_rms = 0.1
+        self.cross_fade_duration = 0.15
+        self.nfe_step = 32
+        self.cfg_strength = 2.0
+        self.sway_sampling_coef = -1.0
+        self.speed = 1.0
+        self.fix_duration = None
+
+    def _load_checkpoint(self, model, ckpt_path, dtype=None, use_ema=True):
+        """
+        Load model checkpoint with proper handling of different checkpoint formats.
+        
+        Args:
+            model: The model to load weights into
+            ckpt_path: Path to the checkpoint file
+            dtype: Data type for model weights
+            use_ema: Whether to use EMA weights from the checkpoint
+        
+        Returns:
+            Loaded model
+        """
+        if dtype is None:
+            dtype = (
+                torch.float16
+                if "cuda" in self.device
+                and torch.cuda.get_device_properties(self.device).major >= 7
+                and not torch.cuda.get_device_name().endswith("[ZLUDA]")
+                else torch.float32
+            )
+        model = model.to(dtype)
+
+        ckpt_type = ckpt_path.split(".")[-1]
+        if ckpt_type == "safetensors":
+            from safetensors.torch import load_file
+            checkpoint = load_file(ckpt_path, device=self.device)
+        else:
+            checkpoint = torch.load(ckpt_path, map_location=self.device, weights_only=True)
+
+        if use_ema:
+            if ckpt_type == "safetensors":
+                checkpoint = {"ema_model_state_dict": checkpoint}
+            checkpoint["model_state_dict"] = {
+                k.replace("ema_model.", ""): v
+                for k, v in checkpoint["ema_model_state_dict"].items()
+                if k not in ["initted", "step"]
+            }
+
+            # patch for backward compatibility
+            for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
+                if key in checkpoint["model_state_dict"]:
+                    del checkpoint["model_state_dict"][key]
+
+            model.load_state_dict(checkpoint["model_state_dict"])
+        else:
+            if ckpt_type == "safetensors":
+                checkpoint = {"model_state_dict": checkpoint}
+            model.load_state_dict(checkpoint["model_state_dict"])
+
+        del checkpoint
+        torch.cuda.empty_cache()
+
+        return model.to(self.device)
+    
+    def preprocess_reference(self, ref_audio_path: str, ref_text: str = "", clip_short: bool = True):
+        """
+        Preprocess the reference audio and text, storing them for later use.
+        
+        Args:
+            ref_audio_path: Path to the reference audio file
+            ref_text: Text transcript of reference audio. If empty, will auto-transcribe.
+            clip_short: Whether to clip long audio to shorter segments
+            
+        Returns:
+            Tuple of processed audio and text
+        """
+        print("Converting audio...")
+        # Load audio file
+        aseg = AudioSegment.from_file(ref_audio_path)
+        
+        if clip_short:
+            # 1. try to find long silence for clipping
+            non_silent_segs = silence.split_on_silence(
+                aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000, seek_step=10
+            )
+            non_silent_wave = AudioSegment.silent(duration=0)
+            for non_silent_seg in non_silent_segs:
+                if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:
+                    print("Audio is over 12s, clipping short. (1)")
+                    break
+                non_silent_wave += non_silent_seg
+                
+            # 2. try to find short silence for clipping if 1. failed
+            if len(non_silent_wave) > 12000:
+                non_silent_segs = silence.split_on_silence(
+                    aseg, min_silence_len=100, silence_thresh=-40, keep_silence=1000, seek_step=10
+                )
+                non_silent_wave = AudioSegment.silent(duration=0)
+                for non_silent_seg in non_silent_segs:
+                    if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:
+                        print("Audio is over 12s, clipping short. (2)")
+                        break
+                    non_silent_wave += non_silent_seg
+                    
+            aseg = non_silent_wave
+                
+            # 3. if no proper silence found for clipping
+            if len(aseg) > 12000:
+                aseg = aseg[:12000]
+                print("Audio is over 12s, clipping short. (3)")
+            
+        # Remove silence edges
+        aseg = self._remove_silence_edges(aseg) + AudioSegment.silent(duration=50)
+        
+        # Export to temporary file and load as tensor
+        import tempfile
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
+            aseg.export(tmp_file.name, format="wav")
+            processed_audio_path = tmp_file.name
+            
+        # Transcribe if needed
+        if not ref_text.strip():
+            print("No reference text provided, transcribing reference audio...")
+            ref_text = transcribe(processed_audio_path)
+        else:
+            print("Using custom reference text...")
+            
+        # Ensure ref_text ends with proper punctuation
+        if not ref_text.endswith(". ") and not ref_text.endswith("。"):
+            if ref_text.endswith("."):
+                ref_text += " "
+            else:
+                ref_text += ". "
+                
+        print("\nReference text:", ref_text)
+        
+        # Load and process audio
+        audio, sr = torchaudio.load(processed_audio_path)
+        if audio.shape[0] > 1:  # Convert stereo to mono
+            audio = torch.mean(audio, dim=0, keepdim=True)
+            
+        # Normalize volume
+        rms = torch.sqrt(torch.mean(torch.square(audio)))
+        if rms < self.target_rms:
+            audio = audio * self.target_rms / rms
+            
+        # Resample if needed
+        if sr != self.target_sample_rate:
+            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
+            audio = resampler(audio)
+            
+        # Move to device
+        audio = audio.to(self.device)
+        
+        # Store reference data
+        self.ref_audio_processed = audio
+        self.ref_text = ref_text
+        self.ref_audio_len = audio.shape[-1] // self.hop_length
+        
+        # Remove temporary file
+        os.unlink(processed_audio_path)
+        
+        return audio, ref_text
+    
+    def _remove_silence_edges(self, audio, silence_threshold=-42):
+        """
+        Remove silence from the start and end of audio.
+
+        Args:
+            audio: AudioSegment to process
+            silence_threshold: dB threshold to consider as silence
+
+        Returns:
+            Processed AudioSegment
+        """
+        # Remove silence from the start
+        non_silent_start_idx = silence.detect_leading_silence(audio, silence_threshold=silence_threshold)
+        audio = audio[non_silent_start_idx:]
+
+        # Remove silence from the end
+        non_silent_end_duration = audio.duration_seconds
+        for ms in reversed(audio):
+            if ms.dBFS > silence_threshold:
+                break
+            non_silent_end_duration -= 0.001
+        trimmed_audio = audio[: int(non_silent_end_duration * 1000)]
+
+        return trimmed_audio
+    
+    def generate(
+        self, 
+        text: str,
+        output_path: Optional[str] = None,
+        nfe_step: Optional[int] = None,
+        cfg_strength: Optional[float] = None,
+        sway_sampling_coef: Optional[float] = None,
+        speed: Optional[float] = None,
+        fix_duration: Optional[float] = None,
+        cross_fade_duration: Optional[float] = None,
+        return_numpy: bool = False,
+        return_spectrogram: bool = False,
+    ) -> Union[str, Tuple[np.ndarray, int], Tuple[np.ndarray, int, np.ndarray]]:
+        """
+        Generate speech for the given text using the stored reference audio.
+        
+        Args:
+            text: Text to synthesize
+            output_path: Path to save the generated audio. If None, won't save.
+            nfe_step: Number of function evaluation steps
+            cfg_strength: Classifier-free guidance strength
+            sway_sampling_coef: Sway sampling coefficient
+            speed: Speed of generated audio
+            fix_duration: Fixed duration in seconds
+            cross_fade_duration: Duration of cross-fade between segments
+            return_numpy: If True, returns the audio as a numpy array
+            return_spectrogram: If True, also returns the spectrogram
+            
+        Returns:
+            If output_path provided: path to output file
+            If return_numpy=True: tuple of (audio_array, sample_rate)
+            If return_spectrogram=True: tuple of (audio_array, sample_rate, spectrogram)
+        """
+        if self.ref_audio_processed is None or self.ref_text is None:
+            raise ValueError("Reference audio not preprocessed. Call preprocess_reference() first.")
+            
+        # Use default values if not specified
+        nfe_step = nfe_step if nfe_step is not None else self.nfe_step
+        cfg_strength = cfg_strength if cfg_strength is not None else self.cfg_strength
+        sway_sampling_coef = sway_sampling_coef if sway_sampling_coef is not None else self.sway_sampling_coef
+        speed = speed if speed is not None else self.speed
+        fix_duration = fix_duration if fix_duration is not None else self.fix_duration
+        cross_fade_duration = cross_fade_duration if cross_fade_duration is not None else self.cross_fade_duration
+        
+        # Split the input text into batches
+        audio_len = self.ref_audio_processed.shape[-1] / self.target_sample_rate
+        max_chars = int(len(self.ref_text.encode("utf-8")) / audio_len * (22 - audio_len))
+        text_batches = chunk_text(text, max_chars=max_chars)
+        
+        for i, text_batch in enumerate(text_batches):
+            print(f"Text batch {i}: {text_batch}")
+        print("\n")
+        
+        # Generate audio for each batch
+        generated_waves = []
+        spectrograms = []
+        
+        for text_batch in text_batches:
+            # Adjust speed for very short texts
+            local_speed = speed
+            if len(text_batch.encode("utf-8")) < 10:
+                local_speed = 0.3
+                
+            # Prepare the text
+            text_list = [self.ref_text + text_batch]
+            final_text_list = convert_char_to_pinyin(text_list)
+            
+            # Calculate duration
+            if fix_duration is not None:
+                duration = int(fix_duration * self.target_sample_rate / self.hop_length)
+            else:
+                # Calculate duration based on text length
+                ref_text_len = len(self.ref_text.encode("utf-8"))
+                gen_text_len = len(text_batch.encode("utf-8"))
+                duration = self.ref_audio_len + int(self.ref_audio_len / ref_text_len * gen_text_len / local_speed)
+                
+            # Generate audio
+            with torch.inference_mode():
+                generated, _ = self.model.sample(
+                    cond=self.ref_audio_processed,
+                    text=final_text_list,
+                    duration=duration,
+                    steps=nfe_step,
+                    cfg_strength=cfg_strength,
+                    sway_sampling_coef=sway_sampling_coef,
+                )
+                
+                # Process the generated mel spectrogram
+                generated = generated.to(torch.float32)
+                generated = generated[:, self.ref_audio_len:, :]
+                generated = generated.permute(0, 2, 1)
+                
+                # Convert to audio
+                if self.mel_spec_type == "vocos":
+                    generated_wave = self.vocoder.decode(generated)
+                elif self.mel_spec_type == "bigvgan":
+                    generated_wave = self.vocoder(generated)
+                    
+                # Normalize volume if needed
+                rms = torch.sqrt(torch.mean(torch.square(self.ref_audio_processed)))
+                if rms < self.target_rms:
+                    generated_wave = generated_wave * rms / self.target_rms
+                    
+                # Convert to numpy and append to list
+                generated_wave = generated_wave.squeeze().cpu().numpy()
+                generated_waves.append(generated_wave)
+                
+                # Store spectrogram if needed
+                if return_spectrogram or output_path is not None:
+                    spectrograms.append(generated.squeeze().cpu().numpy())
+        
+        # Combine all segments
+        if generated_waves:
+            if cross_fade_duration <= 0:
+                # Simply concatenate
+                final_wave = np.concatenate(generated_waves)
+            else:
+                # Cross-fade between segments
+                final_wave = generated_waves[0]
+                for i in range(1, len(generated_waves)):
+                    prev_wave = final_wave
+                    next_wave = generated_waves[i]
+                    
+                    # Calculate cross-fade samples
+                    cross_fade_samples = int(cross_fade_duration * self.target_sample_rate)
+                    cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))
+                    
+                    if cross_fade_samples <= 0:
+                        # No overlap possible, concatenate
+                        final_wave = np.concatenate([prev_wave, next_wave])
+                        continue
+                        
+                    # Create cross-fade
+                    prev_overlap = prev_wave[-cross_fade_samples:]
+                    next_overlap = next_wave[:cross_fade_samples]
+                    
+                    fade_out = np.linspace(1, 0, cross_fade_samples)
+                    fade_in = np.linspace(0, 1, cross_fade_samples)
+                    
+                    cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in
+                    
+                    final_wave = np.concatenate([
+                        prev_wave[:-cross_fade_samples], 
+                        cross_faded_overlap, 
+                        next_wave[cross_fade_samples:]
+                    ])
+            
+            # Combine spectrograms if needed
+            if return_spectrogram or output_path is not None:
+                combined_spectrogram = np.concatenate(spectrograms, axis=1)
+                
+            # Save to file if path provided
+            if output_path is not None:
+                output_dir = os.path.dirname(output_path)
+                if output_dir and not os.path.exists(output_dir):
+                    os.makedirs(output_dir)
+                
+                # Save audio
+                torchaudio.save(output_path, 
+                                torch.tensor(final_wave).unsqueeze(0), 
+                                self.target_sample_rate)
+                
+                # Save spectrogram if needed
+                if return_spectrogram:
+                    spectrogram_path = os.path.splitext(output_path)[0] + '_spec.png'
+                    self._save_spectrogram(combined_spectrogram, spectrogram_path)
+                
+                if not return_numpy:
+                    return output_path
+            
+            # Return as requested
+            if return_spectrogram:
+                return final_wave, self.target_sample_rate, combined_spectrogram
+            else:
+                return final_wave, self.target_sample_rate
+            
+        else:
+            raise RuntimeError("No audio generated")
+    
+    def _save_spectrogram(self, spectrogram, path):
+        """Save spectrogram as image"""
+        import matplotlib.pyplot as plt
+        plt.figure(figsize=(12, 4))
+        plt.imshow(spectrogram, origin="lower", aspect="auto")
+        plt.colorbar()
+        plt.savefig(path)
+        plt.close()
+        
+    def get_current_audio_length(self):
+        """Get the length of the reference audio in seconds"""
+        if self.ref_audio_processed is None:
+            return 0
+        return self.ref_audio_processed.shape[-1] / self.target_sample_rate