import streamlit as st
import torch
from transformers import pipeline
from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
from diffusers.utils import export_to_gif
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from gtts import gTTS
from moviepy.editor import VideoFileClip, AudioFileClip

# Load the text generation model
generator = pipeline('text-generation', model='distilgpt2')

def generate_text(prompt):
    response = generator(prompt, max_length=150, num_return_sequences=1)
    return response[0]['generated_text']

# Text-to-speech conversion
def text_to_speech(text, filename='output_audio.mp3'):
    tts = gTTS(text)
    tts.save(filename)
    return filename

# Generate animation using AnimateDiffPipeline
def create_animation(prompt, output_file='animation.gif'):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    step = 4
    repo = "ByteDance/AnimateDiff-Lightning"
    ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
    base = "emilianJR/epiCRealism"

    # Load adapter and pipeline
    adapter = MotionAdapter().to(device, dtype)
    adapter.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
    pipe = AnimateDiffPipeline.from_pretrained(base, motion_adapter=adapter, torch_dtype=dtype).to(device)
    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")

    # Generate animation based on prompt
    output = pipe(prompt=prompt, guidance_scale=1.0, num_inference_steps=step)
    export_to_gif(output.frames[0], output_file)

    return output_file

# Combine animation and audio into a video
def create_video(animation_file, audio_file, output_file='output_video.mp4'):
    clip = VideoFileClip(animation_file)
    audio = AudioFileClip(audio_file)
    clip = clip.set_audio(audio)
    clip.write_videofile(output_file, fps=24)

def generate_educational_video(prompt):
    # Step 1: Generate text from prompt
    generated_text = generate_text(prompt)

    # Step 2: Convert text to speech
    audio_file = text_to_speech(generated_text)

    # Step 3: Create animation based on prompt
    animation_file = create_animation(prompt)

    # Step 4: Assemble the video
    create_video(animation_file, audio_file)

    # Return the path to the video
    return 'output_video.mp4'

# Streamlit UI
st.title("Educational Video Generator")

# User input for prompt
prompt = st.text_input("Enter your prompt here:")

if st.button("Generate Video"):
    if prompt:
        st.write("Generating video, please wait...")
        
        # Generate the video
        video_path = generate_educational_video(prompt)
        
        # Display the video in Streamlit
        st.video(video_path)
    else:
        st.warning("Please enter a prompt to generate the video.")