Update app.py

app.py

@@ -1,113 +1,64 @@
-import gradio as gr
-import google.generativeai as genai
-import numpy as np
-import re
+import spaces
+from snac import SNAC
 import torch
+import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import snapshot_download
-import logging
-import os
-import spaces
-import warnings
-from snac import SNAC
-
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-warnings.filterwarnings("ignore", category=UserWarning)
-warnings.filterwarnings("ignore", category=RuntimeWarning)
+from dotenv import load_dotenv
+load_dotenv()
 
+# Check if CUDA is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
-logger.info(f"Using device: {device}")
-
-model = None
-tokenizer = None
-snac_model = None
-
-@spaces.GPU()
-def load_model():
-    global model, tokenizer, snac_model
-    
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    print("Loading SNAC model...")
-    snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
-    snac_model = snac_model.to(device)
-
-    print("Loading Orpheus model...")
-    model_name = "canopylabs/orpheus-3b-0.1-ft"
-
-    snapshot_download(
-        repo_id=model_name,
-        allow_patterns=[
-            "config.json",
-            "*.safetensors",
-            "model.safetensors.index.json",
-            "tokenizer.json",
-            "tokenizer_config.json",
-            "special_tokens_map.json",
-            "vocab.json",
-            "merges.txt",
-        ],
-        ignore_patterns=[
-            "optimizer.pt",
-            "pytorch_model.bin",
-            "training_args.bin",
-            "scheduler.pt",
-        ]
-    )
-
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
-    model.to(device)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    print(f"Orpheus model and tokenizer loaded to {device}")
-
-@spaces.GPU()
-def generate_podcast_script(api_key, content, uploaded_file, duration, num_hosts):
-    try:
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
-        
-        combined_content = content or ""
-        if uploaded_file:
-            file_content = uploaded_file.read().decode('utf-8')
-            combined_content += "\n" + file_content if combined_content else file_content
-        
-        prompt = f"""
-        Create a podcast script for {'one person' if num_hosts == 1 else 'two people'} discussing:
-        {combined_content}
-        
-        Duration: {duration}. Include natural speech, humor, and occasional off-topic thoughts.
-        Use speech fillers like um, ah. Vary emotional tone.
-        
-        Format: {'Monologue' if num_hosts == 1 else 'Alternating dialogue'} without speaker labels.
-        Separate {'paragraphs' if num_hosts == 1 else 'lines'} with blank lines.
-        
-        Use emotion tags in angle brackets: <laugh>, <sigh>, <chuckle>, <cough>, <sniffle>, <groan>, <yawn>, <gasp>.
-        
-        Example: "I can't believe I stayed up all night <yawn> only to find out the meeting was canceled <groan>."
-        
-        Ensure content flows naturally and stays on topic. Match the script length to {duration}.
-        """
-        
-        response = model.generate_content(prompt)
-        return re.sub(r'[^a-zA-Z0-9\s.,?!<>]', '', response.text)
-    except Exception as e:
-        logger.error(f"Error generating podcast script: {str(e)}")
-        raise
 
+print("Loading SNAC model...")
+snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
+snac_model = snac_model.to(device)
+
+model_name = "canopylabs/orpheus-3b-0.1-ft"
+
+# Download only model config and safetensors
+snapshot_download(
+    repo_id=model_name,
+    allow_patterns=[
+        "config.json",
+        "*.safetensors",
+        "model.safetensors.index.json",
+    ],
+    ignore_patterns=[
+        "optimizer.pt",
+        "pytorch_model.bin",
+        "training_args.bin",
+        "scheduler.pt",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "special_tokens_map.json",
+        "vocab.json",
+        "merges.txt",
+        "tokenizer.*"
+    ]
+)
+
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+model.to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+print(f"Orpheus model loaded to {device}")
+
+# Process text prompt
 def process_prompt(prompt, voice, tokenizer, device):
     prompt = f"{voice}: {prompt}"
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids
     
