Create cog_sdxl_dataset_and_utils.py

cog_sdxl_dataset_and_utils.py

@@ -0,0 +1,422 @@
+# dataset_and_utils.py file taken from https://github.com/replicate/cog-sdxl/blob/main/dataset_and_utils.py
+import os
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import PIL
+import torch
+import torch.utils.checkpoint
+from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
+from PIL import Image
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer, PretrainedConfig
+
+
+def prepare_image(
+    pil_image: PIL.Image.Image, w: int = 512, h: int = 512
+) -> torch.Tensor:
+    pil_image = pil_image.resize((w, h), resample=Image.BICUBIC, reducing_gap=1)
+    arr = np.array(pil_image.convert("RGB"))
+    arr = arr.astype(np.float32) / 127.5 - 1
+    arr = np.transpose(arr, [2, 0, 1])
+    image = torch.from_numpy(arr).unsqueeze(0)
+    return image
+
+
+def prepare_mask(
+    pil_image: PIL.Image.Image, w: int = 512, h: int = 512
+) -> torch.Tensor:
+    pil_image = pil_image.resize((w, h), resample=Image.BICUBIC, reducing_gap=1)
+    arr = np.array(pil_image.convert("L"))
+    arr = arr.astype(np.float32) / 255.0
+    arr = np.expand_dims(arr, 0)
+    image = torch.from_numpy(arr).unsqueeze(0)
+    return image
+
+
+class PreprocessedDataset(Dataset):
+    def __init__(
+        self,
+        csv_path: str,
+        tokenizer_1,
+        tokenizer_2,
+        vae_encoder,
+        text_encoder_1=None,
+        text_encoder_2=None,
+        do_cache: bool = False,
+        size: int = 512,
+        text_dropout: float = 0.0,
+        scale_vae_latents: bool = True,
+        substitute_caption_map: Dict[str, str] = {},
+    ):
+        super().__init__()
+
+        self.data = pd.read_csv(csv_path)
+        self.csv_path = csv_path
+
+        self.caption = self.data["caption"]
+        # make it lowercase
+        self.caption = self.caption.str.lower()
+        for key, value in substitute_caption_map.items():
+            self.caption = self.caption.str.replace(key.lower(), value)
+
+        self.image_path = self.data["image_path"]
+
+        if "mask_path" not in self.data.columns:
+            self.mask_path = None
+        else:
+            self.mask_path = self.data["mask_path"]
+
+        if text_encoder_1 is None:
+            self.return_text_embeddings = False
+        else:
+            self.text_encoder_1 = text_encoder_1
+            self.text_encoder_2 = text_encoder_2
+            self.return_text_embeddings = True
+            assert (
+                NotImplementedError
+            ), "Preprocessing Text Encoder is not implemented yet"
+
+        self.tokenizer_1 = tokenizer_1
+        self.tokenizer_2 = tokenizer_2
+
+        self.vae_encoder = vae_encoder
+        self.scale_vae_latents = scale_vae_latents
+        self.text_dropout = text_dropout
+
+        self.size = size
+
+        if do_cache:
+            self.vae_latents = []
+            self.tokens_tuple = []
+            self.masks = []
+
+            self.do_cache = True
+
+            print("Captions to train on: ")
+            for idx in range(len(self.data)):
+                token, vae_latent, mask = self._process(idx)
+                self.vae_latents.append(vae_latent)
+                self.tokens_tuple.append(token)
+                self.masks.append(mask)
+
+            del self.vae_encoder
+
+        else:
+            self.do_cache = False
+
+    @torch.no_grad()
+    def _process(
+        self, idx: int
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
+        image_path = self.image_path[idx]
+        image_path = os.path.join(os.path.dirname(self.csv_path), image_path)
+
+        image = PIL.Image.open(image_path).convert("RGB")
+        image = prepare_image(image, self.size, self.size).to(
+            dtype=self.vae_encoder.dtype, device=self.vae_encoder.device
+        )
+
+        caption = self.caption[idx]
+
+        print(caption)
+
+        # tokenizer_1
+        ti1 = self.tokenizer_1(
+            caption,
+            padding="max_length",
+            max_length=77,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        ).input_ids
+
+        ti2 = self.tokenizer_2(
+            caption,
+            padding="max_length",
+            max_length=77,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        ).input_ids
+
+        vae_latent = self.vae_encoder.encode(image).latent_dist.sample()
+
+        if self.scale_vae_latents:
+            vae_latent = vae_latent * self.vae_encoder.config.scaling_factor
+
+        if self.mask_path is None:
+            mask = torch.ones_like(
+                vae_latent, dtype=self.vae_encoder.dtype, device=self.vae_encoder.device
+            )
+
+        else:
+            mask_path = self.mask_path[idx]
+            mask_path = os.path.join(os.path.dirname(self.csv_path), mask_path)
+
+            mask = PIL.Image.open(mask_path)
+            mask = prepare_mask(mask, self.size, self.size).to(
+                dtype=self.vae_encoder.dtype, device=self.vae_encoder.device
+            )
+
+            mask = torch.nn.functional.interpolate(
+                mask, size=(vae_latent.shape[-2], vae_latent.shape[-1]), mode="nearest"
+            )
+            mask = mask.repeat(1, vae_latent.shape[1], 1, 1)
+
+        assert len(mask.shape) == 4 and len(vae_latent.shape) == 4
+
+        return (ti1.squeeze(), ti2.squeeze()), vae_latent.squeeze(), mask.squeeze()
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def atidx(
+        self, idx: int
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
+        if self.do_cache:
+            return self.tokens_tuple[idx], self.vae_latents[idx], self.masks[idx]
+        else:
+            return self._process(idx)
+
+    def __getitem__(
+        self, idx: int
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
+        token, vae_latent, mask = self.atidx(idx)
+        return token, vae_latent, mask
+
+
+def import_model_class_from_model_name_or_path(
+    pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder"
+):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path, subfolder=subfolder, revision=revision
+    )
+    model_class = text_encoder_config.architectures[0]
+
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+
+        return CLIPTextModel
+    elif model_class == "CLIPTextModelWithProjection":
+        from transformers import CLIPTextModelWithProjection
+
+        return CLIPTextModelWithProjection
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+
+
+def load_models(pretrained_model_name_or_path, revision, device, weight_dtype):
+    tokenizer_one = AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="tokenizer",
+        revision=revision,
+        use_fast=False,
+    )
+    tokenizer_two = AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="tokenizer_2",
+        revision=revision,
+        use_fast=False,
+    )
+
+    # Load scheduler and models
+    noise_scheduler = DDPMScheduler.from_pretrained(
+        pretrained_model_name_or_path, subfolder="scheduler"
+    )
+    # import correct text encoder classes
+    text_encoder_cls_one = import_model_class_from_model_name_or_path(
+        pretrained_model_name_or_path, revision
+    )
+    text_encoder_cls_two = import_model_class_from_model_name_or_path(
+        pretrained_model_name_or_path, revision, subfolder="text_encoder_2"
+    )
+    text_encoder_one = text_encoder_cls_one.from_pretrained(
+        pretrained_model_name_or_path, subfolder="text_encoder", revision=revision
+    )
+    text_encoder_two = text_encoder_cls_two.from_pretrained(
+        pretrained_model_name_or_path, subfolder="text_encoder_2", revision=revision
+    )
+
+    vae = AutoencoderKL.from_pretrained(
+        pretrained_model_name_or_path, subfolder="vae", revision=revision
+    )
+    unet = UNet2DConditionModel.from_pretrained(
+        pretrained_model_name_or_path, subfolder="unet", revision=revision
+    )
+
+    vae.requires_grad_(False)
+    text_encoder_one.requires_grad_(False)
+    text_encoder_two.requires_grad_(False)
+
+    unet.to(device, dtype=weight_dtype)
+    vae.to(device, dtype=torch.float32)
+    text_encoder_one.to(device, dtype=weight_dtype)
+    text_encoder_two.to(device, dtype=weight_dtype)
+
+    return (
+        tokenizer_one,
+        tokenizer_two,
+        noise_scheduler,
+        text_encoder_one,
+        text_encoder_two,
+        vae,
+        unet,
+    )
+
+
+def unet_attn_processors_state_dict(unet) -> Dict[str, torch.tensor]:
+    """
+    Returns:
+        a state dict containing just the attention processor parameters.
+    """
+    attn_processors = unet.attn_processors
+
+    attn_processors_state_dict = {}
+
+    for attn_processor_key, attn_processor in attn_processors.items():
+        for parameter_key, parameter in attn_processor.state_dict().items():
+            attn_processors_state_dict[
+                f"{attn_processor_key}.{parameter_key}"
+            ] = parameter
+
+    return attn_processors_state_dict
+
+
+class TokenEmbeddingsHandler:
+    def __init__(self, text_encoders, tokenizers):
+        self.text_encoders = text_encoders
+        self.tokenizers = tokenizers
+
+        self.train_ids: Optional[torch.Tensor] = None
+        self.inserting_toks: Optional[List[str]] = None
+        self.embeddings_settings = {}
+
+    def initialize_new_tokens(self, inserting_toks: List[str]):
+        idx = 0
+        for tokenizer, text_encoder in zip(self.tokenizers, self.text_encoders):
+            assert isinstance(
+                inserting_toks, list
+            ), "inserting_toks should be a list of strings."
+            assert all(
+                isinstance(tok, str) for tok in inserting_toks
+            ), "All elements in inserting_toks should be strings."
+
+            self.inserting_toks = inserting_toks
+            special_tokens_dict = {"additional_special_tokens": self.inserting_toks}
+            tokenizer.add_special_tokens(special_tokens_dict)
+            text_encoder.resize_token_embeddings(len(tokenizer))
+
+            self.train_ids = tokenizer.convert_tokens_to_ids(self.inserting_toks)
+
+            # random initialization of new tokens
+
+            std_token_embedding = (
+                text_encoder.text_model.embeddings.token_embedding.weight.data.std()
+            )
+
+            print(f"{idx} text encodedr's std_token_embedding: {std_token_embedding}")
+
+            text_encoder.text_model.embeddings.token_embedding.weight.data[
+                self.train_ids
+            ] = (
+                torch.randn(
+                    len(self.train_ids), text_encoder.text_model.config.hidden_size
+                )
+                .to(device=self.device)
+                .to(dtype=self.dtype)
+                * std_token_embedding
+            )
+            self.embeddings_settings[
+                f"original_embeddings_{idx}"
+            ] = text_encoder.text_model.embeddings.token_embedding.weight.data.clone()
+            self.embeddings_settings[f"std_token_embedding_{idx}"] = std_token_embedding
+
+            inu = torch.ones((len(tokenizer),), dtype=torch.bool)
+            inu[self.train_ids] = False
+
+            self.embeddings_settings[f"index_no_updates_{idx}"] = inu
+
+            print(self.embeddings_settings[f"index_no_updates_{idx}"].shape)
+
+            idx += 1
+
+    def save_embeddings(self, file_path: str):
+        assert (
+            self.train_ids is not None
+        ), "Initialize new tokens before saving embeddings."
+        tensors = {}
+        for idx, text_encoder in enumerate(self.text_encoders):
+            assert text_encoder.text_model.embeddings.token_embedding.weight.data.shape[
+                0
+            ] == len(self.tokenizers[0]), "Tokenizers should be the same."
+            new_token_embeddings = (
+                text_encoder.text_model.embeddings.token_embedding.weight.data[
+                    self.train_ids
+                ]
+            )
+            tensors[f"text_encoders_{idx}"] = new_token_embeddings
+
+        save_file(tensors, file_path)
+
+    @property
+    def dtype(self):
+        return self.text_encoders[0].dtype
+
+    @property
+    def device(self):
+        return self.text_encoders[0].device
+
+    def _load_embeddings(self, loaded_embeddings, tokenizer, text_encoder):
+        # Assuming new tokens are of the format <s_i>
+        self.inserting_toks = [f"<s{i}>" for i in range(loaded_embeddings.shape[0])]
+        special_tokens_dict = {"additional_special_tokens": self.inserting_toks}
+        tokenizer.add_special_tokens(special_tokens_dict)
+        text_encoder.resize_token_embeddings(len(tokenizer))
+
+        self.train_ids = tokenizer.convert_tokens_to_ids(self.inserting_toks)
+        assert self.train_ids is not None, "New tokens could not be converted to IDs."
+        text_encoder.text_model.embeddings.token_embedding.weight.data[
+            self.train_ids
+        ] = loaded_embeddings.to(device=self.device).to(dtype=self.dtype)
+
+    @torch.no_grad()
+    def retract_embeddings(self):
+        for idx, text_encoder in enumerate(self.text_encoders):
+            index_no_updates = self.embeddings_settings[f"index_no_updates_{idx}"]
+            text_encoder.text_model.embeddings.token_embedding.weight.data[
+                index_no_updates
+            ] = (
+                self.embeddings_settings[f"original_embeddings_{idx}"][index_no_updates]
+                .to(device=text_encoder.device)
+                .to(dtype=text_encoder.dtype)
+            )
+
+            # for the parts that were updated, we need to normalize them
+            # to have the same std as before
+            std_token_embedding = self.embeddings_settings[f"std_token_embedding_{idx}"]
+
+            index_updates = ~index_no_updates
+            new_embeddings = (
+                text_encoder.text_model.embeddings.token_embedding.weight.data[
+                    index_updates
+                ]
+            )
+            off_ratio = std_token_embedding / new_embeddings.std()
+
+            new_embeddings = new_embeddings * (off_ratio**0.1)
+            text_encoder.text_model.embeddings.token_embedding.weight.data[
+                index_updates
+            ] = new_embeddings
+
+    def load_embeddings(self, file_path: str):
+        with safe_open(file_path, framework="pt", device=self.device.type) as f:
+            for idx in range(len(self.text_encoders)):
+                text_encoder = self.text_encoders[idx]
+                tokenizer = self.tokenizers[idx]
+
+                loaded_embeddings = f.get_tensor(f"text_encoders_{idx}")
+                self._load_embeddings(loaded_embeddings, tokenizer, text_encoder)