Update app.py

app.py

@@ -14,7 +14,7 @@ from tqdm import tqdm
 from einops import rearrange
 import numpy as np
 import json
-
+import argparse
 import torch
 import torchaudio
 from torchaudio.transforms import Resample
@@ -83,6 +83,8 @@ MAX_NEW_TOKENS = 15         # Duration slider (in seconds, scaled internally)
 RUN_N_SEGMENTS = 2          # Number of segments to generate
 STAGE2_BATCH_SIZE = 4       # Batch size for stage2 inference
 
+default_args = argparse.Namespace(cuda_idx=4)
+
 # You may change these defaults via Gradio input (see below)
 
 # ----------------------- Device Setup -----------------------
@@ -493,7 +495,7 @@ def generate_music(
                 stage2_result[0],
                 os.path.join(vocoder_stems_dir, "vocal.mp3"),
                 rescale,
-                None,
+                default_args,
                 vocal_decoder,
                 codec_model,
             )
@@ -501,7 +503,7 @@ def generate_music(
                 stage2_result[1],
                 os.path.join(vocoder_stems_dir, "instrumental.mp3"),
                 rescale,
-                None,
+                default_args,
                 inst_decoder,
                 codec_model,
             )