-    start_token = torch.tensor([[128259]], dtype=torch.int64)
-    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)
+    start_token = torch.tensor([[128259]], dtype=torch.int64)  # Start of human
+    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # End of text, End of human
+    
+    modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)  # SOH SOT Text EOT EOH
     
-    modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
+    # No padding needed for single input
     attention_mask = torch.ones_like(modified_input_ids)
     
     return modified_input_ids.to(device), attention_mask.to(device)
 
+# Parse output tokens to audio
 def parse_output(generated_ids):
     token_to_find = 128257
     token_to_remove = 128258
@@ -133,47 +84,45 @@ def parse_output(generated_ids):
         trimmed_row = [t - 128266 for t in trimmed_row]
         code_lists.append(trimmed_row)
         
-    return code_lists[0]
+    return code_lists[0]  # Return just the first one for single sample
 
+# Redistribute codes for audio generation
 def redistribute_codes(code_list, snac_model):
-    try:
-        device = next(snac_model.parameters()).device
-        
-        layer_1, layer_2, layer_3 = [], [], []
-        for i in range((len(code_list)+1)//7):
-            layer_1.append(code_list[7*i])
-            layer_2.append(code_list[7*i+1]-4096)
-            layer_3.append(code_list[7*i+2]-(2*4096))
-            layer_3.append(code_list[7*i+3]-(3*4096))
-            layer_2.append(code_list[7*i+4]-(4*4096))
-            layer_3.append(code_list[7*i+5]-(5*4096))
-            layer_3.append(code_list[7*i+6]-(6*4096))
-        
-        codes = [
-            torch.tensor(layer_1, device=device).unsqueeze(0),
-            torch.tensor(layer_2, device=device).unsqueeze(0),
-            torch.tensor(layer_3, device=device).unsqueeze(0)
-        ]
-        
-        audio_hat = snac_model.decode(codes)
-        return audio_hat.detach().squeeze().cpu().numpy()
-    except Exception as e:
-        logger.error(f"Error in redistribute_codes: {e}", exc_info=True)
-        return None
+    device = next(snac_model.parameters()).device  # Get the device of SNAC model
+    
+    layer_1 = []
+    layer_2 = []
+    layer_3 = []
+    for i in range((len(code_list)+1)//7):
+        layer_1.append(code_list[7*i])
+        layer_2.append(code_list[7*i+1]-4096)
+        layer_3.append(code_list[7*i+2]-(2*4096))
+        layer_3.append(code_list[7*i+3]-(3*4096))
+        layer_2.append(code_list[7*i+4]-(4*4096))
+        layer_3.append(code_list[7*i+5]-(5*4096))
+        layer_3.append(code_list[7*i+6]-(6*4096))
+        
+    # Move tensors to the same device as the SNAC model
+    codes = [
+        torch.tensor(layer_1, device=device).unsqueeze(0),
+        torch.tensor(layer_2, device=device).unsqueeze(0),
+        torch.tensor(layer_3, device=device).unsqueeze(0)
+    ]
+    
+    audio_hat = snac_model.decode(codes)
+    return audio_hat.detach().squeeze().cpu().numpy()  # Always return CPU numpy array
 
+# Main generation function
 @spaces.GPU()
-def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens):
-    global tokenizer, model
-    if tokenizer is None or model is None:
-        print("Model or tokenizer is not initialized. Please ensure the model is properly loaded.")
-        return None
-
+def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens, progress=gr.Progress()):
     if not text.strip():
         return None
     
     try:
+        progress(0.1, "Processing text...")
         input_ids, attention_mask = process_prompt(text, voice, tokenizer, device)
         
+        progress(0.3, "Generating speech tokens...")
         with torch.no_grad():
             generated_ids = model.generate(
                 input_ids=input_ids,
@@ -187,7 +136,10 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
                 eos_token_id=128258,
             )
         
