Update app.py

app.py

@@ -1,80 +1,115 @@
 import gradio as gr
-from transformers import pipeline
 import torch
-from PIL import Image, ImageDraw, ImageFont
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from TTS.api import TTS
+import numpy as np
+import librosa
+import soundfile as sf
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+from mpl_toolkits.mplot3d import Axes3D
+import io
+import base64
 import os
-import random
-
-# Initialize text generation model
-text_generator = pipeline('text-generation', model='gpt2')
-
-# Function to generate comedy script
-def generate_comedy_script(prompt, max_length=100):
-    generated = text_generator(f"Write a short comedy script about {prompt}: ", max_length=max_length, num_return_sequences=1)[0]
-    return generated['generated_text']
-
-# Function to create a simple animation frame
-def create_animation_frame(text, frame_number):
-    img = Image.new('RGB', (400, 200), color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
-    d = ImageDraw.Draw(img)
-    fnt = ImageFont.load_default()
-    d.text((10, 10), text, font=fnt, fill=(255, 255, 0))
-    d.text((10, 180), f"Frame {frame_number}", font=fnt, fill=(255, 255, 255))
-    return img
-
-# Function to create a simple animation
-def create_animation(text, num_frames=5):
-    frames = []
-    for i in range(num_frames):
-        frame = create_animation_frame(text[:50], i+1)  # Use first 50 characters for each frame
-        frames.append(frame)
-    
-    animation_path = 'animation.gif'
-    frames[0].save(animation_path, save_all=True, append_images=frames[1:], duration=500, loop=0)
-    return animation_path
 
-# Function to generate kids music theme
-def generate_kids_music_theme(theme):
-    return f"🎵 {theme} for kids! 🎵\nImagine a catchy tune with lyrics about {theme}."
+# Initialize text generation model (GPT-2)
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+# Initialize TTS model
+tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC")
+
+# Function to generate text using GPT-2
+def generate_text(prompt, max_length=100):
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+# Function to generate speech using TTS
+def generate_speech(text):
+    output_path = "generated_speech.wav"
+    tts.tts_to_file(text=text, file_path=output_path)
+    return output_path
+
+# Function to create a more advanced character animation
+def create_character_animation(text):
+    def update_point(num, data, line):
+        line.set_data(data[:2, :num])
+        line.set_3d_properties(data[2, :num])
+        return line
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+
+    # Generate a more complex 3D path for the character
+    t = np.linspace(0, 20, 100)
+    x = np.sin(t)
+    y = np.cos(t)
+    z = t/10
+
+    data = np.array([x, y, z])
+    line, = ax.plot(data[0, 0:1], data[1, 0:1], data[2, 0:1])
+
+    # Setting the axes properties
+    ax.set_xlim3d([-1.0, 1.0])
+    ax.set_xlabel('X')
+    ax.set_ylim3d([-1.0, 1.0])
+    ax.set_ylabel('Y')
+    ax.set_zlim3d([0.0, 2.0])
+    ax.set_zlabel('Z')
+
+    # Adding text
+    ax.text2D(0.05, 0.95, text[:50], transform=ax.transAxes)
+
+    ani = animation.FuncAnimation(fig, update_point, frames=100, fargs=(data, line), interval=100, blit=False)
+    
+    # Save animation as gif
+    ani.save('character_animation.gif', writer='pillow')
+    
+    return 'character_animation.gif'
 
 # Main function to generate comedy animation
 def generate_comedy_animation(prompt):
-    script = generate_comedy_script(prompt)
-    animation_path = create_animation(script)
-    return script, animation_path
+    script = generate_text(f"Write a short comedy script about {prompt}: ")
+    animation_path = create_character_animation(script)
+    speech_path = generate_speech(script)
+    return script, animation_path, speech_path
 
 # Main function to generate kids music animation
 def generate_kids_music_animation(theme):
-    music_theme = generate_kids_music_theme(theme)
-    animation_path = create_animation(music_theme)
-    return music_theme, animation_path
+    lyrics = generate_text(f"Write a short children's song about {theme}: ")
+    animation_path = create_character_animation(lyrics)
+    speech_path = generate_speech(lyrics)
+    return lyrics, animation_path, speech_path
 
 # Gradio Interface
 with gr.Blocks() as app:
-    gr.Markdown("## Simplified Comedy and Kids Animation Generator")
+    gr.Markdown("## Advanced Character Animation Generator with Open Source Models")
     
     with gr.Tab("Comedy Animation"):
         comedy_prompt = gr.Textbox(label="Enter comedy prompt")
         comedy_generate_btn = gr.Button("Generate Comedy Animation")
         comedy_script = gr.Textbox(label="Generated Comedy Script")
         comedy_animation = gr.Image(label="Comedy Animation")
+        comedy_audio = gr.Audio(label="Comedy Speech")
 
         comedy_generate_btn.click(
             generate_comedy_animation,
             inputs=comedy_prompt,
-            outputs=[comedy_script, comedy_animation]
+            outputs=[comedy_script, comedy_animation, comedy_audio]
         )
     
     with gr.Tab("Kids Music Animation"):
         kids_theme = gr.Textbox(label="Enter kids music theme")
         kids_generate_btn = gr.Button("Generate Kids Music Animation")
-        kids_music_theme = gr.Textbox(label="Generated Music Theme")
+        kids_lyrics = gr.Textbox(label="Generated Lyrics")
         kids_animation = gr.Image(label="Kids Music Animation")
+        kids_audio = gr.Audio(label="Kids Speech")
 
         kids_generate_btn.click(
             generate_kids_music_animation,
             inputs=kids_theme,
-            outputs=[kids_music_theme, kids_animation]
+            outputs=[kids_lyrics, kids_animation, kids_audio]
         )
 
 app.launch()