Update app.py

app.py

@@ -115,12 +115,12 @@ def run_tts(
     return save_path
 
 
-def build_ui(model_dir):
+def build_ui(model_dir, device=0):
     
     global MODEL
     
     # Initialize model with proper device handling
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cuda" if torch.cuda.is_available() and device != "cpu" else "cpu"
     if MODEL is None:
         MODEL = initialize_model(model_dir, device=device)
         if device == "cuda":
@@ -164,7 +164,7 @@ def build_ui(model_dir):
 
     with gr.Blocks() as demo:
         # Use HTML for centered title
-        gr.HTML('<h1 style="text-align: center;">(Official) Spark-TTS by SparkAudio</h1>')
+        gr.HTML('<h1 style="text-align: center;">(Unofficial) Spark-TTS by SparkAudio</h1>')
         with gr.Tabs():
             # Voice Clone Tab
             with gr.TabItem("Voice Clone"):
@@ -260,7 +260,12 @@ def parse_arguments():
         default=None,
         help="Path to the model directory."
     )
-
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help="Device to use (e.g., 'cpu' or 'cuda:0')."
+    )
     parser.add_argument(
         "--server_name",
         type=str,
@@ -279,8 +284,11 @@ if __name__ == "__main__":
     # Parse command-line arguments
     args = parse_arguments()
 
-    # Build the Gradio demo by specifying the model directory
-    demo = build_ui(model_dir=args.model_dir)
+    # Build the Gradio demo by specifying the model directory and GPU device
+    demo = build_ui(
+        model_dir=args.model_dir,
+        device=args.device
+    )
 
     # Launch Gradio with the specified server name and port
     demo.launch(