Update app.py

app.py

@@ -4,8 +4,13 @@ import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import snapshot_download
-from dotenv import load_dotenv
-load_dotenv()
+import google.generativeai as genai
+import re
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Check if CUDA is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -16,49 +21,57 @@ 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
+@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
+
 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)  # Start of human
-    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # End of text, End of human
+    start_token = torch.tensor([[128259]], dtype=torch.int64)
+    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)
     
-    modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)  # SOH SOT Text EOT EOH
-    
-    # No padding needed for single input
+    modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
     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
@@ -84,35 +97,8 @@ def parse_output(generated_ids):
         trimmed_row = [t - 128266 for t in trimmed_row]
         code_lists.append(trimmed_row)
         
-    return code_lists[0]  # Return just the first one for single sample
+    return code_lists[0]
 
-# Redistribute codes for audio generation
-def redistribute_codes(code_list, snac_model):
-    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, progress=gr.Progress()):
     if not text.strip():
@@ -125,13 +111,13 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
         progress(0.3, "Generating speech tokens...")
         with torch.no_grad():
             generated_ids = model.generate(
-                input_ids=input_ids,
+                input_ids,
                 attention_mask=attention_mask,
-                max_new_tokens=max_new_tokens,
                 do_sample=True,
                 temperature=temperature,
                 top_p=top_p,
                 repetition_penalty=repetition_penalty,
+                max_new_tokens=max_new_tokens,
                 num_return_sequences=1,
                 eos_token_id=128258,
             )
@@ -147,87 +133,43 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
         print(f"Error generating speech: {e}")
         return None
 
-# 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.Blocks(title="AI Podcaster") as demo:
     with gr.Row():
-        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.Column(scale=1):
+            gemini_api_key = gr.Textbox(label="Gemini API Key", type="password")
+            content = gr.Textbox(label="Content", lines=5)
+            uploaded_file = gr.File(label="Upload File")
+            duration = gr.Slider(minimum=1, maximum=60, value=5, step=1, label="Duration (minutes)")
+            num_hosts = gr.Radio(["1", "2"], label="Number of Hosts", value="1")
+            generate_script_btn = gr.Button("Generate Podcast Script")
+
+        with gr.Column(scale=2):
+            script_output = gr.Textbox(label="Generated Script", lines=10)
+            text_input = gr.Textbox(label="Text to speak", lines=5)
+            voice = gr.Dropdown(choices=["Narrator", "Male", "Female"], value="Narrator", 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():
+                temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
+                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top P")
             
             with gr.Row():
-                submit_btn = gr.Button("Generate Speech", variant="primary")
-                clear_btn = gr.Button("Clear")
+                repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, value=1.2, step=0.1, label="Repetition Penalty")
+                max_new_tokens = gr.Slider(minimum=100, maximum=1000, value=500, step=50, label="Max New Tokens")
+            
+            submit_btn = gr.Button("Generate Speech")
+            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,
+    # Set up event handlers
+    generate_script_btn.click(
+        fn=generate_podcast_script,
+        inputs=[gemini_api_key, content, uploaded_file, 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],
@@ -242,4 +184,4 @@ with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
 
 # Launch the app
 if __name__ == "__main__":
-    demo.queue().launch(share=False, ssr_mode=False)
+    demo.queue().launch(share=False, ssr_mode=False)