app.py

import spaces  # imported first, as ZeroGPU requires
import torch
import os
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
import gradio as gr
import traceback
import gc
import numpy as np
import librosa
from pydub import AudioSegment
from pydub.effects import normalize
from huggingface_hub import snapshot_download
from tts.infer_cli import MegaTTS3DiTInfer, convert_to_wav, cut_wav

from pyharp.core import ModelCard, build_endpoint
from pyharp.media.audio import save_audio
from pyharp.tags import Channels, SampleRate, Subcategory
from audiotools import AudioSignal


model_card = ModelCard(
    name="MegaTTS 3",
    description=(
        "MegaTTS 3 is a text-to-speech model trained by ByteDance with "
        "exceptional voice cloning capabilities, as 24 kHz mono speech. Please use this Space "
        "responsibly and do not abuse it! This demo is for research and "
        "educational purposes only. Repository: "
        "https://github.com/bytedance/MegaTTS3 Paper: "
        "https://arxiv.org/abs/2502.18924"
    ),
    author="Ziyue Jiang et al.",
    tags=[Subcategory.TEXT_TO_SPEECH, SampleRate(24000), Channels(1), 'voice-cloning', 'domain:speech']
)


def download_weights():
    """Download model weights from HuggingFace if not already present."""
    repo_id = "mrfakename/MegaTTS3-VoiceCloning"
    weights_dir = "checkpoints"
    
    if not os.path.exists(weights_dir):
        print("Downloading model weights from HuggingFace...")
        snapshot_download(
            repo_id=repo_id,
            revision="b3a2c297eea319931ee985751f117f866085d0c5",
            local_dir=weights_dir,
            local_dir_use_symlinks=False
        )
        print("Model weights downloaded successfully!")
    else:
        print("Model weights already exist.")
    
    return weights_dir


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the weights are downloaded and the model built there alone
if __name__ != "__main__":
    download_weights()
    print("Initializing MegaTTS3 model...")
    infer_pipe = MegaTTS3DiTInfer()
    print("Model loaded successfully!")

# PyHARP stops a job that runs past its timeout_s, 15 minutes by default, so a request whose
# reservation would exceed that is refused rather than started. A shorter one is left to
# ZeroGPU, which admits it when the caller's own allowance covers it.
MAX_GPU_SECONDS = 900

# The text is spoken in one pass the GPU holds whole, so it is limited to the longest measured to
# fit. The cost grows with its length and with the timesteps
MAX_TEXT_CHARACTERS = 2500


def gpu_seconds(wav_path, inp_text, infer_timestep, p_w, t_w):
    # Measured at about 15 s for 70 characters and 35 s for 550, both at 32 timesteps, padded
    return 15 + 0.002 * len(inp_text) * int(infer_timestep)


@spaces.GPU(duration=gpu_seconds)
def generate_speech(wav_path, inp_text, infer_timestep, p_w, t_w):
    with open(wav_path, 'rb') as file:
        file_content = file.read()

    try:
        print(f"Generating speech with: {inp_text}...")
        resource_context = infer_pipe.preprocess(file_content)
        fs, wav = infer_pipe.forward(resource_context, inp_text, time_step=infer_timestep, p_w=p_w, t_w=t_w)
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Speech generation failed: {e}") from e
    finally:
        cleanup_memory()

    return fs, wav


def process_fn(inp_audio, inp_text, infer_timestep, p_w, t_w):
    # Checked and prepared before the GPU is reserved, so a request that cannot run spends none of
    # the caller's ZeroGPU allowance. A failure is raised as gr.Error, since HARP never sees a
    # gr.Warning.
    if not inp_audio:
        raise gr.Error("Provide a recording of the voice to clone.")
    if not (inp_text or "").strip():
        raise gr.Error("Enter the text for the cloned voice to speak.")
    if len(inp_text) > MAX_TEXT_CHARACTERS:
        raise gr.Error(
            f"The text is {len(inp_text)} characters long. Shorten it to {MAX_TEXT_CHARACTERS} "
            "or fewer."
        )
    if gpu_seconds(None, inp_text, infer_timestep, p_w, t_w) > MAX_GPU_SECONDS:
        raise gr.Error(
            "This much text at this many timesteps needs more GPU time than one request may "
            "reserve. Shorten the text, or use fewer timesteps."
        )

    try:
        wav_path = preprocess_audio_robust(inp_audio)
        cut_wav(wav_path, max_len=10)
    except Exception as e:
        raise gr.Error(f"Could not use the reference recording: {e}") from e

    fs, wav = generate_speech(wav_path, inp_text, int(infer_timestep), p_w, t_w)
    sig = AudioSignal(wav, sample_rate=fs)
    return save_audio(sig)


