Optimize CUDA device handling - auto-detect GPU and remove unnecessary device transfers

api.py

@@ -25,7 +25,11 @@ except ImportError:
 # --- Model Loading ---
 
 MODEL_ID = "ai4bharat/indic-parler-tts"
-DEVICE = "cpu"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(f"Using device: {DEVICE}")
+if DEVICE == "cuda":
+    print(f"GPU: {torch.cuda.get_device_name(0)}")
 
 print("Loading Indic Parler-TTS model...")
 model = ParlerTTSForConditionalGeneration.from_pretrained(MODEL_ID).to(DEVICE)
@@ -101,7 +105,6 @@ def generate_speech_internal(text, speaker, pitch, rate, temperature, do_sample)
         speaker_name = SPEAKERS.get(speaker, "")
         gender = "female" if "Female" in speaker or speaker in ["Divya", "Rani"] else "male"
         description = build_description(speaker_name, gender, pitch, rate)
-        model.to("cuda")
 
         sentences = split_sentences(text)
         if not sentences:
@@ -120,10 +123,10 @@ def generate_speech_internal(text, speaker, pitch, rate, temperature, do_sample)
 
             with torch.no_grad():
                 generation = model.generate(
-                    input_ids=desc_tokens.input_ids.to("cuda"),
-                    attention_mask=desc_tokens.attention_mask.to("cuda"),
-                    prompt_input_ids=prompt_tokens.input_ids.to("cuda"),
-                    prompt_attention_mask=prompt_tokens.attention_mask.to("cuda"),
+                    input_ids=desc_tokens.input_ids.to(DEVICE),
+                    attention_mask=desc_tokens.attention_mask.to(DEVICE),
+                    prompt_input_ids=prompt_tokens.input_ids.to(DEVICE),
+                    prompt_attention_mask=prompt_tokens.attention_mask.to(DEVICE),
                     do_sample=do_sample,
                     temperature=temperature if do_sample else 1.0,
                     min_new_tokens=10,
@@ -136,8 +139,6 @@ def generate_speech_internal(text, speaker, pitch, rate, temperature, do_sample)
             silence = np.zeros(int(SAMPLE_RATE * 0.3), dtype=np.int16)
             all_audio.append(silence)
 
-        model.to("cpu")
-
         if not all_audio:
             return None
 
@@ -146,7 +147,6 @@ def generate_speech_internal(text, speaker, pitch, rate, temperature, do_sample)
 
     except Exception as e:
         print(f"Error generating speech: {e}")
-        model.to("cpu")
         return None