import gradio as gr
import torch
import os
import numpy as np
import ffmpeg
from moviepy.editor import AudioFileClip, ImageSequenceClip
from diffusers import StableDiffusionPipeline
from tempfile import TemporaryDirectory

# Load AI Image Generation Model (Stable Diffusion)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)

def generate_frames(prompt, num_frames=30):
    """Generate a series of images based on the given prompt."""
    images = []
    for i in range(num_frames):
        image = pipe(prompt).images[0]
        images.append(image)
    return images

def create_music_video(song, prompt):
    """Generate AI-based video from the uploaded song."""
    
    with TemporaryDirectory() as temp_dir:
        audio_path = os.path.join(temp_dir, "song.mp3")
        video_path = os.path.join(temp_dir, "output.mp4")

        # Save uploaded song
        with open(audio_path, "wb") as f:
            f.write(song)

        # Extract duration
        audio_clip = AudioFileClip(audio_path)
        duration = audio_clip.duration
        fps = 10
        num_frames = int(duration * fps)

        # Generate images
        images = generate_frames(prompt, num_frames)

        # Create video from images
        image_sequence_clip = ImageSequenceClip([np.array(img) for img in images], fps=fps)
        image_sequence_clip = image_sequence_clip.set_audio(audio_clip)
        image_sequence_clip.write_videofile(video_path, codec="libx264", audio_codec="aac")

        return video_path

# Gradio UI
interface = gr.Interface(
    fn=create_music_video,
    inputs=[
        gr.Audio(type="file", label="Upload Song"),
        gr.Textbox(label="Describe the Music Video (Prompt)")
    ],
    outputs=gr.Video(label="Generated AI Music Video"),
    title="AI Music Video Generator",
    description="Upload a song and describe the music video you want. AI will generate a video using Stable Diffusion.",
)

# Launc