Update app.py

app.py

@@ -10,13 +10,13 @@ from transformers import (
     MusicgenForConditionalGeneration
 )
 from io import BytesIO
-from streamlit_lottie import st_lottie  # pip install streamlit-lottie
+from streamlit_lottie import st_lottie
 
 # ---------------------------------------------------------------------
 # 1) PAGE CONFIG
 # ---------------------------------------------------------------------
 st.set_page_config(
-    page_title="Radio Imaging AI MVP",
+    page_title="Radio Imaging AI with Llama 3",
     page_icon="🎧",
     layout="wide"
 )
@@ -26,26 +26,24 @@ st.set_page_config(
 # ---------------------------------------------------------------------
 CUSTOM_CSS = """
 <style>
-/* Body styling for a dark, music-app vibe */
+/* Dark background with Spotify-like vibe */
 body {
     background-color: #121212;
     color: #FFFFFF;
     font-family: "Helvetica Neue", sans-serif;
 }
 
-/* Main container width */
 .block-container {
     max-width: 1100px;
     padding: 1rem 1.5rem;
 }
 
-/* Headings with a neon-ish green accent */
 h1, h2, h3 {
-    color: #1DB954; 
+    color: #1DB954;
     margin-bottom: 0.5rem;
 }
 
-/* Buttons: rounded, bright Spotify-like green on hover */
+/* Rounded, bright green button on hover */
 .stButton>button {
     background-color: #1DB954 !important;
     color: #FFFFFF !important;
@@ -59,13 +57,12 @@ h1, h2, h3 {
     background-color: #1ed760 !important;
 }
 
-/* Sidebar: black background, white text */
+/* Sidebar: black background */
 .sidebar .sidebar-content {
     background-color: #000000;
     color: #FFFFFF;
 }
 
-/* Text inputs and text areas */
 textarea, input, select {
     border-radius: 8px !important;
     background-color: #282828 !important;
@@ -73,20 +70,20 @@ textarea, input, select {
     border: 1px solid #3e3e3e;
 }
 
-/* Audio player styling */
+/* Audio styling */
 audio {
     width: 100%;
     margin-top: 1rem;
 }
 
-/* Lottie container styling */
+/* Lottie container */
 .lottie-container {
     display: flex;
     justify-content: center;
     margin-bottom: 20px;
 }
 
-/* Footer styling */
+/* Footer */
 .footer-note {
     text-align: center; 
     font-size: 14px; 
@@ -94,31 +91,83 @@ audio {
     margin-top: 2rem;
 }
 
-/* Hide Streamlit's default branding if desired */
+/* Hide Streamlit branding if you wish */
 #MainMenu, footer {visibility: hidden;}
 </style>
 """
 st.markdown(CUSTOM_CSS, unsafe_allow_html=True)
 
 # ---------------------------------------------------------------------
-# 3) HELPER: LOAD LOTTIE ANIMATION
+# 3) LOAD LOTTIE ANIMATION
 # ---------------------------------------------------------------------
 @st.cache_data
 def load_lottie_url(url: str):
-    """
-    Fetch Lottie JSON for animations.
-    """
     r = requests.get(url)
     if r.status_code != 200:
         return None
     return r.json()
 
-# Example Lottie animation (radio waves / music eq, etc.)
 LOTTIE_URL = "https://assets3.lottiefiles.com/temp/lf20_Q6h5zV.json"
 lottie_animation = load_lottie_url(LOTTIE_URL)
 
 # ---------------------------------------------------------------------
-# 4) SIDEBAR: "LIBRARY" NAVIGATION (MIMICS SPOTIFY)
+# 4) LOAD LLAMA 3 (GATED MODEL) - WITH use_auth_token
+# ---------------------------------------------------------------------
+@st.cache_resource
+def load_llama_pipeline(model_id: str, device: str):
+    """
+    Load the Llama 3 model from Hugging Face.
+    Requires huggingface-cli login if model is gated.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16 if device == "auto" else torch.float32,
+        device_map=device,
+        use_auth_token=True
+    )
+    text_gen_pipeline = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device_map=device
+    )
+    return text_gen_pipeline
+
+# ---------------------------------------------------------------------
+# 5) REFINE SCRIPT (LLAMA)
+# ---------------------------------------------------------------------
+def generate_radio_script(user_input: str, pipeline_llama) -> str:
+    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_input}\nRefined script:"
+
+    result = pipeline_llama(
+        combined_prompt,
+        max_new_tokens=200,
+        do_sample=True,
+        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)"
+    return output_text
+
+# ---------------------------------------------------------------------
+# 6) LOAD MUSICGEN
+# ---------------------------------------------------------------------
+@st.cache_resource
+def load_musicgen_model():
+    mg_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+    mg_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+    return mg_model, mg_processor
+
+# ---------------------------------------------------------------------
+# 7) SIDEBAR
 # ---------------------------------------------------------------------
 with st.sidebar:
     st.header("🎚 Radio Library")
@@ -131,21 +180,19 @@ with st.sidebar:
     st.markdown("<br>", unsafe_allow_html=True)
 
 # ---------------------------------------------------------------------
-# 5) HEADER SECTION WITH LOTS OF FLARE
+# 8) HEADER
 # ---------------------------------------------------------------------
 col1, col2 = st.columns([3, 2], gap="large")
 
 with col1:
-    st.title("AI Radio Imaging MVP")
-    st.subheader("Llama-Driven Promo Scripts, MusicGen Audio")
+    st.title("AI Radio Imaging with Llama 3")
+    st.subheader("Gated Model + MusicGen Audio")
 
     st.markdown(
         """
-        Create **radio imaging promos** and **jingles** with a minimal but creative MVP. 
-        This app:
-        - Uses a (hypothetical) [Llama 3] model for **script generation**.
-        - Uses Meta's [MusicGen](https://github.com/facebookresearch/audiocraft) for **audio**.
-        - Features a Spotify-like UI & Lottie animations for a modern user experience.
+        Create **radio imaging promos** and **jingles** with Llama 3 + MusicGen.  
+        **Note**: You must have access to `"meta-llama/Llama-3-70B-Instruct"` on Hugging Face, 
+        and be logged in via `huggingface-cli login`.
         """
     )
 with col2:
@@ -158,152 +205,86 @@ with col2:
 st.markdown("---")
 
 # ---------------------------------------------------------------------
-# 6) PROMPT INPUT & MODEL SELECTION
+# 9) SCRIPT GENERATION
 # ---------------------------------------------------------------------
-st.subheader("🎙 Step 1: Briefly Describe Your Promo Idea")
+st.subheader("🎙 Step 1: Describe Your Promo Idea")
 
 prompt = st.text_area(
-    "E.g. 'A 15-second upbeat jingle with a catchy hook for a Top 40 morning show'",
+    "Example: 'A 15-second hype jingle for a morning talk show, fun and energetic.'",
     height=120
 )
 
 col_model, col_device = st.columns(2)
 with col_model:
     llama_model_id = st.text_input(
-        "Llama Model (Hugging Face ID)",
-        value="meta-llama/Llama-3.3-70B-Instruct",  # Replace with a real model
-        help="If non-existent, you'll see errors. Try Llama 2 (e.g. meta-llama/Llama-2-7b-chat-hf)."
+        "Llama 3 Model ID",
+        value="meta-llama/Llama-3-70B-Instruct",  # Official ID if you have it
+        help="Use the exact name you see on the Hugging Face model page."
     )
 with col_device:
     device_option = st.selectbox(
-        "Choose Device",
+        "Device (GPU vs CPU)",
         ["auto", "cpu"],
-        help="For GPU usage, pick 'auto'. CPU can be slow for big models."
+        help="If you have GPU, 'auto' tries to use it; CPU might be slow."
     )
 
-# ---------------------------------------------------------------------
-# 7) BUTTON: GENERATE RADIO SCRIPT WITH LLAMA
-# ---------------------------------------------------------------------
 if st.button("📝 Generate Promo Script"):
     if not prompt.strip():
-        st.error("Please enter a radio imaging concept first.")
+        st.error("Please type some concept first.")
     else:
-        with st.spinner("Generating script..."):
+        with st.spinner("Generating script with Llama 3..."):
             try:
-                # Load Llama pipeline
-                pipeline_llama = load_llama_pipeline(llama_model_id, device_option)
-                # Generate refined script
-                refined_text = generate_radio_script(prompt, pipeline_llama)
-                st.session_state["refined_script"] = refined_text
+                llm_pipeline = load_llama_pipeline(llama_model_id, device_option)
+                final_script = generate_radio_script(prompt, llm_pipeline)
+                st.session_state["final_script"] = final_script
                 st.success("Promo script generated!")
-                st.write(refined_text)
+                st.write(final_script)
             except Exception as e:
-                st.error(f"Error during Llama generation: {e}")
+                st.error(f"Llama generation error: {e}")
 
 st.markdown("---")
 
 # ---------------------------------------------------------------------
-# 8) AUDIO GENERATION: MUSICGEN
+# 10) AUDIO GENERATION: MUSICGEN
 # ---------------------------------------------------------------------
-st.subheader("🎶 Step 2: Generate Your Radio Audio")
+st.subheader("🎶 Step 2: Generate Audio")
 
-audio_tokens = st.slider("MusicGen Max Tokens (Track Length)", 128, 1024, 512, 64)
+audio_length = st.slider("MusicGen Max Tokens (approx track length)", 128, 1024, 512, 64)
 
 if st.button("🎧 Create Audio with MusicGen"):
-    # Check if we have a refined script
-    if "refined_script" not in st.session_state:
-        st.error("Please generate a promo script first.")
+    if "final_script" not in st.session_state:
+        st.error("No script found. Please generate a script first.")
     else:
-        with st.spinner("Generating audio..."):
+        with st.spinner("Creating audio..."):
             try:
-                # Load MusicGen
                 mg_model, mg_processor = load_musicgen_model()
-                descriptive_text = st.session_state["refined_script"]
-                
-                # Prepare model input
+                text_for_audio = st.session_state["final_script"]
+
                 inputs = mg_processor(
-                    text=[descriptive_text],
-                    return_tensors="pt",
-                    padding=True
+                    text=[text_for_audio],
+                    padding=True,
+                    return_tensors="pt"
                 )
-                # Generate audio
-                audio_values = mg_model.generate(**inputs, max_new_tokens=audio_tokens)
+                audio_values = mg_model.generate(**inputs, max_new_tokens=audio_length)
                 sr = mg_model.config.audio_encoder.sampling_rate
-                
-                # Save audio to WAV
-                out_filename = "radio_imaging_output.wav"
-                scipy.io.wavfile.write(out_filename, rate=sr, data=audio_values[0,0].numpy())
-
-                st.success("Audio created! Press play to listen:")
-                st.audio(out_filename)
-            except Exception as e:
-                st.error(f"Error generating audio: {e}")
-
-# ---------------------------------------------------------------------
-# 9) HELPER FUNCTIONS
-# ---------------------------------------------------------------------
-@st.cache_resource
-def load_llama_pipeline(model_id: str, device: str):
-    """
-    Load the Llama model & pipeline.
-    """
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16 if device == "auto" else torch.float32,
-        device_map=device
-    )
-    text_gen_pipeline = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device_map=device
-    )
-    return text_gen_pipeline
 
-def generate_radio_script(user_input: str, pipeline_llama) -> str:
-    """
-    Use Llama to refine the user's input into a brief but creative radio imaging script.
-    """
-    system_prompt = (
-        "You are a top-tier radio imaging producer. "
-        "Take the user's concept and craft a short, high-impact promo script. "
-        "Include style, tone, and potential CTA if relevant."
-    )
-    full_prompt = f"{system_prompt}\nUser concept: {user_input}\nRefined script:"
-    
-    output = pipeline_llama(
-        full_prompt,
-        max_new_tokens=200,
-        do_sample=True,
-        temperature=0.9
-    )[0]["generated_text"]
-    
-    # Attempt to isolate the final script portion
-    if "Refined script:" in output:
-        output = output.split("Refined script:", 1)[-1].strip()
-    output += "\n\n(Generated by Llama in Radio Imaging MVP)"
-    
-    return output
+                outfile = "llama3_radio_jingle.wav"
+                scipy.io.wavfile.write(outfile, rate=sr, data=audio_values[0, 0].numpy())
 
-@st.cache_resource
-def load_musicgen_model():
-    """
-    Load MusicGen (small version).
-    """
-    mg_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
-    mg_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
-    return mg_model, mg_processor
+                st.success("Audio generated! Press play below:")
+                st.audio(outfile)
+            except Exception as e:
+                st.error(f"MusicGen error: {e}")
 
 # ---------------------------------------------------------------------
-# 10) FOOTER
+# 11) FOOTER
 # ---------------------------------------------------------------------
 st.markdown("---")
 st.markdown(
     """
     <div class="footer-note">
-    &copy; 2025 Radio Imaging MVP &ndash; Built with Llama & MusicGen. <br>
-    Inspired by Spotify's UI for a sleek, modern experience.
+    © 2025 Radio Imaging with Llama 3 – Built using Hugging Face & Streamlit. <br>
+    Log in via <code>huggingface-cli</code> and ensure access to <strong>meta-llama/Llama-3-70B-Instruct</strong>.
     </div>
     """,
     unsafe_allow_html=True