app.py

import os

# Set a default watermark key to avoid the NoneType error
# Do this BEFORE any imports
if "WATERMARK_KEY" not in os.environ:
    os.environ["WATERMARK_KEY"] = "0 0 0 0"  # Default placeholder

import subprocess
import tempfile

import gradio as gr
import numpy as np
import spaces
import torch
import torchaudio
from generator import Segment, load_csm_1b
from watermarking import watermark

# Simplified environment variables handling
gpu_timeout = int(os.getenv("GPU_TIMEOUT", 60))

SPACE_INTRO_TEXT = """\
# Sesame CSM 1B

Generate from CSM 1B (Conversational Speech Model). 
Code is available on GitHub: [SesameAILabs/csm](https://github.com/SesameAILabs/csm). 
Checkpoint is [hosted on HuggingFace](https://huggingface.co/sesame/csm-1b).

---

"""

CONVO_INTRO_TEXT = """\
## Conversation content

Each line is an utterance in the conversation to generate. Speakers alternate between A and B, starting with speaker A.
"""

DEFAULT_CONVERSATION = """\
Hey how are you doing.
Pretty good, pretty good.
I'm great, so happy to be speaking to you.
Me too, this is some cool stuff huh?
Yeah, I've been reading more about speech generation, and it really seems like context is important.
Definitely.
"""

SPEAKER_PROMPTS = {
    "conversational_a": {
        "text": (
            "like revising for an exam I'd have to try and like keep up the momentum because I'd "
            "start really early I'd be like okay I'm gonna start revising now and then like "
            "you're revising for ages and then I just like start losing steam I didn't do that "
            "for the exam we had recently to be fair that was a more of a last minute scenario "
            "but like yeah I'm trying to like yeah I noticed this yesterday that like Mondays I "
            "sort of start the day with this not like a panic but like a"
        ),
        "audio": "prompts/conversational_a.wav",
    },
    "conversational_b": {
        "text": (
            "like a super Mario level. Like it's very like high detail. And like, once you get "
            "into the park, it just like, everything looks like a computer game and they have all "
            "these, like, you know, if, if there's like a, you know, like in a Mario game, they "
            "will have like a question block. And if you like, you know, punch it, a coin will "
            "come out. So like everyone, when they come into the park, they get like this little "
            "bracelet and then you can go punching question blocks around."
        ),
        "audio": "prompts/conversational_b.wav",
    },
    "read_speech_a": {
        "text": (
            "And Lake turned round upon me, a little abruptly, his odd yellowish eyes, a little "
            "like those of the sea eagle, and the ghost of his smile that flickered on his "
            "singularly pale face, with a stern and insidious look, confronted me."
        ),
        "audio": "prompts/read_speech_a.wav",
    },
    "read_speech_b": {
        "text": (
            "He was such a big boy that he wore high boots and carried a jack knife. He gazed and "
            "gazed at the cap, and could not keep from fingering the blue tassel."
        ),
        "audio": "prompts/read_speech_b.wav",
    },
}

device = "cuda" if torch.cuda.is_available() else "cpu"
generator = load_csm_1b(device=device)


def convert_ebook_to_txt(ebook_path):
    """Convert an ebook file to text using Calibre's ebook-convert."""
    if not ebook_path:
        return None
    
    # Create a temporary file for the output
    with tempfile.NamedTemporaryFile(suffix='.txt', delete=False) as temp_txt:
        txt_path = temp_txt.name
    
    try:
        # Run ebook-convert from Calibre
        subprocess.run(
            ["ebook-convert", ebook_path, txt_path], 
            check=True, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE
        )
        
        # Read the converted text
        with open(txt_path, 'r', encoding='utf-8') as f:
            text_content = f.read()
        
        # Clean up
        os.unlink(txt_path)
        
        # Format the text into alternating lines for conversation
        lines = [line.strip() for line in text_content.split('.') if line.strip()]
        formatted_lines = []
        
        # Take up to 20 sentences to avoid extremely long conversations
        for i, line in enumerate(lines[:20]):
            formatted_lines.append(line + ".")
        
        return "\n".join(formatted_lines)
    
    except Exception as e:
        if os.path.exists(txt_path):
            os.unlink(txt_path)
        raise gr.Error(f"Error converting ebook: {str(e)}")


