Update app.py

app.py

@@ -1,59 +1,84 @@
 import gradio as gr
-import random
-import time
-import os
-import torch
-from pathlib import Path
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from pydub import AudioSegment
 import json
 import uuid
 import edge_tts
 import asyncio
 import aiofiles
+import os
+import time
 import mimetypes
-from typing import List
-
-from PyPDF2 import PdfReader
-from pydub import AudioSegment
+import torch
+from typing import List, Dict
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Define model name clearly
-MODEL_NAME = "unsloth/gemma-3-1b-pt"
+# Constants
+MAX_FILE_SIZE_MB = 20
+MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024  # Convert MB to bytes
 
-# Device setup
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
+MODEL_ID = "HuggingFaceH4/zephyr-7b-alpha"
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
+    MODEL_ID,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-).eval().to(device)
+    device_map="auto"
+).eval()
 
-# Constants
-MAX_FILE_SIZE_MB = 20
-MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024
 
 class PodcastGenerator:
     def __init__(self):
         pass
 
-    async def generate_script(self, prompt: str, language: str, api_key: str, file_obj=None, progress=None):
+    async def generate_script(self, prompt: str, language: str, api_key: str, file_obj=None, progress=None) -> Dict:
         example = """
 {
     "topic": "AGI",
     "podcast": [
-        {"speaker": 2, "line": "So, AGI, huh? Seems like everyone's talking about it these days."},
-        {"speaker": 1, "line": "Yeah, it's definitely having a moment, isn't it?"},
-        {"speaker": 2, "line": "It is and for good reason, right? I mean, you've been digging into this stuff, listening to the podcasts and everything. What really stood out to you? What got you hooked?"},
-        {"speaker": 1, "line": "It's easy to get lost in the noise, for sure."},
-        {"speaker": 2, "line": "Exactly. So how about we try to cut through some of that, shall we?"},
-        {"speaker": 1, "line": "Sounds like a plan."},
-        {"speaker": 2, "line": "It certainly is and on that note, we'll wrap up this deep dive. Thanks for listening, everyone."},
-        {"speaker": 1, "line": "Peace."}
+        {
+            "speaker": 2,
+            "line": "So, AGI, huh? Seems like everyone's talking about it these days."
+        },
+        {
+            "speaker": 1,
+            "line": "Yeah, it's definitely having a moment, isn't it?"
+        },
+        {
+            "speaker": 2,
+            "line": "It is and for good reason, right? I mean, you've been digging into this stuff, listening to the podcasts and everything. What really stood out to you? What got you hooked?"
+        },
+        {
+            "speaker": 1,
+            "line": "I like that. It really is."
+        },
+        {
+            "speaker": 2,
+            "line": "And honestly, that's a responsibility that extends beyond just the researchers and the policymakers."
+        },
+        {
+            "speaker": 1,
+            "line": "100%"
+        },
+        {
+            "speaker": 2,
+            "line": "So to everyone listening out there I'll leave you with this. As AGI continues to develop, what role do you want to play in shaping its future?"
+        },
+        {
+            "speaker": 1,
+            "line": "That's a question worth pondering."
+        },
+        {
+            "speaker": 2,
+            "line": "It certainly is and on that note, we'll wrap up this deep dive. Thanks for listening, everyone."
+        },
+        {
+            "speaker": 1,
+            "line": "Peace."
+        }
     ]
 }
         """
+
         if language == "Auto Detect":
             language_instruction = "- The podcast MUST be in the same language as the user input."
         else:
