Up

Author: tolgacangoz

examples/research_projects/anytext/auxiliary_latent_module.py

@@ -37,9 +37,9 @@ def retrieve_latents(
 class AuxiliaryLatentModule(nn.Module):
     def __init__(self, glyph_channels=1, position_channels=1, model_channels=320, **kwargs):
         super().__init__()
-        self.font = ImageFont.truetype("Arial_Unicode.ttf", 60)
+        self.font = ImageFont.truetype("font/Arial_Unicode.ttf", 60)
         self.use_fp16 = kwargs.get("use_fp16", False)
-        self.device = kwargs.get("device", "cuda")
+        self.device = kwargs.get("device", "cpu")
 
         self.glyph_block = nn.Sequential(
             nn.Conv2d(glyph_channels, 8, 3, padding=1),
@@ -78,17 +78,22 @@ def __init__(self, glyph_channels=1, position_channels=1, model_channels=320, **
             nn.Conv2d(32, 64, 3, padding=1, stride=2),
             nn.SiLU(),
         )
-        self.glyph_block.load_state_dict(load_file("glyph_block.safetensors"))
-        self.position_block.load_state_dict(load_file("position_block.safetensors"))
-        self.glyph_block = self.glyph_block.to(device="cuda", dtype=torch.float16)
-        self.position_block = self.position_block.to(device="cuda", dtype=torch.float16)
 
         self.vae = kwargs.get("vae")
         self.vae.eval()
 
         self.fuse_block = zero_module(nn.Conv2d(256 + 64 + 4, model_channels, 3, padding=1))
-        self.fuse_block.load_state_dict(load_file("fuse_block.safetensors"))
-        self.fuse_block = self.fuse_block.to(device="cuda", dtype=torch.float16)
+
+        self.glyph_block.load_state_dict(
+            load_file("AuxiliaryLatentModule/glyph_block.safetensors", device=self.device)
+        )
+        self.glyph_block = self.glyph_block.to(dtype=torch.float16 if self.use_fp16 else torch.float32)
+        self.position_block.load_state_dict(
+            load_file("AuxiliaryLatentModule/position_block.safetensors", device=self.device)
+        )
+        self.position_block = self.position_block.to(dtype=torch.float16 if self.use_fp16 else torch.float32)
+        self.fuse_block.load_state_dict(load_file("AuxiliaryLatentModule/fuse_block.safetensors", device=self.device))
+        self.fuse_block = self.fuse_block.to(dtype=torch.float16 if self.use_fp16 else torch.float32)
 
     @torch.no_grad()
     def forward(
@@ -121,30 +126,26 @@ def forward(
             edit_image = self.resize_image(
                 edit_image, max_length=768
             )  # make w h multiple of 64, resize if w or h > max_length
-            h, w = edit_image.shape[:2]  # change h, w by input ref_img
 
         # get masked_x
         masked_img = ((edit_image.astype(np.float32) / 127.5) - 1.0) * (1 - np_hint)
         masked_img = np.transpose(masked_img, (2, 0, 1))
         masked_img = torch.from_numpy(masked_img.copy()).float().to(self.device)
         if self.use_fp16:
             masked_img = masked_img.half()
-        masked_x = self.encode_first_stage(masked_img[None, ...]).detach()
+        masked_x = (retrieve_latents(self.vae.encode(masked_img[None, ...])) * self.vae.config.scaling_factor).detach()
         if self.use_fp16:
             masked_x = masked_x.half()
         text_info["masked_x"] = torch.cat([masked_x for _ in range(num_images_per_prompt)], dim=0)
 
         glyphs = torch.cat(text_info["glyphs"], dim=1).sum(dim=1, keepdim=True)
         positions = torch.cat(text_info["positions"], dim=1).sum(dim=1, keepdim=True)
-        enc_glyph = self.glyph_block(glyphs.cuda())
-        enc_pos = self.position_block(positions.cuda())
-        guided_hint = self.fuse_block(torch.cat([enc_glyph, enc_pos, text_info["masked_x"].cuda()], dim=1))
+        enc_glyph = self.glyph_block(glyphs)
+        enc_pos = self.position_block(positions)
+        guided_hint = self.fuse_block(torch.cat([enc_glyph, enc_pos, text_info["masked_x"]], dim=1))
 
         return guided_hint
 
-    def encode_first_stage(self, masked_img):
-        return retrieve_latents(self.vae.encode(masked_img)) * self.vae.config.scaling_factor
-
     def check_channels(self, image):
         channels = image.shape[2] if len(image.shape) == 3 else 1
         if channels == 1:

examples/research_projects/anytext/convert_from_ckpt.py

@@ -34,13 +34,18 @@
     CLIPVisionModelWithProjection,
 )
 
-from ...models import (
+from diffusers.models import (
     AutoencoderKL,
     ControlNetModel,
     PriorTransformer,
     UNet2DConditionModel,
 )
-from ...schedulers import (
+from diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel
+from diffusers.pipelines.paint_by_example import PaintByExampleImageEncoder
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
+from diffusers.schedulers import (
     DDIMScheduler,
     DDPMScheduler,
     DPMSolverMultistepScheduler,
@@ -51,12 +56,7 @@
     PNDMScheduler,
     UnCLIPScheduler,
 )
-from ...utils import is_accelerate_available, logging
-from ..latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel
-from ..paint_by_example import PaintByExampleImageEncoder
-from ..pipeline_utils import DiffusionPipeline
-from .safety_checker import StableDiffusionSafetyChecker
-from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
+from diffusers.utils import is_accelerate_available, logging
 
 
 if is_accelerate_available():

examples/research_projects/anytext/embedding_manager.py

@@ -107,25 +107,17 @@ class EmbeddingManager(nn.Module):
     def __init__(
         self,
         embedder,
-        valid=True,
-        glyph_channels=20,
         position_channels=1,
         placeholder_string="*",
         add_pos=False,
         emb_type="ocr",
-        **kwargs,
+        use_fp16=False,
     ):
         super().__init__()
-        if hasattr(embedder, "tokenizer"):  # using Stable Diffusion's CLIP encoder
-            get_token_for_string = partial(get_clip_token_for_string, embedder.tokenizer)
-            token_dim = 768
-            if hasattr(embedder, "vit"):
-                assert emb_type == "vit"
-                self.get_vision_emb = partial(get_clip_vision_emb, embedder.vit, embedder.processor)
-            self.get_recog_emb = None
-        else:  # using LDM's BERT encoder
-            get_token_for_string = partial(get_bert_token_for_string, embedder.tknz_fn)
-            token_dim = 1280
+        get_token_for_string = partial(get_clip_token_for_string, embedder.tokenizer)
+        token_dim = 768
+        self.get_recog_emb = None
+        token_dim = 1280
         self.token_dim = token_dim
         self.emb_type = emb_type
 
@@ -134,9 +126,7 @@ def __init__(
             self.position_encoder = EncodeNet(position_channels, token_dim)
         if emb_type == "ocr":
             self.proj = nn.Sequential(zero_module(nn.Linear(40 * 64, token_dim)), nn.LayerNorm(token_dim))
-            self.proj = self.proj.to(dtype=torch.float16 if kwargs.get("use_fp16", False) else torch.float32)
-        if emb_type == "conv":
-            self.glyph_encoder = EncodeNet(glyph_channels, token_dim)
+            self.proj = self.proj.to(dtype=torch.float16 if use_fp16 else torch.float32)
 
         self.placeholder_token = get_token_for_string(placeholder_string)
 

examples/research_projects/anytext/frozen_clip_embedder_t3.py

@@ -1,6 +1,6 @@
 import torch
 from torch import nn
-from transformers import AutoProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
+from transformers import CLIPTextModel, CLIPTokenizer
 from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
 
 
@@ -16,14 +16,18 @@ class FrozenCLIPEmbedderT3(AbstractEncoder):
     """Uses the CLIP transformer encoder for text (from Hugging Face)"""
 
     def __init__(
-        self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77, freeze=True, use_vision=False
+        self,
+        version="openai/clip-vit-large-patch14",
+        device="cpu",
+        max_length=77,
+        freeze=True,
+        use_fp16=False,
     ):
         super().__init__()
         self.tokenizer = CLIPTokenizer.from_pretrained(version)
-        self.transformer = CLIPTextModel.from_pretrained(version, torch_dtype=torch.float16).to(device)
-        if use_vision:
-            self.vit = CLIPVisionModelWithProjection.from_pretrained(version)
-            self.processor = AutoProcessor.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(
+            version, use_safetensors=True, torch_dtype=torch.float16 if use_fp16 else torch.float32
+        ).to(device)
         self.device = device
         self.max_length = max_length
         if freeze:

examples/research_projects/anytext/pipeline_anytext.py

@@ -218,10 +218,15 @@ def __init__(
         feature_extractor: CLIPImageProcessor,
         image_encoder: CLIPVisionModelWithProjection = None,
         requires_safety_checker: bool = True,
+        font_path: str = "font/Arial_Unicode.ttf",
     ):
         super().__init__()
-        self.text_embedding_module = TextEmbeddingModule(use_fp16=unet.dtype == torch.float16)
-        self.auxiliary_latent_module = AuxiliaryLatentModule(vae=vae, use_fp16=unet.dtype == torch.float16)
+        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
+        )
 
         if safety_checker is None and requires_safety_checker:
             logger.warning(

examples/research_projects/anytext/recognizer.py

@@ -78,37 +78,30 @@ def crop_image(src_img, mask):
     return result
 
 
-def create_predictor(model_dir=None, model_lang="ch", is_onnx=False):
+def create_predictor(model_dir=None, model_lang="ch", device="cpu", use_fp16=False):
     model_file_path = model_dir
     if model_file_path is not None and not os.path.exists(model_file_path):
         raise ValueError("not find model file path {}".format(model_file_path))
 
-    if is_onnx:
-        import onnxruntime as ort
-
-        sess = ort.InferenceSession(
-            model_file_path, providers=["CPUExecutionProvider"]
-        )  # 'TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'
-        return sess
+    if model_lang == "ch":
+        n_class = 6625
+    elif model_lang == "en":
+        n_class = 97
     else:
-        if model_lang == "ch":
-            n_class = 6625
-        elif model_lang == "en":
-            n_class = 97
-        else:
-            raise ValueError(f"Unsupported OCR recog model_lang: {model_lang}")
-        rec_config = edict(
-            in_channels=3,
-            backbone=edict(type="MobileNetV1Enhance", scale=0.5, last_conv_stride=[1, 2], last_pool_type="avg"),
-            neck=edict(type="SequenceEncoder", encoder_type="svtr", dims=64, depth=2, hidden_dims=120, use_guide=True),
-            head=edict(type="CTCHead", fc_decay=0.00001, out_channels=n_class, return_feats=True),
+        raise ValueError(f"Unsupported OCR recog model_lang: {model_lang}")
+    rec_config = edict(
+        in_channels=3,
+        backbone=edict(type="MobileNetV1Enhance", scale=0.5, last_conv_stride=[1, 2], last_pool_type="avg"),
+        neck=edict(type="SequenceEncoder", encoder_type="svtr", dims=64, depth=2, hidden_dims=120, use_guide=True),
+        head=edict(type="CTCHead", fc_decay=0.00001, out_channels=n_class, return_feats=True),
+    )
+
+    rec_model = RecModel(rec_config)
+    if model_file_path is not None:
+        rec_model.load_state_dict(load_file(model_file_path, device=device)).to(
+            dtype=torch.float16 if use_fp16 else torch.float32
         )
-
-        rec_model = RecModel(rec_config)
-        if model_file_path is not None:
-            rec_model.load_state_dict(load_file(model_file_path))
-            rec_model.eval()
-        return rec_model.eval()
+    return rec_model
 
 
 def _check_image_file(path):

examples/research_projects/anytext/text_embedding_module.py

@@ -16,36 +16,25 @@
 
 
 class TextEmbeddingModule(nn.Module):
-    def __init__(self, use_fp16):
+    def __init__(self, font_path, use_fp16=False, device="cpu"):
         super().__init__()
         self.use_fp16 = use_fp16
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
         # TODO: Learn if the recommended font file is free to use
-        self.font = ImageFont.truetype("Arial_Unicode.ttf", 60)
-        self.frozen_CLIP_embedder_t3 = FrozenCLIPEmbedderT3(device=self.device)
-        self.embedding_manager_config = {
-            "valid": True,
-            "emb_type": "ocr",
-            "glyph_channels": 1,
-            "position_channels": 1,
-            "add_pos": False,
-            "placeholder_string": "*",
-            "use_fp16": self.use_fp16,
-        }
-        self.embedding_manager = EmbeddingManager(self.frozen_CLIP_embedder_t3, **self.embedding_manager_config)
+        self.font = ImageFont.truetype(font_path, 60)
+        self.frozen_CLIP_embedder_t3 = FrozenCLIPEmbedderT3(device=self.device, use_fp16=self.use_fp16)
+        self.embedding_manager = EmbeddingManager(self.frozen_CLIP_embedder_t3, use_fp16=self.use_fp16)
         # TODO: Understand the reason of param.requires_grad = True
         for param in self.embedding_manager.embedding_parameters():
             param.requires_grad = True
-        rec_model_dir = "ppv3_rec.safetensors"
-        self.text_predictor = create_predictor(rec_model_dir).eval()
+        rec_model_dir = "OCR/ppv3_rec.safetensors"
+        self.text_predictor = create_predictor(rec_model_dir, device=self.device, use_fp16=self.use_fp16)
         args = {}
         args["rec_image_shape"] = "3, 48, 320"
         args["rec_batch_num"] = 6
-        args["rec_char_dict_path"] = "ppocr_keys_v1.txt"
-        args["use_fp16"] = True
-        self.cn_recognizer = TextRecognizer(
-            args, self.text_predictor.to(dtype=torch.float16 if use_fp16 else torch.float32)
-        )
+        args["rec_char_dict_path"] = "OCR/ppocr_keys_v1.txt"
+        args["use_fp16"] = self.use_fp16
+        self.cn_recognizer = TextRecognizer(args, self.text_predictor, device=self.device, use_fp16=self.use_fp16)
         for param in self.text_predictor.parameters():
             param.requires_grad = False
         self.embedding_manager.recog = self.cn_recognizer