@spaces.GPU(duration=gpu_timeout)
def infer(
    text_prompt_speaker_a,
    text_prompt_speaker_b,
    audio_prompt_speaker_a,
    audio_prompt_speaker_b,
    gen_conversation_input,
) -> tuple[np.ndarray, int]:
    # Estimate token limit
    if len(gen_conversation_input.strip() + text_prompt_speaker_a.strip() + text_prompt_speaker_b.strip()) >= 2000:
        raise gr.Error("Prompts and conversation too long.", duration=30)

    try:
        return _infer(
            text_prompt_speaker_a,
            text_prompt_speaker_b,
            audio_prompt_speaker_a,
            audio_prompt_speaker_b,
            gen_conversation_input,
        )
    except ValueError as e:
        raise gr.Error(f"Error generating audio: {e}", duration=120)


def _infer(
    text_prompt_speaker_a,
    text_prompt_speaker_b,
    audio_prompt_speaker_a,
    audio_prompt_speaker_b,
    gen_conversation_input,
) -> tuple[np.ndarray, int]:
    audio_prompt_a = prepare_prompt(text_prompt_speaker_a, 0, audio_prompt_speaker_a)
    audio_prompt_b = prepare_prompt(text_prompt_speaker_b, 1, audio_prompt_speaker_b)

    prompt_segments: list[Segment] = [audio_prompt_a, audio_prompt_b]
    generated_segments: list[Segment] = []

    conversation_lines = [line.strip() for line in gen_conversation_input.strip().split("\n") if line.strip()]
    for i, line in enumerate(conversation_lines):
        # Alternating speakers A and B, starting with A
        speaker_id = i % 2

        audio_tensor = generator.generate(
            text=line,
            speaker=speaker_id,
            context=prompt_segments + generated_segments,
            max_audio_length_ms=30_000,
        )
        generated_segments.append(Segment(text=line, speaker=speaker_id, audio=audio_tensor))

    # Concatenate all generations and convert to 16-bit int format
    audio_tensors = [segment.audio for segment in generated_segments]
    audio_tensor = torch.cat(audio_tensors, dim=0)

    # Get the watermark key from environment
    watermark_key = list(map(int, os.getenv("WATERMARK_KEY").split(" ")))
    
    # Watermarking
    audio_tensor, wm_sample_rate = watermark(
        generator._watermarker, audio_tensor, generator.sample_rate, watermark_key
    )
    audio_tensor = torchaudio.functional.resample(
        audio_tensor, orig_freq=wm_sample_rate, new_freq=generator.sample_rate
    )

    audio_array = (audio_tensor * 32768).to(torch.int16).cpu().numpy()

    return generator.sample_rate, audio_array


def prepare_prompt(text: str, speaker: int, audio_path: str) -> Segment:
    audio_tensor, _ = load_prompt_audio(audio_path)
    return Segment(text=text, speaker=speaker, audio=audio_tensor)


def load_prompt_audio(audio_path: str) -> torch.Tensor:
    audio_tensor, sample_rate = torchaudio.load(audio_path)
    audio_tensor = audio_tensor.squeeze(0)
    if sample_rate != generator.sample_rate:
        audio_tensor = torchaudio.functional.resample(
            audio_tensor, orig_freq=sample_rate, new_freq=generator.sample_rate
        )
    return audio_tensor, generator.sample_rate