@@ -71,7 +96,7 @@ You are a professional podcast generator. Your task is to generate a professiona
 Follow this example structure:
 {example}
 """
-        # Build the user prompt
+        # Construct system and user prompt
         if prompt and file_obj:
             user_prompt = f"Please generate a podcast script based on the uploaded file following user input:\n{prompt}"
         elif prompt:
@@ -79,114 +104,186 @@ Follow this example structure:
         else:
             user_prompt = "Please generate a podcast script based on the uploaded file."
 
-        # Handle file content
+        # NOTE: file_obj cannot be passed to a text-only LLM
         if file_obj:
-            file_size = getattr(file_obj, 'size', os.path.getsize(file_obj.name))
-            if file_size > MAX_FILE_SIZE_BYTES:
-                raise Exception(f"File size exceeds the {MAX_FILE_SIZE_MB}MB limit. Please upload a smaller file.")
-            ext = os.path.splitext(file_obj.name)[1].lower()
-            if ext == '.pdf':
-                reader = PdfReader(file_obj)
-                text = "\n\n".join(page.extract_text() or '' for page in reader.pages)
-            else:
-                raw = file_obj.read() if hasattr(file_obj, 'read') else await aiofiles.open(file_obj.name, 'rb').read()
-                text = raw.decode(errors='ignore')
-            user_prompt += f"\n\n―― FILE CONTENT ――\n{text}"
-
-        prompt_text = system_prompt + "\n" + user_prompt
+            print("Warning: Uploaded file is ignored in this version because external LLM does not support file input.")
+
+        # Build prompt
+        full_prompt = f"""{system_prompt}
+
+{user_prompt}
+
+Return the result strictly as a JSON object in the format:
+{{
+  "topic": "{prompt}",
+  "podcast": [
+    {{ "speaker": 1, "line": "..." }},
+    {{ "speaker": 2, "line": "..." }}
+  ]
+}}
+"""
+
         try:
             if progress:
                 progress(0.3, "Generating podcast script...")
-            def hf_generate(p):
-                inputs = tokenizer(p, return_tensors="pt").to(model.device)
-                outs = model.generate(
-                    **inputs,
-                    max_new_tokens=1024,
-                    do_sample=True,
-                    temperature=1.0
-                )
-                return tokenizer.decode(outs[0], skip_special_tokens=True)
-            generated_text = await asyncio.wait_for(
-                asyncio.to_thread(hf_generate, prompt_text),
-                timeout=60
-            )
-        except asyncio.TimeoutError:
-            raise Exception("The script generation request timed out. Please try again later.")
+
+            inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+            output = model.generate(**inputs, max_new_tokens=1024)
+            text = tokenizer.decode(output[0], skip_special_tokens=True)
+
         except Exception as e:
             raise Exception(f"Failed to generate podcast script: {e}")
 
+        print(f"Generated podcast script:\n{text}")
+
         if progress:
             progress(0.4, "Script generated successfully!")
 
-        return json.loads(generated_text)
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            raise Exception("The model did not return valid JSON. Please refine the prompt.")
+
+    
+    async def _read_file_bytes(self, file_obj) -> bytes:
+        """Read file bytes from a file object"""
+        # Check file size before reading
+        if hasattr(file_obj, 'size'):
+            file_size = file_obj.size
+        else:
+            file_size = os.path.getsize(file_obj.name)
+            
+        if file_size > MAX_FILE_SIZE_BYTES:
+            raise Exception(f"File size exceeds the {MAX_FILE_SIZE_MB}MB limit. Please upload a smaller file.")
+            
+        if hasattr(file_obj, 'read'):
+            return file_obj.read()
+        else:
+            async with aiofiles.open(file_obj.name, 'rb') as f:
+                return await f.read()
+    
+    def _get_mime_type(self, filename: str) -> str:
+        """Determine MIME type based on file extension"""
+        ext = os.path.splitext(filename)[1].lower()
+        if ext == '.pdf':
+            return "application/pdf"
+        elif ext == '.txt':
+            return "text/plain"
+        else:
+            # Fallback to the default mime type detector
+            mime_type, _ = mimetypes.guess_type(filename)
+            return mime_type or "application/octet-stream"
 
     async def tts_generate(self, text: str, speaker: int, speaker1: str, speaker2: str) -> str:
         voice = speaker1 if speaker == 1 else speaker2
         speech = edge_tts.Communicate(text, voice)
+        
         temp_filename = f"temp_{uuid.uuid4()}.wav"
         try:
-            await asyncio.wait_for(speech.save(temp_filename), timeout=30)
+            # Add timeout to TTS generation
+            await asyncio.wait_for(speech.save(temp_filename), timeout=30)  # 30 seconds timeout
             return temp_filename
         except asyncio.TimeoutError:
-            if os.path.exists(temp_filename): os.remove(temp_filename)
+            if os.path.exists(temp_filename):
+                os.remove(temp_filename)
             raise Exception("Text-to-speech generation timed out. Please try with a shorter text.")
         except Exception as e:
-            if os.path.exists(temp_filename): os.remove(temp_filename)
+            if os.path.exists(temp_filename):
+                os.remove(temp_filename)
             raise e
 
     async def combine_audio_files(self, audio_files: List[str], progress=None) -> str:
-        if progress: progress(0.9, "Combining audio files...")
+        if progress:
+            progress(0.9, "Combining audio files...")
+            
         combined_audio = AudioSegment.empty()
         for audio_file in audio_files:
             combined_audio += AudioSegment.from_file(audio_file)
-            os.remove(audio_file)
+            os.remove(audio_file)  # Clean up temporary files
+
         output_filename = f"output_{uuid.uuid4()}.wav"
         combined_audio.export(output_filename, format="wav")
-        if progress: progress(1.0, "Podcast generated successfully!")
+        
+        if progress:
+            progress(1.0, "Podcast generated successfully!")
+            
         return output_filename
 
     async def generate_podcast(self, input_text: str, language: str, speaker1: str, speaker2: str, api_key: str, file_obj=None, progress=None) -> str:
         try:
-            if progress: progress(0.1, "Starting podcast generation...")
+            if progress:
+                progress(0.1, "Starting podcast generation...")
+                
+            # Set overall timeout for the entire process
             return await asyncio.wait_for(
                 self._generate_podcast_internal(input_text, language, speaker1, speaker2, api_key, file_obj, progress),
-                timeout=600
+                timeout=600  # 10 minutes total timeout
             )
         except asyncio.TimeoutError:
             raise Exception("The podcast generation process timed out. Please try with shorter text or try again later.")
         except Exception as e:
             raise Exception(f"Error generating podcast: {str(e)}")
-
+    
     async def _generate_podcast_internal(self, input_text: str, language: str, speaker1: str, speaker2: str, api_key: str, file_obj=None, progress=None) -> str:
-        if progress: progress(0.2, "Generating podcast script...")
+        if progress:
+            progress(0.2, "Generating podcast script...")
+            
         podcast_json = await self.generate_script(input_text, language, api_key, file_obj, progress)
-        if progress: progress(0.5, "Converting text to speech...")
+        
+        if progress:
+            progress(0.5, "Converting text to speech...")
+        
+        # Process TTS in batches for concurrent processing
         audio_files = []
         total_lines = len(podcast_json['podcast'])
-        batch_size = 10
+        
+        # Define batch size to control concurrency
+        batch_size = 10  # Adjust based on system resources
+        
+        # Process in batches
         for batch_start in range(0, total_lines, batch_size):
             batch_end = min(batch_start + batch_size, total_lines)
             batch = podcast_json['podcast'][batch_start:batch_end]
-            tts_tasks = [self.tts_generate(item['line'], item['speaker'], speaker1, speaker2) for item in batch]
+            
+            # Create tasks for concurrent processing
+            tts_tasks = []
+            for item in batch:
+                tts_task = self.tts_generate(item['line'], item['speaker'], speaker1, speaker2)
+                tts_tasks.append(tts_task)
+            
             try:
+                # Process batch concurrently
                 batch_results = await asyncio.gather(*tts_tasks, return_exceptions=True)
-                for result in batch_results:
+                
+                # Check for exceptions and handle results
+                for i, result in enumerate(batch_results):
                     if isinstance(result, Exception):
+                        # Clean up any files already created
                         for file in audio_files:
-                            if os.path.exists(file): os.remove(file)
+                            if os.path.exists(file):
+                                os.remove(file)
                         raise Exception(f"Error generating speech: {str(result)}")
-                    audio_files.append(result)
+                    else:
+                        audio_files.append(result)
+                        
+                # Update progress
                 if progress:
-                    progress(0.5 + (0.4 * (batch_end / total_lines)), f"Processed {batch_end}/{total_lines} speech segments...")
+                    current_progress = 0.5 + (0.4 * (batch_end / total_lines))
+                    progress(current_progress, f"Processed {batch_end}/{total_lines} speech segments...")
+            
             except Exception as e:
+                # Clean up any files already created
                 for file in audio_files:
-                    if os.path.exists(file): os.remove(file)
+                    if os.path.exists(file):
+                        os.remove(file)
                 raise Exception(f"Error in batch TTS generation: {str(e)}")
-        combined = await self.combine_audio_files(audio_files, progress)
-        return combined
+        
+        combined_audio = await self.combine_audio_files(audio_files, progress)
+        return combined_audio
 
 async def process_input(input_text: str, input_file, language: str, speaker1: str, speaker2: str, api_key: str = "", progress=None) -> str:
     start_time = time.time()
+
     voice_names = {
         "Andrew - English (United States)": "en-US-AndrewMultilingualNeural",
         "Ava - English (United States)": "en-US-AvaMultilingualNeural",
@@ -197,44 +294,124 @@ async def process_input(input_text: str, input_file, language: str, speaker1: st
         "Remy - French (France)": "fr-FR-RemyMultilingualNeural",
         "Vivienne - French (France)": "fr-FR-VivienneMultilingualNeural"
     }
+
     speaker1 = voice_names[speaker1]
     speaker2 = voice_names[speaker2]
+
     try:
-        if progress: progress(0.05, "Processing input...")
+        if progress:
+            progress(0.05, "Processing input...")
+
         if not api_key:
-            api_key = "saf"
+            api_key = os.getenv("GENAI_API_KEY")
             if not api_key:
                 raise Exception("No API key provided. Please provide a Gemini API key.")
-        generator = PodcastGenerator()
-        output = await generator.generate_podcast(input_text, lan
-                                                  guage, speaker1, speaker2, api_key, input_file, progress)
-        print(f"Total podcast generation time: {time.time() - start_time:.2f} seconds")
-        return output
+
+        podcast_generator = PodcastGenerator()
+        podcast = await podcast_generator.generate_podcast(input_text, language, speaker1, speaker2, api_key, input_file, progress)
+
+        end_time = time.time()
+        print(f"Total podcast generation time: {end_time - start_time:.2f} seconds")
+        return podcast
+        
     except Exception as e:
-        msg = str(e)
-        if "rate limit" in msg.lower():
+        # Ensure we show a user-friendly error
+        error_msg = str(e)
+        if "rate limit" in error_msg.lower():
             raise Exception("Rate limit exceeded. Please try again later or use your own API key.")
-        elif "timeout" in msg.lower():
-            raise Exception("The request timed out... Please try with shorter text.")
+        elif "timeout" in error_msg.lower():
+            raise Exception("The request timed out. This could be due to server load or the length of your input. Please try again with shorter text.")
         else:
-            raise Exception(f"Error: {msg}")
+            raise Exception(f"Error: {error_msg}")
+
 # Gradio UI
-with gr.Blocks(title="PodcastGen 🎙️") as demo:
-    gr.Markdown("""
-# PodcastGen 🎙️
-Generate a 2-speaker podcast from text or PDF!
-""" )
-    with gr.Row():
-        with gr.Column():
-            input_text = gr.Textbox(label="Input Text", lines=10, placeholder="Enter podcast topic or paste text here...", elem_id="input_text")
-            input_file = gr.File(label="Or Upload a PDF or TXT file", file_types=[".pdf", ".txt"] )
-        with gr.Column():
-            language = gr.Dropdown(label="Podcast Language", choices=["Auto Detect","English","German","French","Spanish","Italian","Dutch","Portuguese","Russian","Chinese","Japanese","Korean","Other" ], value="Auto Detect")
-            speaker1 = gr.Dropdown(label="Speaker 1 Voice", choices=["Andrew - English (United States)","Ava - English (United States)","Brian - English (United States)","Emma - English (United States)","Florian - German (Germany)","Seraphina - German (Germany)","Remy - French (France)","Vivienne - French (France)" ], value="Andrew - English (United States)")
-            speaker2 = gr.Dropdown(label="Speaker 2 Voice", choices=["Andrew - English (United States)","Ava - English (United States)","Brian - English (United States)","Emma - English (United States)","Florian - German (Germany)","Seraphina - German (Germany)","Remy - French (France)","Vivienne - French (France)" ], value="Ava - English (United States)")
-            api_key = gr.Textbox(label="Gemini API Key (Optional)", type="password", placeholder="Needed only if you're getting rate limited.")
-    generate_btn = gr.Button("Generate Podcast 🎙️", variant="primary")
-    output_audio = gr.Audio(label="Generated Podcast", type="filepath", format="wav", elem_id="output_audio")
-    generate_btn.click(fn=process_input, inputs=[input_text, input_file, language, speaker1, speaker2, api_key], outputs=output_audio, show_progress=True)
- demo.queue()
- demo.launch(server_name="0.0.0.0", share=True, debug=True)
+def generate_podcast_gradio(input_text, input_file, language, speaker1, speaker2, api_key, progress=gr.Progress()):
+    # Handle the file if uploaded
+    file_obj = None
+    if input_file is not None:
+        file_obj = input_file
+        
+    # Use the progress function from Gradio
+    def progress_callback(value, text):
+        progress(value, text)
+
+    # Run the async function in the event loop
+    result = asyncio.run(process_input(
+        input_text, 
+        file_obj, 
+        language, 
+        speaker1, 
+        speaker2, 
+        api_key,
+        progress_callback
+    ))
+    
+    return result
+
+def main():
+    # Define language options
+    language_options = [
+        "Auto Detect",
+        "Afrikaans", "Albanian", "Amharic", "Arabic", "Armenian", "Azerbaijani",
+        "Bahasa Indonesian", "Bangla", "Basque", "Bengali", "Bosnian", "Bulgarian",
+        "Burmese", "Catalan", "Chinese Cantonese", "Chinese Mandarin",
+        "Chinese Taiwanese", "Croatian", "Czech", "Danish", "Dutch", "English",
+        "Estonian", "Filipino", "Finnish", "French", "Galician", "Georgian",
+        "German", "Greek", "Hebrew", "Hindi", "Hungarian", "Icelandic", "Irish",
+        "Italian", "Japanese", "Javanese", "Kannada", "Kazakh", "Khmer", "Korean",
+        "Lao", "Latvian", "Lithuanian", "Macedonian", "Malay", "Malayalam",
+        "Maltese", "Mongolian", "Nepali", "Norwegian Bokmål", "Pashto", "Persian",
+        "Polish", "Portuguese", "Romanian", "Russian", "Serbian", "Sinhala",
+        "Slovak", "Slovene", "Somali", "Spanish", "Sundanese", "Swahili",
+        "Swedish", "Tamil", "Telugu", "Thai", "Turkish", "Ukrainian", "Urdu",
+        "Uzbek", "Vietnamese", "Welsh", "Zulu"
+    ]
+    
+    # Define voice options
+    voice_options = [
+        "Andrew - English (United States)",
+        "Ava - English (United States)",
+        "Brian - English (United States)",
+        "Emma - English (United States)",
+        "Florian - German (Germany)",
+        "Seraphina - German (Germany)",
+        "Remy - French (France)",
+        "Vivienne - French (France)"
+    ]
+    
+    # Create Gradio interface
+    with gr.Blocks(title="PodcastGen 🎙️") as demo:
+        gr.Markdown("# PodcastGen 🎙️")
+        gr.Markdown("Generate a 2-speaker podcast from text input or documents!")
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                input_text = gr.Textbox(label="Input Text", lines=10, placeholder="Enter text for podcast generation...")
+            
+            with gr.Column(scale=1):
+                input_file = gr.File(label="Or Upload a PDF or TXT file", file_types=[".pdf", ".txt"])
+        
+        with gr.Row():
+            with gr.Column():
+                api_key = gr.Textbox(label="Your Gemini API Key (Optional)", placeholder="Enter API key here if you're getting rate limited", type="password")
+                language = gr.Dropdown(label="Language", choices=language_options, value="Auto Detect")
+
+            with gr.Column():
+                speaker1 = gr.Dropdown(label="Speaker 1 Voice", choices=voice_options, value="Andrew - English (United States)")
+                speaker2 = gr.Dropdown(label="Speaker 2 Voice", choices=voice_options, value="Ava - English (United States)")
+        
+        generate_btn = gr.Button("Generate Podcast", variant="primary")
+        
+        with gr.Row():
+            output_audio = gr.Audio(label="Generated Podcast", type="filepath", format="wav")
+            
+        generate_btn.click(
+            fn=generate_podcast_gradio,
+            inputs=[input_text, input_file, language, speaker1, speaker2, api_key],
+            outputs=[output_audio]
+        )
+    
+    demo.launch()
+
+if __name__ == "__main__":
+    main()