Upload anytext.py

anytext.py

@@ -206,12 +206,13 @@ def get_recog_emb(encoder, img_list):
 class EmbeddingManager(nn.Module):
     def __init__(
         self,
-        embedder,
+        clip_tokenizer,
         placeholder_string="*",
         use_fp16=False,
+        device="cpu",
     ):
         super().__init__()
-        get_token_for_string = partial(get_clip_token_for_string, embedder.clip_tokenizer)
+        get_token_for_string = partial(get_clip_token_for_string, clip_tokenizer)
         token_dim = 768
         self.get_recog_emb = None
         self.token_dim = token_dim
@@ -222,7 +223,7 @@ class EmbeddingManager(nn.Module):
             filename="text_embedding_module/proj.safetensors",
             cache_dir=HF_MODULES_CACHE,
         )
-        self.proj.load_state_dict(load_file(proj_dir, device=str(embedder.device)))
+        self.proj.load_state_dict(load_file(proj_dir, device=str(device)))
         if use_fp16:
             self.proj = self.proj.to(dtype=torch.float16)
 
@@ -533,7 +534,7 @@ class TextEmbeddingModule(nn.Module):
         self.clip_text_model = CLIPTextModel.from_pretrained(version, torch_dtype=torch_dtype).to(device)
         self.max_length = 77  # same as before
 
-        self.embedding_manager = EmbeddingManager(self, use_fp16=use_fp16)
+        self.embedding_manager = EmbeddingManager(self.clip_tokenizer, use_fp16=use_fp16, device=device)
         rec_model_dir = "./text_embedding_module/OCR/ppv3_rec.pth"
         self.text_predictor = create_predictor(rec_model_dir, device=device, use_fp16=use_fp16).eval()
         args = {}