Ojas1024/ChatterboxTEST

1import random2import numpy as np3import torch4from chatterbox.src.chatterbox.tts import ChatterboxTTS5import gradio as gr6import spaces7 8DEVICE = "cpu"9print(f"🚀 Running on device: {DEVICE}")10 11# --- Global Model Initialization ---12MODEL = None13 14def get_or_load_model():15    """Loads the ChatterboxTTS model if it hasn't been loaded already,16    and ensures it's on the correct device."""17    global MODEL18    if MODEL is None:19        print("Model not loaded, initializing...")20        try:21            MODEL = ChatterboxTTS.from_pretrained(DEVICE)22            if hasattr(MODEL, 'to') and str(MODEL.device) != DEVICE:23                MODEL.to(DEVICE)24            print(f"Model loaded successfully. Internal device: {getattr(MODEL, 'device', 'N/A')}")25        except Exception as e:26            print(f"Error loading model: {e}")27            raise28    return MODEL29 30# Attempt to load the model at startup.31try:32    get_or_load_model()33except Exception as e:34    print(f"CRITICAL: Failed to load model on startup. Application may not function. Error: {e}")35 36def set_seed(seed: int):37    """Sets the random seed for reproducibility across torch, numpy, and random."""38    torch.manual_seed(seed)39    if DEVICE == "cuda":40        torch.cuda.manual_seed(seed)41        torch.cuda.manual_seed_all(seed)42    random.seed(seed)43    np.random.seed(seed)44 45@spaces.GPU46def generate_tts_audio(47    text_input: str,48    audio_prompt_path_input: str,49    exaggeration_input: float,50    temperature_input: float,51    seed_num_input: int,52    cfgw_input: float53) -> tuple[int, np.ndarray]:54    """55    Generates TTS audio using the ChatterboxTTS model.56 57    Args:58        text_input: The text to synthesize (max 300 characters).59        audio_prompt_path_input: Path to the reference audio file.60        exaggeration_input: Exaggeration parameter for the model.61        temperature_input: Temperature parameter for the model.62        seed_num_input: Random seed (0 for random).63        cfgw_input: CFG/Pace weight.64 65    Returns:66        A tuple containing the sample rate (int) and the audio waveform (numpy.ndarray).67    """68    current_model = get_or_load_model()69 70    if current_model is None:71        raise RuntimeError("TTS model is not loaded.")72 73    if seed_num_input != 0:74        set_seed(int(seed_num_input))75 76    print(f"Generating audio for text: '{text_input[:50]}...'")77    wav = current_model.generate(78        text_input[:300],  # Truncate text to max chars79        audio_prompt_path=audio_prompt_path_input,80        exaggeration=exaggeration_input,81        temperature=temperature_input,82        cfg_weight=cfgw_input,83    )84    print("Audio generation complete.")85    return (current_model.sr, wav.squeeze(0).numpy())86 87with gr.Blocks() as demo:88    gr.Markdown(89        """90        # Chatterbox TTS Demo91        Generate high-quality speech from text with reference audio styling.92        """93    )94    with gr.Row():95        with gr.Column():96            text = gr.Textbox(97                value="Now let's make my mum's favourite. So three mars bars into the pan. Then we add the tuna and just stir for a bit, just let the chocolate and fish infuse. A sprinkle of olive oil and some tomato ketchup. Now smell that. Oh boy this is going to be incredible.",98                label="Text to synthesize (max chars 300)",99                max_lines=5100            )101            ref_wav = gr.Audio(102                sources=["upload", "microphone"],103                type="filepath",104                label="Reference Audio File (Optional)",105                value="https://storage.googleapis.com/chatterbox-demo-samples/prompts/female_shadowheart4.flac"106            )107            exaggeration = gr.Slider(108                0.25, 2, step=.05, label="Exaggeration (Neutral = 0.5, extreme values can be unstable)", value=.5109            )110            cfg_weight = gr.Slider(111                0.2, 1, step=.05, label="CFG/Pace", value=0.5112            )113 114            with gr.Accordion("More options", open=False):115                seed_num = gr.Number(value=0, label="Random seed (0 for random)")116                temp = gr.Slider(0.05, 5, step=.05, label="Temperature", value=.8)117 118            run_btn = gr.Button("Generate", variant="primary")119 120        with gr.Column():121            audio_output = gr.Audio(label="Output Audio")122 123    run_btn.click(124        fn=generate_tts_audio,125        inputs=[126            text,127            ref_wav,128            exaggeration,129            temp,130            seed_num,131            cfg_weight,132        ],133        outputs=[audio_output],134    )135 136demo.launch()