#!/usr/bin/env python3
# convert.py: safetensors to GGUF for OmniVoice (LM + audio tokenizer).
# Reads from checkpoints/OmniVoice/, writes 2 byte-perfect F32 GGUFs to models/.
# The source checkpoint is 100% F32 (k2-fsa/OmniVoice on Hugging Face),
# so this converter never downcasts : every tensor is written in its native
# source dtype and quantize.sh derives BF16 / Q8_0 from these F32 GGUFs.
#
# omnivoice-base-F32.gguf       Qwen3 0.6B + audio_embeddings + audio_heads + tokenizer
# omnivoice-tokenizer-F32.gguf  HuBERT + DAC encoder/decoder + RVQ + fc/fc2

import json
import os
import struct
import sys

import numpy as np
import gguf

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CHECKPOINT_DIR = os.path.join(SCRIPT_DIR, "checkpoints", "OmniVoice")
OUTPUT_DIR = os.path.join(SCRIPT_DIR, "models")

def log(tag, msg):
    print("[%s] %s" % (tag, msg), file=sys.stderr, flush=True)

# safetensors header reader
def read_sf_header(path):
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        meta = json.loads(f.read(n))
    meta.pop("__metadata__", None)
    return meta, 8 + n

# read a tensor from safetensors as a numpy view of the right dtype
def read_sf_tensor(sf_path, hdr_size, info):
    off0, off1 = info["data_offsets"]
    nbytes = off1 - off0
    with open(sf_path, "rb") as f:
        f.seek(hdr_size + off0)
        raw = f.read(nbytes)
    dtype_str = info["dtype"]
    shape = info["shape"]
    if dtype_str == "F32":
        return np.frombuffer(raw, dtype=np.float32).reshape(shape).copy()
    if dtype_str == "BF16":
        u16 = np.frombuffer(raw, dtype=np.uint16).reshape(shape)
        return (u16.astype(np.uint32) << 16).view(np.float32).copy()
    if dtype_str == "F16":
        return np.frombuffer(raw, dtype=np.float16).reshape(shape).astype(np.float32).copy()
    raise RuntimeError("unsupported dtype %s" % dtype_str)

# write a tensor to GGUF preserving the source dtype : F32 stays F32, BF16 stays
# BF16, F16 stays F16. No cast, no precision invented or destroyed. The OmniVoice
# checkpoint is integrally F32 in the reference release ; the BF16 / F16 branches
# stay as defensive code in case a future release switches dtype upstream.
def add_tensor_passthrough(w, name, sf_path, hdr_size, info):
    dtype_str = info["dtype"]
    shape = info["shape"]
    off0, off1 = info["data_offsets"]
    nbytes = off1 - off0
    with open(sf_path, "rb") as f:
        f.seek(hdr_size + off0)
        raw = f.read(nbytes)
    if dtype_str == "F32":
        arr = np.frombuffer(raw, dtype=np.float32).reshape(shape)
        w.add_tensor(name, arr)
        return nbytes
    if dtype_str == "BF16":
        arr = np.frombuffer(raw, dtype=np.uint16).reshape(shape)
        w.add_tensor(name, arr, raw_dtype=gguf.GGMLQuantizationType.BF16)
        return nbytes
    if dtype_str == "F16":
        arr = np.frombuffer(raw, dtype=np.float16).reshape(shape)
        w.add_tensor(name, arr)
        return nbytes
    raise RuntimeError("unsupported dtype %s for %s" % (dtype_str, name))

