Merge pull request #60 from okotaku/feat/pre-compute-embs

[Feature] Pre compute Text Embeddings

Author: okotaku

.devcontainer/devcontainer.json

@@ -25,7 +25,8 @@
           "ms-vscode-remote.remote-ssh-edit",
           "ms-vscode.remote-explorer",
           "wayou.vscode-todo-highlight",
-          "Gruntfuggly.todo-tree"
+          "Gruntfuggly.todo-tree",
+          "streetsidesoftware.code-spell-checker"
         ]
       }
     }

configs/_base_/datasets/pokemon_blip_xl_pre_compute.py

@@ -0,0 +1,39 @@
+train_pipeline = [
+    dict(type='SaveImageShape'),
+    dict(type='torchvision/Resize', size=1024, interpolation='bilinear'),
+    dict(type='RandomCrop', size=1024),
+    dict(type='RandomHorizontalFlip', p=0.5),
+    dict(type='ComputeTimeIds'),
+    dict(type='torchvision/ToTensor'),
+    dict(type='torchvision/Normalize', mean=[0.5], std=[0.5]),
+    dict(
+        type='PackInputs',
+        input_keys=[
+            'img', 'time_ids', 'prompt_embeds', 'pooled_prompt_embeds'
+        ]),
+]
+train_dataloader = dict(
+    batch_size=2,
+    num_workers=2,
+    dataset=dict(
+        type='HFDatasetPreComputeEmbs',
+        dataset='lambdalabs/pokemon-blip-captions',
+        text_hasher='text_pokemon_blip',
+        model='stabilityai/stable-diffusion-xl-base-1.0',
+        pipeline=train_pipeline),
+    sampler=dict(type='DefaultSampler', shuffle=True),
+)
+
+val_dataloader = None
+val_evaluator = None
+test_dataloader = val_dataloader
+test_evaluator = val_evaluator
+
+custom_hooks = [
+    dict(
+        type='VisualizationHook',
+        prompt=['yoda pokemon'] * 4,
+        height=1024,
+        width=1024),
+    dict(type='SDCheckpointHook')
+]

configs/stable_diffusion_xl/README.md

