Upload anytext.py

text_embedding_module/anytext.py

@@ -35,6 +35,7 @@ import torch
 import torch.nn.functional as F
 from bert_tokenizer import BasicTokenizer
 from easydict import EasyDict as edict
+from diffusers.utils.constants import HF_MODULES_CACHE
 from frozen_clip_embedder_t3 import FrozenCLIPEmbedderT3
 from ocr_recog.RecModel import RecModel
 from PIL import Image, ImageDraw, ImageFont
@@ -271,18 +272,12 @@ def crop_image(src_img, mask):
 
 def create_predictor(model_dir=None, model_lang="ch", device="cpu", use_fp16=False):
     if model_dir is None or not os.path.exists(model_dir):
-        try:
-            # Use the repo id from which the pipeline was loaded
-            model_dir = hf_hub_download(
-                repo_id="tolgacangoz/anytext",
-                filename="text_embedding_module/OCR/ppv3_rec.pth",
-                local_dir=".cache/diffusers",
-                local_dir_use_symlinks=True
-            )
-        except Exception as e:
-            raise ValueError(f"Could not download the model file: {e}")
-
-    if model_dir is not None and not os.path.exists(model_dir):
+        model_dir = hf_hub_download(
+            repo_id="tolgacangoz/anytext",
+            filename="text_embedding_module/OCR/ppv3_rec.pth",
+            cache_dir=HF_MODULES_CACHE
+        )
+    if not os.path.exists(model_dir):
         raise ValueError("not find model file path {}".format(model_dir))
 
     if model_lang == "ch":
@@ -476,7 +471,7 @@ class TextEmbeddingModule(nn.Module):
         args["rec_image_shape"] = "3, 48, 320"
         args["rec_batch_num"] = 6
         args["rec_char_dict_path"] = "./text_embedding_module/OCR/ppocr_keys_v1.txt"
-        args["use_fp16"] = self.use_fp16
+        args["use_fp16"] = use_fp16
         self.embedding_manager.recog = TextRecognizer(args, self.text_predictor)
 
     @torch.no_grad()