import gradio as gr
from transformers import pipeline
import torch
from PIL import Image, ImageDraw, ImageFont
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}."

# 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

# 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

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("## Simplified Comedy and Kids 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_generate_btn.click(
            generate_comedy_animation,
            inputs=comedy_prompt,
            outputs=[comedy_script, comedy_animation]
        )
    
    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_animation = gr.Image(label="Kids Music Animation")

        kids_generate_btn.click(
            generate_kids_music_animation,
            inputs=kids_theme,
            outputs=[kids_music_theme, kids_animation]
        )

app.launch()