@@ -76,3 +76,7 @@ You can see more details on [`docs/source/run_guides/run_xl.md`](../../docs/sour
 #### stable_diffusion_xl_pokemon_blip
 
 ![example1](https://github.com/okotaku/diffengine/assets/24734142/dd04fb22-64fb-4c4f-8164-b8391d94abab)
+
+#### stable_diffusion_xl_pokemon_blip_pre_compute
+
+![example2](https://github.com/okotaku/diffengine/assets/24734142/5da59a56-ce36-48cc-b113-007f8b9faeba)

configs/stable_diffusion_xl/stable_diffusion_xl_pokemon_blip_pre_compute.py

@@ -0,0 +1,12 @@
+_base_ = [
+    '../_base_/models/stable_diffusion_xl.py',
+    '../_base_/datasets/pokemon_blip_xl_pre_compute.py',
+    '../_base_/schedules/stable_diffusion_xl_50e.py',
+    '../_base_/default_runtime.py'
+]
+
+model = dict(pre_compute_text_embeddings=True)
+
+train_dataloader = dict(batch_size=1)
+
+optim_wrapper_cfg = dict(accumulative_counts=4)  # update every four times

diffengine/datasets/__init__.py

@@ -1,7 +1,10 @@
 from .hf_controlnet_datasets import HFControlNetDataset
-from .hf_datasets import HFDataset
+from .hf_datasets import HFDataset, HFDatasetPreComputeEmbs
 from .hf_dreambooth_datasets import HFDreamBoothDataset
 from .samplers import *  # noqa: F401, F403
 from .transforms import *  # noqa: F401, F403
 
-__all__ = ['HFDataset', 'HFDreamBoothDataset', 'HFControlNetDataset']
+__all__ = [
+    'HFDataset', 'HFDreamBoothDataset', 'HFControlNetDataset',
+    'HFDatasetPreComputeEmbs'
+]

diffengine/datasets/hf_datasets.py

@@ -1,14 +1,22 @@
+import functools
+import gc
 import os
 import random
 from pathlib import Path
 from typing import Optional, Sequence
 
 import numpy as np
+import torch
 from datasets import load_dataset
+from datasets.fingerprint import Hasher
 from mmengine.dataset.base_dataset import Compose
 from PIL import Image
 from torch.utils.data import Dataset
+from transformers import AutoTokenizer
 
+from diffengine.datasets.utils import encode_prompt_sdxl
+from diffengine.models.editors.stable_diffusion_xl.stable_diffusion_xl import \
+    import_model_class_from_model_name_or_path
 from diffengine.registry import DATASETS
 
 Image.MAX_IMAGE_PIXELS = 1000000000
@@ -88,3 +96,84 @@ def __getitem__(self, idx: int) -> dict:
         result = self.pipeline(result)
 
         return result
+
+
+@DATASETS.register_module()
+class HFDatasetPreComputeEmbs(HFDataset):
+    """Dataset for huggingface datasets.
+
+    The difference from HFDataset is
+        1. pre-compute Text Encoder embeddings to save memory.
+
+    Args:
+        model (str): pretrained model name of stable diffusion xl.
+            Defaults to 'stabilityai/stable-diffusion-xl-base-1.0'.
+        text_hasher (str): Text embeddings hasher name. Defaults to 'text'.
+        device (str): Device used to compute embeddings. Defaults to 'cuda'.
+        proportion_empty_prompts (float): The probabilities to replace empty
+            text. Defaults to 0.9.
+    """
+
+    def __init__(self,
+                 *args,
+                 model: str = 'stabilityai/stable-diffusion-xl-base-1.0',
+                 text_hasher: str = 'text',
+                 device: str = 'cuda',
+                 proportion_empty_prompts: float = 0.0,
+                 **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+        tokenizer_one = AutoTokenizer.from_pretrained(
+            model, subfolder='tokenizer', use_fast=False)
+        tokenizer_two = AutoTokenizer.from_pretrained(
+            model, subfolder='tokenizer_2', use_fast=False)
+
+        text_encoder_cls_one = import_model_class_from_model_name_or_path(
+            model)
+        text_encoder_cls_two = import_model_class_from_model_name_or_path(
+            model, subfolder='text_encoder_2')
+        text_encoder_one = text_encoder_cls_one.from_pretrained(
+            model, subfolder='text_encoder').to(device)
+        text_encoder_two = text_encoder_cls_two.from_pretrained(
+            model, subfolder='text_encoder_2').to(device)
+
+        new_fingerprint = Hasher.hash(text_hasher)
+        compute_embeddings_fn = functools.partial(
+            encode_prompt_sdxl,
+            text_encoders=[text_encoder_one, text_encoder_two],
+            tokenizers=[tokenizer_one, tokenizer_two],
+            proportion_empty_prompts=proportion_empty_prompts,
+            caption_column=self.caption_column,
+        )
+        self.dataset = self.dataset.map(
+            compute_embeddings_fn,
+            batched=True,
+            new_fingerprint=new_fingerprint)
+
+        del text_encoder_one, text_encoder_two, tokenizer_one, tokenizer_two
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def __getitem__(self, idx: int) -> dict:
+        """Get the idx-th image and data information of dataset after
+        ``self.train_transforms`.
+
+        Args:
+            idx (int): The index of self.data_list.
+
+        Returns:
+            dict: The idx-th image and data information of dataset after
+            ``self.train_transforms``.
+        """
+        data_info = self.dataset[idx]
+        image = data_info[self.image_column]
+        if type(image) == str:
+            image = Image.open(os.path.join(self.dataset_name, image))
+        image = image.convert('RGB')
+        result = dict(
+            img=image,
+            prompt_embeds=data_info['prompt_embeds'],
+            pooled_prompt_embeds=data_info['pooled_prompt_embeds'])
+        result = self.pipeline(result)
+
+        return result

diffengine/datasets/utils.py

@@ -0,0 +1,56 @@
+import random
+from typing import Dict
+
+import numpy as np
+import torch
+
+
+def encode_prompt_sdxl(batch,
+                       text_encoders,
+                       tokenizers,
+                       proportion_empty_prompts,
+                       caption_column,
+                       is_train: bool = True) -> Dict[str, torch.Tensor]:
+    # Adapted from pipelines.StableDiffusionXLPipeline.encode_prompt
+    prompt_embeds_list = []
+    prompt_batch = batch[caption_column]
+
+    captions = []
+    for caption in prompt_batch:
+        if random.random() < proportion_empty_prompts:
+            captions.append('')
+        elif isinstance(caption, str):
+            captions.append(caption)
+        elif isinstance(caption, (list, np.ndarray)):
+            # take a random caption if there are multiple
+            captions.append(random.choice(caption) if is_train else caption[0])
+
+    with torch.no_grad():
+        for tokenizer, text_encoder in zip(tokenizers, text_encoders):
+            text_inputs = tokenizer(
+                captions,
+                padding='max_length',
+                max_length=tokenizer.model_max_length,
+                truncation=True,
+                return_tensors='pt',
+            )
+            text_input_ids = text_inputs.input_ids
+            prompt_embeds = text_encoder(
+                text_input_ids.to(text_encoder.device),
+                output_hidden_states=True,
+            )
+
+            # We are only ALWAYS interested in the pooled output of the final
+            # text encoder
+            pooled_prompt_embeds = prompt_embeds[0]
+            prompt_embeds = prompt_embeds.hidden_states[-2]
+            bs_embed, seq_len, _ = prompt_embeds.shape
+            prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
+            prompt_embeds_list.append(prompt_embeds)
+
+    prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
+    pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
+    return {
+        'prompt_embeds': prompt_embeds.cpu(),
+        'pooled_prompt_embeds': pooled_prompt_embeds.cpu()
+    }

diffengine/models/editors/stable_diffusion_xl/sdxl_data_preprocessor.py

@@ -34,4 +34,11 @@ def forward(self, data: dict, training: bool = False) -> Union[dict, list]:
 
         data['inputs']['img'] = torch.stack(data['inputs']['img'])
         data['inputs']['time_ids'] = torch.stack(data['inputs']['time_ids'])
+        # pre-compute text embeddings
+        if 'prompt_embeds' in data['inputs']:
+            data['inputs']['prompt_embeds'] = torch.stack(
+                data['inputs']['prompt_embeds'])
+        if 'pooled_prompt_embeds' in data['inputs']:
+            data['inputs']['pooled_prompt_embeds'] = torch.stack(
+                data['inputs']['pooled_prompt_embeds'])
         return super().forward(data)  # type: ignore