Export preprocessor, pre_encode, and encode methods. Args: model: The SortformerEncLabelModel to export. backend: Target backend ("xnnpack" or "portable").
(model, backend: Optional[str] = None)
| 191 | |
| 192 | |
| 193 | def export_all(model, backend: Optional[str] = None): |
| 194 | """Export preprocessor, pre_encode, and encode methods. |
| 195 | |
| 196 | Args: |
| 197 | model: The SortformerEncLabelModel to export. |
| 198 | backend: Target backend ("xnnpack" or "portable"). |
| 199 | """ |
| 200 | programs = {} |
| 201 | |
| 202 | sample_rate = model.preprocessor._cfg.sample_rate |
| 203 | window_stride = float(model.preprocessor._cfg.window_stride) |
| 204 | subsampling_factor = int(model.encoder.subsampling_factor) |
| 205 | |
| 206 | prepare_for_export(model) |
| 207 | |
| 208 | # --- Method 1: preprocessor --- |
| 209 | preprocessor_wrapper = PreprocessorWrapper(model.preprocessor) |
| 210 | preprocessor_wrapper.eval() |
| 211 | |
| 212 | max_audio_samples = int(sample_rate * 120) # 120 seconds max |
| 213 | sample_audio = torch.randn(max_audio_samples, dtype=torch.float) |
| 214 | sample_length = torch.tensor([sample_audio.shape[0]], dtype=torch.int64) |
| 215 | |
| 216 | # Force CPU path to avoid data-dependent CUDA conditionals in preprocessor |
| 217 | old_cuda_is_available = torch.cuda.is_available |
| 218 | torch.cuda.is_available = lambda: False |
| 219 | |
| 220 | print(" Exporting preprocessor...") |
| 221 | programs["preprocessor"] = export( |
| 222 | preprocessor_wrapper, |
| 223 | (sample_audio, sample_length), |
| 224 | dynamic_shapes={ |
| 225 | # min=10 frames = 0.1 sec @ 16kHz, max is one 120s runner chunk. |
| 226 | "audio": {0: Dim.AUTO(min=1600, max=max_audio_samples)}, |
| 227 | "length": {}, |
| 228 | }, |
| 229 | strict=False, |
| 230 | ) |
| 231 | |
| 232 | torch.cuda.is_available = old_cuda_is_available |
| 233 | |
| 234 | # --- Method 2: pre_encode --- |
| 235 | pre_encode_wrapper = PreEncodeWrapper(model.encoder.pre_encode) |
| 236 | pre_encode_wrapper.eval() |
| 237 | |
| 238 | feat_in = getattr(model.encoder, "_feat_in", 128) |
| 239 | max_chunk_mel = 4000 |
| 240 | sample_chunk = torch.randn(1, max_chunk_mel, feat_in, dtype=torch.float) |
| 241 | sample_chunk_len = torch.tensor([max_chunk_mel], dtype=torch.int64) |
| 242 | |
| 243 | print(" Exporting pre_encode...") |
| 244 | programs["pre_encode"] = export( |
| 245 | pre_encode_wrapper, |
| 246 | (sample_chunk, sample_chunk_len), |
| 247 | # Static shapes: conv-derived symbolic expression 1+((L-1)//8) creates |
| 248 | # an unsolvable guard when hitting nn.Linear. Static shapes are practical |
| 249 | # since streaming chunk sizes are fixed per config. |
| 250 | strict=False, |
no test coverage detected