Spaces:
Sleeping
Sleeping
import gradio as gr | |
import os | |
import torch | |
from transformers import ( | |
AutoTokenizer, | |
AutoModelForCausalLM, | |
pipeline, | |
AutoProcessor, | |
MusicgenForConditionalGeneration, | |
) | |
from scipy.io.wavfile import write | |
from pydub import AudioSegment | |
from dotenv import load_dotenv | |
import tempfile | |
import spaces | |
# Load environment variables | |
load_dotenv() | |
hf_token = os.getenv("HF_TOKEN") | |
# --------------------------------------------------------------------- | |
# Script Generation Function | |
# --------------------------------------------------------------------- | |
def generate_script(user_prompt: str, model_id: str, token: str, duration: int): | |
try: | |
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_id, | |
use_auth_token=token, | |
torch_dtype=torch.float16, | |
device_map="auto", | |
trust_remote_code=True, | |
) | |
llama_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer) | |
system_prompt = ( | |
"You are an expert radio imaging producer specializing in sound design and music.\n" | |
"---\n" | |
f"Based on the user's concept and the selected duration of {duration} seconds, craft a concise, engaging promo script.\n" | |
"---\n" | |
"Ensure the script fits within the time limit and suggest a matching music style that complements the theme." | |
) | |
combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nRefined script and music suggestion:" | |
result = llama_pipeline(combined_prompt, max_new_tokens=500, do_sample=True, temperature=0.9) | |
generated_text = result[0]["generated_text"] | |
if "Refined script and music suggestion:" in generated_text: | |
parts = generated_text.split("Refined script and music suggestion:", 1)[-1].strip() | |
if "Music Style:" in parts: | |
script, music_suggestion = parts.split("Music Style:", 1) | |
return script.strip(), music_suggestion.strip() | |
else: | |
return parts.strip(), "No specific music suggestion found." | |
return "Error: Could not parse the script.", None | |
except Exception as e: | |
return f"Error generating script: {e}", None | |
# --------------------------------------------------------------------- | |
# Voice-Over Generation Function | |
# --------------------------------------------------------------------- | |
def generate_voice(script: str): | |
try: | |
tts_model = "coqui/xtts-en-ljspeech-v2" | |
processor = AutoProcessor.from_pretrained(tts_model) | |
model = AutoModelForCausalLM.from_pretrained(tts_model) | |
inputs = processor(script, return_tensors="pt") | |
speech = model.generate(**inputs) | |
output_path = f"{tempfile.gettempdir()}/generated_voice.wav" | |
write(output_path, 22050, speech.cpu().numpy()) | |
return output_path | |
except Exception as e: | |
return f"Error generating voice-over: {e}" | |
# --------------------------------------------------------------------- | |
# Music Generation Function | |
# --------------------------------------------------------------------- | |
def generate_music(prompt: str, audio_length: int): | |
try: | |
musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small") | |
musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small") | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
musicgen_model.to(device) | |
inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device) | |
outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length) | |
audio_data = outputs[0, 0].cpu().numpy() | |
normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16") | |
output_path = f"{tempfile.gettempdir()}/generated_music.wav" | |
write(output_path, 44100, normalized_audio) | |
return output_path | |
except Exception as e: | |
return f"Error generating music: {e}" | |
# --------------------------------------------------------------------- | |
# Audio Blending Function with Ducking | |
# --------------------------------------------------------------------- | |
def blend_audio(voice_path: str, music_path: str, ducking: bool): | |
try: | |
voice = AudioSegment.from_file(voice_path) | |
music = AudioSegment.from_file(music_path) | |
if ducking: | |
music = music - 10 # Lower music volume for ducking | |
combined = music.overlay(voice) | |
output_path = f"{tempfile.gettempdir()}/final_promo.wav" | |
combined.export(output_path, format="wav") | |
return output_path | |
except Exception as e: | |
return f"Error blending audio: {e}" | |
# --------------------------------------------------------------------- | |
# Gradio Interface | |
# --------------------------------------------------------------------- | |
with gr.Blocks() as demo: | |
gr.Markdown(""" | |
# 🎧 AI Promo Studio with Pages 🚀 | |
Follow a step-by-step process to create amazing promos with AI. | |
""") | |
with gr.Tabs(): | |
# Step 1: Script Generation | |
with gr.Tab("Step 1: Generate Script"): | |
user_prompt = gr.Textbox(label="Promo Idea", placeholder="E.g., A 30-second promo for a morning show.") | |
llama_model_id = gr.Textbox(label="Llama Model ID", value="meta-llama/Meta-Llama-3-8B-Instruct") | |
duration = gr.Slider(label="Duration (seconds)", minimum=15, maximum=60, step=15, value=30) | |
generate_script_button = gr.Button("Generate Script") | |
script_output = gr.Textbox(label="Generated Script") | |
music_suggestion_output = gr.Textbox(label="Music Suggestion") | |
generate_script_button.click( | |
fn=generate_script, | |
inputs=[user_prompt, llama_model_id, hf_token, duration], | |
outputs=[script_output, music_suggestion_output], | |
) | |
# Step 2: Voice Generation | |
with gr.Tab("Step 2: Generate Voice"): | |
script_input = gr.Textbox(label="Script for Voice", interactive=False) | |
generate_voice_button = gr.Button("Generate Voice") | |
voice_output = gr.Audio(label="Generated Voice", type="filepath") | |
generate_voice_button.click( | |
fn=generate_voice, | |
inputs=[script_input], | |
outputs=[voice_output], | |
) | |
# Step 3: Music Generation | |
with gr.Tab("Step 3: Generate Music"): | |
music_prompt_input = gr.Textbox(label="Music Suggestion Prompt", interactive=False) | |
audio_length = gr.Slider(label="Music Length (tokens)", minimum=128, maximum=1024, step=64, value=512) | |
generate_music_button = gr.Button("Generate Music") | |
music_output = gr.Audio(label="Generated Music", type="filepath") | |
generate_music_button.click( | |
fn=generate_music, | |
inputs=[music_prompt_input, audio_length], | |
outputs=[music_output], | |
) | |
# Step 4: Blend Audio | |
with gr.Tab("Step 4: Blend Audio"): | |
voice_path = gr.Audio(label="Voice File", type="filepath") | |
music_path = gr.Audio(label="Music File", type="filepath") | |
ducking = gr.Checkbox(label="Enable Ducking", value=True) | |
blend_button = gr.Button("Blend Audio") | |
final_output = gr.Audio(label="Final Promo Audio", type="filepath") | |
blend_button.click( | |
fn=blend_audio, | |
inputs=[voice_path, music_path, ducking], | |
outputs=[final_output], | |
) | |
gr.Markdown(""" | |
<hr> | |
<p style="text-align: center; font-size: 0.9em;"> | |
Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a> | |
</p> | |
""") | |
demo.launch(debug=True) | |