Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import requests
 import torch
-import scipy.io.wavfile
+import scipy.io.wavfile as wav
 import streamlit as st
 from io import BytesIO
 from transformers import (
@@ -14,85 +14,55 @@ from transformers import (
 from streamlit_lottie import st_lottie
 
 # ---------------------------------------------------------------------
-# 1) PAGE CONFIG
+# 1) PAGE CONFIGURATION
 # ---------------------------------------------------------------------
 st.set_page_config(
-    page_title="Radio Imaging AI with Llama 3",
-    page_icon="🎧",
+    page_title="AI Radio Imaging with Llama 3",
+    page_icon="\ud83c\udfa7",
     layout="wide"
 )
 
 # ---------------------------------------------------------------------
-# 2) CUSTOM CSS / SPOTIFY-LIKE UI
+# 2) CUSTOM CSS / UI DESIGN
 # ---------------------------------------------------------------------
 CUSTOM_CSS = """
 <style>
-/* Dark background with Spotify-like vibe */
 body {
     background-color: #121212;
     color: #FFFFFF;
     font-family: "Helvetica Neue", sans-serif;
 }
-
 .block-container {
     max-width: 1100px;
     padding: 1rem 1.5rem;
 }
-
 h1, h2, h3 {
     color: #1DB954;
-    margin-bottom: 0.5rem;
 }
-
-/* Rounded, bright green button on hover */
 .stButton>button {
     background-color: #1DB954 !important;
     color: #FFFFFF !important;
     border-radius: 24px;
-    border: none;
-    font-size: 16px !important;
-    padding: 0.6rem 1.2rem !important;
-    transition: background-color 0.3s ease;
+    padding: 0.6rem 1.2rem;
 }
 .stButton>button:hover {
     background-color: #1ed760 !important;
 }
-
-/* Sidebar: black background */
-.sidebar .sidebar-content {
-    background-color: #000000;
-    color: #FFFFFF;
-}
-
 textarea, input, select {
     border-radius: 8px !important;
     background-color: #282828 !important;
     color: #FFFFFF !important;
-    border: 1px solid #3e3e3e;
 }
-
-/* Audio styling */
 audio {
     width: 100%;
     margin-top: 1rem;
 }
-
-/* Lottie container */
-.lottie-container {
-    display: flex;
-    justify-content: center;
-    margin-bottom: 20px;
-}
-
-/* Footer */
 .footer-note {
     text-align: center; 
     font-size: 14px; 
     opacity: 0.7;
     margin-top: 2rem;
 }
-
-/* Hide Streamlit branding if you wish */
 #MainMenu, footer {visibility: hidden;}
 </style>
 """
@@ -112,18 +82,11 @@ LOTTIE_URL = "https://assets3.lottiefiles.com/temp/lf20_Q6h5zV.json"
 lottie_animation = load_lottie_url(LOTTIE_URL)
 
 # ---------------------------------------------------------------------
-# 4) LOAD LLAMA 3 (GATED MODEL) - WITH use_auth_token
+# 4) LOAD LLAMA 3 (GATED MODEL)
 # ---------------------------------------------------------------------
 @st.cache_resource
 def load_llama_pipeline(model_id: str, device: str, token: str):
-    """
-    Load the Llama 3 model from Hugging Face with a user token.
-    token: The HF access token from environment or secrets.
-    """
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_id, 
-        use_auth_token=token
-    )
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         use_auth_token=token,
@@ -139,7 +102,7 @@ def load_llama_pipeline(model_id: str, device: str, token: str):
     return text_gen_pipeline
 
 # ---------------------------------------------------------------------
-# 5) REFINE SCRIPT (LLAMA)
+# 5) GENERATE RADIO SCRIPT
 # ---------------------------------------------------------------------
 def generate_radio_script(user_input: str, pipeline_llama) -> str:
     system_prompt = (
@@ -155,7 +118,6 @@ def generate_radio_script(user_input: str, pipeline_llama) -> str:
         temperature=0.9
     )
     output_text = result[0]["generated_text"]
-
     if "Refined script:" in output_text:
         output_text = output_text.split("Refined script:", 1)[-1].strip()
     output_text += "\n\n(Generated by Llama 3 - Radio Imaging)"
@@ -171,49 +133,22 @@ def load_musicgen_model():
     return mg_model, mg_processor
 
 # ---------------------------------------------------------------------
-# 7) SIDEBAR
-# ---------------------------------------------------------------------
-with st.sidebar:
-    st.header("🎚 Radio Library")
-    st.write("**My Stations**")
-    st.write("- Favorites")
-    st.write("- Recently Generated")
-    st.write("- Top Hits")
-    st.write("---")
-    st.write("**Settings**")
-    st.markdown("<br>", unsafe_allow_html=True)
-
-# ---------------------------------------------------------------------
-# 8) HEADER
+# 7) HEADER
 # ---------------------------------------------------------------------
-col1, col2 = st.columns([3, 2], gap="large")
+st.title("\ud83c\udfa7 AI Radio Imaging with Llama 3")
+st.subheader("Create engaging radio promos with Llama 3 + MusicGen")
+st.markdown("""Create **radio imaging promos** and **jingles** easily. Ensure you have access to 
+**meta-llama/Meta-Llama-3-70B** on Hugging Face and provide your token below.""")
 
-with col1:
-    st.title("AI Radio Imaging with Llama 3")
-    st.subheader("Gated Model + MusicGen Audio")
-
-    st.markdown(
-        """
-        Create **radio imaging promos** and **jingles** with Llama 3 + MusicGen.  
-        **Note**: 
-        - You must have access to `"meta-llama/Meta-Llama-3-70B"` on Hugging Face.
-        - You must provide your HF token in the environment (e.g., HF_TOKEN).
-        """
-    )
-with col2:
-    if lottie_animation:
-        with st.container():
-            st_lottie(lottie_animation, height=180, loop=True, key="radio_lottie")
-    else:
-        st.write("*No animation loaded.*")
+if lottie_animation:
+    st_lottie(lottie_animation, height=180, loop=True, key="radio_lottie")
 
 st.markdown("---")
 
 # ---------------------------------------------------------------------
-# 9) SCRIPT GENERATION
+# 8) USER INPUT
 # ---------------------------------------------------------------------
-st.subheader("🎙 Step 1: Describe Your Promo Idea")
-
+st.subheader("\ud83c\udfa4 Step 1: Describe Your Promo Idea")
 prompt = st.text_area(
     "Example: 'A 15-second hype jingle for a morning talk show, fun and energetic.'",
     height=120
@@ -223,79 +158,75 @@ col_model, col_device = st.columns(2)
 with col_model:
     llama_model_id = st.text_input(
         "Llama 3 Model ID",
-        value="meta-llama/Meta-Llama-3-70B",  
-        help="Use the exact name you see on the Hugging Face model page."
+        value="meta-llama/Meta-Llama-3-70B",
+        help="Enter the exact model ID from Hugging Face."
     )
 with col_device:
     device_option = st.selectbox(
-        "Device (GPU vs CPU)",
+        "Device",
         ["auto", "cpu"],
-        help="If you have GPU, 'auto' tries to use it; CPU might be slow."
+        help="Choose GPU (auto) or CPU."
     )
 
-# Grab your token from environment
-my_token = os.getenv("HF_TOKEN")
-if not my_token:
-    st.error("No HF_TOKEN found. Please set it in your HF Space secrets or environment variables.")
+hf_token = os.getenv("HF_TOKEN")
+if not hf_token:
+    st.error("No HF_TOKEN found. Please set it in your environment.")
     st.stop()
 
-if st.button("📝 Generate Promo Script"):
+if st.button("\u270d Generate Promo Script"):
     if not prompt.strip():
-        st.error("Please type some concept first.")
+        st.error("Please provide a concept first.")
     else:
-        with st.spinner("Generating script with Llama 3..."):
+        with st.spinner("Generating script..."):
             try:
-                llm_pipeline = load_llama_pipeline(llama_model_id, device_option, my_token)
-                final_script = generate_radio_script(prompt, llm_pipeline)
-                st.session_state["final_script"] = final_script
+                llama_pipeline = load_llama_pipeline(llama_model_id, device_option, hf_token)
+                final_script = generate_radio_script(prompt, llama_pipeline)
                 st.success("Promo script generated!")
-                st.write(final_script)
+                st.text_area("Generated Script", value=final_script, height=200)
             except Exception as e:
                 st.error(f"Llama generation error: {e}")
 
 st.markdown("---")
 
 # ---------------------------------------------------------------------
-# 10) AUDIO GENERATION: MUSICGEN
+# 9) GENERATE AUDIO WITH MUSICGEN
 # ---------------------------------------------------------------------
-st.subheader("🎶 Step 2: Generate Audio")
-
-audio_length = st.slider("MusicGen Max Tokens (approx track length)", 128, 1024, 512, 64)
+st.subheader("\ud83c\udfb5 Step 2: Generate Audio")
+audio_length = st.slider("Track Length (tokens)", 128, 1024, 512, 64)
 
-if st.button("🎧 Create Audio with MusicGen"):
+if st.button("\ud83c\udfa7 Create Audio"):
     if "final_script" not in st.session_state:
-        st.error("No script found. Please generate a script first.")
+        st.error("Please generate a script first.")
     else:
-        with st.spinner("Creating audio..."):
+        with st.spinner("Generating audio..."):
             try:
                 mg_model, mg_processor = load_musicgen_model()
-                text_for_audio = st.session_state["final_script"]
-
                 inputs = mg_processor(
-                    text=[text_for_audio],
+                    text=[st.session_state["final_script"]],
                     padding=True,
                     return_tensors="pt"
                 )
                 audio_values = mg_model.generate(**inputs, max_new_tokens=audio_length)
                 sr = mg_model.config.audio_encoder.sampling_rate
+                output_file = "radio_jingle.wav"
 
-                outfile = "llama3_radio_jingle.wav"
-                scipy.io.wavfile.write(outfile, rate=sr, data=audio_values[0, 0].numpy())
+                audio_data = audio_values[0, 0].cpu().numpy()
+                normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
+                wav.write(output_file, rate=sr, data=normalized_audio)
 
-                st.success("Audio generated! Press play below:")
-                st.audio(outfile)
+                st.success("Audio generated! Play it below:")
+                st.audio(output_file)
             except Exception as e:
                 st.error(f"MusicGen error: {e}")
 
 # ---------------------------------------------------------------------
-# 11) FOOTER
+# 10) FOOTER
 # ---------------------------------------------------------------------
 st.markdown("---")
 st.markdown(
     """
     <div class="footer-note">
-    © 2025 Radio Imaging with Llama 3 – Built using Hugging Face & Streamlit. <br>
-    Log in or provide <code>HF_TOKEN</code> and ensure access to <strong>meta-llama/Llama-3-70B-Instruct</strong>.
+    © 2025 AI Radio Imaging – Built with Hugging Face & Streamlit
     </div>
     """,
     unsafe_allow_html=True