import gradio as gr
from diffusers import StableDiffusionPipeline
import moviepy.editor as mp
import torch
import os
from datetime import datetime

def generate_video(prompt, duration, frame_rate, music_file):
    # Initialize the pipeline
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16).to(device)
    else:
        # Use float32 when running on CPU
        pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)

    # Generate frames
    num_frames = int(duration * frame_rate)
    temp_dir = f"/tmp/{datetime.now().strftime('%Y%m%d%H%M%S')}"
    os.makedirs(temp_dir, exist_ok=True)

    for i in range(num_frames):
        frame = pipeline(prompt).images[0]
        frame_path = os.path.join(temp_dir, f"frame_{i:04d}.png")
        frame.save(frame_path)

    # Create video from frames
    video_path = os.path.join(temp_dir, "video.mp4")
    video_clip = mp.ImageSequenceClip([os.path.join(temp_dir, f"frame_{i:04d}.png") for i in range(num_frames)], fps=frame_rate)

    if music_file:
        audio_clip = mp.AudioFileClip(music_file)
        audio_clip = audio_clip.set_duration(video_clip.duration)
        video_clip = video_clip.set_audio(audio_clip)

    video_clip.write_videofile(video_path, codec="libx264")

    return video_path

with gr.Blocks() as demo:
    gr.Markdown("# AI Dreams & Visions Video Generator")
    gr.Markdown("Generate stunning videos from text prompts using AI technology. For inquiries, contact us at [aidreams@aidreams.company](mailto:aidreams@aidreams.company). Follow us on X: [@TheKingofJewelz](https://x.com/TheKingofJewelz).")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Text Prompt", placeholder="Enter your video description here...")
            duration = gr.Slider(label="Duration (seconds)", minimum=1, maximum=30, step=1, value=5)
            frame_rate = gr.Slider(label="Frame Rate", minimum=1, maximum=60, step=1, value=24)
            music_file = gr.Audio(label="Music File (Optional)", type="file")
            generate_btn = gr.Button("Generate Video")

        with gr.Column():
            video_output = gr.Video(label="Generated Video")
            download_link = gr.File(label="Download Video")

    generate_btn.click(
        generate_and_display_video,
        inputs=[prompt, duration, frame_rate, music_file],
        outputs=[video_output, download_link]
    )

if __name__ == "__main__":
    demo.launch(share=True)