Update app.py

app.py

@@ -3,14 +3,38 @@ from parler_tts import ParlerTTSForConditionalGeneration
 from transformers import AutoTokenizer
 import gradio as gr
 import numpy as np
+import time
+import requests
 
 # Set device to GPU if available, else CPU
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# Load the TTS model and tokenizers
-model = ParlerTTSForConditionalGeneration.from_pretrained("ai4bharat/indic-parler-tts").to(device)
-tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indic-parler-tts")
-description_tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder._name_or_path)
+# Function to load the model with retry logic and increased timeout
+def load_model_with_retry(model_id, retries=3, timeout=60):
+    for attempt in range(1, retries + 1):
+        try:
+            model = ParlerTTSForConditionalGeneration.from_pretrained(model_id, timeout=timeout).to(device)
+            return model
+        except requests.exceptions.ReadTimeout:
+            print(f"Timeout when loading model. Attempt {attempt} of {retries}. Retrying in 5 seconds...")
+            time.sleep(5)
+    raise Exception("Failed to load model after multiple retries.")
+
+# Function to load a tokenizer with retry logic and increased timeout
+def load_tokenizer_with_retry(tokenizer_id, retries=3, timeout=60):
+    for attempt in range(1, retries + 1):
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, timeout=timeout)
+            return tokenizer
+        except requests.exceptions.ReadTimeout:
+            print(f"Timeout when loading tokenizer. Attempt {attempt} of {retries}. Retrying in 5 seconds...")
+            time.sleep(5)
+    raise Exception("Failed to load tokenizer after multiple retries.")
+
+# Load the TTS model and tokenizers using the retry functions
+model = load_model_with_retry("ai4bharat/indic-parler-tts")
+tokenizer = load_tokenizer_with_retry("ai4bharat/indic-parler-tts")
+description_tokenizer = load_tokenizer_with_retry(model.config.text_encoder._name_or_path)
 
 def generate_audio(text: str):
     """
@@ -23,8 +47,10 @@ def generate_audio(text: str):
         tuple: A tuple containing the audio numpy array and the sampling rate.
     """
     # Set a default voice description
-    default_description = ("Divya's voice is monotone yet slightly fast in delivery, with a very close recording "
-                           "that almost has no background noise.")
+    default_description = (
+        "Divya's voice is monotone yet slightly fast in delivery, with a very close recording "
+        "that almost has no background noise."
+    )
     
     # Tokenize the default description and the input text
     description_tokens = description_tokenizer(default_description, return_tensors="pt").to(device)