+        progress(0.6, "Processing speech tokens...")
         code_list = parse_output(generated_ids)
+        
+        progress(0.8, "Converting to audio...")
         audio_samples = redistribute_codes(code_list, snac_model)
         
         return (24000, audio_samples)  # Return sample rate and audio
@@ -195,71 +147,99 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
         print(f"Error generating speech: {e}")
         return None
 
-@spaces.GPU()
-def render_podcast(api_key, script, voice1, voice2, num_hosts):
-    try:
-        lines = [line for line in script.split('\n') if line.strip()]
-        audio_segments = []
-
-        for i, line in enumerate(lines):
-            voice = voice1 if num_hosts == 1 or i % 2 == 0 else voice2
-            result = generate_speech(line, voice, temperature=0.6, top_p=0.95, repetition_penalty=1.1, max_new_tokens=1200)
-            if result is not None:
-                sample_rate, audio = result
-                audio_segments.append(audio)
-
-        if not audio_segments:
-            logger.warning("No valid audio segments were generated.")
-            return (24000, np.zeros(24000, dtype=np.float32))
-
-        podcast_audio = np.concatenate(audio_segments)
-        podcast_audio = np.clip(podcast_audio, -1, 1)
-        podcast_audio = (podcast_audio * 32767).astype(np.int16)
-        
-        return (24000, podcast_audio)
-    except Exception as e:
-        logger.error(f"Error rendering podcast: {str(e)}")
-        raise
-
-with gr.Blocks() as demo:
-    gr.Markdown("# AI Podcast Generator")
-    
-    api_key_input = gr.Textbox(label="Enter your Gemini API Key", type="password")
+# Examples for the UI
+examples = [
+    ["Hey there my name is Tara, <chuckle> and I'm a speech generation model that can sound like a person.", "tara", 0.6, 0.95, 1.1, 1200],
+    ["I've also been taught to understand and produce paralinguistic things <sigh> like sighing, or <laugh> laughing, or <yawn> yawning!", "dan", 0.7, 0.95, 1.1, 1200],
+    ["I live in San Francisco, and have, uhm let's see, 3 billion 7 hundred ... <gasp> well, lets just say a lot of parameters.", "leah", 0.6, 0.9, 1.2, 1200],
+    ["Sometimes when I talk too much, I need to <cough> excuse myself. <sniffle> The weather has been quite cold lately.", "leo", 0.65, 0.9, 1.1, 1200],
+    ["Public speaking can be challenging. <groan> But with enough practice, anyone can become better at it.", "jess", 0.7, 0.95, 1.1, 1200],
+    ["The hike was exhausting but the view from the top was absolutely breathtaking! <sigh> It was totally worth it.", "mia", 0.65, 0.9, 1.15, 1200],
+    ["Did you hear that joke? <laugh> I couldn't stop laughing when I first heard it. <chuckle> It's still funny.", "zac", 0.7, 0.95, 1.1, 1200],
+    ["After running the marathon, I was so tired <yawn> and needed a long rest. <sigh> But I felt accomplished.", "zoe", 0.6, 0.95, 1.1, 1200]
+]
+
+# Available voices
+VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]
+
+# Available Emotive Tags
+EMOTIVE_TAGS = ["`<laugh>`", "`<chuckle>`", "`<sigh>`", "`<cough>`", "`<sniffle>`", "`<groan>`", "`<yawn>`", "`<gasp>`"]
+
+# Create Gradio interface
+with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
+    gr.Markdown(f"""
+    # 🎵 [Orpheus Text-to-Speech](https://github.com/canopyai/Orpheus-TTS)
+    Enter your text below and hear it converted to natural-sounding speech with the Orpheus TTS model.
     
+    ## Tips for better prompts:
+    - Add paralinguistic elements like {", ".join(EMOTIVE_TAGS)} or `uhm` for more human-like speech.
+    - Longer text prompts generally work better than very short phrases
+    - Increasing `repetition_penalty` and `temperature` makes the model speak faster.
+    """)    
     with gr.Row():
