Upload anytext.py

anytext.py

@@ -204,30 +204,32 @@ def get_recog_emb(encoder, img_list):
     return preds_neck
 
 
-class EmbeddingManager(nn.Module):
+class EmbeddingManager(ModelMixin, ConfigMixin):
+    @register_to_config
     def __init__(
         self,
         embedder,
         placeholder_string="*",
         use_fp16=False,
+        token_dim = 768,
+        get_recog_emb = None,
     ):
         super().__init__()
         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
 
-        self.proj = nn.Linear(40 * 64, token_dim)
+        proj = nn.Linear(40 * 64, token_dim)
         proj_dir = hf_hub_download(
             repo_id="tolgacangoz/anytext",
             filename="text_embedding_module/proj.safetensors",
             cache_dir=HF_MODULES_CACHE,
         )
-        self.proj.load_state_dict(load_file(proj_dir, device=str(embedder.device)))
+        proj.load_state_dict(load_file(proj_dir, device=str(embedder.device)))
         if use_fp16:
-            self.proj = self.proj.to(dtype=torch.float16)
+            proj = proj.to(dtype=torch.float16)
 
-        self.placeholder_token = get_token_for_string(placeholder_string)
+        self.register_parameter("proj", proj)
+        placeholder_token = get_token_for_string(placeholder_string)
+        self.register_config(placeholder_token=placeholder_token)
 
     @torch.no_grad()
     def encode_text(self, text_info):
@@ -1024,10 +1026,10 @@ class TextEmbeddingModule(ModelMixin, ConfigMixin):
             new_string += char + " " * nSpace
         return new_string[:-nSpace]
 
-    def to(self, *args, **kwargs):
-        self.frozen_CLIP_embedder_t3 = self.frozen_CLIP_embedder_t3.to(*args, **kwargs)
-        self.embedding_manager = self.embedding_manager.to(*args, **kwargs)
-        return self
+    # def to(self, *args, **kwargs):
+    #     self.frozen_CLIP_embedder_t3 = self.frozen_CLIP_embedder_t3.to(*args, **kwargs)
+    #     self.embedding_manager = self.embedding_manager.to(*args, **kwargs)
+    #     return self
 
 
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents