Merge pull request #111 from okotaku/feat/inpaint

[Feature] Support Inpaint
okotaku · Dec 10, 2023 · 62913d9 · 62913d9
2 parents 79a7a50 + f03364b
commit 62913d9
Show file tree

Hide file tree

Showing 24 changed files with 2,095 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -194,6 +194,7 @@ For detailed user guides and advanced guides, please refer to our [Documentation
             <li><a href="configs/stable_diffusion_dreambooth/README.md">DreamBooth (CVPR'2023)</a></li>
             <li><a href="configs/stable_diffusion_lora/README.md">LoRA (ICLR'2022)</a></li>
             <li><a href="configs/distill_sd_dreambooth/README.md">Distill SD DreamBooth (2023)</a></li>
+            <li><a href="configs/stable_diffusion_inpaint/README.md">Inpaint</a></li>
       </ul>
       </td>
       <td>
@@ -212,6 +213,7 @@ For detailed user guides and advanced guides, please refer to our [Documentation
           <li><a href="configs/lokr/README.md">LoKr (2022)</a></li>
           <li><a href="configs/oft/README.md">OFT (NeurIPS'2023)</a></li>
           <li><a href="configs/stable_diffusion_xl_controlnetxs/README.md">ControlNet-XS (2023)</a></li>
+          <li><a href="configs/stable_diffusion_xl_inpaint/README.md">Inpaint</a></li>
         </ul>
       </td>
       <td>

diff --git a/configs/_base_/datasets/dog_inpaint.py b/configs/_base_/datasets/dog_inpaint.py
@@ -0,0 +1,48 @@
+train_pipeline = [
+    dict(type="torchvision/Resize", size=512, interpolation="bilinear"),
+    dict(type="RandomCrop", size=512),
+    dict(type="RandomHorizontalFlip", p=0.5),
+    dict(
+        type="LoadMask",
+        mask_mode="bbox",
+        mask_config=dict(
+            max_bbox_shape=(256, 256),
+            max_bbox_delta=40,
+            min_margin=20)),
+    dict(type="torchvision/ToTensor"),
+    dict(type="MaskToTensor"),
+    dict(type="DumpImage", max_imgs=10, dump_dir="work_dirs/dump"),
+    dict(type="torchvision/Normalize", mean=[0.5], std=[0.5]),
+    dict(type="GetMaskedImage"),
+    dict(type="PackInputs",
+         input_keys=["img", "mask", "masked_image", "text"]),
+]
+train_dataloader = dict(
+    batch_size=4,
+    num_workers=4,
+    dataset=dict(
+        type="HFDreamBoothDataset",
+        dataset="diffusers/dog-example",
+        instance_prompt="a photo of sks dog",
+        pipeline=train_pipeline,
+        class_prompt=None),
+    sampler=dict(type="InfiniteSampler", shuffle=True),
+)
+
+val_dataloader = None
+val_evaluator = None
+test_dataloader = val_dataloader
+test_evaluator = val_evaluator
+
+custom_hooks = [
+    dict(
+        type="VisualizationHook",
+        prompt=["a photo of sks dog"] * 4,
+        image=["https://github.com/okotaku/diffengine/assets/24734142/8e02bd0e-9dcc-49b6-94b0-86ab3b40bc2b"] * 4,  # noqa
+        mask=["https://github.com/okotaku/diffengine/assets/24734142/d0de4fb9-9183-418a-970d-582e9324f05d"] * 4,  # noqa
+        by_epoch=False,
+        width=512,
+        height=512,
+        interval=100),
+    dict(type="SDCheckpointHook"),
+]
diff --git a/configs/_base_/datasets/dog_inpaint_xl.py b/configs/_base_/datasets/dog_inpaint_xl.py
@@ -0,0 +1,50 @@
+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="LoadMask",
+        mask_mode="bbox",
+        mask_config=dict(
+            max_bbox_shape=(512, 512),
+            max_bbox_delta=80,
+            min_margin=40)),
+    dict(type="torchvision/ToTensor"),
+    dict(type="MaskToTensor"),
+    dict(type="DumpImage", max_imgs=10, dump_dir="work_dirs/dump"),
+    dict(type="torchvision/Normalize", mean=[0.5], std=[0.5]),
+    dict(type="GetMaskedImage"),
+    dict(type="PackInputs",
+         input_keys=["img", "mask", "masked_image", "text", "time_ids"]),
+]
+train_dataloader = dict(
+    batch_size=1,
+    num_workers=2,
+    dataset=dict(
+        type="HFDreamBoothDataset",
+        dataset="diffusers/dog-example",
+        instance_prompt="a photo of sks dog",
+        pipeline=train_pipeline,
+        class_prompt=None),
+    sampler=dict(type="InfiniteSampler", shuffle=True),
+)
+
+val_dataloader = None
+val_evaluator = None
+test_dataloader = val_dataloader
+test_evaluator = val_evaluator
+
+custom_hooks = [
+    dict(
+        type="VisualizationHook",
+        prompt=["a photo of sks dog"] * 4,
+        image=["https://github.com/okotaku/diffengine/assets/24734142/8e02bd0e-9dcc-49b6-94b0-86ab3b40bc2b"] * 4,  # noqa
+        mask=["https://github.com/okotaku/diffengine/assets/24734142/d0de4fb9-9183-418a-970d-582e9324f05d"] * 4,  # noqa
+        by_epoch=False,
+        width=1024,
+        height=1024,
+        interval=100),
+    dict(type="SDCheckpointHook"),
+]
diff --git a/configs/_base_/models/stable_diffusion_inpaint.py b/configs/_base_/models/stable_diffusion_inpaint.py
@@ -0,0 +1,2 @@
+model = dict(type="StableDiffusionInpaint",
+             model="runwayml/stable-diffusion-inpainting")
diff --git a/configs/_base_/models/stable_diffusion_xl_inpaint.py b/configs/_base_/models/stable_diffusion_xl_inpaint.py
@@ -0,0 +1,3 @@
+model = dict(type="StableDiffusionXLInpaint",
+             model="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+             gradient_checkpointing=True)
diff --git a/configs/stable_diffusion_inpaint/README.md b/configs/stable_diffusion_inpaint/README.md
@@ -0,0 +1,89 @@
+# Stable Diffusion Inpaint
+
+[High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)
+
+## Abstract
+
+By decomposing the image formation process into a sequential application of denoising autoencoders, diffusion models (DMs) achieve state-of-the-art synthesis results on image data and beyond. Additionally, their formulation allows for a guiding mechanism to control the image generation process without retraining. However, since these models typically operate directly in pixel space, optimization of powerful DMs often consumes hundreds of GPU days and inference is expensive due to sequential evaluations. To enable DM training on limited computational resources while retaining their quality and flexibility, we apply them in the latent space of powerful pretrained autoencoders. In contrast to previous work, training diffusion models on such a representation allows for the first time to reach a near-optimal point between complexity reduction and detail preservation, greatly boosting visual fidelity. By introducing cross-attention layers into the model architecture, we turn diffusion models into powerful and flexible generators for general conditioning inputs such as text or bounding boxes and high-resolution synthesis becomes possible in a convolutional manner. Our latent diffusion models (LDMs) achieve a new state of the art for image inpainting and highly competitive performance on various tasks, including unconditional image generation, semantic scene synthesis, and super-resolution, while significantly reducing computational requirements compared to pixel-based DMs.
+
+<div align=center>
+<img src="https://github.com/okotaku/diffengine/assets/24734142/d874dc07-af7b-4a88-9629-f3be7c1d2d60"/>
+</div>
+
+## Citation
+
+```
+@InProceedings{Rombach_2022_CVPR,
+    author    = {Rombach, Robin and Blattmann, Andreas and Lorenz, Dominik and Esser, Patrick and Ommer, Bj\"orn},
+    title     = {High-Resolution Image Synthesis With Latent Diffusion Models},
+    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    month     = {June},
+    year      = {2022},
+    pages     = {10684-10695}
+}
+```
+
+## Run Training
+
+Run Training
+
+```
+# single gpu
+$ mim train diffengine ${CONFIG_FILE}
+# multi gpus
+$ mim train diffengine ${CONFIG_FILE} --gpus 2 --launcher pytorch
+
+# Example.
+$ mim train diffengine configs/stable_diffusion_inpaint/stable_diffusion_inpaint_dog.py
+```
+
+## Inference with diffusers
+
+Once you have trained a model, specify the path to the saved model and utilize it for inference using the `diffusers.pipeline` module.
+
+Before inferencing, we should convert weights for diffusers format,
+
+```bash
+$ mim run diffengine publish_model2diffusers ${CONFIG_FILE} ${INPUT_FILENAME} ${OUTPUT_DIR} --save-keys ${SAVE_KEYS}
+# Example
+$ mim run diffengine publish_model2diffusers configs/stable_diffusion_inpaint/stable_diffusion_inpaint_dog.py work_dirs/stable_diffusion_inpaint_dog/iter_1000.pth work_dirs/stable_diffusion_inpaint_dog --save-keys unet
+```
+
+Then we can run inference.
+
+```py
+import torch
+from diffusers import StableDiffusionInpaintPipeline, UNet2DConditionModel
+from diffusers.utils import load_image
+
+prompt = 'a photo of sks dog'
+img = 'https://github.com/okotaku/diffengine/assets/24734142/8e02bd0e-9dcc-49b6-94b0-86ab3b40bc2b'
+mask = 'https://github.com/okotaku/diffengine/assets/24734142/d0de4fb9-9183-418a-970d-582e9324f05d'
+checkpoint = 'work_dirs/stable_diffusion_inpaint_dog'
+
+unet = UNet2DConditionModel.from_pretrained(
+    checkpoint, subfolder='unet', torch_dtype=torch.float16)
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    'runwayml/stable-diffusion-inpainting', unet=unet, torch_dtype=torch.float16)
+pipe.to('cuda')
+
+image = pipe(
+    prompt,
+    load_image(img).convert("RGB"),
+    load_image(mask).convert("L"),
+    num_inference_steps=50,
+).images[0]
+image.save('demo.png')
+```
+
+You can see more details on [`docs/source/run_guides/run_sd.md`](../../docs/source/run_guides/run_sd.md#inference-with-diffusers).
+
+## Results Example
+
+#### stable_diffusion_inpaint_dog
+
+![input](https://github.com/okotaku/diffengine/assets/24734142/8e02bd0e-9dcc-49b6-94b0-86ab3b40bc2b)
+
+![mask](https://github.com/okotaku/diffengine/assets/24734142/d0de4fb9-9183-418a-970d-582e9324f05d)
+
+![example](https://github.com/okotaku/diffengine/assets/24734142/f9ec820b-af75-4c74-8c0b-6558a0a19b95)
diff --git a/configs/stable_diffusion_inpaint/stable_diffusion_inpaint_dog.py b/configs/stable_diffusion_inpaint/stable_diffusion_inpaint_dog.py
@@ -0,0 +1,6 @@
+_base_ = [
+    "../_base_/models/stable_diffusion_inpaint.py",
+    "../_base_/datasets/dog_inpaint.py",
+    "../_base_/schedules/stable_diffusion_1k.py",
+    "../_base_/default_runtime.py",
+]
diff --git a/configs/stable_diffusion_xl_inpaint/README.md b/configs/stable_diffusion_xl_inpaint/README.md
@@ -0,0 +1,123 @@
+# Stable Diffusion XL Inpaint
+
+[Stable Diffusion XL](https://arxiv.org/abs/2307.01952)
+
+## Abstract
+
+We present SDXL, a latent diffusion model for text-to-image synthesis. Compared to previous versions of Stable Diffusion, SDXL leverages a three times larger UNet backbone: The increase of model parameters is mainly due to more attention blocks and a larger cross-attention context as SDXL uses a second text encoder. We design multiple novel conditioning schemes and train SDXL on multiple aspect ratios. We also introduce a refinement model which is used to improve the visual fidelity of samples generated by SDXL using a post-hoc image-to-image technique. We demonstrate that SDXL shows drastically improved performance compared the previous versions of Stable Diffusion and achieves results competitive with those of black-box state-of-the-art image generators.
+
+<div align=center>
+<img src="https://github.com/okotaku/diffengine/assets/24734142/27d4ebad-5705-4500-826f-41f425a08c0d"/>
+</div>
+
+## Citation
+
+```
+```
+
+## Run Training
+
+Run Training
+
+```
+# single gpu
+$ mim train diffengine ${CONFIG_FILE}
+# multi gpus
+$ mim train diffengine ${CONFIG_FILE} --gpus 2 --launcher pytorch
+
+# Example.
+$ mim train diffengine configs/stable_diffusion_xl_inpaint/stable_diffusion_xl_inpaint_dog.py
+```
+
+## Training Speed
+
+#### Single GPU
+
+Environment:
+
+- A6000 Single GPU
+- nvcr.io/nvidia/pytorch:23.10-py3
+
+Settings:
+
+- 1epoch training.
+
+|                   Model                   | total time |
+| :---------------------------------------: | :--------: |
+|  stable_diffusion_xl_pokemon_blip (fp16)  | 12 m 37 s  |
+| stable_diffusion_xl_pokemon_blip_xformers |  10 m 6 s  |
+|   stable_diffusion_xl_pokemon_blip_fast   |  9 m 47 s  |
+
+Note that `stable_diffusion_xl_pokemon_blip_fast` took a few minutes to compile. We will disregard it.
+
+#### Multiple GPUs
+
+Environment:
+
+- A100 x 4 GPUs
+- nvcr.io/nvidia/pytorch:23.11-py3
+
+Settings:
+
+- 1epoch training.
+
+|                           Model                           | total time |
+| :-------------------------------------------------------: | :--------: |
+|       stable_diffusion_xl_pokemon_blip_fast (BS=4)        |  1 m 6 s   |
+| stable_diffusion_xl_pokemon_blip_deepspeed_stage3 (BS=8)  |  1 m 5 s   |
+| stable_diffusion_xl_pokemon_blip_deepspeed_stage2 (BS=8)  |    58 s    |
+| stable_diffusion_xl_pokemon_blip_colossal (stage=2, BS=8) |    58s     |
+
+## Inference with diffusers
+
+Once you have trained a model, specify the path to the saved model and utilize it for inference using the `diffusers.pipeline` module.
+
+Before inferencing, we should convert weights for diffusers format,
+
+```bash
+$ mim run diffengine publish_model2diffusers ${CONFIG_FILE} ${INPUT_FILENAME} ${OUTPUT_DIR} --save-keys ${SAVE_KEYS}
+# Example
+$ mim run diffengine publish_model2diffusers configs/stable_diffusion_xl_inpaint/stable_diffusion_xl_inpaint_dog.py work_dirs/stable_diffusion_xl_inpaint_dog/iter_1000.pth work_dirs/stable_diffusion_xl_inpaint_dog --save-keys unet
+```
+
+Then we can run inference.
+
+```py
+import torch
+from diffusers import AutoPipelineForInpainting, UNet2DConditionModel
+from diffusers.utils import load_image
+
+prompt = 'a photo of sks dog'
+img = 'https://github.com/okotaku/diffengine/assets/24734142/8e02bd0e-9dcc-49b6-94b0-86ab3b40bc2b'
+mask = 'https://github.com/okotaku/diffengine/assets/24734142/d0de4fb9-9183-418a-970d-582e9324f05d'
+checkpoint = 'work_dirs/stable_diffusion_xl_inpaint_dog'
+
+unet = UNet2DConditionModel.from_pretrained(
+    checkpoint, subfolder='unet', torch_dtype=torch.float16)
+pipe = AutoPipelineForInpainting.from_pretrained(
+    'diffusers/stable-diffusion-xl-1.0-inpainting-0.1', unet=unet, torch_dtype=torch.float16)
+pipe.to('cuda')
+
+image = pipe(
+    prompt,
+    prompt,
+    load_image(img).convert("RGB"),
+    load_image(mask).convert("L"),
+    num_inference_steps=50,
+    width=1024,
+    height=1024,
+).images[0]
+image.save('demo.png')
+```
+
+You can see more details on [`docs/source/run_guides/run_xl.md`](../../docs/source/run_guides/run_xl.md#inference-with-diffusers).
+
+## Results Example
+
+#### stable_diffusion_xl_inpaint_dog
+
+![input](https://github.com/okotaku/diffengine/assets/24734142/8e02bd0e-9dcc-49b6-94b0-86ab3b40bc2b)
+
+![mask](https://github.com/okotaku/diffengine/assets/24734142/d0de4fb9-9183-418a-970d-582e9324f05d)
+
+![example](https://github.com/okotaku/diffengine/assets/24734142/a2d20952-bd9f-4893-b5da-4171e24f22e2)
diff --git a/configs/stable_diffusion_xl_inpaint/stable_diffusion_xl_inpaint_dog.py b/configs/stable_diffusion_xl_inpaint/stable_diffusion_xl_inpaint_dog.py
@@ -0,0 +1,18 @@
+_base_ = [
+    "../_base_/models/stable_diffusion_xl_inpaint.py",
+    "../_base_/datasets/dog_inpaint_xl.py",
+    "../_base_/schedules/stable_diffusion_1k.py",
+    "../_base_/default_runtime.py",
+]
+
+optim_wrapper = dict(
+    _delete_=True,
+    type="AmpOptimWrapper",
+    dtype="bfloat16",
+    optimizer=dict(
+        type="Adafactor",
+        lr=1e-5,
+        weight_decay=1e-2,
+        scale_parameter=False,
+        relative_step=False),
+    clip_grad=dict(max_norm=1.0))
diff --git a/diffengine/datasets/transforms/__init__.py b/diffengine/datasets/transforms/__init__.py
@@ -1,12 +1,15 @@
 from .base import BaseTransform
 from .dump_image import DumpImage
 from .formatting import PackInputs
+from .loading import LoadMask
 from .processing import (
     TRANSFORMS,
     CenterCrop,
     CLIPImageProcessor,
     ComputePixArtImgInfo,
     ComputeTimeIds,
+    GetMaskedImage,
+    MaskToTensor,
     MultiAspectRatioResizeCenterCrop,
     RandomCrop,
     RandomHorizontalFlip,
@@ -30,4 +33,7 @@
     "RandomTextDrop",
     "ComputePixArtImgInfo",
     "T5TextPreprocess",
+    "LoadMask",
+    "MaskToTensor",
+    "GetMaskedImage",
 ]
diff --git a/diffengine/datasets/transforms/dump_image.py b/diffengine/datasets/transforms/dump_image.py
@@ -61,4 +61,12 @@ def __call__(self, results) -> dict:
                 cv2.imwrite(cond_out_file,
                             condition_img.numpy().astype(np.uint8)[..., ::-1])
 
+            if "mask" in results:
+                mask = results["mask"]
+                if mask.shape[0] in [1, 3]:
+                    mask = mask.permute(1, 2, 0) * 255
+                mask_out_file = osp.join(self.dump_dir, f"{dump_id}_mask.png")
+                cv2.imwrite(mask_out_file,
+                            mask.numpy().astype(np.uint8))
+
         return results