Consolidated release (history squashed to remove internal filesystem paths from prior revisions)
245508e | #!/usr/bin/env python3 | |
| """Load the ft8 Sortformer diarizer from config.yaml + model.safetensors — fp32, lossless, no .nemo tar. | |
| Native path that works (NeMo 2.6.2, torch 2.10): | |
| 1. cfg = OmegaConf.load("config.yaml") # the .nemo model_config.yaml IS the | |
| model-level cfg (top-level keys: encoder, sortformer_modules, ... — NOT nested | |
| under 'model'). Identical to restore_from(..., return_config=True). | |
| 2. Null out train_ds / validation_ds / test_ds: ModelPT.__init__ otherwise tries to | |
| build dataloaders from cluster manifest paths that don't exist at serve time. | |
| 3. model = SortformerEncLabelModel(cfg=cfg) # direct instantiation works | |
| 4. model.load_state_dict(safetensors_sd, strict=True) | |
| fp16 quirks (ONLY relevant if you re-quantize to fp16 yourself; the shipped weights are | |
| fp32 and the loader skips all of this — kept so the community can quantize freely): | |
| * preprocessor.* (STFT window + mel fb) is kept fp32 in the safetensors; after | |
| model.half() we re-float the preprocessor and cast its output features to fp16. | |
| Halving the STFT itself degrades mel features / breaks torch.stft dtype paths. | |
| * Sortformer's STREAMING path creates fp32 state tensors internally | |
| (sortformer_modules.init_streaming_state / streaming_update use torch.zeros | |
| without dtype). torch.cat promotes the fp16 chunk to fp32 and the fp16 encoder | |
| then throws "expected scalar type Float but found Half". Fix: run forward with | |
| torch.set_default_dtype(torch.float16) so those internal states are fp16 too. | |
| * forward returns fp32 preds (cast at the end) so NeMo's CPU post-processing | |
| (ts_vad_post_processing) never sees Half tensors. | |
| Usage: | |
| from load_diarizer import load_diarizer | |
| model = load_diarizer() # dtype auto-detected (fp32 shipped) | |
| segs = model.diarize(audio=["x.wav"], batch_size=1, | |
| postprocessing_yaml=pp_yaml, verbose=False) | |
| """ | |
| import os | |
| import torch | |
| from omegaconf import OmegaConf, open_dict | |
| _HERE = os.path.dirname(os.path.abspath(__file__)) | |
| def load_diarizer(artifact_dir: str = _HERE, device: str = "cuda"): | |
| from nemo.collections.asr.models import SortformerEncLabelModel | |
| from safetensors.torch import load_file | |
| cfg = OmegaConf.load(os.path.join(artifact_dir, "config.yaml")) | |
| if "model" in cfg and "encoder" not in cfg: # nested checkpoint (not ft8; safety) | |
| cfg = cfg.model | |
| with open_dict(cfg): | |
| for k in ("train_ds", "validation_ds", "test_ds"): | |
| if k in cfg: | |
| cfg[k] = None | |
| model = SortformerEncLabelModel(cfg=cfg) | |
| sd = load_file(os.path.join(artifact_dir, "model.safetensors")) | |
| fp16 = any(v.dtype == torch.float16 for v in sd.values()) | |
| if fp16: | |
| model = model.half() | |
| model.preprocessor.float() # STFT/mel stays fp32 (matches fp32 sd keys) | |
| info = model.load_state_dict(sd, strict=True) | |
| assert not info.missing_keys and not info.unexpected_keys | |
| if fp16: | |
| # bridge fp32 mel features -> fp16 encoder | |
| _orig_pre = model.preprocessor.forward | |
| def _cast_pre(*a, **kw): | |
| out = _orig_pre(*a, **kw) | |
| if isinstance(out, tuple): | |
| return (out[0].half(),) + tuple(out[1:]) | |
| return out.half() | |
| model.preprocessor.forward = _cast_pre | |
| # make internally-created streaming state (spkcache/fifo zeros) fp16 as well, | |
| # and hand fp32 preds back to NeMo's CPU post-processing | |
| _orig_fwd = model.forward | |
| def _half_fwd(*a, **kw): | |
| prev = torch.get_default_dtype() | |
| torch.set_default_dtype(torch.float16) | |
| try: | |
| out = _orig_fwd(*a, **kw) | |
| finally: | |
| torch.set_default_dtype(prev) | |
| if torch.is_tensor(out) and out.is_floating_point(): | |
| return out.float() | |
| return out | |
| model.forward = _half_fwd | |
| model = model.to(device).eval() | |
| return model | |
| if __name__ == "__main__": | |
| m = load_diarizer() | |
| n = sum(p.numel() for p in m.parameters()) | |
| print(f"loaded OK: {n/1e6:.1f}M params, encoder dtype={next(m.encoder.parameters()).dtype}") | |