# BPE tokenizer: extract vocab + merges + added_tokens from tokenizer.json
def add_bpe_tokenizer(w, tokenizer_json_path, tag):
    with open(tokenizer_json_path, "r", encoding="utf-8") as f:
        tj = json.load(f)

    model = tj.get("model", {})
    if model.get("type") != "BPE":
        raise RuntimeError("expected BPE tokenizer, got %s" % model.get("type"))

    vocab_dict = model.get("vocab", {})
    added_tokens = tj.get("added_tokens", [])

    max_id = max(
        max(vocab_dict.values(), default=-1),
        max((a["id"] for a in added_tokens), default=-1),
    )
    tokens = [""] * (max_id + 1)
    for tok_str, tok_id in vocab_dict.items():
        if 0 <= tok_id <= max_id:
            tokens[tok_id] = tok_str
    for a in added_tokens:
        tokens[a["id"]] = a["content"]

    merges = []
    for m in model.get("merges", []):
        if isinstance(m, list):
            merges.append(" ".join(m))
        else:
            merges.append(m)

    w.add_tokenizer_model("gpt2")
    w.add_token_list(tokens)
    w.add_token_merges(merges)

    log(tag, "tokenizer: %d vocab, %d merges, %d added" %
        (len(tokens), len(merges), len(added_tokens)))

# Tensors to skip in model.safetensors (recomputable at runtime).
def is_skip_base_tensor(name):
    # codebook_layer_offsets = arange(num_audio_codebook) * audio_vocab_size,
    # trivially recomputed C++ side from the audio metadata.
    if name == "codebook_layer_offsets":
        return True
    return False

# omnivoice-base: LLM + audio_embeddings + audio_heads + tokenizer
def convert_base(ckpt_dir, output_path):
    tag = "BASE"
    cfg_path = os.path.join(ckpt_dir, "config.json")
    sf_path = os.path.join(ckpt_dir, "model.safetensors")
    tok_path = os.path.join(ckpt_dir, "tokenizer.json")

    with open(cfg_path, "r", encoding="utf-8") as f:
        cfg = json.load(f)
    llm_cfg = cfg["llm_config"]

    log(tag, "writing %s" % os.path.basename(output_path))
    w = gguf.GGUFWriter(output_path, "omnivoice-lm", use_temp_file=True)
    w.add_name("OmniVoice")

    # standard LLM metadata
    w.add_block_count(llm_cfg["num_hidden_layers"])
    w.add_embedding_length(llm_cfg["hidden_size"])
    w.add_feed_forward_length(llm_cfg["intermediate_size"])
    w.add_head_count(llm_cfg["num_attention_heads"])
    w.add_head_count_kv(llm_cfg["num_key_value_heads"])
    w.add_key_length(llm_cfg["head_dim"])
    w.add_value_length(llm_cfg["head_dim"])
    w.add_vocab_size(llm_cfg["vocab_size"])
    w.add_context_length(llm_cfg["max_position_embeddings"])
    w.add_layer_norm_rms_eps(llm_cfg["rms_norm_eps"])
    w.add_rope_freq_base(float(llm_cfg["rope_parameters"]["rope_theta"]))
    w.add_bool("omnivoice.tie_word_embeddings", bool(llm_cfg.get("tie_word_embeddings", False)))

    # audio IO (LLM-side custom)
    w.add_uint32("omnivoice.num_audio_codebook", cfg["num_audio_codebook"])
    w.add_uint32("omnivoice.audio_vocab_size", cfg["audio_vocab_size"])
    w.add_uint32("omnivoice.audio_mask_id", cfg["audio_mask_id"])
    w.add_array("omnivoice.audio_codebook_weights",
                [int(x) for x in cfg["audio_codebook_weights"]])

    # special token IDs from tokenizer.json (denoise / lang / instruct / text bracketing)
    with open(tok_path, "r", encoding="utf-8") as f:
        tj = json.load(f)
    special_map = {a["content"]: a["id"] for a in tj.get("added_tokens", [])}
    for tag_name in [
        "<|denoise|>",
        "<|lang_start|>", "<|lang_end|>",
        "<|instruct_start|>", "<|instruct_end|>",
        "<|text_start|>", "<|text_end|>",
    ]:
        if tag_name in special_map:
            key = "omnivoice.special." + tag_name.strip("<|>").replace("|", "_")
            w.add_uint32(key, special_map[tag_name])

    # tokenizer
    add_bpe_tokenizer(w, tok_path, tag)

    # tensors
    meta, hdr_size = read_sf_header(sf_path)
    n_tensors, n_bytes = 0, 0
    n_skip = 0
    for name in sorted(meta.keys()):
        if is_skip_base_tensor(name):
            n_skip += 1
            continue
        nb = add_tensor_passthrough(w, name, sf_path, hdr_size, meta[name])
        n_tensors += 1
        n_bytes += nb
    log(tag, "total: %d tensors, %.1f MB (%d skipped)" %
        (n_tensors, n_bytes / (1 << 20), n_skip))

    w.write_header_to_file()
    w.write_kv_data_to_file()
    w.write_tensors_to_file(progress=True)
    w.close()

    out_mb = os.path.getsize(output_path) / (1 << 20)
    log(tag, "wrote %.0f MB -> %s" % (out_mb, output_path))

