from transformers import AutoModel, AutoTokenizer
import torch

repo_name = "src/models/Nemotron-Diffusion-Final"
ckpt_path = "src/models/Ministral-3-8B-Instruct-2512/state_dict.pt"
save_path = "src/models/Nemotron-Diffusion-Final-merged"

tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
model = AutoModel.from_pretrained(repo_name, trust_remote_code=True)

# ---- 加载 state_dict.pt 中缺失的 vision_tower 和 multi_modal_projector 权重 ----
print("Loading checkpoint from", ckpt_path)
ckpt = torch.load(ckpt_path, map_location="cpu")

# 映射: model.vision_tower.* -> encoder.vision_tower.*
#        model.multi_modal_projector.* -> encoder.multi_modal_projector.*
mapped_state_dict = {}
for key, value in ckpt.items():
    if key.startswith("model.vision_tower."):
        new_key = key.replace("model.vision_tower.", "encoder.vision_tower.", 1)
        mapped_state_dict[new_key] = value
    elif key.startswith("model.multi_modal_projector."):
        new_key = key.replace("model.multi_modal_projector.", "encoder.multi_modal_projector.", 1)
        mapped_state_dict[new_key] = value

print(f"Mapped {len(mapped_state_dict)} keys from checkpoint")
missing, unexpected = model.load_state_dict(mapped_state_dict, strict=False)

# missing 中应该不再包含 encoder.vision_tower 和 encoder.multi_modal_projector 的 key
encoder_missing = [k for k in missing if "vision_tower" in k or "multi_modal_projector" in k]
if encoder_missing:
    print(f"WARNING: still missing encoder keys: {encoder_missing}")
else:
    print("All encoder vision_tower and multi_modal_projector weights loaded successfully!")

# ---- 保存合并后的完整模型 ----
print(f"Saving merged model to {save_path}")
model.save_pretrained(save_path)
tokenizer.save_pretrained(save_path)
print("Done! Merged model saved.")

# ---- 推理测试 ----
# model = model.cuda().to(torch.bfloat16)

# history = []

# user_input = input("User: ").strip()
# history.append({"role": "user", "content": user_input})

# prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
# prompt_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device='cuda')
# out_ids, nfe = model.generate(prompt_ids, max_new_tokens=128, steps=128, block_length=32, shift_logits=False, causal_context=True, threshold=0.9)

# tokenized_out = tokenizer.batch_decode(out_ids[:, prompt_ids.shape[1]:], skip_special_tokens=True)[0]
# print(f"Model: {tokenized_out}")
# print(f"[Num Function Eval (NFE)={nfe}]")
