app.py

import gradio as gr
import torch
import librosa
import soundfile as sf
from transformers import AutoProcessor, AutoModelForCTC, pipeline
from pydub import AudioSegment
import os
import numpy as np

# --- 1. Load the Voice Conversion Model and Processor ---
# This part runs once when the app starts.
# We'll use a model designed for voice processing.
# For example, facebook/hidef-vc or a speech-to-text model that can generate embeddings.
# NOTE: Directly using hidef-vc is complex as it's not a simple pipeline.
# Let's try a more general approach with a pipeline that allows speaker conditioning,
# or simulate based on available models.

# As of my knowledge cut-off, a direct "Voice Changer" pipeline like this is not standard
# in `transformers` library for arbitrary source/target.
# Most HFace models are TTS (Text-to-Speech) with speaker embeddings for different voices,
# or ASR (Automatic Speech Recognition).
# Voice Conversion often involves separate components: Speaker Encoder, Content Encoder, Decoder.

# Let's pivot to a more realistic (but still simplified) approach:
# Using a ASR model to get content and a speaker embedding model (like RawNet, ECAPA-TDNN)
# and then a Custom Decoder if possible.
# This gets very complex.

# A more practical pre-trained model for Voice Conversion might be outside `transformers`
# directly, e.g., from `ESPnet`, `Coqui-TTS`, `PaddleSpeech`.

# Let's try to mimic with a text-to-speech model if speaker embeddings are easily extracted
# and inputable, allowing a "style transfer" through speaker.

# Placeholder Warning: The following "model" is illustrative.
# Real Voice Conversion requires dedicated models.
# As of now, `transformers` does not have a single "voice-to-voice" pipeline for arbitrary inputs.

# --- Realistic (but still simplified) Approach: Using a Speaker Verification model
# to get an embedding, and then (hypothetically) use that with a TTS model to
# "synthesize" the target audio with the source speaker's voice.
# This assumes we can transcript the target audio and feed it to a TTS model with a speaker ID.
# This is NOT direct voice conversion, but Text-to-Speech with speaker adaptation.

# ---- NEW PLAN: Use a specific Voice Conversion Library/Project ----
# The `transformers` library itself doesn't have a direct "Voice Changer" pipeline
# that takes two audio files and merges their characteristics easily.
# For proper voice conversion, projects like `so-vits-svc`, `RVC (Retrieval-based Voice Conversion)`,
# or deep models from `Coqui-TTS` (`YourTTS`, `VITS`) are used.

# Since the request is for "Hugging Face", I need to find a suitable model on HF Hub.
# `facebook/musicgen` or `facebook/audiocraft` are for music generation, not voice.
# `speechbrain/spkrec-ecapa-voxceleb` for speaker verification (embeddings).
# `suno/bark` for text-to-speech with expressive voice.

# Let's try to use `suno/bark` as it can handle multiple speakers.
# The idea would be:
# 1. Transcribe the target audio.
# 2. Extract prompts (or audio history) from the source audio using Bark's capabilities.
# 3. Generate the transcribed text with the source audio's "voice".
# This is still TTS with speaker conditioning, not direct voice-to-voice.

try:
    from transformers import pipeline, AutoProcessor, AutoModelForSpeechSeq2Seq, SpeechT5ForTextToSpeech, SpeechT5Processor, SpeechT5HifiGan
    import torchaudio
    # Using a general ASR model for transcription (e.g., Whisper)
    # Using SpeechT5 for TTS with speaker embeddings, if suitable models are on HF Hub.

    # Load ASR model (Whisper) for transcription
    asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny.en", device=0 if torch.cuda.is_available() else -1)

    # Load SpeechT5 components for TTS and speaker embeddings
    # SpeechT5 can use speaker embeddings. We need a way to extract embedding from source audio.
    # A pre-trained speaker embedding model:
    speaker_encoder_model_name = "Matthijs/speecht5_vc_melgan" # Not a direct speaker encoder for arbitrary audio. This is a VC model.
    # Let's use a standard speaker embedding from SpeechBrain:
    from speechbrain.pretrained import EncoderClassifier
    speaker_encoder = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb", run_opts={"device":"cuda" if torch.cuda.is_available() else "cpu"})

    tts_processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
    tts_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")

    voice_conversion_enabled = True
    print("Hugging Face models loaded successfully for Voice Conversion attempt.")

except ImportError as e:
    voice_conversion_enabled = False
    print(f"Failed to load Hugging Face models for advanced voice conversion: {e}")
    print("Falling back to simplified (less accurate) logic. Please ensure you have `transformers`, `torch`, `speechbrain`, etc. installed.")
    print("For best results, a CUDA-enabled GPU is highly recommended.")

# --- Helper functions for audio processing ---
def convert_to_wav(audio_file_path):
    try:
        audio = AudioSegment.from_file(audio_file_path)
        # Convert to a common sample rate for consistency (e.g., 16kHz)
        audio = audio.set_frame_rate(16000).set_channels(1)
        wav_file_path = audio_file_path + ".wav"
        audio.export(wav_file_path, format="wav")
        return wav_file_path
    except Exception as e:
        raise gr.Error(f"Error converting audio to WAV: {e}")

