Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import torch | |
| from transformers import ( | |
| AutoTokenizer, | |
| AutoModelForCausalLM, | |
| pipeline, | |
| AutoProcessor, | |
| MusicgenForConditionalGeneration, | |
| ) | |
| from diffusers import StableDiffusionPipeline | |
| from scipy.io.wavfile import write | |
| from pydub import AudioSegment | |
| from dotenv import load_dotenv | |
| import tempfile | |
| import spaces | |
| from TTS.api import TTS | |
| from TTS.utils.synthesizer import Synthesizer | |
| # 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 with clear structure instructions | |
| system_prompt = ( | |
| f"You are an expert radio imaging producer specializing in sound design and music. " | |
| f"Based on the user's concept and the selected duration of {duration} seconds, produce the following: " | |
| f"1. A concise voice-over script. Prefix this section with 'Voice-Over Script:'.\n" | |
| f"2. Suggestions for sound design. Prefix this section with 'Sound Design Suggestions:'.\n" | |
| f"3. Music styles or track recommendations. Prefix this section with 'Music Suggestions:'." | |
| ) | |
| combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nOutput:" | |
| result = llama_pipeline(combined_prompt, max_new_tokens=300, do_sample=True, temperature=0.8) | |
| # Parsing output | |
| generated_text = result[0]["generated_text"].split("Output:")[-1].strip() | |
| # Extract sections based on prefixes | |
| voice_script = generated_text.split("Voice-Over Script:")[1].split("Sound Design Suggestions:")[0].strip() if "Voice-Over Script:" in generated_text else "No voice-over script found." | |
| sound_design = generated_text.split("Sound Design Suggestions:")[1].split("Music Suggestions:")[0].strip() if "Sound Design Suggestions:" in generated_text else "No sound design suggestions found." | |
| music_suggestions = generated_text.split("Music Suggestions:")[1].strip() if "Music Suggestions:" in generated_text else "No music suggestions found." | |
| return voice_script, sound_design, music_suggestions | |
| except Exception as e: | |
| return f"Error generating script: {e}", "", "" | |
| # --------------------------------------------------------------------- | |
| # Voice-Over Generation Function | |
| # --------------------------------------------------------------------- | |
| def generate_voice(script: str, speaker: str = "default"): | |
| try: | |
| # Load TTS model | |
| tts_model_path = "tts_models/en/ljspeech/tacotron2-DDC" | |
| vocoder_model_path = "vocoder_models/en/ljspeech/hifigan_v2" | |
| synthesizer = Synthesizer(tts_model_path, vocoder_model_path) | |
| # Generate audio | |
| wav = synthesizer.tts(script) | |
| # Save output to a file | |
| output_path = f"{tempfile.gettempdir()}/generated_voice.wav" | |
| synthesizer.save_wav(wav, output_path) | |
| 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, model_choice: str): | |
| try: | |
| if model_choice == "Stable Audio Open 1.0": | |
| stable_pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-audio-open-1.0") | |
| stable_pipeline.to("cuda" if torch.cuda.is_available() else "cpu") | |
| audio = stable_pipeline(prompt, num_inference_steps=50, guidance_scale=7.5) | |
| output_path = f"{tempfile.gettempdir()}/stable_generated_music.wav" | |
| write(output_path, 44100, audio["sample"].cpu().numpy()) | |
| return output_path | |
| elif model_choice == "MusicGen": | |
| 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()}/musicgen_generated_music.wav" | |
| write(output_path, 44100, normalized_audio) | |
| return output_path | |
| else: | |
| return "Invalid model choice!" | |
| 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 🚀 | |
| Generate scripts, sound design, and music suggestions with ease. | |
| """) | |
| with gr.Tabs(): | |
| # Step 1: Generate Script | |
| with gr.Tab("Step 1: Generate Script"): | |
| with gr.Row(): | |
| 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 Voice-Over Script", lines=5) | |
| sound_design_output = gr.Textbox(label="Sound Design Ideas", lines=3) | |
| music_suggestion_output = gr.Textbox(label="Music Suggestions", lines=3) | |
| generate_script_button.click( | |
| fn=lambda user_prompt, model_id, duration: generate_script(user_prompt, model_id, hf_token, duration), | |
| inputs=[user_prompt, llama_model_id, duration], | |
| outputs=[script_output, sound_design_output, music_suggestion_output], | |
| ) | |
| # Step 2: Generate Voice | |
| with gr.Tab("Step 2: Generate Voice"): | |
| with gr.Row(): | |
| speaker = gr.Textbox(label="Voice Style (optional)", placeholder="E.g., male, female, or neutral.") | |
| generate_voice_button = gr.Button("Generate Voice") | |
| voice_output = gr.Audio(label="Generated Voice", type="filepath") | |
| generate_voice_button.click( | |
| fn=lambda script, speaker: generate_voice(script, speaker), | |
| inputs=[script_output, speaker], | |
| outputs=[voice_output], | |
| ) | |
| # Step 3: Generate Music | |
| with gr.Tab("Step 3: Generate Music"): | |
| with gr.Row(): | |
| audio_length = gr.Slider(label="Music Length (tokens)", minimum=128, maximum=1024, step=64, value=512) | |
| model_choice = gr.Dropdown( | |
| label="Select Music Generation Model", | |
| choices=["Stable Audio Open 1.0", "MusicGen"], | |
| value="Stable Audio Open 1.0" | |
| ) | |
| generate_music_button = gr.Button("Generate Music") | |
| music_output = gr.Audio(label="Generated Music", type="filepath") | |
| generate_music_button.click( | |
| fn=lambda music_suggestion, audio_length, model_choice: generate_music(music_suggestion, audio_length, model_choice), | |
| inputs=[music_suggestion_output, audio_length, model_choice], | |
| outputs=[music_output], | |
| ) | |
| with gr.Tab("Step 4: Blend Audio"): | |
| with gr.Row(): | |
| 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=lambda voice_path, music_path, ducking: blend_audio(voice_path, music_path, ducking), | |
| inputs=[voice_output, music_output, 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) | |