xcczach commited on
Commit
f0abd45
1 Parent(s): c3f10f2

Update s2s.py

Browse files
Files changed (1) hide show
  1. s2s.py +15 -8
s2s.py CHANGED
@@ -141,9 +141,16 @@ def generate_from_wav(
141
  return audio_hat, output_text
142
 
143
 
 
 
 
 
 
144
  def generate(
145
  wav_path: str, progress_callback: Callable[[str], None] | None = None
146
  ) -> tuple[np.ndarray, int | float]:
 
 
147
  config = OmegaConf.structured(InferenceConfig())
148
  train_config, model_config, dataset_config, decode_config = (
149
  config.train_config,
@@ -156,14 +163,14 @@ def generate(
156
  torch.manual_seed(train_config.seed)
157
  random.seed(train_config.seed)
158
 
159
- update_progress(progress_callback, "Loading model")
160
-
161
- model_factory = get_custom_model_factory(model_config)
162
- model, _ = model_factory(train_config, model_config, CKPT_PATH)
163
- codec_decoder = model.codec_decoder
164
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
165
- model.to(device)
166
- model.eval()
167
 
168
  update_progress(progress_callback, "Generating")
169
  output_wav, output_text = generate_from_wav(
 
141
  return audio_hat, output_text
142
 
143
 
144
+ model = None
145
+ codec_decoder = None
146
+ device = None
147
+
148
+
149
  def generate(
150
  wav_path: str, progress_callback: Callable[[str], None] | None = None
151
  ) -> tuple[np.ndarray, int | float]:
152
+ global model, codec_decoder, device
153
+
154
  config = OmegaConf.structured(InferenceConfig())
155
  train_config, model_config, dataset_config, decode_config = (
156
  config.train_config,
 
163
  torch.manual_seed(train_config.seed)
164
  random.seed(train_config.seed)
165
 
166
+ if model is None or codec_decoder is None or device is None:
167
+ update_progress(progress_callback, "Loading model")
168
+ model_factory = get_custom_model_factory(model_config)
169
+ model, _ = model_factory(train_config, model_config, CKPT_PATH)
170
+ codec_decoder = model.codec_decoder
171
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
172
+ model.to(device)
173
+ model.eval()
174
 
175
  update_progress(progress_callback, "Generating")
176
  output_wav, output_text = generate_from_wav(