Update app.py

app.py

@@ -1,17 +1,18 @@
 import gradio as gr
 import os
 import torch
+import time
 from transformers import (
-    AutoTokenizer, 
-    AutoModelForCausalLM, 
+    AutoTokenizer,
+    AutoModelForCausalLM,
     pipeline,
-    AutoProcessor, 
-    MusicgenForConditionalGeneration
+    AutoProcessor,
+    MusicgenForConditionalGeneration,
 )
 from scipy.io.wavfile import write
 import tempfile
 from dotenv import load_dotenv
-import spaces  # Assumes Hugging Face Spaces library supports `@spaces.GPU`
+import spaces  # Hugging Face Spaces library for ZeroGPU support
 
 # Load environment variables (e.g., Hugging Face token)
 load_dotenv()
@@ -22,10 +23,31 @@ llama_pipeline = None
 musicgen_model = None
 musicgen_processor = None
 
+# ---------------------------------------------------------------------
+# Helper: Safe Model Loader with Retry Logic
+# ---------------------------------------------------------------------
+def safe_load_model(model_id, token, retries=3, delay=5):
+    for attempt in range(retries):
+        try:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                use_auth_token=token,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                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=300)  # Increased duration to 300 seconds
+@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:
@@ -33,13 +55,7 @@ def load_llama_pipeline_zero_gpu(model_id: str, token: str):
             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
-            )
+            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.")
@@ -66,7 +82,7 @@ def generate_script(user_input: str, pipeline_llama):
 # ---------------------------------------------------------------------
 # Load MusicGen Model (Lazy Loading)
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=600)
 def load_musicgen_model():
     global musicgen_model, musicgen_processor
     if musicgen_model is None or musicgen_processor is None:
@@ -83,7 +99,7 @@ def load_musicgen_model():
 # ---------------------------------------------------------------------
 # Generate Audio
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=600)
 def generate_audio(prompt: str, audio_length: int):
     global musicgen_model, musicgen_processor
     if musicgen_model is None or musicgen_processor is None:
@@ -132,7 +148,7 @@ with gr.Blocks() as demo:
 
     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")
+        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():