Upload anytext.py

anytext.py

@@ -69,6 +69,7 @@ from diffusers.utils import (
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor
 from diffusers.configuration_utils import register_to_config, ConfigMixin
 from diffusers.models.modeling_utils import ModelMixin
+from huggingface_hub import hf_hub_download
 
 
 checker = BasicTokenizer()
@@ -269,9 +270,20 @@ def crop_image(src_img, mask):
 
 
 def create_predictor(model_dir=None, model_lang="ch", device="cpu", use_fp16=False):
-    model_file_path = model_dir
-    if model_file_path is not None and not os.path.exists(model_file_path):
-        raise ValueError("not find model file path {}".format(model_file_path))
+    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):
+        raise ValueError("not find model file path {}".format(model_dir))
 
     if model_lang == "ch":
         n_class = 6625
@@ -287,8 +299,8 @@ def create_predictor(model_dir=None, model_lang="ch", device="cpu", use_fp16=Fal
     )
 
     rec_model = RecModel(rec_config)
-    if model_file_path is not None:
-        rec_model.load_state_dict(torch.load(model_file_path, map_location=device))
+    state_dict = torch.load(model_dir, map_location=device)
+    rec_model.load_state_dict(state_dict)
     return rec_model