#!/usr/bin/env python3
"""Merge the name-binding LoRA into full weights on CPU -> models/cosmos_namebind_merged.

Same recipe as the 2026-06-25 rebirth merge (which went to the old E: drive — gone now),
targeting D: instead. CPU-only on purpose: the 4GB card can't hold fp16 1.5B + adapters,
and CPU merge is a few minutes for this size.

Usage:  py scripts/cosmos_merge_namebind.py
Env:    HF_HOME should point at the cache that holds the base (C:\\Users\\corys\\hf_cache_local)
"""
import os
import sys
import time

ADAPTER = os.getenv("COSMOS_NB_ADAPTER", r"D:\Cosmos\models\cosmos_namebind_lora")
OUT = os.getenv("COSMOS_NB_MERGED", r"D:\Cosmos\models\cosmos_namebind_merged")
BASE = os.getenv("COSMOS_FT_BASE", "Qwen/Qwen2.5-1.5B-Instruct")


def main() -> int:
    t0 = time.time()
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from peft import PeftModel

    print(f"[MERGE] base={BASE}")
    print(f"[MERGE] adapter={ADAPTER}")
    model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.float16, device_map="cpu")
    model = PeftModel.from_pretrained(model, ADAPTER, device_map="cpu")
    print(f"[MERGE] merging adapters into base (cpu)... ({time.time()-t0:.0f}s)")
    model = model.merge_and_unload()
    os.makedirs(OUT, exist_ok=True)
    model.save_pretrained(OUT, safe_serialization=True)
    AutoTokenizer.from_pretrained(BASE).save_pretrained(OUT)
    print(f"[MERGE] DONE in {time.time()-t0:.0f}s -> {OUT}")
    return 0


if __name__ == "__main__":
    sys.exit(main())