# Tensors to skip in audio_tokenizer (training-only, not used at inference)
def is_skip_tokenizer_tensor(name):
    # decoder_semantic + fc1: auxiliary HuBERT-feature reconstruction loss
    if name.startswith("decoder_semantic.") or name.startswith("fc1."):
        return True
    # RVQ codebook EMA buffers (cluster_size, embed_avg, inited): training only
    if "quantizer.quantizers." in name:
        if name.endswith(".cluster_size") or name.endswith(".embed_avg") or name.endswith(".inited"):
            return True
    # parametrizations.weight.original{0,1}: replaced by folded weight_norm below
    if "parametrizations.weight.original" in name:
        return True
    return False

# Fold HuBERT pos_conv weight_norm: weight = v * g / ||v||_{dim=(0,1)}
# Convention: torch.nn.utils.weight_norm(conv, dim=2) for grouped conv1d.
# Verified bit-identical to torch._weight_norm(v, g, dim=2) up to FP32 noise.
def fold_pos_conv_weight_norm(meta, hdr_size, sf_path):
    g_name = "semantic_model.encoder.pos_conv_embed.conv.parametrizations.weight.original0"
    v_name = "semantic_model.encoder.pos_conv_embed.conv.parametrizations.weight.original1"
    target = "semantic_model.encoder.pos_conv_embed.conv.weight"

    g = read_sf_tensor(sf_path, hdr_size, meta[g_name])  # (1, 1, 128)
    v = read_sf_tensor(sf_path, hdr_size, meta[v_name])  # (768, 48, 128)
    norm = np.sqrt(np.sum(v * v, axis=(0, 1), keepdims=True))
    weight = v * g / norm
    return target, weight.astype(np.float32, copy=False)

