Upload anytext.py

anytext.py

@@ -547,8 +547,7 @@ class FrozenCLIPEmbedderT3(AbstractEncoder):
     ):
         super().__init__()
         self.tokenizer = CLIPTokenizer.from_pretrained("tolgacangoz/anytext", subfolder="tokenizer")
-        self.transformer = CLIPTextModel.from_pretrained("tolgacangoz/anytext", subfolder="text_encoder"
-                                                         ).to(device)
+        self.transformer = CLIPTextModel.from_pretrained("tolgacangoz/anytext", subfolder="text_encoder")
         self.device = device
         self.max_length = max_length
         if freeze:
@@ -739,22 +738,28 @@ class TextEmbeddingModule(ModelMixin, ConfigMixin):
     @register_to_config
     def __init__(self, font_path, use_fp16=False, device="cpu"):
         super().__init__()
-        self.font = ImageFont.truetype(font_path, 60)
+        font = ImageFont.truetype(font_path, 60)
+
         # self.use_fp16 = use_fp16
         # self.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)
-        self.text_predictor = create_predictor(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"] = hf_hub_download(
-            repo_id="tolgacangoz/anytext",
-            filename="text_embedding_module/OCR/ppocr_keys_v1.txt",
-            cache_dir=HF_MODULES_CACHE,
+        frozen_CLIP_embedder_t3 = FrozenCLIPEmbedderT3()#device=device, use_fp16=use_fp16)
+        embedding_manager = EmbeddingManager(frozen_CLIP_embedder_t3)#, use_fp16=use_fp16)
+        text_predictor = create_predictor(device=device, use_fp16=use_fp16).eval()
+        args = {"rec_image_shape": "3, 48, 320",
+                "rec_batch_num": 6,
+                "rec_char_dict_path": hf_hub_download(
+                    repo_id="tolgacangoz/anytext",
+                    filename="text_embedding_module/OCR/ppocr_keys_v1.txt",
+                    cache_dir=HF_MODULES_CACHE,
+                ),
+                "use_fp16": use_fp16}
+        embedding_manager.recog = TextRecognizer(args, text_predictor)
+
+        self.register_modules(
+            frozen_CLIP_embedder_t3=frozen_CLIP_embedder_t3,
+            embedding_manager=embedding_manager,
         )
-        args["use_fp16"] = use_fp16
-        self.embedding_manager.recog = TextRecognizer(args, self.text_predictor)
+        self.register_to_config(font=font)
 
     @torch.no_grad()
     def forward(
@@ -1046,17 +1051,14 @@ class AuxiliaryLatentModule(ModelMixin, ConfigMixin):
     @register_to_config
     def __init__(
         self,
-        font_path,
-        vae=None,
+        # font_path,
+        vae,
         device="cpu",
         use_fp16=False,
     ):
         super().__init__()
-        self.font = ImageFont.truetype(font_path, 60)
-        self.use_fp16 = use_fp16
-        self.device = device
-
-        self.vae = vae.eval() if vae is not None else None
+        # self.font = ImageFont.truetype(font_path, 60)
+        # self.vae = vae.eval() if vae is not None else None
 
     @torch.no_grad()
     def forward(
@@ -1276,7 +1278,7 @@ class AnyTextPipeline(
             # use_fp16=unet.dtype == torch.float16, device=unet.device,
         )
         auxiliary_latent_module = AuxiliaryLatentModule(
-            font_path=font_path,
+            # font_path=font_path,
             vae=vae,
             # use_fp16=unet.dtype == torch.float16, device=unet.device,
         )