def cleanup_memory():
    """Clean up GPU and system memory."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

def preprocess_audio_robust(audio_path, target_sr=22050, max_duration=30):
    """Robustly preprocess audio to prevent CUDA errors."""
    try:
        # Load with pydub for robust format handling
        audio = AudioSegment.from_file(audio_path)
        
        # Convert to mono if stereo
        if audio.channels > 1:
            audio = audio.set_channels(1)
        
        # Limit duration to prevent memory issues
        if len(audio) > max_duration * 1000:  # pydub uses milliseconds
            audio = audio[:max_duration * 1000]
        
        # Normalize audio to prevent clipping
        audio = normalize(audio)
        
        # Convert to target sample rate
        audio = audio.set_frame_rate(target_sr)
        
        # Export to temporary WAV file with specific parameters
        temp_path = audio_path.replace(os.path.splitext(audio_path)[1], '_processed.wav')
        audio.export(
            temp_path,
            format="wav",
            parameters=["-acodec", "pcm_s16le", "-ac", "1", "-ar", str(target_sr)]
        )
        
        # Validate the audio with librosa
        wav, sr = librosa.load(temp_path, sr=target_sr, mono=True)
        
        # Check for invalid values
        if np.any(np.isnan(wav)) or np.any(np.isinf(wav)):
            raise ValueError("Audio contains NaN or infinite values")
        
        # Ensure reasonable amplitude range
        if np.max(np.abs(wav)) < 1e-6:
            raise ValueError("Audio signal is too quiet")
        
        # Re-save the validated audio
        import soundfile as sf
        sf.write(temp_path, wav, sr)
        
        return temp_path
        
    except Exception as e:
        print(f"Audio preprocessing failed: {e}")
        raise ValueError(f"Failed to process audio: {str(e)}")


if __name__ == "__main__":
    with gr.Blocks(title="MegaTTS3 Voice Cloning") as demo:
        gr.Markdown("# MegaTTS 3 Voice Cloning")
        # gr.Markdown("MegaTTS 3 is a text-to-speech model trained by ByteDance with exceptional voice cloning capabilities. The original authors did not release the WavVAE encoder, so voice cloning was not publicly available; however, thanks to [@ACoderPassBy](https://modelscope.cn/models/ACoderPassBy/MegaTTS-SFT)'s WavVAE encoder, we can now clone voices with MegaTTS 3!")
        # gr.Markdown("This is by no means the best voice cloning solution, but it works pretty well for some specific use-cases. Try out multiple and see which one works best for you.")
        # gr.Markdown("**Please use this Space responsibly and do not abuse it!** This demo is for research and educational purposes only!")
        # gr.Markdown("h/t to MysteryShack on Discord for the info about the unofficial WavVAE encoder!")
        # gr.Markdown("Upload a reference audio clip and enter text to generate speech with the cloned voice.")
    
        reference_audio = gr.Audio(
            label="Reference Audio",
            type="filepath"
        ).set_info((
            "A clean recording of the voice to clone. A few seconds of speech "
            "is enough, and background noise or room tone will be cloned along "
            "with the voice."
        ))
        text_input = gr.Textbox(
            label="Text to Generate",
            placeholder="Enter the text you want to synthesize...",
            info="Text for the cloned voice to speak, up to 2500 characters."
        )
    
        infer_timestep = gr.Slider(
            label="Inference Timesteps",
            value=32,
            minimum=1,
            maximum=100,
            step=1,
            info=(
                "Higher values (e.g., 50+) give clearer audio but take longer to "
                "generate."
            )
        )
        p_w = gr.Slider(
            label="Intelligibility Weight",
            value=1.4,
            minimum=0.1,
            maximum=5.0,
            step=0.1,
            info="Increase this if the speech sounds mumbled or unclear."
        )
        t_w = gr.Slider(
            label="Similarity Weight", 
            value=3.0,
            minimum=0.1,
            maximum=10.0,
            step=0.1,
            info="Increase this for a stronger likeness to the reference voice."
        )

        output_audio = gr.Audio(type="filepath", label="Synthesized Speech").set_info((
            "The text spoken in the reference voice."
        ))
    
        _ = build_endpoint(
            model_card=model_card,
            input_components=[reference_audio, text_input, infer_timestep, p_w, t_w],
            output_components=[output_audio],
            process_fn=process_fn
        )

    demo.queue().launch(pwa=True, show_error=True, server_name='0.0.0.0', server_port=7860)