Spaces:
Sleeping
Sleeping
Update s2s.py
Browse files
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 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
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(
|