Update app.py

app.py

@@ -4,10 +4,11 @@ from parler_tts import ParlerTTSForConditionalGeneration
 from transformers import AutoTokenizer
 import soundfile as sf
 import tempfile
+import spaces  # Import the spaces module for ZeroGPU compatibility
 
 # Load the model and tokenizers
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-model = ParlerTTSForConditionalGeneration.from_pretrained("ai4bharat/indic-parler-tts").to(device)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = ParlerTTSForConditionalGeneration.from_pretrained("ai4bharat/indic-parler-tts")
 tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indic-parler-tts")
 description_tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder._name_or_path)
 
@@ -35,11 +36,15 @@ def generate_description(language, gender, emotion, noise, reverb, expressivity,
     )
     return description
 
-# Generate audio function
+# Generate audio function with GPU allocation
+@spaces.GPU  # Allocate GPU for the duration of this function
 def generate_audio(text, description):
+    # Move model to GPU
+    model.to("cuda")
+
     # Prepare model inputs
-    input_ids = description_tokenizer(description, return_tensors="pt").input_ids.to(device)
-    prompt_input_ids = tokenizer(text, return_tensors="pt").input_ids.to(device)
+    input_ids = description_tokenizer(description, return_tensors="pt").input_ids.to("cuda")
+    prompt_input_ids = tokenizer(text, return_tensors="pt").input_ids.to("cuda")
 
     # Generate audio
     generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
@@ -50,6 +55,9 @@ def generate_audio(text, description):
         sf.write(f.name, audio_arr, model.config.sampling_rate)
         audio_path = f.name
 
+    # Move model back to CPU
+    model.to("cpu")
+
     return audio_path
 
 # Gradio Interface
@@ -146,4 +154,4 @@ def app():
     return demo
 
 # Run the app
-app().launch()
+app().launch()