Spaces:
Sleeping
Sleeping
File size: 556 Bytes
1b58092 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
import argparse
import os.path
import torch
from internvl.model.internvl_chat import InternVLChatModel
argparse = argparse.ArgumentParser()
argparse.add_argument('model_path', type=str, default='')
argparse.add_argument('output_path', type=str, default='')
args = argparse.parse_args()
model = InternVLChatModel.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
model = model.mlp1.to(torch.bfloat16)
ckpt = model.state_dict()
output_path = os.path.join(args.output_path, 'mlp_projector.pth')
torch.save(ckpt, output_path)
print('finished')
|