Upload anytext.py

anytext.py

@@ -334,13 +334,12 @@ def crop_image(src_img, mask):
     return result
 
 
-def create_predictor(model_dir=None, model_lang="ch", device="cpu", use_fp16=False):
-    if model_dir is None or not os.path.exists(model_dir):
-        model_dir = hf_hub_download(
-            repo_id="tolgacangoz/anytext",
-            filename="text_embedding_module/OCR/ppv3_rec.pth",
-            cache_dir=HF_MODULES_CACHE,
-        )
+def create_predictor(model_lang="ch", device="cpu", use_fp16=False):
+    model_dir = hf_hub_download(
+        repo_id="tolgacangoz/anytext",
+        filename="text_embedding_module/OCR/ppv3_rec.pth",
+        cache_dir=HF_MODULES_CACHE,
+    )
     if not os.path.exists(model_dir):
         raise ValueError("not find model file path {}".format(model_dir))
 
@@ -540,16 +539,17 @@ class FrozenCLIPEmbedderT3(AbstractEncoder):
 
     def __init__(
         self,
-        version="openai/clip-vit-large-patch14",
         device="cpu",
         max_length=77,
         freeze=True,
         use_fp16=False,
+        variant="fp32",
     ):
         super().__init__()
-        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.tokenizer = CLIPTokenizer.from_pretrained("tolgacangoz/anytext", subfolder="tokenizer")
         self.transformer = CLIPTextModel.from_pretrained(
-            version, use_safetensors=True, torch_dtype=torch.float16 if use_fp16 else torch.float32
+            "tolgacangoz/anytext", subfolder="text_encoder", use_safetensors=True,
+            torch_dtype=torch.float16 if use_fp16 else torch.float32, variant=variant,
         ).to(device)
         self.device = device
         self.max_length = max_length
@@ -746,8 +746,7 @@ class TextEmbeddingModule(ModelMixin, ConfigMixin):
         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)
-        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()
+        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
@@ -1045,7 +1044,8 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class AuxiliaryLatentModule(nn.Module):
+class AuxiliaryLatentModule(ModelMixin, ConfigMixin):
+    @register_to_config
     def __init__(
         self,
         font_path,
@@ -1229,7 +1229,7 @@ class AnyTextPipeline(
     Args:
         vae ([`AutoencoderKL`]):
             Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
-        text_encoder ([`~anytext.TextEmbeddingModule`]):
+        text_encoder ([`~transformers.CLIPTextModel`]):
             Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
         tokenizer ([`~transformers.CLIPTokenizer`]):
             A `CLIPTokenizer` to tokenize text.
@@ -1259,7 +1259,7 @@ class AnyTextPipeline(
         self,
         font_path: str,
         vae: AutoencoderKL,
-        text_encoder: TextEmbeddingModule,
+        text_encoder: CLIPTextModel,
         tokenizer: CLIPTokenizer,
         unet: UNet2DConditionModel,
         controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
@@ -1267,17 +1267,21 @@ class AnyTextPipeline(
         safety_checker: StableDiffusionSafetyChecker,
         feature_extractor: CLIPImageProcessor,
         auxiliary_latent_module: AuxiliaryLatentModule,
+        text_embedding_module: TextEmbeddingModule,
         trust_remote_code: bool = False,
         image_encoder: CLIPVisionModelWithProjection = None,
         requires_safety_checker: bool = True,
     ):
         super().__init__()
-        # self.text_embedding_module = TextEmbeddingModule(
-        #     use_fp16=unet.dtype == torch.float16, device=unet.device, font_path=font_path
-        # )
-        # self.auxiliary_latent_module = AuxiliaryLatentModule(
-        #     vae=vae, use_fp16=unet.dtype == torch.float16, device=unet.device, font_path=font_path
-        # )
+        self.text_embedding_module = TextEmbeddingModule(
+            font_path=font_path,
+            # use_fp16=unet.dtype == torch.float16, device=unet.device,
+        )
+        self.auxiliary_latent_module = AuxiliaryLatentModule(
+            font_path=font_path,
+            vae=vae,
+            # use_fp16=unet.dtype == torch.float16, device=unet.device,
+        )
 
         if safety_checker is None and requires_safety_checker:
             logger.warning(
@@ -1308,7 +1312,7 @@ class AnyTextPipeline(
             safety_checker=safety_checker,
             feature_extractor=feature_extractor,
             image_encoder=image_encoder,
-            # text_embedding_module=self.text_embedding_module,
+            text_embedding_module=text_embedding_module,
             auxiliary_latent_module=auxiliary_latent_module,
         )
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
@@ -2177,7 +2181,7 @@ class AnyTextPipeline(
             self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
         )
         draw_pos = draw_pos.to(device=device) if isinstance(draw_pos, torch.Tensor) else draw_pos
-        prompt_embeds, negative_prompt_embeds, text_info, np_hint = self.text_encoder(
+        prompt_embeds, negative_prompt_embeds, text_info, np_hint = self.text_embedding_module(
             prompt,
             texts,
             negative_prompt,
@@ -2419,6 +2423,6 @@ class AnyTextPipeline(
 
     def to(self, *args, **kwargs):
         super().to(*args, **kwargs)
-        # self.text_embedding_module.to(*args, **kwargs)
+        self.text_embedding_module.to(*args, **kwargs)
         self.auxiliary_latent_module.to(*args, **kwargs)
         return self