Update app.py

app.py

@@ -11,19 +11,19 @@ from transformers import (
 import scipy.io.wavfile as wav
 
 # ---------------------------------------------------------------------
-# Load Llama 3 Model
+# Load Llama 3 Model with Zero GPU
 # ---------------------------------------------------------------------
-def load_llama_pipeline(model_id: str, token: str, device: str = "cpu"):
+def load_llama_pipeline_zero_gpu(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 if device == "cuda" else torch.float32,
-            device_map="auto" if device == "cuda" else None,
-            low_cpu_mem_usage=True
+            torch_dtype=torch.float16,
+            device_map="auto",  # Use device map to offload computations
+            trust_remote_code=True  # Enables execution of remote code for Zero GPU
         )
-        return pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if device == "cuda" else -1)
+        return pipeline("text-generation", model=model, tokenizer=tokenizer)
     except Exception as e:
         return str(e)
 
@@ -73,8 +73,8 @@ def generate_audio(prompt: str, audio_length: int, mg_model, mg_processor):
 # Gradio Interface
 # ---------------------------------------------------------------------
 def radio_imaging_app(user_prompt, llama_model_id, hf_token, audio_length):
-    # Load Llama 3 Pipeline
-    pipeline_llama = load_llama_pipeline(llama_model_id, hf_token, device="cuda" if torch.cuda.is_available() else "cpu")
+    # 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
 
@@ -97,7 +97,7 @@ def radio_imaging_app(user_prompt, llama_model_id, hf_token, audio_length):
 # Interface
 # ---------------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎧 AI Radio Imaging with Llama 3 + MusicGen")
+    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, fun and energetic.")
         llama_model_id = gr.Textbox(label="Llama 3 Model ID", value="meta-llama/Meta-Llama-3-70B")