maybthingsyouneverknow's picture
Generator.py
b0dae94 verified
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