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
@@ -17,81 +17,44 @@ import spaces  # Assumes Hugging Face Spaces library supports `@spaces.GPU`
 load_dotenv()
 hf_token = os.getenv("HF_TOKEN")
 
-# Globals for lazy loading
-llama_pipeline = None
-musicgen_model = None
-musicgen_processor = None
 
 # ---------------------------------------------------------------------
-# Load Llama 3 Model with Zero GPU (Lazy Loading) - Smaller Model
+# Load Llama 3 Pipeline with Zero GPU (Encapsulated)
 # ---------------------------------------------------------------------
 @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
-            )
-            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
-
-# ---------------------------------------------------------------------
-# Generate Radio Script
-# ---------------------------------------------------------------------
-def generate_script(user_input: str, pipeline_llama):
+def generate_script(user_prompt: str, model_id: str, token: str):
     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",
+            trust_remote_code=True,
+        )
+        llama_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
         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)
-        return result[0]['generated_text'].split("Refined script:")[-1].strip()
+        combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nRefined script:"
+        result = llama_pipeline(combined_prompt, max_new_tokens=200, do_sample=True, temperature=0.9)
+        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=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
+# Load MusicGen Model (Encapsulated)
 # ---------------------------------------------------------------------
 @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:
-        musicgen_model, musicgen_processor = load_musicgen_model()
-        if isinstance(musicgen_model, str):
-            return musicgen_model
     try:
-        musicgen_model.to("cuda")  # Move the model to GPU
+        musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+        musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+
+        musicgen_model.to("cuda")
         inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt")
         outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
         musicgen_model.to("cpu")  # Return the model to CPU
@@ -106,21 +69,17 @@ def generate_audio(prompt: str, audio_length: int):
     except Exception as e:
         return f"Error generating audio: {e}"
 
+
 # ---------------------------------------------------------------------
 # Gradio Interface
 # ---------------------------------------------------------------------
-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, None
+def interface_generate_script(user_prompt, llama_model_id):
+    return generate_script(user_prompt, llama_model_id, hf_token)
 
-    # Generate Script
-    script = generate_script(user_prompt, pipeline_llama)
 
-    # Generate Audio
-    audio_data = generate_audio(script, audio_length)
-    return script, audio_data
+def interface_generate_audio(script, audio_length):
+    return generate_audio(script, audio_length)
+
 
 # ---------------------------------------------------------------------
 # Interface
@@ -129,8 +88,13 @@ 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-Instruct")  # Smaller Model
+        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"
+        )
         audio_length = gr.Slider(label="Audio Length (tokens)", minimum=128, maximum=1024, step=64, value=512)
 
     generate_script_button = gr.Button("Generate Script")
@@ -139,15 +103,15 @@ with gr.Blocks() as demo:
     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],
+        fn=interface_generate_script,
         inputs=[user_prompt, llama_model_id],
-        outputs=script_output
+        outputs=script_output,
     )
 
     generate_audio_button.click(
-        fn=lambda script_output, audio_length: generate_audio(script_output, audio_length),
+        fn=interface_generate_audio,
         inputs=[script_output, audio_length],
-        outputs=audio_output
+        outputs=audio_output,
     )
 
 # ---------------------------------------------------------------------