-        content_input = gr.Textbox(label="Paste your content (optional)", lines=4)
-        document_upload = gr.File(label="Upload Document (optional)")
-    
-    duration = gr.Radio(["1-5 min", "5-10 min", "10-15 min"], label="Estimated podcast duration", value="1-5 min")
-    num_hosts = gr.Radio([1, 2], label="Number of podcast hosts", value=2)
-    
-    voice_options = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]
-    voice1_select = gr.Dropdown(label="Select Voice 1", choices=voice_options, value="tara")
-    voice2_select = gr.Dropdown(label="Select Voice 2", choices=voice_options, value="leo")
-    
-    generate_btn = gr.Button("Generate Script")
-    script_output = gr.Textbox(label="Generated Script", lines=10)
-    
-    render_btn = gr.Button("Render Podcast")
-    audio_output = gr.Audio(label="Generated Podcast")
+        with gr.Column(scale=3):
+            text_input = gr.Textbox(
+                label="Text to speak", 
+                placeholder="Enter your text here...",
+                lines=5
+            )
+            voice = gr.Dropdown(
+                choices=VOICES, 
+                value="tara", 
+                label="Voice"
+            )
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                temperature = gr.Slider(
+                    minimum=0.1, maximum=1.5, value=0.6, step=0.05,
+                    label="Temperature", 
+                    info="Higher values (0.7-1.0) create more expressive but less stable speech"
+                )
+                top_p = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=0.95, step=0.05,
+                    label="Top P", 
+                    info="Nucleus sampling threshold"
+                )
+                repetition_penalty = gr.Slider(
+                    minimum=1.0, maximum=2.0, value=1.1, step=0.05,
+                    label="Repetition Penalty", 
+                    info="Higher values discourage repetitive patterns"
+                )
+                max_new_tokens = gr.Slider(
+                    minimum=100, maximum=2000, value=1200, step=100,
+                    label="Max Length", 
+                    info="Maximum length of generated audio (in tokens)"
+                )
+            
+            with gr.Row():
+                submit_btn = gr.Button("Generate Speech", variant="primary")
+                clear_btn = gr.Button("Clear")
+                
+        with gr.Column(scale=2):
+            audio_output = gr.Audio(label="Generated Speech", type="numpy")
+            
+    # Set up examples
+    gr.Examples(
+        examples=examples,
+        inputs=[text_input, voice, temperature, top_p, repetition_penalty, max_new_tokens],
+        outputs=audio_output,
+        fn=generate_speech,
+        cache_examples=True,
+    )
     
-    generate_btn.click(generate_podcast_script, 
-                       inputs=[api_key_input, content_input, document_upload, duration, num_hosts], 
-                       outputs=script_output)
+    # Set up event handlers
+    submit_btn.click(
+        fn=generate_speech,
+        inputs=[text_input, voice, temperature, top_p, repetition_penalty, max_new_tokens],
+        outputs=audio_output
+    )
     
-    render_btn.click(render_podcast, 
-                     inputs=[api_key_input, script_output, voice1_select, voice2_select, num_hosts], 
-                     outputs=audio_output)
-
-    num_hosts.change(lambda x: gr.update(visible=x == 2), 
-                     inputs=[num_hosts], 
-                     outputs=[voice2_select])
+    clear_btn.click(
+        fn=lambda: (None, None),
+        inputs=[],
+        outputs=[text_input, audio_output]
+    )
 
+# Launch the app
 if __name__ == "__main__":
-    try:
-        print("Loading models...")
-        load_model()  # This function should be defined to load all necessary models
-        print("Models loaded successfully. Launching the interface...")
-        demo.queue().launch(share=False, ssr_mode=False)
-    except Exception as e:
-        print(f"Error during startup: {e}")
+    demo.queue().launch(share=False, ssr_mode=False)