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)