Update app.py

app.py

@@ -1,8 +1,7 @@
 import streamlit as st
+import requests
 import torch
 import scipy.io.wavfile
-import requests
-from io import BytesIO
 from transformers import (
     AutoTokenizer, 
     AutoModelForCausalLM, 
@@ -10,200 +9,243 @@ from transformers import (
     AutoProcessor, 
     MusicgenForConditionalGeneration
 )
+from io import BytesIO
 from streamlit_lottie import st_lottie  # pip install streamlit-lottie
 
 # ---------------------------------------------------------------------
-# 1) Page Configuration
+# 1) PAGE CONFIG
 # ---------------------------------------------------------------------
 st.set_page_config(
-    page_title="Modern Radio Imaging Generator - Llama 3 & MusicGen",
+    page_title="Radio Imaging AI MVP",
     page_icon="🎧",
     layout="wide"
 )
 
 # ---------------------------------------------------------------------
-# 2) Custom CSS for a Sleek, Modern Look
+# 2) CUSTOM CSS / SPOTIFY-LIKE UI
 # ---------------------------------------------------------------------
-MODERN_CSS = """
+CUSTOM_CSS = """
 <style>
-/* Body styling */
+/* Body styling for a dark, music-app vibe */
 body {
-    background: linear-gradient(to bottom right, #ffffff, #f3f4f6);
-    font-family: 'Helvetica Neue', Arial, sans-serif;
-    color: #1F2937;
+    background-color: #121212;
+    color: #FFFFFF;
+    font-family: "Helvetica Neue", sans-serif;
 }
 
-/* Make the container narrower for a sleek look */
+/* Main container width */
 .block-container {
     max-width: 1100px;
+    padding: 1rem 1.5rem;
 }
 
-/* Heading style */
-h1, h2, h3, h4, h5, h6 {
-    color: #3B82F6;
-    margin-bottom: 0.5em;
+/* Headings with a neon-ish green accent */
+h1, h2, h3 {
+    color: #1DB954; 
+    margin-bottom: 0.5rem;
 }
 
-/* Buttons */
+/* Buttons: rounded, bright Spotify-like green on hover */
 .stButton>button {
-    background-color: #3B82F6 !important;
+    background-color: #1DB954 !important;
     color: #FFFFFF !important;
-    border-radius: 0.8rem !important;
-    font-size: 1rem !important;
+    border-radius: 24px;
+    border: none;
+    font-size: 16px !important;
     padding: 0.6rem 1.2rem !important;
+    transition: background-color 0.3s ease;
+}
+.stButton>button:hover {
+    background-color: #1ed760 !important;
 }
 
-/* Sidebar customization */
+/* Sidebar: black background, white text */
 .sidebar .sidebar-content {
-    background: #E0F2FE;
+    background-color: #000000;
+    color: #FFFFFF;
 }
 
-/* Text input areas */
+/* Text inputs and text areas */
 textarea, input, select {
-    border-radius: 0.5rem !important;
+    border-radius: 8px !important;
+    background-color: #282828 !important;
+    color: #FFFFFF !important;
+    border: 1px solid #3e3e3e;
 }
 
-/* Animate some elements on hover (just an example) */
-.stButton>button:hover {
-    background-color: #2563EB !important;
-    transition: background-color 0.3s ease-in-out;
+/* Audio player styling */
+audio {
+    width: 100%;
+    margin-top: 1rem;
 }
 
-/* Lottie container style */
+/* Lottie container styling */
 .lottie-container {
     display: flex;
     justify-content: center;
-    margin: 1rem 0;
+    margin-bottom: 20px;
 }
 
-/* Footer note */
+/* Footer styling */
 .footer-note {
     text-align: center; 
-    opacity: 0.7; 
     font-size: 14px; 
+    opacity: 0.7;
     margin-top: 2rem;
 }
 
-/* Hide default Streamlit branding if desired */
+/* Hide Streamlit's default branding if desired */
 #MainMenu, footer {visibility: hidden;}
 </style>
 """
-st.markdown(MODERN_CSS, unsafe_allow_html=True)
+st.markdown(CUSTOM_CSS, unsafe_allow_html=True)
 
 # ---------------------------------------------------------------------
-# 3) Lottie Animation Loader
+# 3) HELPER: LOAD LOTTIE ANIMATION
 # ---------------------------------------------------------------------
 @st.cache_data
 def load_lottie_url(url: str):
     """
-    Loads a Lottie animation JSON from a given URL.
+    Fetch Lottie JSON for animations.
     """
     r = requests.get(url)
     if r.status_code != 200:
         return None
     return r.json()
 
-# Example Lottie animations (feel free to replace with your own):
-LOTTIE_URL_HEADER = "https://assets1.lottiefiles.com/packages/lf20_amhnytsm.json"  # music-themed animation
-lottie_music = load_lottie_url(LOTTIE_URL_HEADER)
+# 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)
+# ---------------------------------------------------------------------
+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)
 
 # ---------------------------------------------------------------------
-# 4) Header & Intro with a Lottie Animation
+# 5) HEADER SECTION WITH LOTS OF FLARE
 # ---------------------------------------------------------------------
-col_header1, col_header2 = st.columns([3, 2], gap="medium")
+col1, col2 = st.columns([3, 2], gap="large")
+
+with col1:
+    st.title("AI Radio Imaging MVP")
+    st.subheader("Llama-Driven Promo Scripts, MusicGen Audio")
 
-with col_header1:
     st.markdown(
         """
-        <h1>🎙 Radio Imaging Generator (Beta)</h1>
-        <p style='font-size:18px;'>
-            Create catchy radio promos, ads, and station jingles with 
-            a modern UI, Llama 3 text generation, and MusicGen audio!
-        </p>
-        """,
-        unsafe_allow_html=True
+        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.
+        """
     )
-with col_header2:
-    if lottie_music:
+with col2:
+    if lottie_animation:
         with st.container():
-            st_lottie(lottie_music, height=180, key="header_lottie")
+            st_lottie(lottie_animation, height=180, loop=True, key="radio_lottie")
     else:
-        # Fallback if Lottie fails to load
-        st.markdown("*(Animation unavailable)*")
+        st.write("*No animation loaded.*")
 
 st.markdown("---")
 
 # ---------------------------------------------------------------------
-# 5) Explanation in an Expander
+# 6) PROMPT INPUT & MODEL SELECTION
 # ---------------------------------------------------------------------
-with st.expander("📘 How to Use This App"):
-    st.markdown(
-        """
-        **Steps**:
-        1. **Model & Language**: In the sidebar, choose the Llama model ID (e.g. a real Llama 2) and the device.
-        2. **Enter Concept**: Provide a short description of the ad or jingle you want.
-        3. **Refine**: Click on "Refine with Llama 3" to get a polished script in your chosen language or style.
-        4. **Generate Audio**: Use MusicGen to create a short audio snippet from that refined script.
-        5. **Listen & Download**: Enjoy or download the result as a WAV file.
+st.subheader("🎙 Step 1: Briefly Describe Your Promo Idea")
 
-        **Note**: 
-        - If "Llama 3.3" doesn't exist, you'll get errors. Use a real model from [Hugging Face](https://huggingface.co/models) 
-          like `meta-llama/Llama-2-7b-chat-hf`. 
-        - Some large models require GPU (or specialized hardware) for feasible speeds.
-        - This example uses [streamlit-lottie](https://github.com/andfanilo/streamlit-lottie) for animation.
-        """
-    )
+prompt = st.text_area(
+    "E.g. 'A 15-second upbeat jingle with a catchy hook for a Top 40 morning show'",
+    height=120
+)
 
-# ---------------------------------------------------------------------
-# 6) Sidebar Configuration
-# ---------------------------------------------------------------------
-with st.sidebar:
-    st.header("🔧 Llama 3 & Audio Settings")
-    
-    # Model input
+col_model, col_device = st.columns(2)
+with col_model:
     llama_model_id = st.text_input(
-        "Llama Model ID",
-        value="meta-llama/Llama-3.3-70B-Instruct",  # Fictitious, please replace with a real model
-        help="Replace with a real model, e.g. meta-llama/Llama-2-7b-chat-hf"
+        "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)."
     )
-    
+with col_device:
     device_option = st.selectbox(
-        "Hardware Device",
+        "Choose Device",
         ["auto", "cpu"],
-        index=0,
-        help="If local GPU is available, choose 'auto'. CPU might be slow for large models."
+        help="For GPU usage, pick 'auto'. CPU can be slow for big models."
     )
-    
-    # Multi-language or style
-    language_choice = st.selectbox(
-        "Choose Language",
-        ["English", "Spanish", "French", "German", "Other (describe in prompt)"]
-    )
-    
-    # Music style & max tokens
-    music_style = st.selectbox(
-        "Preferred Music Style",
-        ["Pop", "Rock", "Electronic", "Classical", "Hip-Hop", "Reggae", "Ambient", "Other"]
-    )
-    audio_tokens = st.slider("MusicGen Max Tokens (Track Length)", 128, 1024, 512, 64)
 
 # ---------------------------------------------------------------------
-# 7) Prompt for the Radio Imaging Concept
+# 7) BUTTON: GENERATE RADIO SCRIPT WITH LLAMA
 # ---------------------------------------------------------------------
-st.markdown("## ✍️ Your Radio Concept")
-prompt = st.text_area(
-    "Describe the theme, audience, length, energy level, etc.",
-    placeholder="E.g. 'A high-energy 10-second pop jingle for a morning radio show...'"
-)
+if st.button("📝 Generate Promo Script"):
+    if not prompt.strip():
+        st.error("Please enter a radio imaging concept first.")
+    else:
+        with st.spinner("Generating script..."):
+            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
+                st.success("Promo script generated!")
+                st.write(refined_text)
+            except Exception as e:
+                st.error(f"Error during Llama generation: {e}")
+
+st.markdown("---")
 
 # ---------------------------------------------------------------------
-# 8) Load Llama Pipeline
+# 8) AUDIO GENERATION: MUSICGEN
+# ---------------------------------------------------------------------
+st.subheader("🎶 Step 2: Generate Your Radio Audio")
+
+audio_tokens = st.slider("MusicGen Max Tokens (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.")
+    else:
+        with st.spinner("Generating audio..."):
+            try:
+                # Load MusicGen
+                mg_model, mg_processor = load_musicgen_model()
+                descriptive_text = st.session_state["refined_script"]
+                
+                # Prepare model input
+                inputs = mg_processor(
+                    text=[descriptive_text],
+                    return_tensors="pt",
+                    padding=True
+                )
+                # Generate audio
+                audio_values = mg_model.generate(**inputs, max_new_tokens=audio_tokens)
+                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):
     """
-    Loads the specified Llama or other HF model as a text-generation pipeline.
-    This references a hypothetical Llama 3.3.
+    Load the Llama model & pipeline.
     """
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     model = AutoModelForCausalLM.from_pretrained(
@@ -211,127 +253,57 @@ def load_llama_pipeline(model_id: str, device: str):
         torch_dtype=torch.float16 if device == "auto" else torch.float32,
         device_map=device
     )
-    pipe = pipeline(
+    text_gen_pipeline = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
         device_map=device
     )
-    return pipe
+    return text_gen_pipeline
 
-def refine_description_with_llama(user_prompt: str, pipeline_llama, lang: str):
+def generate_radio_script(user_input: str, pipeline_llama) -> str:
     """
-    Create a polished script using Llama. 
-    Incorporate a language preference or style instructions.
+    Use Llama to refine the user's input into a brief but creative radio imaging script.
     """
-    system_msg = (
-        "You are an expert radio imaging script writer. "
-        "Refine the user's concept into a concise, compelling piece. "
-        "Ensure to reflect any language or style requests."
+    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."
     )
-    combined = f"{system_msg}\nLanguage: {lang}\nUser Concept: {user_prompt}\nRefined Script:"
+    full_prompt = f"{system_prompt}\nUser concept: {user_input}\nRefined script:"
     
-    result = pipeline_llama(
-        combined,
-        max_new_tokens=300,
+    output = pipeline_llama(
+        full_prompt,
+        max_new_tokens=200,
         do_sample=True,
-        temperature=0.8
-    )
-    text = result[0]["generated_text"]
+        temperature=0.9
+    )[0]["generated_text"]
     
-    # Attempt to isolate the final portion
-    if "Refined Script:" in text:
-        text = text.split("Refined Script:")[-1].strip()
+    # 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)"
     
-    text += "\n\n(Generated with Llama 3 - Modern Radio Generator)"
-    return text
-
-# ---------------------------------------------------------------------
-# 9) Buttons & Outputs
-# ---------------------------------------------------------------------
-col_gen1, col_gen2 = st.columns(2)
+    return output
 
-with col_gen1:
-    if st.button("📄 Refine with Llama 3"):
-        if not prompt.strip():
-            st.error("Please provide a brief concept first.")
-        else:
-            with st.spinner("Refining your script..."):
-                try:
-                    pipeline_llama = load_llama_pipeline(llama_model_id, device_option)
-                    refined_text = refine_description_with_llama(prompt, pipeline_llama, language_choice)
-                    st.session_state['refined_prompt'] = refined_text
-                    st.success("Refined text generated!")
-                    st.write(refined_text)
-                    st.download_button(
-                        "💾 Download Script",
-                        refined_text,
-                        file_name="refined_jingle_script.txt"
-                    )
-                except Exception as e:
-                    st.error(f"Error: {e}")
-
-with col_gen2:
-    if st.button("▶ Generate Audio with MusicGen"):
-        if 'refined_prompt' not in st.session_state or not st.session_state['refined_prompt']:
-            st.error("No refined prompt found. Please generate/refine your script first.")
-        else:
-            final_text_for_music = st.session_state['refined_prompt']
-            final_text_for_music += f"\nPreferred style: {music_style}"
-            with st.spinner("Generating audio..."):
-                try:
-                    mg_model, mg_processor = None, None
-
-                    # Load MusicGen model once
-                    mg_model, mg_processor = load_musicgen_model()
-                    
-                    inputs = mg_processor(
-                        text=[final_text_for_music],
-                        padding=True,
-                        return_tensors="pt"
-                    )
-                    audio_output = mg_model.generate(**inputs, max_new_tokens=audio_tokens)
-                    sr = mg_model.config.audio_encoder.sampling_rate
-
-                    audio_filename = f"radio_imaging_{music_style.lower()}.wav"
-                    scipy.io.wavfile.write(
-                        audio_filename,
-                        rate=sr,
-                        data=audio_output[0, 0].numpy()
-                    )
-                    st.success("Audio generated! Listen below:")
-                    st.audio(audio_filename)
-                    
-                    # Optional Save/Upload prompt
-                    if st.checkbox("Upload this WAV to a cloud (demo)?"):
-                        with st.spinner("Uploading..."):
-                            # Placeholder for your own S3 or cloud logic
-                            st.success("Uploaded (placeholder).")
-                except Exception as e:
-                    st.error(f"Error generating audio: {e}")
-
-# ---------------------------------------------------------------------
-# 10) Load & Cache MusicGen
-# ---------------------------------------------------------------------
 @st.cache_resource
 def load_musicgen_model():
     """
-    Load and cache the MusicGen model & processor.
-    Using 'facebook/musicgen-small' as example.
+    Load MusicGen (small version).
     """
-    mgm = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
-    mgp = AutoProcessor.from_pretrained("facebook/musicgen-small")
-    return mgm, mgp
+    mg_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+    mg_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+    return mg_model, mg_processor
 
 # ---------------------------------------------------------------------
-# 11) Footer
+# 10) FOOTER
 # ---------------------------------------------------------------------
 st.markdown("---")
 st.markdown(
     """
-    <div class='footer-note'>
-        © 2025 Modern Radio Generator - Built with Llama & MusicGen | 
-        <a href='https://example.com' target='_blank'>YourCompany</a>
+    <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.
     </div>
     """,
     unsafe_allow_html=True