Spaces:
Running
Running
import gradio as gr | |
import os | |
import torch | |
from transformers import ( | |
AutoTokenizer, | |
AutoModelForCausalLM, | |
pipeline, | |
AutoProcessor, | |
MusicgenForConditionalGeneration, | |
) | |
from scipy.io.wavfile import write | |
import tempfile | |
from dotenv import load_dotenv | |
import spaces # Assumes Hugging Face Spaces library supports `@spaces.GPU` | |
# Load environment variables (e.g., Hugging Face token) | |
load_dotenv() | |
hf_token = os.getenv("HF_TOKEN") | |
# --------------------------------------------------------------------- | |
# Load Llama 3 Pipeline with Zero GPU (Encapsulated) | |
# --------------------------------------------------------------------- | |
# Adjust GPU allocation duration | |
def generate_script(user_prompt: str, model_id: str, token: str): | |
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 a top-tier radio imaging producer using Llama 3. " | |
"Take the user's concept and craft a short, creative promo script." | |
) | |
combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nRefined script:" | |
result = llama_pipeline(combined_prompt, max_new_tokens=200, do_sample=True, temperature=0.9) | |
return result[0]["generated_text"].split("Refined script:")[-1].strip() | |
except Exception as e: | |
return f"Error generating script: {e}" | |
# --------------------------------------------------------------------- | |
# Load MusicGen Model (Encapsulated) | |
# --------------------------------------------------------------------- | |
def generate_audio(prompt: str, audio_length: int): | |
try: | |
musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small") | |
musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small") | |
musicgen_model.to("cuda") | |
inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt") | |
outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length) | |
musicgen_model.to("cpu") # Return the model to CPU | |
sr = musicgen_model.config.audio_encoder.sampling_rate | |
audio_data = outputs[0, 0].cpu().numpy() | |
normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16") | |
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav: | |
write(temp_wav.name, sr, normalized_audio) | |
return temp_wav.name | |
except Exception as e: | |
return f"Error generating audio: {e}" | |
# --------------------------------------------------------------------- | |
# Gradio Interface | |
# --------------------------------------------------------------------- | |
def interface_generate_script(user_prompt, llama_model_id): | |
return generate_script(user_prompt, llama_model_id, hf_token) | |
def interface_generate_audio(script, audio_length): | |
return generate_audio(script, audio_length) | |
# --------------------------------------------------------------------- | |
# Interface | |
# --------------------------------------------------------------------- | |
with gr.Blocks() as demo: | |
gr.Markdown("# 🎧 AI Radio Imaging with Llama 3 + MusicGen (Zero GPU)") | |
with gr.Row(): | |
user_prompt = gr.Textbox( | |
label="Enter your promo idea", | |
placeholder="E.g., A 15-second hype jingle for a morning talk show.", | |
) | |
llama_model_id = gr.Textbox( | |
label="Llama 3 Model ID", value="meta-llama/Meta-Llama-3-8B-Instruct" | |
) | |
audio_length = gr.Slider(label="Audio Length (tokens)", minimum=128, maximum=1024, step=64, value=512) | |
generate_script_button = gr.Button("Generate Script") | |
script_output = gr.Textbox(label="Generated Script") | |
generate_audio_button = gr.Button("Generate Audio") | |
audio_output = gr.Audio(label="Generated Audio", type="filepath") | |
generate_script_button.click( | |
fn=interface_generate_script, | |
inputs=[user_prompt, llama_model_id], | |
outputs=script_output, | |
) | |
generate_audio_button.click( | |
fn=interface_generate_audio, | |
inputs=[script_output, audio_length], | |
outputs=audio_output, | |
) | |
# --------------------------------------------------------------------- | |
# Launch App | |
# --------------------------------------------------------------------- | |
demo.launch(debug=True) | |