add config and readme

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](<>) and save to `data/`
+
+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.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/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/textual_inversion.py

@@ -18,6 +18,45 @@
 
 @MODELS.register_module()
 class TextualInversion(StableDiffusion):
+    """Implementation of `An Image is Worth One Word: Personalizing Text-to-
+    Image Generation using Textual Inversion.
+
+    <https://arxiv.org/abs/2208.01618>`_ (Textual Inversion).
+
+    Args:
+        vae (Union[dict, nn.Module]): The config or module for VAE model.
+        text_encoder (Union[dict, nn.Module]): The config or module for text
+            encoder.
+        tokenizer (str): The **name** for CLIP tokenizer.
+        unet (Union[dict, nn.Module]): The config or module for Unet model.
+        schedule (Union[dict, nn.Module]): The config or module for diffusion
+            scheduler.
+        test_scheduler (Union[dict, nn.Module], optional): The config or
+            module for diffusion scheduler in test stage (`self.infer`). If not
+            passed, will use the same scheduler as `schedule`. Defaults to
+            None.
+        dtype (str, optional): The dtype for the model. Defaults to 'fp16'.
+        enable_xformers (bool, optional): Whether to use xformers.
+            Defaults to True.
+        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.
+        initialize_token (str, optional): The initialization token for textual
+            embedding to train. Defaults to None.
+        num_vefctor_per_token (int): The length of the learnable embedding.
+            Defaults to 1.
+        val_prompts (Union[str, List[str]], optional): The prompts for
+            validation. Defaults to None.
+        data_preprocessor (dict, optional): The pre-process config of
+            :class:`BaseDataPreprocessor`. Defaults to
+                dict(type='DataPreprocessor').
+        init_cfg (dict, optional): The weight initialized config for
+            :class:`BaseModule`. Defaults to None/
+    """
 
     def __init__(self,
                  placeholder_token: str,

model-index.yml

@@ -49,6 +49,7 @@ Import:
 - configs/styleganv3/metafile.yml
 - configs/swinir/metafile.yml
 - configs/tdan/metafile.yml
+- configs/textual_inversion/metafile.yml
 - configs/tof/metafile.yml
 - configs/ttsr/metafile.yml
 - configs/wgan-gp/metafile.yml