code update for duration of ZeroGPU

tts/gradio_api.py

@@ -20,14 +20,16 @@ import gradio as gr
 import traceback
 from tts.infer_cli import MegaTTS3DiTInfer, convert_to_wav, cut_wav
 
+CUDA_AVAILABLE = torch.cuda.is_available()
+infer_pipe = MegaTTS3DiTInfer(device='cuda' if CUDA_AVAILABLE else 'cpu')
 
-def model_worker(input_queue, output_queue, device_id):
-    device = None
-    if device_id is not None:
-        device = torch.device(f'cuda:{device_id}')
-    infer_pipe = MegaTTS3DiTInfer(device=device)
-    os.system(f'pkill -f "voidgpu{device_id}"')
+@spaces.GPU(duration=120)
+def forward_gpu(file_content, latent_file, inp_text, time_step, p_w, t_w):
+    resource_context = infer_pipe.preprocess(file_content, latent_file)
+    wav_bytes = infer_pipe.forward(resource_context, inp_text, time_step=time_step, p_w=p_w, t_w=t_w)
+    return wav_bytes
 
+def model_worker(input_queue, output_queue, device_id):
     while True:
         task = input_queue.get()
         inp_audio_path, inp_npy_path, inp_text, infer_timestep, p_w, t_w = task
@@ -37,9 +39,7 @@ def model_worker(input_queue, output_queue, device_id):
             cut_wav(wav_path, max_len=28)
             with open(wav_path, 'rb') as file:
                 file_content = file.read()
-            wav_bytes = infer_pipe.forward_zerogpu(file_content, inp_npy_path, inp_text, time_step=infer_timestep, p_w=p_w, t_w=t_w)
-            # resource_context = infer_pipe.preprocess(file_content, latent_file=inp_npy_path)
-            # wav_bytes = infer_pipe.forward(resource_context, inp_text, time_step=infer_timestep, p_w=p_w, t_w=t_w)
+            wav_bytes = forward_gpu(file_content, inp_npy_path, inp_text, time_step=infer_timestep, p_w=p_w, t_w=t_w)
             output_queue.put(wav_bytes)
         except Exception as e:
             traceback.print_exc()
@@ -63,15 +63,9 @@ if __name__ == '__main__':
 
     mp.set_start_method('spawn', force=True)
     mp_manager = mp.Manager()
-    devices = os.environ.get('CUDA_VISIBLE_DEVICES', '')
-    if devices != '':
-        devices = os.environ.get('CUDA_VISIBLE_DEVICES', '').split(",")
-        for d in devices:
-            os.system(f'pkill -f "voidgpu{d}"')
-    else:
-        devices = None
     
     num_workers = 1
+    devices = [0]
     input_queue = mp_manager.Queue()
     output_queue = mp_manager.Queue()
     processes = []