Update app.py

app.py

@@ -1,181 +1,249 @@
 import streamlit as st
 from pathlib import Path
 import torch
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-from diffusers import StableDiffusionPipeline
-from TTS.api import TTS
-import cv2
-import numpy as np
-from PIL import Image
+from transformers import pipeline
+from PIL import Image, ImageDraw, ImageFont
 import tempfile
 import os
 from moviepy.editor import *
-import base64
+import numpy as np
+from gtts import gTTS
+import textwrap
 
-class VideoGenerator:
+class EnhancedVideoGenerator:
     def __init__(self):
-        # Initialize text generation model
-        self.text_model = AutoModelForCausalLM.from_pretrained(
-            "facebook/opt-1.3b",
-            torch_dtype=torch.float16,
-            device_map="auto"
+        # Use device agnostic setup
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Initialize text generation with a faster model
+        self.text_generator = pipeline(
+            'text-generation',
+            model='distilgpt2',  # Smaller, faster model
+            device=0 if self.device == "cuda" else -1
         )
-        self.text_tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
-
-        # Initialize image generation model
-        self.image_generator = StableDiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5",
-            torch_dtype=torch.float16
-        ).to("cuda")
-
-        # Initialize TTS model
-        self.tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False)
-
+        
         # Create temp directory
         self.temp_dir = Path(tempfile.mkdtemp())
