llava-surgery.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. import argparse
  2. import glob
  3. import os
  4. import torch
  5. ap = argparse.ArgumentParser()
  6. ap.add_argument("-m", "--model", help="Path to LLaVA v1.5 model")
  7. args = ap.parse_args()
  8. # find the model part that includes the the multimodal projector weights
  9. path = sorted(glob.glob(f"{args.model}/pytorch_model*.bin"))[-1]
  10. checkpoint = torch.load(path)
  11. # get a list of mm tensor names
  12. mm_tensors = [k for k, v in checkpoint.items() if k.startswith("model.mm_projector")]
  13. # store these tensors in a new dictionary and torch.save them
  14. projector = {name: checkpoint[name].float() for name in mm_tensors}
  15. torch.save(projector, f"{args.model}/llava.projector")
  16. # remove these tensors from the checkpoint and save it again
  17. for name in mm_tensors:
  18. del checkpoint[name]
  19. # BakLLaVA models contain CLIP tensors in it
  20. clip_tensors = [k for k, v in checkpoint.items() if k.startswith("model.vision_tower")]
  21. if len(clip_tensors) > 0:
  22. clip = {name.replace("vision_tower.vision_tower.", ""): checkpoint[name].float() for name in clip_tensors}
  23. torch.save(clip, f"{args.model}/llava.clip")
  24. # remove these tensors
  25. for name in clip_tensors:
  26. del checkpoint[name]
  27. # added tokens should be removed to be able to convert Mistral models
  28. if os.path.exists(f"{args.model}/added_tokens.json"):
  29. with open(f"{args.model}/added_tokens.json", "w") as f:
  30. f.write("{}\n")
  31. torch.save(checkpoint, path)
  32. print("Done!")
  33. print(f"Now you can convert {args.model} to a a regular LLaMA GGUF file.")
  34. print(f"Also, use {args.model}/llava.projector to prepare a llava-encoder.gguf file.")