[Feature] Support Textual Inversion (#1822)

* support dataset for textual inversion

* support trainer for textual inversin

* solve conflict

* add tomesd to textual inversion

* add noise offset arg in textual inversion

* add config and readme

* remove reversion from init

* remove reversion dataset

* revise readme as comment

configs/textual_inversion/README.md

@@ -0,0 +1,109 @@
+# Textual Inversion (2022)
+
+> [An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion](https://arxiv.org/abs/2208.01618)
+
+> **Task**: Text2Image
+
+<!-- [ALGORITHM] -->
+
+## Abstract
+
+<!-- [ABSTRACT] -->
+
+Text-to-image models offer unprecedented freedom to guide creation through natural language. Yet, it is unclear how such freedom can be exercised to generate images of specific unique concepts, modify their appearance, or compose them in new roles and novel scenes. In other words, we ask: how can we use language-guided models to turn our cat into a painting, or imagine a new product based on our favorite toy? Here we present a simple approach that allows such creative freedom. Using only 3-5 images of a user-provided concept, like an object or a style, we learn to represent it through new "words" in the embedding space of a frozen text-to-image model. These "words" can be composed into natural language sentences, guiding personalized creation in an intuitive way. Notably, we find evidence that a single word embedding is sufficient for capturing unique and varied concepts. We compare our approach to a wide range of baselines, and demonstrate that it can more faithfully portray the concepts across a range of applications and tasks.
+
+<!-- [IMAGE] -->
+
+<div align=center>
+<img src="https://github.com/open-mmlab/mmagic/assets/28132635/b2dac6f1-5151-4199-bcc2-71b5b1523a16">
+</div>
+
+## Configs
+
+|                    Model                    | Dataset | Download |
+| :-----------------------------------------: | :-----: | :------: |
+| [Textual Inversion](./textual_inversion.py) |    -    |    -     |
+
+## Quick Start
+
+1. Download [data](https://drive.google.com/drive/folders/1fmJMs25nxS_rSNqS5hTcRdLem_YQXbq5) and save to `data/cat_toy`
+
+The file structure will be like this:
+
+```text
+data
+└── cat_toy
+    ├── 1.jpeg
+    ├── 2.jpeg
+    ├── 3.jpeg
+    ├── 3.jpeg
+    ├── 4.jpeg
+    ├── 6.jpeg
+    └── 7.jpeg
+```
+
+2. Start training with the following command:
+
+```bash
+bash tools/dist_train.sh configs/textual_inversion/textual_inversion.py 1
+```
+
+<div align="center">
+  <img src="https://github.com/open-mmlab/mmagic/assets/28132635/635a336c-fd6c-4c6f-b2c1-c1621420b9b9" width="400"/>
+  <br/>
+</div>
+
+3. Inference with trained textual embedding:
+
+```python
+import torch
+from mmengine import Config
+
+from mmagic.registry import MODELS
+from mmagic.utils import register_all_modules
+
+register_all_modules()
+
+
+def process_state_dict(state_dict):
+    new_state_dict = dict()
+    for k, v in state_dict.items():
+        new_k = k.replace('module.', '')
+        new_state_dict[new_k] = v
+
+    return new_state_dict
+
+
+cfg = Config.fromfile('configs/textual_inversion/textual_inversion.py')
+checkpoint = torch.load('work_dirs/textual_inversion/iter_3000.pth')
+state_dict = process_state_dict(checkpoint['state_dict'])
+model = MODELS.build(cfg.model)
+model.load_state_dict(state_dict)
+
+model = model.cuda()
+with torch.no_grad():
+    sample = model.infer('a <cat-toy> bag')['samples'][0]
+
+sample.save('cat-toy-bag.png')
+```
+
+## Comments
+
+Our codebase for the stable diffusion models builds heavily on [diffusers codebase](https://github.com/huggingface/diffusers) and the model weights are from [stable-diffusion-1.5](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_controlnet.py).
+
+Thanks for the efforts of the community!
+
+## Citation
+
+```bibtex
+@misc{gal2022textual,
+      doi = {10.48550/ARXIV.2208.01618},
+      url = {https://arxiv.org/abs/2208.01618},
+      author = {Gal, Rinon and Alaluf, Yuval and Atzmon, Yuval and Patashnik, Or and Bermano, Amit H. and Chechik, Gal and Cohen-Or, Daniel},
+      title = {An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion},
+      publisher = {arXiv},
+      year = {2022},
+      primaryClass={cs.CV}
+}
+
+```

configs/textual_inversion/metafile.yml

@@ -0,0 +1,18 @@
+Collections:
+- Name: Textual Inversion
+  Paper:
+    Title: 'An Image is Worth One Word: Personalizing Text-to-Image Generation using
+      Textual Inversion'
+    URL: https://arxiv.org/abs/2208.01618
+  README: configs/textual_inversion/README.md
+  Task:
+  - text2image
+  Year: 2022
+Models:
+- Config: configs/textual_inversion/textual_inversion.py
+  In Collection: Textual Inversion
+  Name: textual_inversion
+  Results:
+  - Dataset: '-'
+    Metrics: {}
+    Task: Text2Image

configs/textual_inversion/textual_inversion.py

@@ -0,0 +1,85 @@
+_base_ = '../_base_/gen_default_runtime.py'
+
+# config for model
+dtype = 'fp16'
+stable_diffusion_v15_url = 'runwayml/stable-diffusion-v1-5'
+
+placeholder_token = '<cat-toy>'
+initialize_token = 'toy'
+num_vectors_per_token = 1
+val_prompts = [
+    'a <cat-toy> on packbag', 'a <cat-toy> on sofa',
+    'a <cat-toy> in swimming pool', 'a <cat-toy>'
+]
+
+model = dict(
+    type='TextualInversion',
+    placeholder_token=placeholder_token,
+    vae=dict(
+        type='AutoencoderKL',
+        from_pretrained=stable_diffusion_v15_url,
+        subfolder='vae'),
+    unet=dict(
+        type='UNet2DConditionModel',
+        from_pretrained=stable_diffusion_v15_url,
+        subfolder='unet'),
+    text_encoder=dict(
+        type='ClipWrapper',
+        clip_type='huggingface',
+        pretrained_model_name_or_path=stable_diffusion_v15_url,
+        subfolder='text_encoder'),
+    tokenizer=stable_diffusion_v15_url,
+    initialize_token=initialize_token,
+    num_vectors_per_token=num_vectors_per_token,
+    val_prompts=val_prompts,
+    scheduler=dict(
+        type='DDPMScheduler',
+        from_pretrained=stable_diffusion_v15_url,
+        subfolder='scheduler'),
+    test_scheduler=dict(
+        type='DDIMScheduler',
+        from_pretrained=stable_diffusion_v15_url,
+        subfolder='scheduler'),
+    data_preprocessor=dict(type='DataPreprocessor', data_keys=None))
+
+train_cfg = dict(max_iters=3000)
+
+optim_wrapper = dict(
+    modules='.*trainable_embeddings',
+    optimizer=dict(type='AdamW', lr=5e-4),
+    accumulative_counts=1)
+
+pipeline = [
+    dict(type='LoadImageFromFile', key='img', channel_order='rgb'),
+    dict(type='Resize', scale=(512, 512)),
+    dict(type='PackInputs')
+]
+
+dataset = dict(
+    type='TextualInversionDataset',
+    data_root='./data/',
+    concept_dir='cat_toy',
+    placeholder=placeholder_token,
+    pipeline=pipeline)
+
+train_dataloader = dict(
+    dataset=dataset,
+    num_workers=16,
+    sampler=dict(type='InfiniteSampler', shuffle=True),
+    persistent_workers=True,
+    batch_size=1)
+val_cfg = val_evaluator = val_dataloader = None
+test_cfg = test_evaluator = test_dataloader = None
+
+default_hooks = dict(
+    logger=dict(interval=10),
+    checkpoint=dict(type='CheckpointHook', interval=10))
+custom_hooks = [
+    dict(
+        type='VisualizationHook',
+        interval=50,
+        fixed_input=True,
+        # visualize train dataset
+        vis_kwargs_list=dict(type='Data', name='fake_img'),
+        n_samples=1)
+]

mmagic/datasets/__init__.py

@@ -11,12 +11,13 @@
 from .mscoco_dataset import MSCoCoDataset
 from .paired_image_dataset import PairedImageDataset
 from .singan_dataset import SinGANDataset
+from .textual_inversion_dataset import TextualInversionDataset
 from .unpaired_image_dataset import UnpairedImageDataset
 
 __all__ = [
     'AdobeComp1kDataset', 'BasicImageDataset', 'BasicFramesDataset',
     'BasicConditionalDataset', 'UnpairedImageDataset', 'PairedImageDataset',
     'ImageNet', 'CIFAR10', 'GrowScaleImgDataset', 'SinGANDataset',
     'MSCoCoDataset', 'ControlNetDataset', 'DreamBoothDataset',
-    'ControlNetDataset', 'SDFinetuneDataset'
+    'ControlNetDataset', 'SDFinetuneDataset', 'TextualInversionDataset'
 ]

mmagic/datasets/textual_inversion_dataset.py

@@ -0,0 +1,122 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+from random import choice
+from typing import Callable, List, Union
+
+from mmengine import FileClient
+from mmengine.dataset import BaseDataset
+
+from mmagic.registry import DATASETS
+
+imagenet_templates_small = [
+    'a photo of a {}',
+    'a rendering of a {}',
+    'a cropped photo of the {}',
+    'the photo of a {}',
+    'a photo of a clean {}',
+    'a photo of a dirty {}',
+    'a dark photo of the {}',
+    'a photo of my {}',
+    'a photo of the cool {}',
+    'a close-up photo of a {}',
+    'a bright photo of the {}',
+    'a cropped photo of a {}',
+    'a photo of the {}',
+    'a good photo of the {}',
+    'a photo of one {}',
+    'a close-up photo of the {}',
+    'a rendition of the {}',
+    'a photo of the clean {}',
+    'a rendition of a {}',
+    'a photo of a nice {}',
+    'a good photo of a {}',
+    'a photo of the nice {}',
+    'a photo of the small {}',
+    'a photo of the weird {}',
+    'a photo of the large {}',
+    'a photo of a cool {}',
+    'a photo of a small {}',
+]
+
+imagenet_style_templates_small = [
+    'a painting in the style of {}',
+    'a rendering in the style of {}',
+    'a cropped painting in the style of {}',
+    'the painting in the style of {}',
+    'a clean painting in the style of {}',
+    'a dirty painting in the style of {}',
+    'a dark painting in the style of {}',
+    'a picture in the style of {}',
+    'a cool painting in the style of {}',
+    'a close-up painting in the style of {}',
+    'a bright painting in the style of {}',
+    'a cropped painting in the style of {}',
+    'a good painting in the style of {}',
+    'a close-up painting in the style of {}',
+    'a rendition in the style of {}',
+    'a nice painting in the style of {}',
+    'a small painting in the style of {}',
+    'a weird painting in the style of {}',
+    'a large painting in the style of {}',
+]
+
+
+@DATASETS.register_module()
+class TextualInversionDataset(BaseDataset):
+    """Dataset for DreamBooth.
+
+    Args:
+        data_root (str): Path to the data root.
+        concept_dir (str): Path to the concept images.
+        is_style (bool)
+        prompt (str): Prompt of the concept.
+        pipeline (list[dict | callable]): A sequence of data transforms.
+    """
+
+    def __init__(self,
+                 data_root: str,
+                 concept_dir: str,
+                 placeholder: str,
+                 is_style: bool = False,
+                 pipeline: List[Union[dict, Callable]] = []):
+
+        data_prefix = dict(img_path=concept_dir)
+
+        self.placeholder = placeholder
+        if is_style:
+            self.template = imagenet_style_templates_small
+        else:
+            self.template = imagenet_templates_small
+
+        super().__init__(
+            data_root=data_root, data_prefix=data_prefix, pipeline=pipeline)
+
+    def load_data_list(self) -> list:
+        """Load data list from concept_dir and class_dir."""
+        data_list = []
+
+        img_dir = self.data_prefix['img_path']
+        file_client = FileClient.infer_client(uri=img_dir)
+        img_dir = osp.abspath(img_dir)
+
+        for data_name in file_client.list_dir_or_file(img_dir, list_dir=False):
+            data_info = dict(
+                img_path=file_client.join_path(img_dir, data_name))
+            data_list.append(data_info)
+        return data_list
+
+    def prepare_data(self, idx):
+        """Get data processed by ``self.pipeline``.
+
+        Args:
+            idx (int): The index of ``data_info``.
+
+        Returns:
+            Any: Depends on ``self.pipeline``.
+        """
+        data_info = self.get_data_info(idx)
+        # load random template
+        selected_template = choice(self.template)
+        prompt = selected_template.format(self.placeholder)
+        data_info['prompt'] = prompt
+        return self.pipeline(data_info)

mmagic/models/editors/__init__.py

@@ -58,6 +58,7 @@
 from .stylegan3 import StyleGAN3, StyleGAN3Generator
 from .swinir import SwinIRNet
 from .tdan import TDAN, TDANNet
+from .textual_inversion import TextualInversion
 from .tof import TOFlowVFINet, TOFlowVSRNet, ToFResBlock
 from .ttsr import LTE, TTSR, SearchTransformer, TTSRDiscriminator, TTSRNet
 from .wgan_gp import WGANGP
@@ -88,5 +89,5 @@
     'StyleGAN3Generator', 'InstColorization', 'NAFBaseline',
     'NAFBaselineLocal', 'NAFNet', 'NAFNetLocal', 'DenoisingUnet',
     'ClipWrapper', 'EG3D', 'Restormer', 'SwinIRNet', 'StableDiffusion',
-    'ControlStableDiffusion', 'DreamBooth'
+    'ControlStableDiffusion', 'DreamBooth', 'TextualInversion'
 ]

mmagic/models/editors/dreambooth/dreambooth.py

@@ -29,8 +29,6 @@ class DreamBooth(StableDiffusion):
             encoder.
         tokenizer (str): The **name** for CLIP tokenizer.
         unet (Union[dict, nn.Module]): The config or module for Unet model.
-        controlnet (Union[dict, nn.Module]): The config or module for
-            ControlNet.
         schedule (Union[dict, nn.Module]): The config or module for diffusion
             scheduler.
         test_scheduler (Union[dict, nn.Module], optional): The config or
@@ -54,6 +52,10 @@ class DreamBooth(StableDiffusion):
         noise_offset_weight (bool, optional): The weight of noise offset
             introduced in https://www.crosslabs.org/blog/diffusion-with-offset-noise  # noqa
             Defaults to 0.
+        tomesd_cfg (dict, optional): The config for TOMESD. Please refers to
+            https://github.com/dbolya/tomesd and
+            https://github.com/open-mmlab/mmagic/blob/main/mmagic/models/utils/tome_utils.py for detail.  # noqa
+            Defaults to None.
         data_preprocessor (dict, optional): The pre-process config of
             :class:`BaseDataPreprocessor`. Defaults to
                 dict(type='DataPreprocessor').

mmagic/models/editors/textual_inversion/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .textual_inversion import TextualInversion
+
+__all__ = ['TextualInversion']