Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import os
 import torch
-import time
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
@@ -12,7 +11,7 @@ from transformers import (
 from scipy.io.wavfile import write
 import tempfile
 from dotenv import load_dotenv
-import spaces  # Hugging Face Spaces library for ZeroGPU support
+import spaces  # Assumes Hugging Face Spaces library supports `@spaces.GPU`
 
 # Load environment variables (e.g., Hugging Face token)
 load_dotenv()
@@ -24,43 +23,23 @@ musicgen_model = None
 musicgen_processor = None
 
 # ---------------------------------------------------------------------
-# Helper: Safe Model Loader with Retry Logic
+# Load Llama 3 Model with Zero GPU (Lazy Loading)
 # ---------------------------------------------------------------------
-def safe_load_model(model_id, token, retries=3, delay=5):
-    for attempt in range(retries):
+@spaces.GPU(duration=300)  # Increased duration to 300 seconds
+def load_llama_pipeline_zero_gpu(model_id: str, token: str):
+    global llama_pipeline
+    if llama_pipeline is None:
         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",
+                device_map="auto",  # Automatically handles GPU allocation
                 trust_remote_code=True,
-                offload_folder="/tmp",  # Stream shards
-                cache_dir="/tmp"        # Cache directory for shard downloads
             )
-            return model
-        except Exception as e:
-            print(f"Attempt {attempt + 1} failed: {e}")
-            time.sleep(delay)
-    raise RuntimeError(f"Failed to load model {model_id} after {retries} attempts")
-
-# ---------------------------------------------------------------------
-# Load Llama 3 Model with Zero GPU (Lazy Loading)
-# ---------------------------------------------------------------------
-@spaces.GPU(duration=600)  # Increased duration to handle large models
-def load_llama_pipeline_zero_gpu(model_id: str, token: str):
-    global llama_pipeline
-    if llama_pipeline is None:
-        try:
-            print("Starting model loading...")
-            tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
-            print("Tokenizer loaded.")
-            model = safe_load_model(model_id, token)
-            print("Model loaded. Initializing pipeline...")
             llama_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
-            print("Pipeline initialized successfully.")
         except Exception as e:
-            print(f"Error loading Llama pipeline: {e}")
             return str(e)
     return llama_pipeline
 
@@ -75,31 +54,28 @@ def generate_script(user_input: str, pipeline_llama):
         )
         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)
-        return result[0]['generated_text'].split("Refined script:")[-1].strip()
+        return result[0]["generated_text"].split("Refined script:")[-1].strip()
     except Exception as e:
         return f"Error generating script: {e}"
 
 # ---------------------------------------------------------------------
 # Load MusicGen Model (Lazy Loading)
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=600)
+@spaces.GPU(duration=300)
 def load_musicgen_model():
     global musicgen_model, musicgen_processor
     if musicgen_model is None or musicgen_processor is None:
         try:
-            print("Loading MusicGen model...")
             musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
             musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
-            print("MusicGen model loaded successfully.")
         except Exception as e:
-            print(f"Error loading MusicGen model: {e}")
             return None, str(e)
     return musicgen_model, musicgen_processor
 
 # ---------------------------------------------------------------------
 # Generate Audio
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=600)
+@spaces.GPU(duration=300)
 def generate_audio(prompt: str, audio_length: int):
     global musicgen_model, musicgen_processor
     if musicgen_model is None or musicgen_processor is None:
@@ -125,51 +101,75 @@ def generate_audio(prompt: str, audio_length: int):
 # ---------------------------------------------------------------------
 # Gradio Interface
 # ---------------------------------------------------------------------
-def generate_script_interface(user_prompt, llama_model_id):
+def radio_imaging_app(user_prompt, llama_model_id, audio_length):
     # Load Llama 3 Pipeline with Zero GPU
     pipeline_llama = load_llama_pipeline_zero_gpu(llama_model_id, hf_token)
     if isinstance(pipeline_llama, str):
-        return pipeline_llama
+        return pipeline_llama, None
 
     # Generate Script
     script = generate_script(user_prompt, pipeline_llama)
-    return script
 
-def generate_audio_interface(script, audio_length):
     # Generate Audio
     audio_data = generate_audio(script, audio_length)
-    return audio_data
+    return script, audio_data
+
 
 # ---------------------------------------------------------------------
 # 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")  # Using a smaller model for better compatibility
-        audio_length = gr.Slider(label="Audio Length (tokens)", minimum=128, maximum=1024, step=64, value=512)
-
-    with gr.Row():
-        generate_script_button = gr.Button("Generate Promo Script")
-        script_output = gr.Textbox(label="Generated Script", interactive=False)
-
-    with gr.Row():
-        generate_audio_button = gr.Button("Generate Audio")
-        audio_output = gr.Audio(label="Generated Audio", type="filepath")
-
-    generate_script_button.click(
-        generate_script_interface, 
-        inputs=[user_prompt, llama_model_id], 
-        outputs=script_output
-    )
-
-    generate_audio_button.click(
-        generate_audio_interface, 
-        inputs=[script_output, audio_length], 
-        outputs=audio_output
-    )
+with gr.Blocks(css="""
+#app-title {
+    text-align: center;
+    font-size: 2rem;
+    font-weight: bold;
+    color: #4CAF50;
+}
+#subsection {
+    margin: 20px 0;
+    font-size: 1.2rem;
+    color: #333;
+    text-align: center;
+}
+""") as demo:
+    gr.Markdown('<div id="app-title">🎧 AI Radio Imaging with Llama 3 + MusicGen (Zero GPU)</div>')
+
+    with gr.Tab("Step 1: Generate Promo Script"):
+        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-70B"
+            )
+
+        generate_script_button = gr.Button("Generate Script")
+        script_output = gr.Textbox(label="Generated Promo Script", interactive=False)
+
+        generate_script_button.click(
+            fn=radio_imaging_app,
+            inputs=[user_prompt, llama_model_id, gr.State(0)],
+            outputs=[script_output, None],
+        )
+
+    with gr.Tab("Step 2: Generate Audio"):
+        with gr.Row():
+            audio_length = gr.Slider(
+                label="Audio Length (tokens)",
+                minimum=128,
+                maximum=1024,
+                step=64,
+                value=512,
+            )
+            generate_audio_button = gr.Button("Generate Audio")
+            audio_output = gr.Audio(label="Generated Audio", type="filepath")
+
+        generate_audio_button.click(
+            fn=generate_audio,
+            inputs=[script_output, audio_length],
+            outputs=audio_output,
+        )
 
 # ---------------------------------------------------------------------
 # Launch App