Spaces:
Running
on
A100
Running
on
A100
| # 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 | |
| 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) | |
| def dtype(self): | |
| return self.text_encoders[0].dtype | |
| 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) | |
| 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) |