# Function to clear temporary files
def cleanup_temp_files(*paths):
    for path in paths:
        if path and os.path.exists(path):
            os.remove(path)

# --- The Voice Changer Function (using Hugging Face models) ---
def advanced_voice_changer(source_audio_path, target_audio_path):
    if not voice_conversion_enabled:
        raise gr.Error("Hugging Face models for voice conversion could not be loaded. Please check your installations and environment.")

    if source_audio_path is None or target_audio_path is None:
        raise gr.Error("Please upload both source and target audio files.")

    source_wav_path = None
    target_wav_path = None
    output_file_path = "output_voice_changed.wav"

    try:
        source_wav_path = convert_to_wav(source_audio_path)
        target_wav_path = convert_to_wav(target_audio_path)

        # 1. Load source audio for speaker embedding
        source_audio_for_embedding, sr_source_embedding = torchaudio.load(source_wav_path)
        # Ensure it's 16kHz for ECAPA-TDNN
        if sr_source_embedding != 16000:
            source_audio_for_embedding = torchaudio.transforms.Resample(orig_freq=sr_source_embedding, new_freq=16000)(source_audio_for_embedding)

        # Ensure mono
        if source_audio_for_embedding.shape[0] > 1:
            source_audio_for_embedding = torch.mean(source_audio_for_embedding, dim=0, keepdim=True)

        # Get speaker embedding from source audio
        # Pass the audio directly to the speaker_encoder model as a tensor
        # The input for speaker_encoder is usually [batch_size, num_samples]
        embeddings = speaker_encoder.encode_batch(source_audio_for_embedding)
        # Take the first embedding from the batch (assuming batch_size=1)
        speaker_embeddings_source = embeddings[0].squeeze().cpu().numpy()

        # 2. Transcribe the target audio to get its content (text)
        # Load audio for ASR at its native sample rate first, then ASR resamples it
        y_target_asr, sr_target_asr = librosa.load(target_wav_path, sr=None)
        # ASR pipeline expects numpy array
        transcription = asr_pipe(y_target_asr)['text']
        print(f"Transcibed Target Audio: '{transcription}'")

        if not transcription:
            raise gr.Error("Could not transcribe target audio. Please ensure it contains clear speech.")

        # 3. Generate speech using SpeechT5 with the extracted speaker embedding
        inputs = tts_processor(text=transcription, return_tensors="pt")
        
        # Reshape speaker_embeddings to [1, 512] as expected by SpeechT5
        # The ECAPA-TDNN embedding is 512, which SpeechT5 expects.
        speaker_embeddings_tensor = torch.tensor(speaker_embeddings_source).unsqueeze(0)

        with torch.no_grad(): # Disable gradient calculation for inference
            speech = tts_model.generate_speech(inputs["input_ids"], speaker_embeddings_tensor, vocoder=vocoder)

        # Convert the generated speech tensor to a numpy array and save
        sf.write(output_file_path, speech.cpu().numpy(), samplerate=16000) # SpeechT5 outputs at 16kHz

        return output_file_path

    except Exception as e:
        import traceback
        traceback.print_exc() # Print full traceback for debugging
        raise gr.Error(f"An error occurred during voice processing: {e}")
    finally:
        cleanup_temp_files(source_wav_path, target_wav_path)


# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Hugging Face Voice Changer (SpeechT5 + Ecapa-TDNN)
        This tool attempts to generate the content of the **Target Audio** in the voice/style of the **Source Audio**.

        **How it works:**
        1. It extracts the speaker's identity from the **Source Audio** using a speaker embedding model (ECAPA-TDNN).
        2. It transcribes the speech content from the **Target Audio** using a Speech-to-Text model (Whisper).
        3. It then synthesizes the transcribed text using a Text-to-Speech model (SpeechT5), conditioned on the speaker's identity extracted from the Source Audio.

        **Important Notes:**
        *   This is **Text-to-Speech with Speaker Adaptation**, not a direct sample-level voice-to-voice conversion like a true voice changer (e.g., changing singing voice).
        *   Quality depends heavily on the clarity of the input audios and the capabilities of the underlying models.
        *   Transcription errors from the Target Audio will affect the output.
        *   **A GPU (CUDA) is highly recommended for faster processing.**
        """
    )
    if not voice_conversion_enabled:
        gr.Markdown("<p style='color:red;'><strong>Warning: Hugging Face models could not be loaded. This app will not function correctly. Please check your Python environment and required packages.</strong></p>")

    with gr.Row():
        source_audio_input = gr.Audio(type="filepath", label="Source Speaker Audio (Voice to imitate)", sources=["upload"])
        target_audio_input = gr.Audio(type="filepath", label="Target Content Audio (What to say)", sources=["upload"])

    output_audio = gr.Audio(label="Generated Audio (Target content in Source speaker's voice)", key="output_audio")

    voice_changer_button = gr.Button("Generate Voice")

    voice_changer_button.click(
        fn=advanced_voice_changer,
        inputs=[source_audio_input, target_audio_input],
        outputs=output_audio
    )

if __name__ == "__main__":
    demo.launch()