-
-    def generate_script(self, prompt):
-        """Generate detailed script with facts and scenes"""
-        input_ids = self.text_tokenizer(
-            f"Generate a detailed video script with facts about: {prompt}. Include scene descriptions.",
-            return_tensors="pt"
-        ).input_ids.to("cuda")
-        
-        outputs = self.text_model.generate(
-            input_ids,
-            max_length=500,
-            temperature=0.7,
-            num_return_sequences=1
+        
+        # Store theme colors
+        self.themes = {
+            'Professional': {
+                'bg': (245, 245, 245),
+                'text': (33, 33, 33),
+                'accent': (0, 102, 204)
+            },
+            'Creative': {
+                'bg': (255, 240, 245),
+                'text': (51, 51, 51),
+                'accent': (255, 64, 129)
+            },
+            'Educational': {
+                'bg': (240, 249, 255),
+                'text': (25, 25, 25),
+                'accent': (0, 151, 167)
+            }
+        }
+
+    @st.cache_data
+    def generate_script(self, prompt, style, length):
+        """Generate engaging script based on style and length"""
+        style_prompts = {
+            'Professional': "Write a clear, professional script about:",
+            'Creative': "Write an engaging, creative script about:",
+            'Educational': "Write an informative educational script about:"
+        }
+        
+        max_length = int(length * 2.5)  # Approximate words needed for duration
+        
+        output = self.text_generator(
+            f"{style_prompts[style]} {prompt}. Keep it concise and engaging.",
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.7
         )
         
-        script = self.text_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        script = output[0]['generated_text']
+        # Clean up the generated text
+        script = script.replace(style_prompts[style], '').strip()
         return script
 
-    def generate_scene_images(self, scene_descriptions):
-        """Generate images for each scene using Stable Diffusion"""
-        image_paths = []
-        for i, desc in enumerate(scene_descriptions):
-            image = self.image_generator(desc).images[0]
-            path = self.temp_dir / f"scene_{i}.png"
-            image.save(path)
-            image_paths.append(path)
-        return image_paths
+    def create_styled_frame(self, text, theme, frame_number, total_frames):
+        """Create a visually appealing frame with text and progress bar"""
+        # Create base image
+        img = Image.new('RGB', (1280, 720), color=theme['bg'])
+        draw = ImageDraw.Draw(img)
+        
+        # Wrap text for better presentation
+        wrapped_text = textwrap.fill(text, width=50)
+        
+        # Add text
+        font_size = 40
+        try:
+            font = ImageFont.truetype("arial.ttf", font_size)
+        except:
+            font = ImageFont.load_default()
+            
+        # Calculate text position for center alignment
+        text_bbox = draw.textbbox((0, 0), wrapped_text, font=font)
+        text_width = text_bbox[2] - text_bbox[0]
+        text_height = text_bbox[3] - text_bbox[1]
+        text_x = (1280 - text_width) // 2
+        text_y = (720 - text_height) // 2
+        
+        # Draw text with subtle shadow for depth
+        draw.text((text_x+2, text_y+2), wrapped_text, fill=(0,0,0,128), font=font)
+        draw.text((text_x, text_y), wrapped_text, fill=theme['text'], font=font)
+        
+        # Add progress bar
+        progress = frame_number / total_frames
+        bar_width = int(1000 * progress)
+        draw.rectangle([140, 650, 1140, 660], fill=(200,200,200))
+        draw.rectangle([140, 650, 140+bar_width, 660], fill=theme['accent'])
+        
+        return np.array(img)
 
-    def generate_voiceover(self, script):
-        """Generate voice narration using TTS"""
-        audio_path = self.temp_dir / "voiceover.wav"
-        self.tts.tts_to_file(script, file_path=str(audio_path))
+    def generate_voiceover(self, script, lang='en'):
+        """Generate voice narration using gTTS"""
+        audio_path = self.temp_dir / "voiceover.mp3"
+        tts = gTTS(text=script, lang=lang, slow=False)
+        tts.save(str(audio_path))
         return audio_path
 
-    def create_video(self, image_paths, audio_path, duration_per_image=5):
-        """Combine images and audio into video"""
-        clips = []
-        for img_path in image_paths:
-            clip = ImageClip(str(img_path)).set_duration(duration_per_image)
-            clips.append(clip)
-        
-        video = concatenate_videoclips(clips)
+    def create_video(self, script, theme, duration=30):
+        """Create video with animated text and voiceover"""
+        # Generate frames
+        fps = 24
+        total_frames = duration * fps
+        
+        def make_frame(t):
+            frame_number = int(t * fps)
+            return self.create_styled_frame(
+                script, 
+                self.themes[theme], 
+                frame_number,
+                total_frames
+            )
+        
+        # Create video clip
+        clip = VideoClip(make_frame, duration=duration)
+        
+        # Generate and add voiceover
+        audio_path = self.generate_voiceover(script)
         audio = AudioFileClip(str(audio_path))
         
-        # Adjust video duration to match audio
-        video = video.set_duration(audio.duration)
-        final_video = video.set_audio(audio)
+        # Combine video and audio
+        final_clip = clip.set_audio(audio)
         
+        # Save video
         output_path = self.temp_dir / "output_video.mp4"
-        final_video.write_videofile(str(output_path), fps=24)
+        final_clip.write_videofile(str(output_path), fps=fps, codec='libx264')
         return output_path
 
 def main():
-    st.set_page_config(page_title="AI Video Generator", layout="wide")
-    st.title("🎬 AI Text-to-Video Generator")
+    st.set_page_config(
+        page_title="✨ Smart Video Generator",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+
+    # Custom CSS for better UI
+    st.markdown("""
+        <style>
+        .stButton>button {
+            width: 100%;
+            height: 3em;
+            margin-top: 2em;
+        }
+        .stProgress .st-bo {
+            height: 20px;
+        }
+        </style>
+    """, unsafe_allow_html=True)
 
     # Initialize session state
     if 'video_generator' not in st.session_state:
-        st.session_state.video_generator = VideoGenerator()
+        with st.spinner("🚀 Initializing... (just a moment)"):
+            st.session_state.video_generator = EnhancedVideoGenerator()
+
+    # Sidebar for settings
+    with st.sidebar:
+        st.title("🎮 Video Settings")
+        theme = st.selectbox(
+            "Choose Theme",
+            ["Professional", "Creative", "Educational"],
+            help="Select the style for your video"
+        )
+        
+        duration = st.slider(
+            "Video Duration (seconds)",
+            min_value=15,
+            max_value=90,
+            value=30,
+            step=15,
+            help="Choose the length of your video"
+        )
 
-    # Input section
-    st.header("Enter Your Topic")
+    # Main content
+    st.title("✨ Smart Video Generator")
+    st.markdown("""
+    Transform your ideas into engaging videos with AI-powered content generation.
+    Just enter your topic below! 🎥
+    """)
+
+    # Input section with example
     text_input = st.text_area(
-        "What would you like to create a video about?",
+        "What's your video about?",
         height=100,
-        placeholder="Example: Explain the process of photosynthesis in plants..."
+        placeholder="Example: Explain how rainbows form in the sky..."
     )
 
-    # Generation settings
-    st.header("Video Settings")
-    col1, col2 = st.columns(2)
-    with col1:
-        video_length = st.slider("Approximate video length (seconds)", 30, 300, 60)
-    with col2:
-        style = st.selectbox(
-            "Video style",
-            ["Educational", "Documentary", "Engaging", "Professional"]
-        )
-
-    # Generate button
-    if st.button("🎥 Generate Video"):
+    # Generate button with loading animation
+    if st.button("🎬 Generate Video", use_container_width=True):
         if text_input:
-            with st.spinner("🤖 Generating your video..."):
-                try:
-                    # Progress bar
-                    progress_bar = st.progress(0)
-                    progress_text = st.empty()
-
-                    # Generate script
-                    progress_text.text("Generating script...")
-                    script = st.session_state.video_generator.generate_script(text_input)
-                    progress_bar.progress(25)
-
-                    # Extract scene descriptions
-                    progress_text.text("Processing scenes...")
-                    scenes = [s.strip() for s in script.split("Scene:") if s.strip()]
-                    progress_bar.progress(40)
-
-                    # Generate images
-                    progress_text.text("Creating visuals...")
-                    image_paths = st.session_state.video_generator.generate_scene_images(scenes)
-                    progress_bar.progress(60)
-
-                    # Generate voiceover
-                    progress_text.text("Generating voiceover...")
-                    audio_path = st.session_state.video_generator.generate_voiceover(script)
-                    progress_bar.progress(80)
-
-                    # Create video
-                    progress_text.text("Composing final video...")
-                    video_path = st.session_state.video_generator.create_video(
-                        image_paths, 
-                        audio_path,
-                        duration_per_image=video_length/len(scenes)
-                    )
-                    progress_bar.progress(100)
-                    progress_text.text("Video generation complete!")
-
-                    # Display results
-                    st.header("Generated Content")
-                    
-                    # Show script
-                    with st.expander("📝 Generated Script"):
-                        st.write(script)
-                    
-                    # Show video
-                    st.header("🎥 Your Video")
-                    video_file = open(str(video_path), 'rb')
-                    video_bytes = video_file.read()
-                    st.video(video_bytes)
+            try:
+                progress_bar = st.progress(0)
+                status_text = st.empty()
+                
+                # Generate script
+                status_text.text("✍️ Creating engaging script...")
+                script = st.session_state.video_generator.generate_script(
+                    text_input, theme, duration
+                )
+                progress_bar.progress(30)
+                
+                # Create video
+                status_text.text("🎨 Crafting visual elements...")
+                progress_bar.progress(60)
+                
+                video_path = st.session_state.video_generator.create_video(
+                    script, theme, duration
+                )
+                progress_bar.progress(100)
+                status_text.text("✨ Done! Your video is ready!")
+                
+                # Display results in tabs
+                tab1, tab2 = st.tabs(["📽️ Video", "📝 Script"])
+                
+                with tab1:
+                    st.video(str(video_path))
                     
                     # Download button
-                    st.download_button(
-                        label="Download Video",
-                        data=video_bytes,
-                        file_name="generated_video.mp4",
-                        mime="video/mp4"
-                    )
-
-                except Exception as e:
-                    st.error(f"An error occurred: {str(e)}")
+                    with open(str(video_path), 'rb') as f:
+                        st.download_button(
+                            "⬇️ Download Video",
+                            f,
+                            file_name="generated_video.mp4",
+                            mime="video/mp4"
+                        )
+                
+                with tab2:
+                    st.markdown("### Generated Script")
+                    st.write(script)
+                    
+            except Exception as e:
+                st.error(f"💥 Oops! Something went wrong: {str(e)}")
         else:
-            st.warning("Please enter some text to generate a video!")
+            st.warning("🎯 Please enter a topic to generate a video!")
 
 if __name__ == "__main__":
     main()