Update app.py

app.py

@@ -8,87 +8,162 @@ from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 from torchvision.io import write_video
 import os
 import groq
+import logging
+from pathlib import Path
 
-# Initialize Groq client
-groq_client = groq.Groq()
-API_KEY = os.getenv("GROQ_API_KEY")  
-groq_client.api_key = API_KEY
-
-# Initialize TTS model
-tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC")
-
-# Initialize Stable Diffusion pipeline for CPU
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-pipe = pipe.to("cpu")
-
-def generate_text_with_groq(prompt, max_tokens=200):
-    chat_completion = groq_client.chat.completions.create(
-        messages=[
-            {
-                "role": "system",
-                "content": "You are a professional comedy writer skilled in creating short, witty scripts."
-            },
-            {
-                "role": "user",
-                "content": prompt
-            }
-        ],
-        model="mixtral-8x7b-32768",
-        max_tokens=max_tokens,
-        temperature=0.7,
-    )
-    return chat_completion.choices[0].message.content
-
-def generate_speech(text):
-    output_path = "generated_speech.wav"
-    tts.tts_to_file(text=text, file_path=output_path)
-    return output_path
-
-def generate_video_frames(prompt, num_frames=10):
-    frames = []
-    for i in range(num_frames):
-        frame_prompt = f"{prompt}, frame {i+1} of {num_frames}"
-        with torch.no_grad():
-            image = pipe(frame_prompt, num_inference_steps=20).images[0]
-        frames.append(np.array(image))
-    return frames
-
-def create_video_from_frames(frames, output_path="output_video.mp4", fps=5):
-    frames_tensor = torch.from_numpy(np.array(frames)).permute(0, 3, 1, 2)
-    write_video(output_path, frames_tensor, fps=fps)
-    return output_path
-
-def generate_comedy_animation(prompt):
-    script_prompt = f"""Write a short, witty comedy script with two characters about {prompt}. 
-    Use the format 'Character: Dialogue or Action' for each line. 
-    Include clever wordplay, unexpected twists, and snappy dialogue. 
-    Keep it concise, around 5-8 exchanges. Make it genuinely funny!"""
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Environment setup and validation
+if not (API_KEY := os.getenv("GROQ_API_KEY")):
+    raise ValueError("GROQ_API_KEY not found in environment variables")
+
+# Initialize clients and models with error handling
+try:
+    groq_client = groq.Groq(api_key=API_KEY)
     
-    script = generate_text_with_groq(script_prompt)
-    video_prompt = f"A comedic scene with two characters: {prompt}"
-    frames = generate_video_frames(video_prompt)
-    video_path = create_video_from_frames(frames)
-    speech_path = generate_speech(script)
-    return script, video_path, speech_path
-
-def generate_kids_music_animation(theme):
-    lyrics_prompt = f"""Write short, catchy, and simple lyrics for a children's song about {theme}. 
-    Each line should be on a new line. Don't include 'Verse' or 'Chorus' labels. 
-    Make it educational, fun, and easy to remember. Include a repeating chorus."""
+    # Initialize TTS model
+    tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC")
     
-    lyrics = generate_text_with_groq(lyrics_prompt)
-    video_prompt = f"A colorful, animated music video for children about {theme}"
-    frames = generate_video_frames(video_prompt)
-    video_path = create_video_from_frames(frames)
-    speech_path = generate_speech(lyrics)
-    return lyrics, video_path, speech_path
+    # Initialize Stable Diffusion with optimizations
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4", 
+        torch_dtype=torch.float32
+    )
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe = pipe.to("cpu")
+    pipe.enable_attention_slicing()  # Memory optimization
+except Exception as e:
+    logger.error(f"Error initializing models: {str(e)}")
+    raise
+
+class ContentGenerator:
+    def __init__(self):
+        self.output_dir = Path("generated_content")
+        self.output_dir.mkdir(exist_ok=True)
+        
+    def cleanup_old_files(self):
+        """Clean up previously generated files"""
+        for file in self.output_dir.glob("*"):
+            try:
+                file.unlink()
+            except Exception as e:
+                logger.warning(f"Could not delete {file}: {e}")
+
+    def generate_text_with_groq(self, prompt, max_tokens=200):
+        """Generate text with error handling"""
+        try:
+            chat_completion = groq_client.chat.completions.create(
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a professional comedy writer skilled in creating short, witty scripts."
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt
+                    }
+                ],
+                model="mixtral-8x7b-32768",
+                max_tokens=max_tokens,
+                temperature=0.7,
+            )
+            return chat_completion.choices[0].message.content
+        except Exception as e:
+            logger.error(f"Error generating text: {str(e)}")
+            raise
+
+    def generate_speech(self, text):
+        """Generate speech with unique filenames"""
+        try:
+            output_path = self.output_dir / f"speech_{hash(text)}.wav"
+            tts.tts_to_file(text=text, file_path=str(output_path))
+            return str(output_path)
+        except Exception as e:
+            logger.error(f"Error generating speech: {str(e)}")
+            raise
 
