Skip to content

Commit

Permalink
convert_dinov2: ignore pyright errors
Browse files Browse the repository at this point in the history
And save converted weights into safetensors instead of pickle
  • Loading branch information
deltheil committed Dec 14, 2023
1 parent 9337d65 commit e978b36
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions scripts/conversion/convert_dinov2.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

import torch

from refiners.fluxion.utils import save_to_safetensors


def convert_dinov2_facebook(weights: dict[str, torch.Tensor]) -> None:
"""Convert a DINOv2 weights from facebook to refiners."""
Expand Down Expand Up @@ -126,9 +128,9 @@ def main() -> None:
parser.add_argument("--output_path", type=str, required=True)
args = parser.parse_args()

weights = torch.load(args.weights_path)
weights = torch.load(args.weights_path) # type: ignore
convert_dinov2_facebook(weights)
torch.save(weights, args.output_path)
save_to_safetensors(path=args.output_path, tensors=weights)


if __name__ == "__main__":
Expand Down

0 comments on commit e978b36

Please sign in to comment.