.

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import os
 import torch
 import json
+import time
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Check if CUDA is available and set the device accordingly
@@ -37,31 +38,43 @@ def create_lyrics_prompt(classification_results):
     """
     return prompt
 
-def generate_lyrics(prompt):
-    """Generate lyrics using the Janus model"""
-    try:
-        response = requests.post(
-            JANUS_API_URL,
-            headers=headers,
-            json={
-                "inputs": prompt,
-                "parameters": {
-                    "max_new_tokens": 200,
-                    "temperature": 0.7,
-                    "top_p": 0.9,
-                    "return_full_text": False
+def generate_lyrics_with_retry(prompt, max_retries=5, initial_wait=2):
+    """Generate lyrics using the Janus model with retry logic"""
+    wait_time = initial_wait
+    
+    for attempt in range(max_retries):
+        try:
+            response = requests.post(
+                JANUS_API_URL,
+                headers=headers,
+                json={
+                    "inputs": prompt,
+                    "parameters": {
+                        "max_new_tokens": 200,
+                        "temperature": 0.7,
+                        "top_p": 0.9,
+                        "return_full_text": False
+                    }
                 }
-            }
-        )
-        
-        if response.status_code == 200:
-            return response.json()[0]["generated_text"]
-        elif response.status_code == 503:
-            return "Model is loading. Please try again in a few seconds."
-        else:
-            return f"Error generating lyrics: {response.text}"
-    except Exception as e:
-        return f"Error: {str(e)}"
+            )
+            
+            if response.status_code == 200:
+                return response.json()[0]["generated_text"]
+            elif response.status_code == 503:
+                print(f"Model loading, attempt {attempt + 1}/{max_retries}. Waiting {wait_time} seconds...")
+                time.sleep(wait_time)
+                wait_time *= 1.5  # Increase wait time for next attempt
+                continue
+            else:
+                return f"Error generating lyrics: {response.text}"
+                
+        except Exception as e:
+            if attempt == max_retries - 1:  # Last attempt
+                return f"Error after {max_retries} attempts: {str(e)}"
+            time.sleep(wait_time)
+            wait_time *= 1.5
+    
+    return "Failed to generate lyrics after multiple attempts. Please try again."
 
 def format_results(classification_results, lyrics, prompt):
     """Format the results for display"""
@@ -107,10 +120,10 @@ def classify_and_generate(audio_file):
                     'score': f"{result['score']*100:.2f}%"
                 })
             
-            # Generate lyrics based on classification
+            # Generate lyrics based on classification with retry logic
             print("Generating lyrics based on classification...")
             prompt = create_lyrics_prompt(formatted_results)
-            lyrics = generate_lyrics(prompt)
+            lyrics = generate_lyrics_with_retry(prompt)
             
             # Format and return results
             return format_results(formatted_results, lyrics, prompt)