import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from diffusers import StableDiffusionPipeline
from moviepy.editor import ImageSequenceClip, AudioFileClip, concatenate_videoclips
from gtts import gTTS
import os

# Function to generate a script from a text prompt
def generate_script(prompt):
    model_name = "gpt2"
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_name)

    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(inputs, max_length=500, num_return_sequences=1)

    script = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return script

# Function to convert text to speech using gTTS
def text_to_speech(text, output_path="output.mp3"):
    tts = gTTS(text=text, lang='en')
    tts.save(output_path)
    return output_path

# Function to generate images from a text prompt using Stable Diffusion (CPU)
def generate_images_from_prompt(prompt, num_images=3, num_inference_steps=20):
    model_id = "runwayml/stable-diffusion-v1-5"
    pipe = StableDiffusionPipeline.from_pretrained(model_id)
    pipe = pipe.to("cpu")  # Ensure we are using CPU
    images = []
    for i in range(num_images):
        image = pipe(prompt, num_inference_steps=num_inference_steps).images[0]
        image_path = f"image_{i}.png"
        image.save(image_path)
        images.append(image_path)
    return images

# Function to create a video from images and synthesized speech
def create_video(images, audio_path, output_video_path="output.mp4"):
    # Load the images
    image_clips = [ImageSequenceClip(images, fps=1)]
    
    # Concatenate the image clips
    video = concatenate_videoclips(image_clips)
    
    # Load the audio clip
    audio_clip = AudioFileClip(audio_path)

    # Set the audio to the video
    video = video.set_audio(audio_clip)

    # Export the video
    video.write_videofile(output_video_path, fps=24)

# Main function to generate video from prompt
def generate_video_from_prompt(prompt):
    script = generate_script(prompt)
    audio_path = text_to_speech(script)
    images = generate_images_from_prompt(prompt)
    create_video(images, audio_path)
    return "output.mp4"

# Gradio interface
iface = gr.Interface(
    fn=generate_video_from_prompt,
    inputs="text",
    outputs="video",
    title="Text to Video Generator",
    description="Enter a prompt to generate a high-definition video."
)

# Launch the interface
iface.launch()