Update app.py

app.py

@@ -1,16 +1,18 @@
 import gradio as gr
 import torch
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
 from TTS.api import TTS
 import numpy as np
 from PIL import Image
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 from torchvision.io import write_video
 import os
+import groq
 
-# Initialize text generation model (GPT-2)
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-model = AutoModelForCausalLM.from_pretrained("gpt2")
+# Initialize Groq client
+groq_client = groq.Groq()
+API_KEY = "GROQ_API_KEY"  
+groq_client.api_key = API_KEY
 
 # Initialize TTS model
 tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC")
@@ -20,11 +22,23 @@ pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4",
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 pipe = pipe.to("cpu")
 
-def generate_text(prompt, max_length=200):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    attention_mask = torch.ones_like(input_ids)
-    output = model.generate(input_ids, attention_mask=attention_mask, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2, pad_token_id=tokenizer.eos_token_id)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+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"
@@ -46,7 +60,12 @@ def create_video_from_frames(frames, output_path="output_video.mp4", fps=5):
     return output_path
 
 def generate_comedy_animation(prompt):
-    script = generate_text(f"Write a short comedy script with two characters about {prompt}. Use the format 'Character: Action' for each line of dialogue or action.")
+    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 = 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)
@@ -54,7 +73,11 @@ def generate_comedy_animation(prompt):
     return script, video_path, speech_path
 
 def generate_kids_music_animation(theme):
-    lyrics = generate_text(f"Write short 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.")
+    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 = 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)
@@ -63,7 +86,7 @@ def generate_kids_music_animation(theme):
 
 # Gradio Interface
 with gr.Blocks() as app:
-    gr.Markdown("## AI-Generated Video and Audio Content (CPU Version)")
+    gr.Markdown("## AI-Generated Video and Audio Content (Optimized CPU Version with Groq API)")
     
     with gr.Tab("Comedy Animation"):
         comedy_prompt = gr.Textbox(label="Enter comedy prompt")