-my_theme='ysharma/steampunk'
+    def generate_video_frames(self, prompt, num_frames=15):
+        """Generate video frames with progress tracking"""
+        frames = []
+        try:
+            for i in range(num_frames):
+                frame_prompt = f"{prompt}, frame {i+1} of {num_frames}"
+                with torch.no_grad():
+                    image = pipe(
+                        frame_prompt, 
+                        num_inference_steps=20,
+                        guidance_scale=7.5
+                    ).images[0]
+                frames.append(np.array(image))
+                logger.info(f"Generated frame {i+1}/{num_frames}")
+        except Exception as e:
+            logger.error(f"Error generating frames: {str(e)}")
+            raise
+        return frames
+
+    def create_video_from_frames(self, frames, prompt):
+        """Create video with unique filenames"""
+        try:
+            output_path = self.output_dir / f"video_{hash(prompt)}.mp4"
+            frames_tensor = torch.from_numpy(np.array(frames)).permute(0, 3, 1, 2)
+            write_video(str(output_path), frames_tensor, fps=8)
+            return str(output_path)
+        except Exception as e:
+            logger.error(f"Error creating video: {str(e)}")
+            raise
+
+    def generate_comedy_animation(self, prompt):
+        """Generate comedy animation with error handling"""
+        try:
+            self.cleanup_old_files()
+            
+            script_prompt = f"""Write a short, witty comedy script with two characters about {prompt}. 
+            Use the format 'Character: Dialogue or Action' for each line. 
+            Include clever wordplay, unexpected twists, and snappy dialogue. 
+            Keep it concise, around 5-8 exchanges. Make it genuinely funny!"""
+            
+            script = self.generate_text_with_groq(script_prompt)
+            video_prompt = f"A comedic scene with two characters: {prompt}"
+            frames = self.generate_video_frames(video_prompt)
+            video_path = self.create_video_from_frames(frames, video_prompt)
+            speech_path = self.generate_speech(script)
+            
+            return script, video_path, speech_path
+        except Exception as e:
+            logger.error(f"Error in comedy animation generation: {str(e)}")
+            return "Error generating content", None, None
+
+    def generate_kids_music_animation(self, theme):
+        """Generate kids music animation with error handling"""
+        try:
+            self.cleanup_old_files()
+            
+            lyrics_prompt = f"""Write short, catchy, and simple lyrics for a children's song about {theme}. 
+            Each line should be on a new line. Don't include 'Verse' or 'Chorus' labels. 
+            Make it educational, fun, and easy to remember. Include a repeating chorus."""
+            
+            lyrics = self.generate_text_with_groq(lyrics_prompt)
+            video_prompt = f"A colorful, animated music video for children about {theme}"
+            frames = self.generate_video_frames(video_prompt)
+            video_path = self.create_video_from_frames(frames, video_prompt)
+            speech_path = self.generate_speech(lyrics)
+            
+            return lyrics, video_path, speech_path
+        except Exception as e:
+            logger.error(f"Error in kids music animation generation: {str(e)}")
+            return "Error generating content", None, None
+
+# Initialize content generator
+generator = ContentGenerator()
 
 # Gradio Interface
-with gr.Blocks(theme=my_theme) as app:
-    gr.Markdown("## AI-Generated Video and Audio Content (Optimized CPU Version with Groq API)")
+with gr.Blocks(theme='ysharma/steampunk') as app:
+    gr.Markdown("## AI-Generated Video and Audio Content")
+    
+    # Status message for errors
+    status_msg = gr.Textbox(label="Status", visible=False)
     
     with gr.Tab("Comedy Animation"):
         comedy_prompt = gr.Textbox(label="Enter comedy prompt")
@@ -97,8 +172,16 @@ with gr.Blocks(theme=my_theme) as app:
         comedy_animation = gr.Video(label="Comedy Animation")
         comedy_audio = gr.Audio(label="Comedy Speech")
 
+        def comedy_wrapper(prompt):
+            status_msg.visible = True
+            try:
+                return generator.generate_comedy_animation(prompt)
+            except Exception as e:
+                status_msg.value = f"Error: {str(e)}"
+                return None, None, None
+
         comedy_generate_btn.click(
-            generate_comedy_animation,
+            comedy_wrapper,
             inputs=comedy_prompt,
             outputs=[comedy_script, comedy_animation, comedy_audio]
         )
@@ -110,11 +193,20 @@ with gr.Blocks(theme=my_theme) as app:
         music_animation = gr.Video(label="Music Animation")
         music_audio = gr.Audio(label="Music Audio")
 
+        def music_wrapper(theme):
+            status_msg.visible = True
+            try:
+                return generator.generate_kids_music_animation(theme)
+            except Exception as e:
+                status_msg.value = f"Error: {str(e)}"
+                return None, None, None
+
         music_generate_btn.click(
-            generate_kids_music_animation,
+            music_wrapper,
             inputs=music_theme,
             outputs=[music_lyrics, music_animation, music_audio]
         )
 
-app.launch()
+if __name__ == "__main__":
+    app.launch()