feat: Add safetensors module for loading model file

Author: tolgacangoz

examples/research_projects/anytext/recognizer.py

@@ -13,6 +13,7 @@
 import torch.nn.functional as F
 from easydict import EasyDict as edict
 from ocr_recog.RecModel import RecModel
+from safetensors.torch import load_file
 from skimage.transform._geometric import _umeyama as get_sym_mat
 
 
@@ -105,7 +106,7 @@ def create_predictor(model_dir=None, model_lang="ch", is_onnx=False):
 
         rec_model = RecModel(rec_config)
         if model_file_path is not None:
-            rec_model.load_state_dict(torch.load(model_file_path, map_location="cpu"))
+            rec_model.load_state_dict(load_file(model_file_path))
             rec_model.eval()
         return rec_model.eval()
 

examples/research_projects/anytext/text_embedding_module.py

@@ -35,7 +35,7 @@ def __init__(self, use_fp16):
         # TODO: Understand the reason of param.requires_grad = True
         for param in self.embedding_manager.embedding_parameters():
             param.requires_grad = True
-        rec_model_dir = "ppv3_rec.pth"
+        rec_model_dir = "ppv3_rec.safetensors"
         self.text_predictor = create_predictor(rec_model_dir).eval()
         args = {}
         args["rec_image_shape"] = "3, 48, 320"