Update cog_sdxl_dataset_and_utils.py

cog_sdxl_dataset_and_utils.py

@@ -1,4 +1,4 @@
-# dataset_and_utils.py file taken from https://github.com/replicate/cog-sdxl/blob/main/dataset_and_utils.py
+# dataset_and_utils.py - Optimized and Improved Version
 import os
 from typing import Dict, List, Optional, Tuple
 
@@ -15,26 +15,22 @@ 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_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_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
+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)
 
 
 class PreprocessedDataset(Dataset):
@@ -50,373 +46,132 @@ class PreprocessedDataset(Dataset):
         size: int = 512,
         text_dropout: float = 0.0,
         scale_vae_latents: bool = True,
-        substitute_caption_map: Dict[str, str] = {},
+        substitute_caption_map: Dict[str, str] = None,
     ):
+        """
+        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"]
-        # 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.caption = self.data["caption"].str.lower()
 
-        self.image_path = self.data["image_path"]
+        if substitute_caption_map:
+            for key, value in substitute_caption_map.items():
+                self.caption = self.caption.str.replace(key.lower(), value)
 
-        if "mask_path" not in self.data.columns:
-            self.mask_path = None
-        else:
-            self.mask_path = self.data["mask_path"]
+        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 is None:
-            self.return_text_embeddings = False
-        else:
+        if text_encoder_1:
             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
+            raise NotImplementedError("Preprocessing for text encoder is not implemented yet.")
+        else:
+            self.return_text_embeddings = False
 
-        if do_cache:
+        if self.do_cache:
             self.vae_latents = []
             self.tokens_tuple = []
             self.masks = []
-
-            self.do_cache = True
-
-            print("Captions to train on: ")
+            print("Caching dataset...")
             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.vae_latents.append(vae_latent)
                 self.masks.append(mask)
-
-            del self.vae_encoder
-
-        else:
-            self.do_cache = False
+            del self.vae_encoder  # Free up memory
 
     @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(
+    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(
             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
+        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
 
         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
+            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 = 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(
+            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(
                 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 len(mask.shape) == 4 and len(vae_latent.shape) == 4
+        assert mask.shape == vae_latent.shape, "Mask and latent dimensions must match."
 
         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]:
+        return self.atidx(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 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 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]
+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]
 
     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.")
+        raise ValueError(f"Unsupported model class: {model_class}")
 
 
 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.
+    Loads required models from a given pretrained path.
     """
-    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
+    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)
 
-        self.train_ids: Optional[torch.Tensor] = None
-        self.inserting_toks: Optional[List[str]] = None
-        self.embeddings_settings = {}
+    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
 
-    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."
+    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")
 
-            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))
+    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)
 
-            self.train_ids = tokenizer.convert_tokens_to_ids(self.inserting_toks)
+    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)
 
-            # random initialization of new tokens
+    for model in [vae, text_encoder_1, text_encoder_2]:
+        model.requires_grad_(False)
+        model.to(device, dtype=weight_dtype)
 
-            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]
+    unet.to(device, dtype=weight_dtype)
 
-                loaded_embeddings = f.get_tensor(f"text_encoders_{idx}")
-                self._load_embeddings(loaded_embeddings, tokenizer, text_encoder)
+    return tokenizer_1, tokenizer_2, noise_scheduler, text_encoder_1, text_encoder_2, vae, unet