Update cog_sdxl_dataset_and_utils.py

cog_sdxl_dataset_and_utils.py

@@ -1,4 +1,4 @@
-# dataset_and_utils.py - Optimized and Improved Version
+# 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
 
@@ -15,29 +15,28 @@ from torch.utils.data import Dataset
 from transformers import AutoTokenizer, PretrainedConfig
 
 
-def prepare_image(image: PIL.Image.Image, width: int = 512, height: int = 512) -> torch.Tensor:
-    """
-    Prepares an image for model input by resizing and normalizing it.
-    """
-    image = image.resize((width, height), resample=Image.BICUBIC, reducing_gap=1)
-    arr = np.array(image.convert("RGB"), dtype=np.float32) / 127.5 - 1
-    return torch.from_numpy(np.transpose(arr, (2, 0, 1))).unsqueeze(0)
+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(mask: PIL.Image.Image, width: int = 512, height: int = 512) -> torch.Tensor:
-    """
-    Prepares a mask image for model input by resizing and normalizing it.
-    """
-    mask = mask.resize((width, height), resample=Image.BICUBIC, reducing_gap=1)
-    arr = np.array(mask.convert("L"), dtype=np.float32) / 255.0
-    return torch.from_numpy(np.expand_dims(arr, 0)).unsqueeze(0)
+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 TokenEmbeddingsHandler:
-    def __init__(self, text_encoders, tokenizers):
-        self.text_encoders = text_encoders
-        self.tokenizers = tokenizers
-
 class PreprocessedDataset(Dataset):
     def __init__(
         self,
@@ -51,133 +50,373 @@ class PreprocessedDataset(Dataset):
         size: int = 512,
         text_dropout: float = 0.0,
         scale_vae_latents: bool = True,
-        substitute_caption_map: Dict[str, str] = None,
+        substitute_caption_map: Dict[str, str] = {},
     ):
-        """
-        Dataset class that pre-processes images, masks, and text data for training.
-        """
         super().__init__()
+
         self.data = pd.read_csv(csv_path)
-        self.size = size
-        self.scale_vae_latents = scale_vae_latents
-        self.text_dropout = text_dropout
         self.csv_path = csv_path
-        self.tokenizer_1 = tokenizer_1
-        self.tokenizer_2 = tokenizer_2
-        self.vae_encoder = vae_encoder
-        self.do_cache = do_cache
 
-        self.caption = self.data["caption"].str.lower()
-
-        if substitute_caption_map:
-            for key, value in substitute_caption_map.items():
-                self.caption = self.caption.str.replace(key.lower(), value)
+        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"]
-        self.mask_path = self.data["mask_path"] if "mask_path" in self.data.columns else None
 
-        if text_encoder_1:
+        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
-            raise NotImplementedError("Preprocessing for text encoder is not implemented yet.")
-        else:
-            self.return_text_embeddings = False
+            assert (
+                NotImplementedError
+            ), "Preprocessing Text Encoder is not implemented yet"
 
-        if self.do_cache:
+        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 = []
-            print("Caching dataset...")
+
+            self.do_cache = True
+
+            print("Captions to train on: ")
             for idx in range(len(self.data)):
                 token, vae_latent, mask = self._process(idx)
-                self.tokens_tuple.append(token)
                 self.vae_latents.append(vae_latent)
+                self.tokens_tuple.append(token)
                 self.masks.append(mask)
-            del self.vae_encoder  # Free up memory
+
+            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]:
-        """
-        Internal function to process images, text, and masks for a given index.
-        """
-        image_path = os.path.join(os.path.dirname(self.csv_path), self.image_path[idx])
-        image = prepare_image(Image.open(image_path).convert("RGB"), self.size, self.size).to(
+    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]
-        ti1 = self.tokenizer_1(caption, padding="max_length", max_length=77, truncation=True, return_tensors="pt").input_ids
-        ti2 = self.tokenizer_2(caption, padding="max_length", max_length=77, truncation=True, return_tensors="pt").input_ids
+
+        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 *= self.vae_encoder.config.scaling_factor
+            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)
+            mask = torch.ones_like(
+                vae_latent, dtype=self.vae_encoder.dtype, device=self.vae_encoder.device
+            )
+
         else:
-            mask_path = os.path.join(os.path.dirname(self.csv_path), self.mask_path[idx])
-            mask = prepare_mask(Image.open(mask_path), self.size, self.size).to(
+            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 = 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 mask.shape == vae_latent.shape, "Mask and latent dimensions must match."
+        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 __getitem__(self, idx: int) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
-        return self.atidx(idx)
+    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 atidx(self, idx: int) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
-        return self._process(idx) if not self.do_cache else (self.tokens_tuple[idx], self.vae_latents[idx], self.masks[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"):
-    """
-    Dynamically imports a model class based on configuration.
-    """
-    config = PretrainedConfig.from_pretrained(pretrained_model_name_or_path, subfolder=subfolder, revision=revision)
-    model_class = config.architectures[0]
+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"Unsupported model class: {model_class}")
+        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]:
     """
-    Loads required models from a given pretrained path.
+    Returns:
+        a state dict containing just the attention processor parameters.
     """
-    tokenizer_1 = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer", revision=revision, use_fast=False)
-    tokenizer_2 = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer_2", revision=revision, use_fast=False)
+    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
 
-    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+    return attn_processors_state_dict
 
-    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_1 = text_encoder_cls_one.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder", revision=revision)
-    text_encoder_2 = text_encoder_cls_two.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder_2", revision=revision)
+class TokenEmbeddingsHandler:
+    def __init__(self, text_encoders, tokenizers):
+        self.text_encoders = text_encoders
+        self.tokenizers = tokenizers
 
-    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)
+        self.train_ids: Optional[torch.Tensor] = None
+        self.inserting_toks: Optional[List[str]] = None
+        self.embeddings_settings = {}
 
-    for model in [vae, text_encoder_1, text_encoder_2]:
-        model.requires_grad_(False)
-        model.to(device, dtype=weight_dtype)
+    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."
 
-    unet.to(device, dtype=weight_dtype)
+            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]
 
-    return tokenizer_1, tokenizer_2, noise_scheduler, text_encoder_1, text_encoder_2, vae, unet
-    
+                loaded_embeddings = f.get_tensor(f"text_encoders_{idx}")
+                self._load_embeddings(loaded_embeddings, tokenizer, text_encoder)