#!/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}")