import gradio as gr
from transformers import pipeline
from gtts import gTTS
from io import BytesIO
from PIL import Image

# Initialize transformers for each task
text_generation_pipe = pipeline("text-generation", model="gpt2")
text_to_speech_pipe = pipeline("text-to-speech", model="google")
image_generation_pipe = pipeline("image-generation", model="stable-diffusion")

# Function for generating stories, audio, and illustrations
def generate_story_and_illustration(prompt, character_name):
    # Generate story text
    story_text = text_generation_pipe(prompt=f"Once upon a time, {character_name}... {prompt}", max_length=150)[0]["generated_text"]

    # Generate audio from story text
    tts_audio = gTTS(text=story_text, lang="en")
    audio_buffer = BytesIO()
    tts_audio.save(audio_buffer)
    audio_buffer.seek(0)

    # Generate image based on story
    image = image_generation_pipe(prompt=story_text)

    # Return generated story, audio, and image
    return story_text, audio_buffer.read(), image

# Gradio interface for user interaction
interface = gr.Interface(
    fn=generate_story_and_illustration,
    inputs=[
        gr.Textbox(label="Start your story with..."),
        gr.Textbox(label="Give your character a name:"),
    ],
    outputs=[
        gr.Textbox(label="Story"),
        gr.Audio(label="Listen to the story"),
        gr.Image(label="See the story come alive"),
    ],
    title="Storyteller Playground",
    description="Create amazing stories with the help of AI!",
    theme="kids",
)

# Launch the app
interface.launch()