code update for duration of ZeroGPU

tts/gradio_api.py

@@ -20,9 +20,7 @@ import gradio as gr
 import traceback
 from tts.infer_cli import MegaTTS3DiTInfer, convert_to_wav, cut_wav
 
-import spaces
 
-@spaces.GPU(duration=120)
 def model_worker(input_queue, output_queue, device_id):
     device = None
     if device_id is not None:
@@ -39,8 +37,9 @@ 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()
-            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 = 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)
             output_queue.put(wav_bytes)
         except Exception as e:
             traceback.print_exc()

tts/infer_cli.py

@@ -252,6 +252,12 @@ class MegaTTS3DiTInfer():
             wav_pred = combine_audio_segments(wav_pred_, sr=self.sr).astype(float)
             return to_wav_bytes(wav_pred, self.sr)
 
+        @spaces.GPU(duration=120)
+        def forward_zerogpu(self, file_content, latent_file, inp_text, time_step, p_w, t_w):
+            resource_context = self.preprocess(file_content, latent_file)
+            wav_bytes = self.forward(resource_context, inp_text, time_step=time_step, p_w=p_w, t_w=t_w)
+            return wav_bytes
+
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()