#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["torch", "transformers>=4.40.0", "peft>=0.10.0", "accelerate", "huggingface_hub>=0.21.0", "pillow", "qwen-vl-utils"]
# ///
import os, torch
from huggingface_hub import login, HfApi
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from peft import PeftModel
from pathlib import Path

login(token=os.environ.get("HF_TOKEN"))
api = HfApi()

print("Loading base model...")
base = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-2B-Instruct", torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", trust_remote_code=True)

print("Loading Stage 4 adapter...")
model = PeftModel.from_pretrained(base, "mmrech/pitvqa-qwen2vl-unified-v2", adapter_name="stage4", subfolder="stage4")

print("Merging...")
merged = model.merge_and_unload()

print("Saving...")
out = Path("./pitvqa-merged")
out.mkdir(exist_ok=True)
merged.save_pretrained(out)
processor.save_pretrained(out)

print("Uploading...")
api.create_repo("mmrech/pitvqa-qwen2vl-merged", exist_ok=True)
api.upload_folder(folder_path=str(out), repo_id="mmrech/pitvqa-qwen2vl-merged", repo_type="model")
print("Done! https://huggingface.co/mmrech/pitvqa-qwen2vl-merged")