# omnivoice-tokenizer: HuBERT + DAC + RVQ + fc/fc2
def convert_tokenizer(ckpt_dir, output_path):
    tag = "TOK"
    audio_dir = os.path.join(ckpt_dir, "audio_tokenizer")
    cfg_path = os.path.join(audio_dir, "config.json")
    sf_path = os.path.join(audio_dir, "model.safetensors")

    with open(cfg_path, "r", encoding="utf-8") as f:
        cfg = json.load(f)

    log(tag, "writing %s" % os.path.basename(output_path))
    w = gguf.GGUFWriter(output_path, "omnivoice-tokenizer", use_temp_file=True)
    w.add_name("OmniVoice-tokenizer")

    # global audio config
    w.add_uint32("omnivoice.sample_rate", cfg["sample_rate"])
    w.add_uint32("omnivoice.semantic_sample_rate", cfg["semantic_sample_rate"])
    w.add_uint32("omnivoice.downsample_factor", cfg["downsample_factor"])
    w.add_uint32("omnivoice.codebook_size", cfg["codebook_size"])
    w.add_uint32("omnivoice.codebook_dim", cfg["codebook_dim"])

    # acoustic (DAC) config
    ac = cfg["acoustic_model_config"]
    w.add_uint32("omnivoice.acoustic.encoder_hidden_size", ac["encoder_hidden_size"])
    w.add_uint32("omnivoice.acoustic.decoder_hidden_size", ac["decoder_hidden_size"])
    w.add_uint32("omnivoice.acoustic.hidden_size", ac["hidden_size"])
    w.add_uint32("omnivoice.acoustic.n_codebooks", ac["n_codebooks"])
    w.add_uint32("omnivoice.acoustic.hop_length", ac["hop_length"])
    w.add_array("omnivoice.acoustic.upsampling_ratios",
                [int(x) for x in ac["upsampling_ratios"]])
    w.add_array("omnivoice.acoustic.downsampling_ratios",
                [int(x) for x in ac["downsampling_ratios"]])

    # semantic (HuBERT) config
    sm = cfg["semantic_model_config"]
    w.add_uint32("omnivoice.semantic.hidden_size", sm["hidden_size"])
    w.add_uint32("omnivoice.semantic.intermediate_size", sm["intermediate_size"])
    w.add_uint32("omnivoice.semantic.num_attention_heads", sm["num_attention_heads"])
    w.add_uint32("omnivoice.semantic.num_hidden_layers", sm["num_hidden_layers"])
    w.add_uint32("omnivoice.semantic.num_feat_extract_layers", sm["num_feat_extract_layers"])
    w.add_array("omnivoice.semantic.conv_dim", [int(x) for x in sm["conv_dim"]])
    w.add_array("omnivoice.semantic.conv_kernel", [int(x) for x in sm["conv_kernel"]])
    w.add_array("omnivoice.semantic.conv_stride", [int(x) for x in sm["conv_stride"]])
    w.add_uint32("omnivoice.semantic.num_conv_pos_embeddings", sm["num_conv_pos_embeddings"])
    w.add_uint32("omnivoice.semantic.num_conv_pos_embedding_groups", sm["num_conv_pos_embedding_groups"])
    w.add_float32("omnivoice.semantic.layer_norm_eps", float(sm["layer_norm_eps"]))

    # tensors
    meta, hdr_size = read_sf_header(sf_path)

    # fold pos_conv weight_norm before iterating, so the folded weight replaces
    # both parametrizations entries in the output GGUF
    folded_name, folded_arr = fold_pos_conv_weight_norm(meta, hdr_size, sf_path)

    n_tensors, n_bytes = 0, 0
    n_skip = 0
    for name in sorted(meta.keys()):
        if is_skip_tokenizer_tensor(name):
            n_skip += 1
            continue
        nb = add_tensor_passthrough(w, name, sf_path, hdr_size, meta[name])
        n_tensors += 1
        n_bytes += nb

    # folded pos_conv weight is computed F32 from F32 source params, write F32
    w.add_tensor(folded_name, folded_arr)
    n_tensors += 1
    n_bytes += folded_arr.nbytes

    log(tag, "total: %d tensors, %.1f MB (%d skipped: training-only)" %
        (n_tensors, n_bytes / (1 << 20), n_skip))

    w.write_header_to_file()
    w.write_kv_data_to_file()
    w.write_tensors_to_file(progress=True)
    w.close()

    out_mb = os.path.getsize(output_path) / (1 << 20)
    log(tag, "wrote %.0f MB -> %s" % (out_mb, output_path))

def main():
    if not os.path.isdir(CHECKPOINT_DIR):
        log("GGUF", "checkpoints/OmniVoice not found, run checkpoints.sh first")
        sys.exit(1)

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    base_path = os.path.join(OUTPUT_DIR, "omnivoice-base-F32.gguf")
    tok_path = os.path.join(OUTPUT_DIR, "omnivoice-tokenizer-F32.gguf")

    if os.path.exists(base_path):
        log("BASE", "skip: %s exists" % os.path.basename(base_path))
    else:
        convert_base(CHECKPOINT_DIR, base_path)

    if os.path.exists(tok_path):
        log("TOK", "skip: %s exists" % os.path.basename(tok_path))
    else:
        convert_tokenizer(CHECKPOINT_DIR, tok_path)

    log("GGUF", "done -> %s" % OUTPUT_DIR)

if __name__ == "__main__":
    main()