import streamlit as st
from gtts import gTTS
from moviepy.editor import ImageSequenceClip, AudioFileClip
from diffusers import StableDiffusionPipeline
from PIL import Image, ImageFilter, ImageEnhance
import os

# Function to generate audio using gTTS
def generate_audio(text, output_path="audio.mp3"):
    tts = gTTS(text=text, lang='en')
    tts.save(output_path)
    return output_path

# Function to generate images using Stable Diffusion
def generate_images(prompt, num_images=5):
    # Load Stable Diffusion pipeline (ensure you have the model locally)
    pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
    pipeline = pipeline.to("cpu")  # Use "cuda" for GPU if available
    image_paths = []
    
    for i in range(num_images):
        image = pipeline(prompt).images[0]  # Generate an image
        image = apply_cartoon_filter(image)  # Apply cartoon effect
        image_path = f"image_{i}.png"
        image.save(image_path)
        image_paths.append(image_path)
    return image_paths

# Function to apply a cartoon effect to an image
def apply_cartoon_filter(image):
    # Convert to Pillow format if needed
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    # Enhance colors and edges
    image = ImageEnhance.Color(image).enhance(1.5)  # Increase saturation
    image = image.filter(ImageFilter.CONTOUR)  # Add cartoon-style edges
    return image

# Function to create a video using MoviePy
def create_video(image_paths, audio_path, output_path="output_video.mp4"):
    clip = ImageSequenceClip(image_paths, fps=1)  # 1 FPS for a slow animation
    audio = AudioFileClip(audio_path)
    final_clip = clip.set_audio(audio)
    final_clip.write_videofile(output_path, fps=24)
    return output_path

# Streamlit App
st.title("Cartoon Video Generator")
st.write("Enter a prompt, and we'll create a cartoon video with scenes and voices!")

# Input prompt
prompt = st.text_area("Enter your prompt:", "A cartoon story about a cat and a dog playing in the park.")

# Generate Button
if st.button("Generate Video"):
    with st.spinner("Generating cartoon video..."):
        # Generate Audio
        audio_file = generate_audio(prompt)

        # Generate Images
        st.info("Generating cartoon images...")
        images = generate_images(prompt)

        # Create Video
        st.info("Creating video...")
        video_file = create_video(images, audio_file)

        st.success("Video generated successfully!")

        # Show video preview
        st.video(video_file)

        # Provide download link
        with open(video_file, "rb") as file:
            st.download_button(
                label="Download Video",
                data=file,
                file_name="cartoon_video.mp4",
                mime="video/mp4"
            )