app.py

import gradio as gr
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)
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 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))

# Function to create a character animation based on script keywords
def create_character_animation(script):
    keywords = extract_keywords(script)
    
    def update_point(num, data, line, ax, keywords):
        line.set_data(data[:2, :num])
        line.set_3d_properties(data[2, :num])
        
        # Update text with cycling keywords
        keyword = keywords[num % len(keywords)]
        ax.texts.clear()
        ax.text2D(0.05, 0.95, keyword, transform=ax.transAxes)
        
        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])

    # 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')

    ani = animation.FuncAnimation(fig, update_point, frames=100, fargs=(data, line, ax, keywords), 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_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):
    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]
        )
    
    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]
        )

app.launch()