def create_speaker_prompt_ui(speaker_name: str):
    speaker_dropdown = gr.Dropdown(
        choices=list(SPEAKER_PROMPTS.keys()), label="Select a predefined speaker", value=speaker_name
    )
    with gr.Accordion("Or add your own voice prompt", open=False):
        text_prompt_speaker = gr.Textbox(label="Speaker prompt", lines=4, value=SPEAKER_PROMPTS[speaker_name]["text"])
        audio_prompt_speaker = gr.Audio(
            label="Speaker prompt", type="filepath", value=SPEAKER_PROMPTS[speaker_name]["audio"]
        )

    return speaker_dropdown, text_prompt_speaker, audio_prompt_speaker


def process_ebook(ebook_file):
    if ebook_file is None:
        return None
    text_content = convert_ebook_to_txt(ebook_file)
    return text_content


def update_input_method(choice):
    if choice == "text_input":
        return gr.update(visible=True), gr.update(visible=False), None
    else:
        return gr.update(visible=False), gr.update(visible=True), None


with gr.Blocks() as app:
    gr.Markdown(SPACE_INTRO_TEXT)
    gr.Markdown("## Voices")
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Speaker A")
            speaker_a_dropdown, text_prompt_speaker_a, audio_prompt_speaker_a = create_speaker_prompt_ui(
                "conversational_a"
            )

        with gr.Column():
            gr.Markdown("### Speaker B")
            speaker_b_dropdown, text_prompt_speaker_b, audio_prompt_speaker_b = create_speaker_prompt_ui(
                "conversational_b"
            )

    def update_audio(speaker):
        if speaker in SPEAKER_PROMPTS:
            return SPEAKER_PROMPTS[speaker]["audio"]
        return None

    def update_text(speaker):
        if speaker in SPEAKER_PROMPTS:
            return SPEAKER_PROMPTS[speaker]["text"]
        return None

    speaker_a_dropdown.change(fn=update_audio, inputs=[speaker_a_dropdown], outputs=[audio_prompt_speaker_a])
    speaker_b_dropdown.change(fn=update_audio, inputs=[speaker_b_dropdown], outputs=[audio_prompt_speaker_b])

    speaker_a_dropdown.change(fn=update_text, inputs=[speaker_a_dropdown], outputs=[text_prompt_speaker_a])
    speaker_b_dropdown.change(fn=update_text, inputs=[speaker_b_dropdown], outputs=[text_prompt_speaker_b])

    gr.Markdown(CONVO_INTRO_TEXT)
    
    # Radio button for selecting input method
    input_method = gr.Radio(
        ["Direct text input", "Upload ebook file"], 
        label="Choose input method", 
        value="Direct text input"
    )
    
    # Container for text input method
    with gr.Group(visible=True) as text_input_group:
        gen_conversation_input = gr.TextArea(label="Conversation", lines=20, value=DEFAULT_CONVERSATION)
    
    # Container for ebook upload method
    with gr.Group(visible=False) as ebook_input_group:
        ebook_file = gr.File(label="Upload ebook file (will be converted using Calibre)", file_types=[".epub", ".mobi", ".azw", ".azw3", ".fb2", ".pdf"])
        process_ebook_btn = gr.Button("Process Ebook")
    
    input_method.change(
        fn=lambda choice: update_input_method("text_input" if choice == "Direct text input" else "ebook"),
        inputs=[input_method],
        outputs=[text_input_group, ebook_input_group, gen_conversation_input]
    )
    
    process_ebook_btn.click(
        fn=process_ebook,
        inputs=[ebook_file],
        outputs=[gen_conversation_input]
    )
    
    generate_btn = gr.Button("Generate conversation", variant="primary")
    gr.Markdown("GPU time limited to 3 minutes, for longer usage duplicate the space.")
    audio_output = gr.Audio(label="Synthesized audio")

    generate_btn.click(
        infer,
        inputs=[
            text_prompt_speaker_a,
            text_prompt_speaker_b,
            audio_prompt_speaker_a,
            audio_prompt_speaker_b,
            gen_conversation_input,
        ],
        outputs=[audio_output],
    )

app.launch(ssr_mode=True)