Update app.py

app.py

@@ -6,7 +6,7 @@ from transformers import (
     AutoModelForCausalLM,
     pipeline,
     AutoProcessor,
-    MusicgenForConditionalGeneration,
+    MusicgenForConditionalGeneration
 )
 from scipy.io.wavfile import write
 import tempfile
@@ -23,23 +23,28 @@ musicgen_model = None
 musicgen_processor = None
 
 # ---------------------------------------------------------------------
-# Load Llama 3 Model with Zero GPU (Lazy Loading)
+# Load Llama 3 Model with Zero GPU (Lazy Loading) - Smaller Model
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=300)  # Adjust GPU allocation duration
 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 = AutoModelForCausalLM.from_pretrained(
                 model_id,
                 use_auth_token=token,
                 torch_dtype=torch.float16,
                 device_map="auto",  # Automatically handles GPU allocation
-                trust_remote_code=True,
+                trust_remote_code=True
             )
+            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
 
@@ -54,7 +59,7 @@ 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}"
 
@@ -66,9 +71,12 @@ 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
 
@@ -101,19 +109,18 @@ def generate_audio(prompt: str, audio_length: int):
 # ---------------------------------------------------------------------
 # Gradio Interface
 # ---------------------------------------------------------------------
-def radio_imaging_app(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
-    return generate_script(user_prompt, pipeline_llama)
-
-
-def generate_audio_from_script(script, audio_length):
-    return generate_audio(script, audio_length)
+    script = generate_script(user_prompt, pipeline_llama)
 
+    # Generate Audio
+    audio_data = generate_audio(script, audio_length)
+    return script, audio_data
 
 # ---------------------------------------------------------------------
 # Interface
@@ -121,42 +128,27 @@ def generate_audio_from_script(script, audio_length):
 with gr.Blocks() as demo:
     gr.Markdown("# 🎧 AI Radio Imaging with Llama 3 + MusicGen (Zero GPU)")
 
-    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],
-            outputs=script_output,
-        )
-
-    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_from_script,
-            inputs=[script_output, audio_length],
-            outputs=audio_output,
-        )
+    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")  # Smaller Model
+        audio_length = gr.Slider(label="Audio Length (tokens)", minimum=128, maximum=1024, step=64, value=512)
+
+    generate_script_button = gr.Button("Generate Script")
+    generate_audio_button = gr.Button("Generate Audio")
+    script_output = gr.Textbox(label="Generated Script")
+    audio_output = gr.Audio(label="Generated Audio", type="filepath")
+
+    generate_script_button.click(
+        fn=lambda user_prompt, llama_model_id: radio_imaging_app(user_prompt, llama_model_id, None)[0],
+        inputs=[user_prompt, llama_model_id],
+        outputs=script_output
+    )
+
+    generate_audio_button.click(
+        fn=lambda script_output, audio_length: generate_audio(script_output, audio_length),
+        inputs=[script_output, audio_length],
+        outputs=audio_output
+    )
 
 # ---------------------------------------------------------------------
 # Launch App