Update app.py

app.py

@@ -3,14 +3,8 @@ import torch
 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 re
 
 # Initialize text generation model (GPT-2)
@@ -20,109 +14,120 @@ 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):
+def generate_text(prompt, max_length=200):
     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 extract keywords from the script
-def extract_keywords(script):
-    # Simple keyword extraction (you might want to use a more sophisticated method)
-    keywords = re.findall(r'\b\w+\b', script.lower())
-    return list(set(keywords))
-
-def create_character_animation(script):
-    keywords = extract_keywords(script)
+def parse_script(script):
+    lines = script.split('\n')
+    scenes = []
+    current_scene = {"characters": set(), "actions": []}
+    
+    for line in lines:
+        if line.strip():
+            if ':' in line:
+                character, action = line.split(':', 1)
+                current_scene["characters"].add(character.strip())
+                current_scene["actions"].append((character.strip(), action.strip()))
+            else:
+                if current_scene["actions"]:
+                    scenes.append(current_scene)
+                    current_scene = {"characters": set(), "actions": []}
+    
+    if current_scene["actions"]:
+        scenes.append(current_scene)
+    
+    return scenes
+
+def create_stick_figure(ax, x, y, color):
+    circle = plt.Circle((x, y+0.1), 0.1, fc=color)
+    line = plt.Line2D([x, x], [y-0.3, y], color=color)
+    left_arm = plt.Line2D([x-0.2, x], [y, y-0.1], color=color)
+    right_arm = plt.Line2D([x, x+0.2], [y-0.1, y], color=color)
+    left_leg = plt.Line2D([x-0.1, x], [y-0.5, y-0.3], color=color)
+    right_leg = plt.Line2D([x, x+0.1], [y-0.3, y-0.5], color=color)
     
-    def update_point(num, data, line, ax, keywords):
-        line.set_data(data[:2, :num])
-        line.set_3d_properties(data[2, :num])
+    ax.add_artist(circle)
+    ax.add_artist(line)
+    ax.add_artist(left_arm)
+    ax.add_artist(right_arm)
+    ax.add_artist(left_leg)
+    ax.add_artist(right_leg)
+
+def animate_scene(scene, ax):
+    characters = list(scene["characters"])
+    colors = plt.cm.get_cmap('Set3')(np.linspace(0, 1, len(characters)))
+    character_positions = {char: (i/(len(characters)-1) if len(characters) > 1 else 0.5, 0.5) for i, char in enumerate(characters)}
+    
+    def init():
+        ax.clear()
+        ax.set_xlim(0, 1)
+        ax.set_ylim(0, 1)
+        ax.axis('off')
+        for char, (x, y) in character_positions.items():
+            create_stick_figure(ax, x, y, colors[characters.index(char)])
+        return []
+
+    def animate(frame):
+        ax.clear()
+        ax.set_xlim(0, 1)
+        ax.set_ylim(0, 1)
+        ax.axis('off')
         
-        # Update text with cycling keywords
-        keyword = keywords[num % len(keywords)]
-        for txt in ax.texts:
-            txt.remove()
-        ax.text2D(0.05, 0.95, keyword, transform=ax.transAxes)
+        action = scene["actions"][frame % len(scene["actions"])]
+        speaking_char, text = action
         
-        return line,
-
-    fig = plt.figure()
-    ax = fig.add_subplot(111, projection='3d')
-
-    # Generate a 3D path influenced by the number of keywords
-    t = np.linspace(0, 2*np.pi*len(keywords)/10, 100)
-    x = np.sin(t) * (1 + len(keywords)/20)
-    y = np.cos(t) * (1 + len(keywords)/20)
-    z = t/10
-
-    data = np.array([x, y, z])
-    line, = ax.plot(data[0, 0:1], data[1, 0:1], data[2, 0:1])
+        for char, (x, y) in character_positions.items():
+            if char == speaking_char:
+                y += 0.05 * np.sin(frame * 0.5)  # Make the speaking character bounce
+            create_stick_figure(ax, x, y, colors[characters.index(char)])
+        
+        ax.text(0.5, 0.9, text, ha='center', va='center', wrap=True)
+        return []
 
-    # Setting the axes properties
-    ax.set_xlim3d([-2.0, 2.0])
-    ax.set_xlabel('X')
-    ax.set_ylim3d([-2.0, 2.0])
-    ax.set_ylabel('Y')
-    ax.set_zlim3d([0.0, 4.0])
-    ax.set_zlabel('Z')
+    return animation.FuncAnimation(fig, animate, init_func=init, frames=len(scene["actions"])*5, interval=1000, blit=True)
 
-    ani = animation.FuncAnimation(fig, update_point, frames=100, fargs=(data, line, ax, keywords), interval=100, blit=False)
+def create_character_animation(script):
+    scenes = parse_script(script)
+    
+    fig, ax = plt.subplots(figsize=(10, 6))
+    animations = [animate_scene(scene, ax) for scene in scenes]
+    
+    # Combine all animations
+    combined_animation = animation.ArtistAnimation(fig, sum([anim._framedata for anim in animations], []), interval=1000, blit=True, repeat_delay=1000)
     
     # Save animation as gif
-    ani.save('character_animation.gif', writer='pillow')
+    combined_animation.save('character_animation.gif', writer='pillow')
     
     return 'character_animation.gif'
 
-# Main function to generate comedy animation
 def generate_comedy_animation(prompt):
-    script = generate_text(f"Write a short comedy script about {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.")
     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):
-    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("## Script-based Character Animation Generator")
-    
-    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, comedy_audio]
-        )
+    gr.Markdown("## Character-based Animation Generator")
     
-    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_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_lyrics, kids_animation, kids_audio]
-        )
+    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, comedy_audio]
+    )
 
 app.launch()