Upload anytext.py

anytext.py

@@ -35,6 +35,7 @@ import PIL.Image
 import torch
 import torch.nn.functional as F
 from easydict import EasyDict as edict
+from frozen_clip_embedder_t3 import FrozenCLIPEmbedderT3
 from huggingface_hub import hf_hub_download
 from ocr_recog.RecModel import RecModel
 from PIL import Image, ImageDraw, ImageFont
@@ -206,13 +207,12 @@ def get_recog_emb(encoder, img_list):
 class EmbeddingManager(nn.Module):
     def __init__(
         self,
-        clip_tokenizer,
+        embedder,
         placeholder_string="*",
         use_fp16=False,
-        device="cpu",
     ):
         super().__init__()
-        get_token_for_string = partial(get_clip_token_for_string, clip_tokenizer)
+        get_token_for_string = partial(get_clip_token_for_string, embedder.tokenizer)
         token_dim = 768
         self.get_recog_emb = None
         self.token_dim = token_dim
@@ -223,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(device)))
+        self.proj.load_state_dict(load_file(proj_dir, device=str(embedder.device)))
         if use_fp16:
             self.proj = self.proj.to(dtype=torch.float16)
 
@@ -526,20 +526,14 @@ class TextEmbeddingModule(nn.Module):
         self.font = ImageFont.truetype(font_path, 60)
         self.use_fp16 = use_fp16
         self.device = device
-
-        # Replace instantiation of frozen_CLIP_embedder_t3
-        version = "openai/clip-vit-large-patch14"
-        torch_dtype = torch.float16 if use_fp16 else torch.float32
-        self.clip_tokenizer = CLIPTokenizer.from_pretrained(version)
-        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.clip_tokenizer, use_fp16=use_fp16, device=device)
+        self.frozen_CLIP_embedder_t3 = FrozenCLIPEmbedderT3(device=device, use_fp16=use_fp16)
+        self.embedding_manager = EmbeddingManager(self.frozen_CLIP_embedder_t3, use_fp16=use_fp16)
         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 = {}
         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["rec_char_dict_path"] = hf_hub_download(
             repo_id="tolgacangoz/anytext",
             filename="text_embedding_module/OCR/ppocr_keys_v1.txt",
@@ -548,50 +542,6 @@ class TextEmbeddingModule(nn.Module):
         args["use_fp16"] = use_fp16
         self.embedding_manager.recog = TextRecognizer(args, self.text_predictor)
 
-    # New helper method to mimic old encode() functionality with chunk splitting
-    def _encode_text(self, texts, embedding_manager=None, **kwargs):
-        batch_encoding = self.clip_tokenizer(
-            texts,
-            truncation=False,
-            max_length=self.max_length,
-            padding="longest",
-            return_tensors="pt",
-        )
-        input_ids = batch_encoding["input_ids"]
-        tokens_list = self._split_chunks(input_ids)
-        embeds_list = []
-        for tokens in tokens_list:
-            tokens = tokens.to(self.device)
-            outputs = self.clip_text_model(input_ids=tokens, **kwargs)
-            # use last_hidden_state as in the old version
-            embeds_list.append(outputs.last_hidden_state)
-        return torch.cat(embeds_list, dim=1)
-
-    # New helper for splitting tokens (mimicking split_chunks behavior)
-    def _split_chunks(self, input_ids, chunk_size=75):
-        tokens_list = []
-        bs, n = input_ids.shape
-        id_start = input_ids[:, 0].unsqueeze(1)
-        id_end = input_ids[:, -1].unsqueeze(1)
-        if n == 2:  # empty caption
-            tokens_list.append(torch.cat((id_start,) + (id_end,) * (chunk_size + 1), dim=1))
-            return tokens_list
-
-        trimmed = input_ids[:, 1:-1]
-        num_full = (n - 2) // chunk_size
-        for i in range(num_full):
-            group = trimmed[:, i*chunk_size:(i+1)*chunk_size]
-            group_pad = torch.cat((id_start, group, id_end), dim=1)
-            tokens_list.append(group_pad)
-        rem = (n - 2) % chunk_size
-        if rem > 0:
-            group = trimmed[:, -rem:]
-            pad_cols = chunk_size - group.shape[1]
-            padding = id_end.expand(bs, pad_cols)
-            group_pad = torch.cat((id_start, group, padding, id_end), dim=1)
-            tokens_list.append(group_pad)
-        return tokens_list
-
     @torch.no_grad()
     def forward(
         self,
@@ -704,9 +654,10 @@ class TextEmbeddingModule(nn.Module):
         # hint = self.arr2tensor(np_hint, len(prompt))
 
         self.embedding_manager.encode_text(text_info)
-        prompt_embeds = self._encode_text([prompt], embedding_manager=self.embedding_manager)
+        prompt_embeds = self.frozen_CLIP_embedder_t3.encode([prompt], embedding_manager=self.embedding_manager)
+
         self.embedding_manager.encode_text(text_info)
-        negative_prompt_embeds = self._encode_text(
+        negative_prompt_embeds = self.frozen_CLIP_embedder_t3.encode(
             [negative_prompt or ""], embedding_manager=self.embedding_manager
         )
 
@@ -856,11 +807,10 @@ class TextEmbeddingModule(nn.Module):
         return new_string[:-nSpace]
 
     def to(self, *args, **kwargs):
-        self.clip_text_model = self.clip_text_model.to(*args, **kwargs)
-        self.device = self.clip_text_model.device
+        self.frozen_CLIP_embedder_t3 = self.frozen_CLIP_embedder_t3.to(*args, **kwargs)
         self.embedding_manager = self.embedding_manager.to(*args, **kwargs)
         self.text_predictor = self.text_predictor.to(*args, **kwargs)
-        self.device = self.clip_text_model.device
+        self.device = self.frozen_CLIP_embedder_t3.device
         return self