├── LICENSE
├── README.md
├── __assets__
└── github
│ ├── method_arch.png
│ └── teaser.png
├── config
├── ddpm
│ ├── v1.yaml
│ └── v2-upsample.yaml
├── encoders
│ ├── clip.yaml
│ └── openclip.yaml
├── unet
│ ├── inpainting
│ │ ├── v1.yaml
│ │ └── v2.yaml
│ ├── upsample
│ │ └── v2.yaml
│ ├── v
│ │ └── v2.yaml
│ ├── v1.yaml
│ └── v2.yaml
├── vae-upsample.yaml
└── vae.yaml
├── data
├── masks
│ ├── 0_rgb.png
│ ├── 1_rgb.png
│ └── 2_rgb.png
├── metadata
│ ├── 0.json
│ ├── 1.json
│ └── 2.json
└── outputs
│ ├── 0_rgb.png
│ ├── 1_rgb.png
│ └── 2_rgb.png
├── requirements.txt
├── src
├── smplfusion
│ ├── __init__.py
│ ├── animation.py
│ ├── api
│ │ └── __init__.py
│ ├── common.py
│ ├── config.py
│ ├── ddim.py
│ ├── libimage
│ │ ├── __init__.py
│ │ ├── __pycache__
│ │ │ ├── __init__.cpython-310.pyc
│ │ │ ├── __init__.cpython-39.pyc
│ │ │ ├── iimage.cpython-310.pyc
│ │ │ ├── iimage.cpython-39.pyc
│ │ │ ├── iimage_gallery.cpython-310.pyc
│ │ │ ├── iimage_gallery.cpython-39.pyc
│ │ │ ├── utils.cpython-310.pyc
│ │ │ └── utils.cpython-39.pyc
│ │ ├── iimage.py
│ │ ├── iimage_gallery.py
│ │ └── utils.py
│ ├── libpath.py
│ ├── models
│ │ ├── __init__.py
│ │ ├── __pycache__
│ │ │ ├── __init__.cpython-39.pyc
│ │ │ ├── unet.cpython-39.pyc
│ │ │ ├── util.cpython-39.pyc
│ │ │ └── vae.cpython-39.pyc
│ │ ├── encoders
│ │ │ ├── __pycache__
│ │ │ │ └── clip_embedder.cpython-39.pyc
│ │ │ ├── clip_embedder.py
│ │ │ ├── clip_image_embedder.py
│ │ │ ├── modules.py
│ │ │ ├── open_clip_embedder.py
│ │ │ ├── open_clip_image_embedder.py
│ │ │ └── t5_mebedder.py
│ │ ├── unet.py
│ │ ├── util.py
│ │ └── vae.py
│ ├── modelzoo
│ │ ├── __init__.py
│ │ ├── dreamshaper8.py
│ │ └── dreamshaper8_inpainting.py
│ ├── modules
│ │ ├── __init__.py
│ │ ├── __pycache__
│ │ │ ├── __init__.cpython-39.pyc
│ │ │ ├── autoencoder.cpython-39.pyc
│ │ │ ├── distributions.cpython-39.pyc
│ │ │ ├── ema.cpython-39.pyc
│ │ │ └── util.cpython-39.pyc
│ │ ├── attention
│ │ │ ├── __init__.py
│ │ │ ├── __pycache__
│ │ │ │ ├── __init__.cpython-39.pyc
│ │ │ │ ├── basic_transformer_block.cpython-39.pyc
│ │ │ │ ├── cross_attention.cpython-39.pyc
│ │ │ │ ├── feed_forward.cpython-39.pyc
│ │ │ │ ├── memory_efficient_cross_attention.cpython-39.pyc
│ │ │ │ └── spatial_transformer.cpython-39.pyc
│ │ │ ├── basic_transformer_block.py
│ │ │ ├── cross_attention.py
│ │ │ ├── feed_forward.py
│ │ │ ├── memory_efficient_cross_attention.py
│ │ │ └── spatial_transformer.py
│ │ ├── autoencoder.py
│ │ ├── distributions.py
│ │ ├── ema.py
│ │ ├── partial_conv2d.py
│ │ └── util.py
│ ├── patches
│ │ ├── __init__.py
│ │ ├── __pycache__
│ │ │ ├── __init__.cpython-39.pyc
│ │ │ └── router.cpython-39.pyc
│ │ ├── attentionpatch
│ │ │ ├── __init__.py
│ │ │ ├── __pycache__
│ │ │ │ ├── __init__.cpython-39.pyc
│ │ │ │ ├── attstore.cpython-39.pyc
│ │ │ │ ├── boosted_maskscore.cpython-39.pyc
│ │ │ │ ├── default.cpython-39.pyc
│ │ │ │ ├── ediff.cpython-39.pyc
│ │ │ │ ├── inpaint.cpython-39.pyc
│ │ │ │ ├── introvert.cpython-39.pyc
│ │ │ │ ├── shuffled.cpython-39.pyc
│ │ │ │ └── zeropaint.cpython-39.pyc
│ │ │ ├── attstore.py
│ │ │ ├── boosted_maskscore copy.py
│ │ │ ├── boosted_maskscore.py
│ │ │ ├── default.py
│ │ │ ├── inpaint.py
│ │ │ ├── introvert.py
│ │ │ ├── maskscore.py
│ │ │ ├── other.py
│ │ │ ├── shuffled.py
│ │ │ └── zeropaint.py
│ │ ├── router.py
│ │ └── transformerpatch
│ │ │ ├── __init__.py
│ │ │ ├── __pycache__
│ │ │ ├── __init__.cpython-39.pyc
│ │ │ ├── default.cpython-39.pyc
│ │ │ ├── guided.cpython-39.pyc
│ │ │ ├── introvert.cpython-39.pyc
│ │ │ └── weighting_versions.cpython-39.pyc
│ │ │ ├── default.py
│ │ │ ├── guided.py
│ │ │ ├── introvert.py
│ │ │ └── weighting_versions.py
│ ├── scheduler.py
│ ├── share.py
│ ├── util.py
│ └── utils
│ │ ├── __init__.py
│ │ ├── __pycache__
│ │ ├── __init__.cpython-39.pyc
│ │ ├── input_image.cpython-39.pyc
│ │ ├── input_mask.cpython-39.pyc
│ │ ├── input_shape.cpython-39.pyc
│ │ └── layer_mask.cpython-39.pyc
│ │ ├── input_image.py
│ │ ├── input_mask.py
│ │ ├── input_shape.py
│ │ └── layer_mask.py
└── zeropainter
│ ├── __init__.py
│ ├── convert_diffusers.py
│ ├── dreamshaper.py
│ ├── generation.py
│ ├── inpainting.py
│ ├── models.py
│ ├── segmentation.py
│ ├── zero_painter_dataset.py
│ └── zero_painter_pipline.py
└── zero_painter.py
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2023 Picsart AI Research (PAIR)
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Zero-Painter: Training-Free Layout Control for Text-to-Image Synthesis [CVPR 2024]
2 |
3 | This repository is the official implementation of [Zero-Painter](https://arxiv.org/abs/2406.04032).
4 |
5 |
6 | **[Zero-Painter: Training-Free Layout Control for Text-to-Image Synthesis](https://arxiv.org/abs/2406.04032)**
7 |
8 | Marianna Ohanyan*,
9 | Hayk Manukyan*,
10 | Zhangyang Wang,
11 | Shant Navasardyan,
12 | [Humphrey Shi](https://www.humphreyshi.com)
13 |
14 |
15 | [Arxiv](https://arxiv.org/abs/2406.04032)
16 |
17 |
18 |
21 | We present Zero-Painter , a novel training-free framework for layout-conditional text-to-image synthesis that facilitates the creation of detailed and controlled imagery from textual prompts. Our method utilizes object masks and individual descriptions, coupled with a global text prompt, to generate images with high fidelity. Zero-Painter employs a two-stage process involving our novel Prompt-Adjusted Cross-Attention (PACA) and Region-Grouped Cross-Attention (ReGCA) blocks, ensuring precise alignment of generated objects with textual prompts and mask shapes. Our extensive experiments demonstrate that Zero-Painter surpasses current state-of-the-art methods in preserving textual details and adhering to mask shapes.
22 |
23 |
24 |
25 |
26 | ## 🔥 News
27 | - [2024.06.6] ZeroPainter paper and code is released.
28 | - [2024.02.27] Paper is accepted to CVPR 2024.
29 |
30 |
31 | ## ⚒️ Installation
32 |
33 |
38 | Install with `pip`:
39 | ```bash
40 | pip3 install -r requirements.txt
41 | ```
42 |
43 | ## 💃 Inference: Generate images with Zero-Painter
44 |
45 | 1. Download [models](https://huggingface.co/PAIR/Zero-Painter) and put them in the `models` folder.
46 | 2. You can use the following script to perform inference on the given mask and prompts pair:
47 | ```
48 | python zero_painter.py \
49 | --mask-path data/masks/1_rgb.png \
50 | --metadata data/metadata/1.json \
51 | --output-dir data/outputs/
52 | ```
53 |
54 | `meatadata` sould be in the following format
55 | ```
56 | [{
57 | "prompt": "Brown gift box beside red candle.",
58 | "color_context_dict": {
59 | "(244, 54, 32)": "Brown gift box",
60 | "(54, 245, 32)": "red candle"
61 | }
62 | }]
63 | ```
64 |
73 |
74 | ## Method
75 |
76 | '))) + [tokenize(prompt + positive_prompt).index('')]
19 |
20 | def load_model_v2_inpainting(folder):
21 | model_config = options.ddpm.v1_yaml
22 |
23 | unet = options.unet.inpainting.v2_yaml.eval().cuda()
24 | unet.load_state_dict(folder.unet_ckpt)
25 | unet = unet.requires_grad_(False)
26 |
27 | vae = options.vae_yaml.eval().cuda()
28 | vae.load_state_dict(folder.vae_ckpt)
29 | vae = vae.requires_grad_(False)
30 |
31 | encoder = options.encoders.openclip_yaml.eval().cuda()
32 | encoder.load_state_dict(folder.encoder_ckpt)
33 | encoder = encoder.requires_grad_(False)
34 |
35 | return model_config, unet, vae, encoder
36 |
37 | def load_model_v15_inpainting(folder):
38 | model_config = options.ddpm.v1_yaml
39 |
40 | unet = options.unet.inpainting.v1_yaml.eval().cuda()
41 | unet.load_state_dict(folder.unet_ckpt)
42 | unet = unet.requires_grad_(False)
43 |
44 | vae = options.vae_yaml.eval().cuda()
45 | vae.load_state_dict(folder.vae_ckpt)
46 | vae = vae.requires_grad_(False)
47 |
48 | encoder = options.encoders.clip_yaml.eval().cuda()
49 | encoder.load_state_dict(folder.encoder_ckpt)
50 | encoder = encoder.requires_grad_(False)
51 |
52 | return model_config, unet, vae, encoder
53 |
54 | def load_model_v2(folder):
55 | model_config = options.ddpm.v1_yaml
56 | unet = options.unet.v2_yaml.eval().cuda()
57 | unet.load_state_dict(folder.unet_ckpt)
58 | unet = unet.requires_grad_(False)
59 | vae = options.vae_yaml.eval().cuda()
60 | vae.load_state_dict(folder.vae_ckpt)
61 | vae = vae.requires_grad_(False)
62 | encoder = options.encoders.openclip_yaml.eval().cuda()
63 | encoder.load_state_dict(folder.encoder_ckpt)
64 | encoder = encoder.requires_grad_(False)
65 | return model_config, unet, vae, encoder
66 |
67 | def load_model_v1(folder):
68 | model_config = options.ddpm.v1_yaml
69 | unet = options.unet.v1_yaml.eval().cuda()
70 | unet.load_state_dict(folder.unet_ckpt)
71 | unet = unet.requires_grad_(False)
72 | vae = options.vae_yaml.eval().cuda()
73 | vae.load_state_dict(folder.vae_ckpt)
74 | vae = vae.requires_grad_(False)
75 | encoder = options.encoders.clip_yaml.eval().cuda()
76 | encoder.load_state_dict(folder.encoder_ckpt)
77 | encoder = encoder.requires_grad_(False)
78 | return model_config, unet, vae, encoder
79 |
80 | def load_unet_v2(folder):
81 | unet = options.unet.v2_yaml.eval().cuda()
82 | unet.load_state_dict(folder.unet_ckpt)
83 | unet = unet.requires_grad_(False)
84 | return unet
--------------------------------------------------------------------------------
/src/smplfusion/config.py:
--------------------------------------------------------------------------------
1 | IMG_THUMBSIZE = None
--------------------------------------------------------------------------------
/src/smplfusion/ddim.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from tqdm.notebook import tqdm
3 | from . import scheduler
4 | from . import share
5 |
6 | from .libimage import IImage
7 |
8 | class DDIM:
9 | def __init__(self, config, vae, encoder, unet):
10 | self.vae = vae
11 | self.encoder = encoder
12 | self.unet = unet
13 | self.config = config
14 | self.schedule = scheduler.linear(1000, config.linear_start, config.linear_end)
15 |
16 | def __call__(
17 | self, prompt = '', dt = 50, shape = (1,4,64,64), seed = None, negative_prompt = '', unet_condition = None,
18 | context = None, verbose = True):
19 | if seed is not None: torch.manual_seed(seed)
20 | if unet_condition is not None:
21 | zT = torch.randn((1,4) + unet_condition.shape[2:]).cuda()
22 | else:
23 | zT = torch.randn(shape).cuda()
24 |
25 | with torch.autocast('cuda'), torch.no_grad():
26 | if context is None: context = self.encoder.encode([negative_prompt, prompt])
27 |
28 | zt = zT
29 | pbar = tqdm(range(999, 0, -dt)) if verbose else range(999, 0, -dt)
30 | for timestep in share.DDIMIterator(pbar):
31 | _zt = zt if unet_condition is None else torch.cat([zt, unet_condition], 1)
32 | eps_uncond, eps = self.unet(
33 | torch.cat([_zt, _zt]),
34 | timesteps = torch.tensor([timestep, timestep]).cuda(),
35 | context = context
36 | ).chunk(2)
37 |
38 | eps = (eps_uncond + 7.5 * (eps - eps_uncond))
39 |
40 | z0 = (zt - self.schedule.sqrt_one_minus_alphas[timestep] * eps) / self.schedule.sqrt_alphas[timestep]
41 | zt = self.schedule.sqrt_alphas[timestep - dt] * z0 + self.schedule.sqrt_one_minus_alphas[timestep - dt] * eps
42 | return IImage(self.vae.decode(z0 / self.config.scale_factor))
43 |
44 | def encode(self, image):
45 | return self.vae.encode(image.padx(64).torch().cuda()).mean * self.config.scale_factor
46 | def decode(self, latent):
47 | return IImage(self.vae.decode(latent / self.config.scale_factor))
48 |
49 | def get_inpainting_condition(self, image, mask):
50 | latent_size = [x//8 for x in image.size]
51 | condition_x0 = self.vae.encode(image.torch().cuda() * ~mask.torch(0).bool().cuda()).mean * self.config.scale_factor
52 | condition_mask = mask.resize(latent_size[::-1]).cuda().torch(0).bool().float()
53 |
54 | # condition_x0 += 0.01 * condition_mask * torch.randn_like(condition_mask)
55 | return torch.cat([condition_mask, condition_x0], 1)
56 |
57 | inpainting_condition = get_inpainting_condition
58 |
59 |
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__init__.py:
--------------------------------------------------------------------------------
1 | from .iimage import IImage
2 | from .iimage_gallery import ImageGallery
3 | from .utils import bytes2html
4 |
5 | import math
6 | import numpy as np
7 | import warnings
8 |
9 | # ========= STATIC FUNCTIONS =============
10 | def find_max_h(images):
11 | return max([x.size[1] for x in images])
12 | def find_max_w(images):
13 | return max([x.size[0] for x in images])
14 | def find_max_size(images):
15 | return find_max_w(images), find_max_h(images)
16 |
17 |
18 | def stack(images, axis = 0):
19 | return IImage(np.concatenate([x.data for x in images], axis))
20 | def tstack(images):
21 | w,h = find_max_size(images)
22 | images = [x.pad2wh(w,h) for x in images]
23 | return IImage(np.concatenate([x.data for x in images], 0))
24 | def hstack(images):
25 | h = find_max_h(images)
26 | images = [x.pad2wh(h = h) for x in images]
27 | return IImage(np.concatenate([x.data for x in images], 2))
28 | def vstack(images):
29 | w = find_max_w(images)
30 | images = [x.pad2wh(w = w) for x in images]
31 | return IImage(np.concatenate([x.data for x in images], 1))
32 |
33 | def grid(images, nrows = None, ncols = None):
34 | combined = stack(images)
35 | if nrows is not None:
36 | ncols = math.ceil(combined.data.shape[0] / nrows)
37 | elif ncols is not None:
38 | nrows = math.ceil(combined.data.shape[0] / ncols)
39 | else:
40 | warnings.warn("No dimensions specified, creating a grid with 5 columns (default)")
41 | ncols = 5
42 | nrows = math.ceil(combined.data.shape[0] / ncols)
43 |
44 | pad = nrows * ncols - combined.data.shape[0]
45 | data = np.pad(combined.data, ((0,pad),(0,0),(0,0),(0,0)))
46 | rows = [np.concatenate(x,1,dtype=np.uint8) for x in np.array_split(data, nrows)]
47 | return IImage(np.concatenate(rows, 0, dtype = np.uint8)[None])
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/__init__.cpython-310.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/__init__.cpython-310.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/__init__.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/__init__.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/iimage.cpython-310.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/iimage.cpython-310.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/iimage.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/iimage.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/iimage_gallery.cpython-310.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/iimage_gallery.cpython-310.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/iimage_gallery.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/iimage_gallery.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/utils.cpython-310.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/utils.cpython-310.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/__pycache__/utils.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/libimage/__pycache__/utils.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/libimage/iimage_gallery.py:
--------------------------------------------------------------------------------
1 | import os
2 | from .utils import bytes2html
3 |
4 | class ImageGallery:
5 | def __init__(self, images, captions = None, size='auto', max_rows = -1, root_path = '/', caption_font_size = '2em'):
6 | self.size = size
7 | self.images = images
8 | self.display_str = None
9 | self.captions = captions if captions is not None else [''] * len(self.images)
10 | self.max_rows = max_rows
11 | self.root_path = root_path
12 | self.caption_font_size = caption_font_size
13 |
14 | def generate_display(self):
15 | """Shows a set of images in a gallery that flexes with the width of the notebook.
16 |
17 | Parameters
18 | ----------
19 | images: list of str or bytes
20 | URLs or bytes of images to display
21 |
22 | row_height: str
23 | CSS height value to assign to all images. Set to 'auto' by default to show images
24 | with their native dimensions. Set to a value e.g. '250px' to make all rows
25 | in the gallery equal height.
26 | """
27 | figures = []
28 | row_figures = 0
29 | for image, caption in zip(self.images, self.captions):
30 | if isinstance(image, str):
31 | with open(image,'rb') as f:
32 | link = os.path.relpath(image, self.root_path)
33 | # src = bytes2html(f.read(), width = self.size)
34 | src = bytes2html(f.read())
35 | src = f'{src} '
36 | else:
37 | if image.display_str is None: image.generate_display()
38 | if image.is_video():
39 | src = image.display_str
40 | else:
41 | # src = _src_from_data(image.display_str)
42 | src = image.to_html(width = "100%", root_path = self.root_path)
43 | if caption != '':
44 | caption = f'{caption} '
45 | figures.append(f'''
46 |
47 | {caption}
48 | {src}
49 |
50 | ''')
51 | row_figures += 1
52 | if row_figures == self.max_rows:
53 | row_figures = 0
54 | figures.append('
')
55 | self.display_str = f'''
56 |
57 | {''.join(figures)}
58 |
59 | '''
60 |
61 | def _repr_html_(self):
62 | if self.display_str is None: self.generate_display()
63 | return self.display_str
64 |
65 | def save(self, path):
66 | if self.display_str is None: self.generate_display()
67 | with open(path, 'w') as f:
68 | f.write(self.display_str)
--------------------------------------------------------------------------------
/src/smplfusion/libimage/utils.py:
--------------------------------------------------------------------------------
1 | from IPython.display import Image as IpyImage
2 |
3 | def bytes2html(data, width='auto'):
4 | img_obj = IpyImage(data=data, format='JPG')
5 | for bundle in img_obj._repr_mimebundle_():
6 | for mimetype, b64value in bundle.items():
7 | if mimetype.startswith('image/'):
8 | return f' (1 word label)
32 | # sim: (16x4096x77); mask: (64x64) -> mask.reshape(-1): (4096)
33 | N = sim.shape[0]//2
34 | sim[N:] += (
35 | share.w
36 | * share.noise_signal_ratio[share.timestep]
37 | * share.input_mask.get_res(q, 'cuda').reshape(1,-1,1)# (1,4096,1)
38 | * sim[N:].amax(dim=qkv_reduce_dims, keepdim = True) # (16x1x1)
39 | ) * (torch.eye(77)[increase_indices].sum(0).cuda()) # (1,1,77)
40 |
41 | del q, k
42 | sim = sim.softmax(dim=-1)
43 |
44 | out = einsum("b i j, b j d -> b i d", sim, v)
45 | out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
46 | return self.to_out(out)
--------------------------------------------------------------------------------
/src/smplfusion/patches/attentionpatch/introvert.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import math
3 | import numbers
4 | import numpy as np
5 | import torch
6 | import torch.nn.functional as F
7 | import torchvision
8 | from torch import nn, einsum
9 | from einops import rearrange, repeat
10 |
11 | from ... import share
12 | from ...libimage import IImage
13 |
14 |
15 | # Default for version 4
16 | introvert_res = [16, 32, 64]
17 | introvert_on = True
18 | token_idx = [1,2]
19 |
20 | # Visualization purpose
21 | viz_image = None
22 | viz_mask = None
23 |
24 | video_frames_selfattn = []
25 | video_frames_crossattn = []
26 | visualize_resolution = 16
27 | visualize_selfattn = False
28 | visualize_crossattn = False
29 |
30 | class GaussianSmoothing(nn.Module):
31 | """
32 | Apply gaussian smoothing on a
33 | 1d, 2d or 3d tensor. Filtering is performed seperately for each channel
34 | in the input using a depthwise convolution.
35 | Arguments:
36 | channels (int, sequence): Number of channels of the input tensors. Output will
37 | have this number of channels as well.
38 | kernel_size (int, sequence): Size of the gaussian kernel.
39 | sigma (float, sequence): Standard deviation of the gaussian kernel.
40 | dim (int, optional): The number of dimensions of the data.
41 | Default value is 2 (spatial).
42 | """
43 | def __init__(self, channels, kernel_size, sigma, dim=2):
44 | super(GaussianSmoothing, self).__init__()
45 | if isinstance(kernel_size, numbers.Number):
46 | kernel_size = [kernel_size] * dim
47 | if isinstance(sigma, numbers.Number):
48 | sigma = [sigma] * dim
49 |
50 | # The gaussian kernel is the product of the
51 | # gaussian function of each dimension.
52 | kernel = 1
53 | meshgrids = torch.meshgrid(
54 | [
55 | torch.arange(size, dtype=torch.float32)
56 | for size in kernel_size
57 | ]
58 | )
59 | for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
60 | mean = (size - 1) / 2
61 | kernel *= 1 / (std * math.sqrt(2 * math.pi)) * \
62 | torch.exp(-((mgrid - mean) / (2 * std)) ** 2)
63 |
64 | # Make sure sum of values in gaussian kernel equals 1.
65 | kernel = kernel / torch.sum(kernel)
66 |
67 | # Reshape to depthwise convolutional weight
68 | kernel = kernel.view(1, 1, *kernel.size())
69 | kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
70 |
71 | self.register_buffer('weight', kernel)
72 | self.groups = channels
73 |
74 | if dim == 1:
75 | self.conv = F.conv1d
76 | elif dim == 2:
77 | self.conv = F.conv2d
78 | elif dim == 3:
79 | self.conv = F.conv3d
80 | else:
81 | raise RuntimeError(
82 | 'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim)
83 | )
84 |
85 | def forward(self, input):
86 | """
87 | Apply gaussian filter to input.
88 | Arguments:
89 | input (torch.Tensor): Input to apply gaussian filter on.
90 | Returns:
91 | filtered (torch.Tensor): Filtered output.
92 | """
93 | return self.conv(input, weight=self.weight.to(input.dtype), groups=self.groups, padding='same')
94 |
95 |
96 | def forward(self, x, context=None, mask=None):
97 | is_cross = context is not None
98 | att_type = "self" if context is None else "cross"
99 |
100 | h = self.heads
101 |
102 | q = self.to_q(x)
103 | context = x if context is None else context
104 | k = self.to_k(context)
105 | v = self.to_v(context)
106 |
107 | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
108 |
109 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
110 | sim_before = sim
111 | del q, k
112 |
113 | if mask is not None:
114 | mask = rearrange(mask, 'b ... -> b (...)')
115 | max_neg_value = -torch.finfo(sim.dtype).max
116 | mask = repeat(mask, 'b j -> (b h) () j', h=h)
117 | sim.masked_fill_(~mask, max_neg_value)
118 |
119 | if hasattr(share, '_crossattn_similarity_res8') and x.shape[1] == share.input_shape.res8 and att_type == 'cross':
120 | share._crossattn_similarity_res8.append(torch.stack(share.reshape(sim).chunk(2))) # Chunk into 2 parts to differentiate the unconditional and conditional parts
121 | if hasattr(share, '_crossattn_similarity_res16') and x.shape[1] == share.input_shape.res16 and att_type == 'cross':
122 | share._crossattn_similarity_res16.append(torch.stack(share.reshape(sim).chunk(2))) # Chunk into 2 parts to differentiate the unconditional and conditional parts
123 | if hasattr(share, '_crossattn_similarity_res32') and x.shape[1] == share.input_shape.res32 and att_type == 'cross':
124 | share._crossattn_similarity_res32.append(torch.stack(share.reshape(sim).chunk(2))) # Chunk into 2 parts to differentiate the unconditional and conditional parts
125 | if hasattr(share, '_crossattn_similarity_res64') and x.shape[1] == share.input_shape.res64 and att_type == 'cross':
126 | share._crossattn_similarity_res64.append(torch.stack(share.reshape(sim).chunk(2))) # Chunk into 2 parts to differentiate the unconditional and conditional parts
127 |
128 | sim = sim.softmax(dim=-1)
129 | out = einsum('b i j, b j d -> b i d', sim, v)
130 | out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
131 |
132 | if is_cross:
133 | return self.to_out(out)
134 |
135 | return self.to_out(out), v, sim_before
136 |
137 |
138 | def introvert_rescale(y, self_v, self_sim, cross_sim, self_h, to_out):
139 | mask = share.introvert_mask.get_res(self_v)
140 | shape = share.introvert_mask.get_shape(self_v)
141 | res = share.introvert_mask.get_res_val(self_v)
142 | # print (res, shape)
143 |
144 | ################# Introvert Attention V4 ################
145 | # Use this with 50% of DDIM steps
146 |
147 | # TODO: LOOK INTO THIS. WHY WITHOUT BINARY WORKS BETTER????
148 | # mask = (mask > 0.5).to(torch.float32)
149 | m = mask.to(self_v.device)
150 | # mask_smoothing = GaussianSmoothing(channels=1, kernel_size=3, sigma=0.5, dim=2).cuda()
151 | # m = mask_smoothing(m) # Smoothing on binary mask also works
152 | m = rearrange(m, 'b c h w -> b (h w) c').contiguous()
153 | mo = m
154 | m = torch.matmul(m, m.permute(0, 2, 1)) + (1-m)
155 |
156 | cross_sim = cross_sim[:, token_idx].sum(dim=1)
157 | cross_sim = cross_sim.reshape(shape)
158 | # TODO: comment this out if it is not neccessary
159 | heatmap_smoothing = GaussianSmoothing(channels=1, kernel_size=3, sigma=0.5, dim=2).cuda()
160 | cross_sim = heatmap_smoothing(cross_sim.unsqueeze(0))[0]
161 | cross_sim = cross_sim.reshape(-1)
162 | cross_sim = ((cross_sim - torch.median(cross_sim.ravel())) / torch.max(cross_sim.ravel())).clip(0, 1)
163 |
164 | # If introvert attention is off, return original SA result
165 | if introvert_on and res in introvert_res:
166 | w = (1 - m) * cross_sim.reshape(1, 1, -1) + m
167 | # On 64 resolution make scaling with constant, as cross_sim doesn't contain semantic meaning
168 | if res == 64: w = m
169 | self_sim = self_sim * w
170 | self_sim_viz = self_sim # Keep for viz purposes
171 | self_sim = self_sim.softmax(dim=-1)
172 | out = einsum('b i j, b j d -> b i d', self_sim, self_v)
173 | out = rearrange(out, '(b h) n d -> b n (h d)', h=self_h)
174 | out = to_out(out)
175 | else:
176 | self_sim_viz = self_sim # Keep for viz purposes
177 | out = y
178 | ################## END OF Introvert Attention V4 ###########################
179 |
180 |
181 | ################# VISUALIZE CROSS ATTENTION ###############################
182 | if visualize_crossattn and res == visualize_resolution:
183 | cross_vis = cross_sim.reshape(shape)
184 | up = (64 // res) * 8
185 | if viz_image is not None:
186 | heatmap = IImage(cross_vis, vmin=0).heatmap((cross_vis.shape[0]*up, cross_vis.shape[1]*up))
187 | video_frames_crossattn.append(viz_image + heatmap)
188 | else:
189 | heatmap = IImage(cross_vis, vmin=0).heatmap((cross_vis.shape[0]*up, cross_vis.shape[1]*up))
190 | video_frames_crossattn.append(heatmap)
191 | ###############################
192 |
193 | ################# VISUALIZE SELF ATTENTION ###############################
194 | if visualize_selfattn and res == visualize_resolution:
195 | selected = []
196 | up = (64 // res) * 8
197 | for i in range(mo.shape[1]):
198 | if mo[0, i, 0]:
199 | selected.append(self_sim_viz[:, i, :])
200 | selected = torch.stack(selected, dim=1)
201 | selected_vis = selected.mean(0).mean(0)
202 |
203 | selected_vis = selected_vis.reshape(shape)
204 |
205 | if viz_image is not None:
206 | heatmap = IImage(selected_vis, vmin=0, vmax=1).heatmap((selected_vis.shape[0]*up, selected_vis.shape[1]*up))
207 | video_frames_selfattn.append(viz_image+heatmap)
208 | else:
209 | heatmap = IImage(selected_vis, vmin=0).heatmap((selected_vis.shape[0]*up, selected_vis.shape[1]*up))
210 | video_frames_selfattn.append(heatmap)
211 | ###############################
212 |
213 | return out
214 |
215 |
--------------------------------------------------------------------------------
/src/smplfusion/patches/attentionpatch/maskscore.py:
--------------------------------------------------------------------------------
1 | from ... import share
2 | import torch
3 |
4 | # import xformers
5 | # import xformers.ops
6 |
7 | import torchvision.transforms.functional as TF
8 |
9 | from torch import nn, einsum
10 | from inspect import isfunction
11 | from einops import rearrange, repeat
12 |
13 | qkv_reduce_dims = [-1]
14 | increase_indices = [1,2]
15 |
16 | def forward(self, x, context=None, mask=None):
17 | h = self.heads
18 | q = self.to_q(x)
19 | is_cross = context is not None
20 | context = x if context is None else context
21 | k = self.to_k(context)
22 | v = self.to_v(context)
23 |
24 | q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
25 |
26 | scale = self.dim_head**-0.5
27 | sim = einsum("b i d, b j d -> b i j", q, k) * scale
28 |
29 | if is_cross:
30 | if q.shape[1] in [share.input_shape.res16]:
31 | # For simplicity, assume token 1 is target, token 2 is (1 word label)
32 | # sim: (16x4096x77); mask: (64x64) -> mask.reshape(-1): (4096)
33 | N = sim.shape[0]//2
34 | sim[N:] += (
35 | share.w
36 | * share.noise_signal_ratio[share.timestep]
37 | * share.input_mask.get_res(q, 'cuda').reshape(1,-1,1)# (1,4096,1)
38 | * sim[N:].amax(dim=qkv_reduce_dims, keepdim = True) # (16x1x1)
39 | ) * (torch.eye(77)[increase_indices].sum(0).cuda()) # (1,1,77)
40 |
41 | del q, k
42 | sim = sim.softmax(dim=-1)
43 |
44 | out = einsum("b i j, b j d -> b i d", sim, v)
45 | out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
46 | return self.to_out(out)
--------------------------------------------------------------------------------
/src/smplfusion/patches/attentionpatch/other.py:
--------------------------------------------------------------------------------
1 | # import xformers
2 | # import xformers.ops
3 |
4 | def forward(self, x, context=None, mask=None):
5 | q = self.to_q(x)
6 | context = x if context is None else context
7 | k = self.to_k(context)
8 | v = self.to_v(context)
9 |
10 | b, _, _ = q.shape
11 | q, k, v = map(
12 | lambda t: t.unsqueeze(3)
13 | .reshape(b, t.shape[1], self.heads, self.dim_head)
14 | .permute(0, 2, 1, 3)
15 | .reshape(b * self.heads, t.shape[1], self.dim_head)
16 | .contiguous(),
17 | (q, k, v),
18 | )
19 |
20 | out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
21 |
22 | if mask is not None:
23 | raise NotImplementedError
24 | out = (
25 | out.unsqueeze(0)
26 | .reshape(b, self.heads, out.shape[1], self.dim_head)
27 | .permute(0, 2, 1, 3)
28 | .reshape(b, out.shape[1], self.heads * self.dim_head)
29 | )
30 | return self.to_out(1 - out)
--------------------------------------------------------------------------------
/src/smplfusion/patches/attentionpatch/shuffled.py:
--------------------------------------------------------------------------------
1 | from ... import share
2 |
3 | # import xformers
4 | # import xformers.ops
5 |
6 | import torch
7 | from torch import nn, einsum
8 | import torchvision.transforms.functional as TF
9 | from einops import rearrange, repeat
10 |
11 | layer_mask = share.LayerMask()
12 |
13 | def forward(self, x, context=None, mask=None):
14 | att_type = "self" if context is None else "cross"
15 |
16 | h = self.heads
17 | q = self.to_q(x)
18 | context = x if context is None else context
19 | k = self.to_k(context)
20 | v = self.to_v(context)
21 |
22 | q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
23 |
24 | sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
25 |
26 | # attention, what we cannot get enough of
27 | sim = sim.softmax(dim=-1)
28 | out = einsum("b i j, b j d -> b i d", sim, v)
29 |
30 | if att_type == 'self' and q.shape[1] in [share.input_shape.res16]:
31 | out = out[:,torch.randperm(out.shape[1]),:]
32 |
33 | out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
34 | return self.to_out(out)
--------------------------------------------------------------------------------
/src/smplfusion/patches/attentionpatch/zeropaint.py:
--------------------------------------------------------------------------------
1 | from ... import share
2 | import torch
3 |
4 | # import xformers
5 | # import xformers.ops
6 |
7 | import torchvision.transforms.functional as TF
8 |
9 | from torch import nn, einsum
10 | from inspect import isfunction
11 | from einops import rearrange, repeat
12 |
13 | qkv_reduce_dims = [-1, -2]
14 |
15 | def forward(self, x, context=None, mask=None):
16 | h = self.heads
17 | q = self.to_q(x)
18 | is_cross = context is not None
19 | context = x if context is None else context
20 | k = self.to_k(context)
21 | v = self.to_v(context)
22 |
23 | q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
24 | scale = self.dim_head**-0.5
25 | sim = einsum("b i d, b j d -> b i j", q, k)
26 | steps_40 = [999, 974, 949, 924, 899, 874, 849, 824, 799, 774, 749, 724, 699, 674, 649, 624, 599, 574, 549, 524, 499, 474, 449, 424, 399, 374, 349, 324, 299, 274, 249, 224, 199, 174, 149, 124, 99, 74, 49, 24]
27 | test_ind = steps_40.index(share.timestep)
28 | if hasattr(share, 'list_of_masks') and is_cross:
29 | if q.shape[1] in [share.input_shape.res16, share.input_shape.res32] or True:
30 | for masked_object in share.list_of_masks:
31 | # sim: (16x4096x77); mask: (64x64) -> mask.reshape(-1): (4096)
32 | zp_condition = (
33 | masked_object['w']
34 | # * share.zp_sigmas[share.timestep]
35 | * share.zp_sigmas[test_ind]
36 | * masked_object['mask'].get_res(q, 'cuda').reshape(1,-1,1)# (1,4096,1)
37 | * sim.amax(dim=qkv_reduce_dims, keepdim = True) # (16x1x1)
38 | ) * (torch.eye(77)[masked_object['token_idx']].sum(0).cuda()) # (1,1,77)
39 |
40 | zp_zeros = torch.zeros_like(zp_condition).cuda()
41 | final_condition = torch.concat([zp_zeros[:8,:,:],zp_condition[8:,:,:]])
42 |
43 | sim = sim+final_condition
44 | # print(sim.shape,"wwwww")
45 | del q, k
46 |
47 | sim = sim*scale
48 | sim = sim.softmax(dim=-1)
49 |
50 | out = einsum("b i j, b j d -> b i d", sim, v)
51 | out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
52 | return self.to_out(out)
--------------------------------------------------------------------------------
/src/smplfusion/patches/router.py:
--------------------------------------------------------------------------------
1 | from . import attentionpatch
2 | from . import transformerpatch
3 |
4 | VERBOSE = False
5 | attention_forward = attentionpatch.default.forward
6 | basic_transformer_forward = transformerpatch.default.forward
7 |
8 | def reset():
9 | global attention_forward, basic_transformer_forward
10 | attention_forward = attentionpatch.default.forward
11 | basic_transformer_forward = transformerpatch.default.forward
12 | if VERBOSE: print ("Resetting Diffusion Model")
13 |
14 | print ("RELOADING ROUTER")
15 | print (attention_forward)
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/__init__.py:
--------------------------------------------------------------------------------
1 | from . import default
2 | from . import guided
3 | from . import weighting_versions
4 | from . import introvert
5 |
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/__pycache__/__init__.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/patches/transformerpatch/__pycache__/__init__.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/__pycache__/default.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/patches/transformerpatch/__pycache__/default.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/__pycache__/guided.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/patches/transformerpatch/__pycache__/guided.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/__pycache__/introvert.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/patches/transformerpatch/__pycache__/introvert.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/__pycache__/weighting_versions.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/patches/transformerpatch/__pycache__/weighting_versions.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/default.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from ... import share
3 |
4 | def set_wsa(value):
5 | global wsa
6 | wsa = torch.tensor([1., value])[:,None,None].cuda()
7 | def set_wca(value):
8 | global wca
9 | wca = torch.tensor([1., value])[:,None,None].cuda()
10 | def set_wff(value):
11 | global wff
12 | wff = torch.tensor([1., value])[:,None,None].cuda()
13 |
14 | # set_wca(1.)
15 | # set_wsa(1.)
16 | # set_wff(1.)
17 |
18 | wca = 1.0
19 | wsa = 1.0
20 | wff = 1.0
21 |
22 | def forward(self, x, context=None):
23 | x = x + wsa * self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) # Self Attn.
24 | x = x + wca * self.attn2(self.norm2(x), context=context) # Cross Attn.
25 | x = x + wff * self.ff(self.norm3(x))
26 | return x
27 |
28 | def forward_and_save(self, x, context=None):
29 | val = [x]
30 | val.append(wsa * self.attn1(self.norm1(x), context=context if self.disable_self_attn else None)) # Self Attn.
31 | x = x + val[-1]
32 | val.append(wca * self.attn2(self.norm2(x), context=context)) # Cross Attn.
33 | x = x + val[-1]
34 | val.append(wff * self.ff(self.norm3(x)))
35 | x = x + val[-1]
36 |
37 | # Save outputs
38 | if hasattr(share, '_basic_transformer_input') and x.shape[1] == share.input_shape.res16:
39 | share._basic_transformer_input.append(share.reshape(val[0]))
40 | if hasattr(share, '_basic_transformer_selfattn') and x.shape[1] == share.input_shape.res16:
41 | share._basic_transformer_selfattn.append(share.reshape(val[1]))
42 | if hasattr(share, '_basic_transformer_crossattn') and x.shape[1] == share.input_shape.res16:
43 | share._basic_transformer_crossattn.append(share.reshape(val[2]))
44 | if hasattr(share, '_basic_transformer_ff') and x.shape[1] == share.input_shape.res16:
45 | share._basic_transformer_ff.append(share.reshape(val[3]))
46 | return x
47 |
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/guided.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from ... import share
3 |
4 | def set_wsa(value):
5 | global wsa
6 | wsa = torch.tensor([1., value])[:,None,None].cuda()
7 | def set_wca(value):
8 | global wca
9 | wca = torch.tensor([1., value])[:,None,None].cuda()
10 | def set_wff(value):
11 | global wff
12 | wff = torch.tensor([1., value])[:,None,None].cuda()
13 |
14 | set_wca(1.)
15 | set_wsa(1.)
16 | set_wff(1.)
17 |
18 | guidance_scale = 7.5
19 |
20 | def forward(self, x, context=None):
21 | # print (x.shape)
22 | mask = share.input_mask.get_res(x).reshape(-1,1).cuda()
23 |
24 | _out_sa = wsa * self.attn1(self.norm1(x), None) # Self Attn.
25 |
26 | # _out_ca_uncond = wca * self.attn2(self.norm2(x), context=context) # Cross Attn. "Unconditional"
27 | _out_ca = wca * self.attn2(self.norm2(x + _out_sa), context=context) # Cross Attn.
28 |
29 | if share.timestep < 0:
30 | # x = x + (1 - mask) * _out_sa + mask * (guidance_scale * _out_ca_uncond + (1 - guidance_scale) * (_out_sa + _out_ca))
31 | x = x + (1 - mask) * (_out_sa + _out_ca) + mask * (guidance_scale * _out_ca_uncond + (1 - guidance_scale) * (_out_sa + _out_ca))
32 | else:
33 | x = x + _out_sa + _out_ca
34 |
35 | _out_ff = wff * self.ff(self.norm3(x))
36 | x = x + _out_ff
37 |
38 | if False:
39 | save([_out_ca, _out_sa, _out_ff])
40 |
41 | return x
42 |
43 | forward_and_save = forward
44 |
45 | def save(val):# Save outputs
46 | if hasattr(share, '_basic_transformer_selfattn'):
47 | share._basic_transformer_selfattn.append(share.reshape(val[0]))
48 | if hasattr(share, '_basic_transformer_crossattn'):
49 | share._basic_transformer_crossattn.append(share.reshape(val[1]))
50 | if hasattr(share, '_basic_transformer_ff'):
51 | share._basic_transformer_ff.append(share.reshape(val[1]))
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/introvert.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn, einsum
3 | from einops import rearrange, repeat
4 | from ... import share
5 | from ..attentionpatch import introvert
6 |
7 | w_ff = 1.
8 | w_sa = 1.
9 | w_ca = 1.
10 |
11 | use_grad = True
12 |
13 |
14 | def forward(self, x, context=None):
15 | # with torch.no_grad():
16 | if use_grad:
17 | y, self_v, self_sim = self.attn1(self.norm1(x), None) # Self Attn.
18 |
19 | x_uncond, x_cond = x.chunk(2)
20 | context_uncond, context_cond = context.chunk(2)
21 |
22 | y_uncond, y_cond = y.chunk(2)
23 | self_sim_uncond, self_sim_cond = self_sim.chunk(2)
24 | self_v_uncond, self_v_cond = self_v.chunk(2)
25 |
26 | # Calculate CA similarities with conditional context
27 | cross_h = self.attn2.heads
28 | cross_q = self.attn2.to_q(self.norm2(x_cond+y_cond))
29 | cross_k = self.attn2.to_k(context_cond)
30 | cross_v = self.attn2.to_v(context_cond)
31 |
32 | cross_q, cross_k, cross_v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=cross_h), (cross_q, cross_k, cross_v))
33 |
34 | with torch.autocast(enabled=False, device_type = 'cuda'):
35 | cross_q, cross_k = cross_q.float(), cross_k.float()
36 | cross_sim = einsum('b i d, b j d -> b i j', cross_q, cross_k) * self.attn2.scale
37 |
38 | del cross_q, cross_k
39 | cross_sim = cross_sim.softmax(dim=-1) # Up to this point cross_sim is regular cross_sim in CA layer
40 |
41 | cross_sim = cross_sim.mean(dim=0) # Calculate mean across heads
42 |
43 | # Introvert Attention rescale heppens here
44 | y_cond = introvert.introvert_rescale(
45 | y_cond, self_v_cond, self_sim_cond, cross_sim, self.attn1.heads, self.attn1.to_out) # Rescale cond
46 | y_uncond = introvert.introvert_rescale(
47 | y_uncond, self_v_uncond, self_sim_uncond, cross_sim, self.attn1.heads, self.attn1.to_out) # Rescale uncond
48 |
49 | y = torch.cat([y_uncond, y_cond], dim=0)
50 |
51 | x = x + w_sa * y
52 | x = x + w_ca * self.attn2(self.norm2(x), context=context) # Cross Attn.
53 | x = x + w_ff * self.ff(self.norm3(x))
54 | return x
--------------------------------------------------------------------------------
/src/smplfusion/patches/transformerpatch/weighting_versions.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from ... import share
3 |
4 | def set_wsa(value):
5 | global wsa
6 | wsa = torch.tensor([1., value])[:,None,None].cuda()
7 | def set_wca(value):
8 | global wca
9 | wca = torch.tensor([1., value])[:,None,None].cuda()
10 | def set_wff(value):
11 | global wff
12 | wff = torch.tensor([1., value])[:,None,None].cuda()
13 |
14 | set_wca(1.)
15 | set_wsa(1.)
16 | set_wff(1.)
17 |
18 | def forward(self, x, context=None):
19 | x = x + wca * self.attn2(self.norm2(x), context=context) # Cross Attn.
20 | x = x + wsa * self.attn1(self.norm1(x), None) # Self Attn.
21 | x = x + wff * self.ff(self.norm3(x))
22 | return x
23 |
24 | def forward_and_save(self, x, context=None):
25 | val = [x]
26 | val.append(wsa * self.attn1(self.norm1(x), None)) # Self Attn.
27 | x = x + val[-1]
28 | val.append(wca * self.attn2(self.norm2(x), context=context)) # Cross Attn.
29 | x = x + val[-1]
30 | val.append(wff * self.ff(self.norm3(x)))
31 | x = x + val[-1]
32 |
33 | # Save outputs
34 | if hasattr(share, '_basic_transformer_input') and x.shape[1] == share.input_shape.res16:
35 | share._basic_transformer_input.append(share.reshape(val[0]))
36 | if hasattr(share, '_basic_transformer_selfattn') and x.shape[1] == share.input_shape.res16:
37 | share._basic_transformer_selfattn.append(share.reshape(val[1]))
38 | if hasattr(share, '_basic_transformer_crossattn') and x.shape[1] == share.input_shape.res16:
39 | share._basic_transformer_crossattn.append(share.reshape(val[2]))
40 | if hasattr(share, '_basic_transformer_ff') and x.shape[1] == share.input_shape.res16:
41 | share._basic_transformer_ff.append(share.reshape(val[3]))
42 | return x
43 |
44 |
45 | # ========= MODIFICATIONS ============= #
46 |
47 |
48 | def forward_and_save3(self, x, context=None):
49 | val = []
50 | val.append(self.attn1(self.norm1(x), None)) # Self Attn.
51 | modify_res = [share.input_mask.res16, share.input_mask.res32, share.input_mask.res64]
52 |
53 | if x.shape[1] in modify_res:
54 | val[-1] = val[-1] + (wsa - 1) * share.input_mask.get_res(x).reshape(-1)[:,None].cuda() * val[-1]
55 | x = x + val[-1]
56 |
57 | val.append(self.attn2(self.norm2(x), context=context)) # Cross Attn.
58 |
59 | if x.shape[1] in modify_res:
60 | val[-1] = val[-1] + (wca - 1) * share.input_mask.get_res(x).reshape(-1)[:,None].cuda() * val[-1]
61 | x = x + val[-1]
62 |
63 | val.append(self.ff(self.norm3(x)))
64 |
65 | if x.shape[1] in modify_res:
66 | val[-1] = val[-1] + (wff - 1) * share.input_mask.get_res(x).reshape(-1)[:,None].cuda() * val[-1]
67 | x = x + val[-1]
68 | if hasattr(share, 'out_basic_transformer_block') and x.shape[1] == share.input_mask.res16:
69 | share.out_basic_transformer_block.append(val)
70 | return x
71 |
72 | def forward_and_save2(self, x, context=None):
73 | val = []
74 | val.append(self.attn1(self.norm1(x), None)) # Self Attn.
75 | x = x + wsa * val[-1]
76 | val.append(self.attn2(self.norm2(x), context=context)) # Cross Attn.
77 | print (val[-1].shape, share.input_mask.val16.shape)
78 | x = x + val[-1] + wca * share.input_mask.val16 * val[-1]
79 | val.append(self.ff(self.norm3(x)))
80 | x = x + wff * val[-1]
81 | if hasattr(share, 'out_basic_transformer_block') and x.shape[1] == share.input_mask.res16:
82 | share.out_basic_transformer_block.append(val)
83 | return x
84 |
85 | def forward_and_reweight(self, x, context=None):
86 | modify_res = [share.input_mask.res16, share.input_mask.res32, share.input_mask.res64]
87 |
88 | _attn1 = self.attn1(self.norm1(x), None) # Self Attn.
89 | _attn2 = self.attn2(self.norm2(x + _attn1), context=context) # Cross Attn.
90 | _ff = self.ff(self.norm3(x + _attn1 + _attn2))
91 |
92 | if x.shape[1] in modify_res:
93 | lm1,lm2,lm3 = wsa - 1, wca - 1, wff - 1
94 | _attn1 *= (1 + lm1 * share.input_mask.get_res(x).reshape(-1)[:,None].cuda())
95 | _attn2 *= (1 + lm2 * share.input_mask.get_res(x).reshape(-1)[:,None].cuda())
96 | _ff *= (1 + lm3 * share.input_mask.get_res(x).reshape(-1)[:,None].cuda())
97 |
98 | if hasattr(share, 'out_basic_transformer_block') and x.shape[1] == share.input_mask.res16:
99 | share.out_basic_transformer_block.append([_attn1, _attn2, _ff])
100 |
101 | return x + _attn1 + _attn2 + _ff
102 | # return x + (w_sa * _attn1 + w_ca * _attn2 + w_ff * _ff)
103 | # return x + (w_sa * _attn1 + w_ca * _attn2 + w_ff * _ff) / ((w_sa + w_ca + w_ff) / 3)
104 |
105 | def forward_and_reweight2(self, x, context=None):
106 | _attn1 = self.attn1(self.norm1(x), None) # Self Attn.
107 | _attn2 = self.attn2(self.norm2(x + _attn1), context=context) # Cross Attn.
108 | _ff = self.ff(self.norm3(x + _attn1 + _attn2))
109 |
110 | _attn1 += (wsa - 1) / ((wsa + wca + wff) / 3) * share.input_mask.get_res(x).reshape(-1)[:,None].cuda() * _attn1
111 | _attn2 += (wca - 1) / ((wsa + wca + wff) / 3) * share.input_mask.get_res(x).reshape(-1)[:,None].cuda() * _attn2
112 | _ff += (wff - 1) / ((wsa + wca + wff) / 3) * share.input_mask.get_res(x).reshape(-1)[:,None].cuda() * _ff
113 |
114 |
115 | if hasattr(share, 'out_basic_transformer_block') and x.shape[1] == share.input_mask.res16:
116 | share.out_basic_transformer_block.append([_attn1, _attn2, _ff])
117 | # share.out_basic_transformer_block.append([
118 | # w_sa / ((w_sa + w_ca + w_ff) / 3) * _attn1,
119 | # w_ca / ((w_sa + w_ca + w_ff) / 3) * _attn2,
120 | # w_ff / ((w_sa + w_ca + w_ff) / 3) * _ff
121 | # ])
122 |
123 | return x + _attn1 + _attn2 + _ff
124 | # return x + (w_sa * _attn1 + w_ca * _attn2 + w_ff * _ff)
125 | # return x + (w_sa * _attn1 + w_ca * _attn2 + w_ff * _ff) / ((w_sa + w_ca + w_ff) / 3)
--------------------------------------------------------------------------------
/src/smplfusion/scheduler.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | def linear(n_timestep = 1000, start = 1e-4, end = 2e-2):
4 | return Schedule(torch.linspace(start ** 0.5, end ** 0.5, n_timestep, dtype = torch.float64) ** 2)
5 |
6 | class Schedule:
7 | def __init__(self, betas):
8 | self.betas = betas
9 | self._alphas = 1 - betas
10 | self.alphas = torch.cumprod(self._alphas, 0)
11 | self.one_minus_alphas = 1 - self.alphas
12 | self.sqrt_alphas = torch.sqrt(self.alphas)
13 | self.sqrt_one_minus_alphas = torch.sqrt(1 - self.alphas)
14 | self.sqrt_noise_signal_ratio = self.sqrt_one_minus_alphas / self.sqrt_alphas
15 | self.noise_signal_ratio = (1 - self.alphas) / self.alphas
16 |
17 | def range(self, dt):
18 | return range(len(self.betas)-1, 0, -dt)
19 |
20 | def sigma(self, t, dt):
21 | return torch.sqrt((1 - self._alphas[t - dt]) / (1 - self._alphas[t]) * (1 - self._alphas[t] / self._alphas[t - dt])) # Like I did initially
22 | # return torch.sqrt((1 - self.alphas[t - dt]) / (1 - self.alphas[t]) * (1 - self.alphas[t] / self.alphas[t - dt])) # Like in diffusers
23 |
24 | def sigma_(self, t, dt):
25 | return torch.sqrt((1 - self.alphas[t - dt]) / (1 - self.alphas[t]) * (1 - self.alphas[t] / self.alphas[t - dt])) # Like in diffusers
--------------------------------------------------------------------------------
/src/smplfusion/share.py:
--------------------------------------------------------------------------------
1 | import torchvision.transforms.functional as TF
2 | import torch
3 | import sys
4 | from .utils import *
5 |
6 | input_mask = None
7 | input_shape = None
8 | timestep = None
9 | timestep_index = None
10 |
11 | class Seed:
12 | def __getitem__(self, idx):
13 | if isinstance(idx, slice):
14 | idx = list(range(*idx.indices(idx.stop)))
15 | if isinstance(idx, list) or isinstance(idx, tuple):
16 | return [self[_idx] for _idx in idx]
17 | return 12345 ** idx % 54321
18 | seed = Seed()
19 |
20 | lock = {}
21 | def get_lock(value):
22 | global lock
23 | if value not in lock:
24 | lock[value] = True
25 | if lock[value]:
26 | lock[value] = False
27 | return True
28 | return False
29 |
30 |
31 | class DDIMIterator:
32 | def __init__(self, iterator):
33 | self.iterator = iterator
34 | def __iter__(self):
35 | self.iterator = iter(self.iterator)
36 | global timestep_index
37 | timestep_index = 0
38 | return self
39 | def __next__(self):
40 | global timestep, timestep_index, lock
41 | for x in lock: lock[x] = True
42 | timestep = next(self.iterator)
43 | timestep_index += 1
44 | return timestep
45 |
46 | def reshape(x):
47 | return input_shape.reshape(x)
48 |
49 | def set_shape(image_or_shape):
50 | global input_shape
51 | if hasattr(image_or_shape, 'size'):
52 | input_shape = InputShape(image_or_shape.size)
53 | if isinstance(image_or_shape, torch.Tensor):
54 | input_shape = InputShape(image_or_shape.shape[-2:][::-1])
55 | elif isinstance(image_or_shape, list) or isinstance(image_or_shape, tuple):
56 | input_shape = InputShape(image_or_shape)
57 |
58 | self = sys.modules[__name__]
59 |
60 | def set_mask(mask):
61 | global input_mask, mask64, mask32, mask16, mask8, introvert_mask
62 | input_mask = InputMask(mask)
63 | introvert_mask = InputMask(mask)
64 |
65 | mask64 = input_mask.val64[0,0]
66 | mask32 = input_mask.val32[0,0]
67 | mask16 = input_mask.val16[0,0]
68 | mask8 = input_mask.val8[0,0]
69 | set_shape(mask)
70 |
71 | def exists(name):
72 | return hasattr(self, name) and getattr(self, name) is not None
--------------------------------------------------------------------------------
/src/smplfusion/util.py:
--------------------------------------------------------------------------------
1 | # TODO: remove everything below as long as it doesn't break anything!
2 |
3 | import importlib
4 |
5 | import torch
6 | from torch import optim
7 | import numpy as np
8 |
9 | from inspect import isfunction
10 | from PIL import Image, ImageDraw, ImageFont
11 |
12 |
13 | def autocast(f):
14 | def do_autocast(*args, **kwargs):
15 | with torch.cuda.amp.autocast(enabled=True,
16 | dtype=torch.get_autocast_gpu_dtype(),
17 | cache_enabled=torch.is_autocast_cache_enabled()):
18 | return f(*args, **kwargs)
19 |
20 | return do_autocast
21 |
22 |
23 | def log_txt_as_img(wh, xc, size=10):
24 | # wh a tuple of (width, height)
25 | # xc a list of captions to plot
26 | b = len(xc)
27 | txts = list()
28 | for bi in range(b):
29 | txt = Image.new("RGB", wh, color="white")
30 | draw = ImageDraw.Draw(txt)
31 | font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
32 | nc = int(40 * (wh[0] / 256))
33 | lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
34 |
35 | try:
36 | draw.text((0, 0), lines, fill="black", font=font)
37 | except UnicodeEncodeError:
38 | print("Cant encode string for logging. Skipping.")
39 |
40 | txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
41 | txts.append(txt)
42 | txts = np.stack(txts)
43 | txts = torch.tensor(txts)
44 | return txts
45 |
46 |
47 | def ismap(x):
48 | if not isinstance(x, torch.Tensor):
49 | return False
50 | return (len(x.shape) == 4) and (x.shape[1] > 3)
51 |
52 |
53 | def isimage(x):
54 | if not isinstance(x,torch.Tensor):
55 | return False
56 | return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
57 |
58 |
59 | def exists(x):
60 | return x is not None
61 |
62 |
63 | def default(val, d):
64 | if exists(val):
65 | return val
66 | return d() if isfunction(d) else d
67 |
68 |
69 | def mean_flat(tensor):
70 | """
71 | https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
72 | Take the mean over all non-batch dimensions.
73 | """
74 | return tensor.mean(dim=list(range(1, len(tensor.shape))))
75 |
76 |
77 | def count_params(model, verbose=False):
78 | total_params = sum(p.numel() for p in model.parameters())
79 | if verbose:
80 | print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
81 | return total_params
82 |
83 |
84 | def instantiate_from_config(config):
85 | if not "target" in config:
86 | if config == '__is_first_stage__':
87 | return None
88 | elif config == "__is_unconditional__":
89 | return None
90 | raise KeyError("Expected key `target` to instantiate.")
91 | return get_obj_from_str(config["target"])(**config.get("params", dict()))
92 |
93 |
94 | def get_obj_from_str(string, reload=False):
95 | module, cls = string.rsplit(".", 1)
96 | if reload:
97 | module_imp = importlib.import_module(module)
98 | importlib.reload(module_imp)
99 | return getattr(importlib.import_module(module, package=None), cls)
100 |
101 |
102 | class AdamWwithEMAandWings(optim.Optimizer):
103 | # credit to https://gist.github.com/crowsonkb/65f7265353f403714fce3b2595e0b298
104 | def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8, # TODO: check hyperparameters before using
105 | weight_decay=1.e-2, amsgrad=False, ema_decay=0.9999, # ema decay to match previous code
106 | ema_power=1., param_names=()):
107 | """AdamW that saves EMA versions of the parameters."""
108 | if not 0.0 <= lr:
109 | raise ValueError("Invalid learning rate: {}".format(lr))
110 | if not 0.0 <= eps:
111 | raise ValueError("Invalid epsilon value: {}".format(eps))
112 | if not 0.0 <= betas[0] < 1.0:
113 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
114 | if not 0.0 <= betas[1] < 1.0:
115 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
116 | if not 0.0 <= weight_decay:
117 | raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
118 | if not 0.0 <= ema_decay <= 1.0:
119 | raise ValueError("Invalid ema_decay value: {}".format(ema_decay))
120 | defaults = dict(lr=lr, betas=betas, eps=eps,
121 | weight_decay=weight_decay, amsgrad=amsgrad, ema_decay=ema_decay,
122 | ema_power=ema_power, param_names=param_names)
123 | super().__init__(params, defaults)
124 |
125 | def __setstate__(self, state):
126 | super().__setstate__(state)
127 | for group in self.param_groups:
128 | group.setdefault('amsgrad', False)
129 |
130 | @torch.no_grad()
131 | def step(self, closure=None):
132 | """Performs a single optimization step.
133 | Args:
134 | closure (callable, optional): A closure that reevaluates the model
135 | and returns the loss.
136 | """
137 | loss = None
138 | if closure is not None:
139 | with torch.enable_grad():
140 | loss = closure()
141 |
142 | for group in self.param_groups:
143 | params_with_grad = []
144 | grads = []
145 | exp_avgs = []
146 | exp_avg_sqs = []
147 | ema_params_with_grad = []
148 | state_sums = []
149 | max_exp_avg_sqs = []
150 | state_steps = []
151 | amsgrad = group['amsgrad']
152 | beta1, beta2 = group['betas']
153 | ema_decay = group['ema_decay']
154 | ema_power = group['ema_power']
155 |
156 | for p in group['params']:
157 | if p.grad is None:
158 | continue
159 | params_with_grad.append(p)
160 | if p.grad.is_sparse:
161 | raise RuntimeError('AdamW does not support sparse gradients')
162 | grads.append(p.grad)
163 |
164 | state = self.state[p]
165 |
166 | # State initialization
167 | if len(state) == 0:
168 | state['step'] = 0
169 | # Exponential moving average of gradient values
170 | state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
171 | # Exponential moving average of squared gradient values
172 | state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
173 | if amsgrad:
174 | # Maintains max of all exp. moving avg. of sq. grad. values
175 | state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
176 | # Exponential moving average of parameter values
177 | state['param_exp_avg'] = p.detach().float().clone()
178 |
179 | exp_avgs.append(state['exp_avg'])
180 | exp_avg_sqs.append(state['exp_avg_sq'])
181 | ema_params_with_grad.append(state['param_exp_avg'])
182 |
183 | if amsgrad:
184 | max_exp_avg_sqs.append(state['max_exp_avg_sq'])
185 |
186 | # update the steps for each param group update
187 | state['step'] += 1
188 | # record the step after step update
189 | state_steps.append(state['step'])
190 |
191 | optim._functional.adamw(params_with_grad,
192 | grads,
193 | exp_avgs,
194 | exp_avg_sqs,
195 | max_exp_avg_sqs,
196 | state_steps,
197 | amsgrad=amsgrad,
198 | beta1=beta1,
199 | beta2=beta2,
200 | lr=group['lr'],
201 | weight_decay=group['weight_decay'],
202 | eps=group['eps'],
203 | maximize=False)
204 |
205 | cur_ema_decay = min(ema_decay, 1 - state['step'] ** -ema_power)
206 | for param, ema_param in zip(params_with_grad, ema_params_with_grad):
207 | ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)
208 |
209 | return loss
--------------------------------------------------------------------------------
/src/smplfusion/utils/__init__.py:
--------------------------------------------------------------------------------
1 | from .input_image import InputImage
2 | from .input_mask import InputMask
3 | from .input_shape import InputShape
4 | from .layer_mask import LayerMask
--------------------------------------------------------------------------------
/src/smplfusion/utils/__pycache__/__init__.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/utils/__pycache__/__init__.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/utils/__pycache__/input_image.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/utils/__pycache__/input_image.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/utils/__pycache__/input_mask.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/utils/__pycache__/input_mask.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/utils/__pycache__/input_shape.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/utils/__pycache__/input_shape.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/utils/__pycache__/layer_mask.cpython-39.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/smplfusion/utils/__pycache__/layer_mask.cpython-39.pyc
--------------------------------------------------------------------------------
/src/smplfusion/utils/input_image.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from ..libimage import IImage
3 |
4 | class InputImage:
5 | def to(self, device): return InputImage(self.image, device = device)
6 | def cuda(self): return InputImage(self.image, device = 'cuda')
7 | def cpu(self): return InputImage(self.image, device = 'cpu')
8 |
9 | def __init__(self, input_image):
10 | '''
11 | args:
12 | input_image: (b,c,h,w) tensor
13 | '''
14 | if hasattr(input_image, 'is_iimage'):
15 | self.image = input_image
16 | self.val512 = self.full = input_image.torch(0)
17 | elif isinstance(input_image, torch.Tensor):
18 | self.val512 = self.full = input_image.clone()
19 | self.image = IImage(input_image,0)
20 |
21 | self.h,self.w = h,w = self.val512.shape[-2:]
22 | self.shape = [self.h, self.w]
23 | self.shape64 = [self.h // 8, self.w // 8]
24 | self.shape32 = [self.h // 16, self.w // 16]
25 | self.shape16 = [self.h // 32, self.w // 32]
26 | self.shape8 = [self.h // 64, self.w // 64]
27 |
28 | self.res = self.h * self.w
29 | self.res64 = self.res // 64
30 | self.res32 = self.res // 64 // 4
31 | self.res16 = self.res // 64 // 16
32 | self.res8 = self.res // 64 // 64
33 |
34 | self.img = self.image
35 | self.img512 = self.image
36 | self.img64 = self.image.resize((h//8,w//8))
37 | self.img32 = self.image.resize((h//16,w//16))
38 | self.img16 = self.image.resize((h//32,w//32))
39 | self.img8 = self.image.resize((h//64,w//64))
40 |
41 | self.val64 = self.img64.torch()
42 | self.val32 = self.img32.torch()
43 | self.val16 = self.img16.torch()
44 | self.val8 = self.img8.torch()
45 |
46 | def get_res(self, q, device = 'cpu'):
47 | if q.shape[1] == self.res64: return self.val64.to(device)
48 | if q.shape[1] == self.res32: return self.val32.to(device)
49 | if q.shape[1] == self.res16: return self.val16.to(device)
50 | if q.shape[1] == self.res8: return self.val8.to(device)
51 |
52 | def get_shape(self, q, device = 'cpu'):
53 | if q.shape[1] == self.res64: return self.shape64
54 | if q.shape[1] == self.res32: return self.shape32
55 | if q.shape[1] == self.res16: return self.shape16
56 | if q.shape[1] == self.res8: return self.shape8
57 |
58 | def get_res_val(self, q, device = 'cpu'):
59 | if q.shape[1] == self.res64: return 64
60 | if q.shape[1] == self.res32: return 32
61 | if q.shape[1] == self.res16: return 16
62 | if q.shape[1] == self.res8: return 8
63 |
64 | def _repr_png_(self):
65 | return self.img16._repr_png_()
--------------------------------------------------------------------------------
/src/smplfusion/utils/input_mask.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from ..libimage import IImage
3 |
4 | class InputMask:
5 | def to(self, device): return InputMask(self.image, device = device)
6 | def cuda(self): return InputMask(self.image, device = 'cuda')
7 | def cpu(self): return InputMask(self.image, device = 'cpu')
8 |
9 | def __init__(self, input_image, device = 'cpu'):
10 | '''
11 | args:
12 | input_image: (b,c,h,w) tensor
13 | '''
14 | if hasattr(input_image, 'is_iimage'):
15 | self.image = input_image
16 | self.val512 = self.full = (input_image.torch(0) > 0.5).float()
17 | elif isinstance(input_image, torch.Tensor):
18 | self.val512 = self.full = input_image.clone()
19 | self.image = IImage(input_image,0)
20 |
21 | self.h,self.w = h,w = self.val512.shape[-2:]
22 | self.shape = [self.h, self.w]
23 | self.shape64 = [self.h // 8, self.w // 8]
24 | self.shape32 = [self.h // 16, self.w // 16]
25 | self.shape16 = [self.h // 32, self.w // 32]
26 | self.shape8 = [self.h // 64, self.w // 64]
27 |
28 | self.res = self.h * self.w
29 | self.res64 = self.res // 64
30 | self.res32 = self.res // 64 // 4
31 | self.res16 = self.res // 64 // 16
32 | self.res8 = self.res // 64 // 64
33 |
34 | self.img = self.image
35 | self.img512 = self.image
36 | self.img64 = self.image.resize((h//8,w//8))
37 | self.img32 = self.image.resize((h//16,w//16))
38 | self.img16 = self.image.resize((h//32,w//32))
39 | self.img8 = self.image.resize((h//64,w//64))
40 |
41 | self.val64 = (self.img64.torch(0) > 0.5).float()
42 | self.val32 = (self.img32.torch(0) > 0.5).float()
43 | self.val16 = (self.img16.torch(0) > 0.5).float()
44 | self.val8 = ( self.img8.torch(0) > 0.5).float()
45 |
46 |
47 | def get_res(self, q, device = 'cpu'):
48 | if q.shape[1] == self.res64: return self.val64.to(device)
49 | if q.shape[1] == self.res32: return self.val32.to(device)
50 | if q.shape[1] == self.res16: return self.val16.to(device)
51 | if q.shape[1] == self.res8: return self.val8.to(device)
52 |
53 | def _repr_png_(self):
54 | return self.img16._repr_png_()
55 |
56 | def get_res(self, q, device = 'cpu'):
57 | if q.shape[1] == self.res64: return self.val64.to(device)
58 | if q.shape[1] == self.res32: return self.val32.to(device)
59 | if q.shape[1] == self.res16: return self.val16.to(device)
60 | if q.shape[1] == self.res8: return self.val8.to(device)
61 |
62 | def get_shape(self, q, device = 'cpu'):
63 | if q.shape[1] == self.res64: return self.shape64
64 | if q.shape[1] == self.res32: return self.shape32
65 | if q.shape[1] == self.res16: return self.shape16
66 | if q.shape[1] == self.res8: return self.shape8
67 |
68 | def get_res_val(self, q, device = 'cpu'):
69 | if q.shape[1] == self.res64: return 64
70 | if q.shape[1] == self.res32: return 32
71 | if q.shape[1] == self.res16: return 16
72 | if q.shape[1] == self.res8: return 8
73 |
74 |
75 | # class InputMask2:
76 | # def to(self, device): return InputMask2(self.image, device = device)
77 | # def cuda(self): return InputMask2(self.image, device = 'cuda')
78 | # def cpu(self): return InputMask2(self.image, device = 'cpu')
79 |
80 | # def __init__(self, input_image, device = 'cpu'):
81 | # '''
82 | # args:
83 | # input_image: (b,c,h,w) tensor
84 | # '''
85 | # if hasattr(input_image, 'is_iimage'):
86 | # self.image = input_image
87 | # self.val512 = self.full = input_image.torch(0).bool().float()
88 | # elif isinstance(input_image, torch.Tensor):
89 | # self.val512 = self.full = input_image.clone()
90 | # self.image = IImage(input_image,0)
91 |
92 | # self.h,self.w = h,w = self.val512.shape[-2:]
93 | # self.shape = [self.h, self.w]
94 | # self.shape64 = [self.h // 8, self.w // 8]
95 | # self.shape32 = [self.h // 16, self.w // 16]
96 | # self.shape16 = [self.h // 32, self.w // 32]
97 | # self.shape8 = [self.h // 64, self.w // 64]
98 |
99 | # self.res = self.h * self.w
100 | # self.res64 = self.res // 64
101 | # self.res32 = self.res // 64 // 4
102 | # self.res16 = self.res // 64 // 16
103 | # self.res8 = self.res // 64 // 64
104 |
105 | # self.img = self.image
106 | # self.img512 = self.image
107 | # self.img64 = self.image.resize((h//8,w//8))
108 | # self.img32 = self.image.resize((h//16,w//16))
109 | # self.img16 = self.image.resize((h//32,w//32)).dilate(1)
110 | # self.img8 = self.image.resize((h//64,w//64)).dilate(1)
111 |
112 | # self.val64 = self.img64.torch(0).bool().float()
113 | # self.val32 = self.img32.torch(0).bool().float()
114 | # self.val16 = self.img16.torch(0).bool().float()
115 | # self.val8 = self.img8.torch(0).bool().float()
116 |
117 |
118 | # def get_res(self, q, device = 'cpu'):
119 | # if q.shape[1] == self.res64: return self.val64.to(device)
120 | # if q.shape[1] == self.res32: return self.val32.to(device)
121 | # if q.shape[1] == self.res16: return self.val16.to(device)
122 | # if q.shape[1] == self.res8: return self.val8.to(device)
123 |
124 | # def _repr_png_(self):
125 | # return self.img16._repr_png_()
126 |
127 | # def get_res(self, q, device = 'cpu'):
128 | # if q.shape[1] == self.res64: return self.val64.to(device)
129 | # if q.shape[1] == self.res32: return self.val32.to(device)
130 | # if q.shape[1] == self.res16: return self.val16.to(device)
131 | # if q.shape[1] == self.res8: return self.val8.to(device)
132 |
133 | # def get_shape(self, q, device = 'cpu'):
134 | # if q.shape[1] == self.res64: return self.shape64
135 | # if q.shape[1] == self.res32: return self.shape32
136 | # if q.shape[1] == self.res16: return self.shape16
137 | # if q.shape[1] == self.res8: return self.shape8
138 |
139 | # def get_res_val(self, q, device = 'cpu'):
140 | # if q.shape[1] == self.res64: return 64
141 | # if q.shape[1] == self.res32: return 32
142 | # if q.shape[1] == self.res16: return 16
143 | # if q.shape[1] == self.res8: return 8
--------------------------------------------------------------------------------
/src/smplfusion/utils/input_shape.py:
--------------------------------------------------------------------------------
1 | class InputShape:
2 | def __init__(self, image_size):
3 | self.h,self.w = image_size[::-1]
4 | self.res = self.h * self.w
5 | self.res64 = self.res // 64
6 | self.res32 = self.res // 64 // 4
7 | self.res16 = self.res // 64 // 16
8 | self.res8 = self.res // 64 // 64
9 | self.shape = [self.h, self.w]
10 | self.shape64 = [self.h // 8, self.w // 8]
11 | self.shape32 = [self.h // 16, self.w // 16]
12 | self.shape16 = [self.h // 32, self.w // 32]
13 | self.shape8 = [self.h // 64, self.w // 64]
14 |
15 | def reshape(self, x):
16 | assert len(x.shape) == 3
17 | if x.shape[1] == self.res64: return x.reshape([x.shape[0]] + self.shape64 + [x.shape[-1]])
18 | if x.shape[1] == self.res32: return x.reshape([x.shape[0]] + self.shape32 + [x.shape[-1]])
19 | if x.shape[1] == self.res16: return x.reshape([x.shape[0]] + self.shape16 + [x.shape[-1]])
20 | if x.shape[1] == self.res8: return x.reshape([x.shape[0]] + self.shape8 + [x.shape[-1]])
21 | raise Exception("Unknown shape")
22 |
23 | def get_res(self, q, device = 'cpu'):
24 | if q.shape[1] == self.res64: return 64
25 | if q.shape[1] == self.res32: return 32
26 | if q.shape[1] == self.res16: return 16
27 | if q.shape[1] == self.res8: return 8
--------------------------------------------------------------------------------
/src/smplfusion/utils/layer_mask.py:
--------------------------------------------------------------------------------
1 | class LayerMask:
2 | def __init__(self, self64 = False, self32 = False, self16 = False, self8 = False, cross64 = False, cross32 = False, cross16 = False, cross8 = False):
3 | self.self64 = self64
4 | self.self32 = self32
5 | self.self16 = self16
6 | self.self8 = self8
7 | self.cross64 = cross64
8 | self.cross32 = cross32
9 | self.cross16 = cross16
10 | self.cross8 = cross8
11 |
--------------------------------------------------------------------------------
/src/zeropainter/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/src/zeropainter/__init__.py
--------------------------------------------------------------------------------
/src/zeropainter/convert_diffusers.py:
--------------------------------------------------------------------------------
1 | # Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
2 | # *Only* converts the UNet, VAE, and Text Encoder.
3 | # Does not convert optimizer state or any other thing.
4 | # Written by jachiam https://github.com/jachiam
5 |
6 | import argparse
7 | import os.path as osp
8 |
9 | import torch
10 |
11 |
12 | # =================#
13 | # UNet Conversion #
14 | # =================#
15 |
16 | unet_conversion_map = [
17 | # (stable-diffusion, HF Diffusers)
18 | ("time_embed.0.weight", "time_embedding.linear_1.weight"),
19 | ("time_embed.0.bias", "time_embedding.linear_1.bias"),
20 | ("time_embed.2.weight", "time_embedding.linear_2.weight"),
21 | ("time_embed.2.bias", "time_embedding.linear_2.bias"),
22 | ("input_blocks.0.0.weight", "conv_in.weight"),
23 | ("input_blocks.0.0.bias", "conv_in.bias"),
24 | ("out.0.weight", "conv_norm_out.weight"),
25 | ("out.0.bias", "conv_norm_out.bias"),
26 | ("out.2.weight", "conv_out.weight"),
27 | ("out.2.bias", "conv_out.bias"),
28 | ]
29 |
30 | unet_conversion_map_resnet = [
31 | # (stable-diffusion, HF Diffusers)
32 | ("in_layers.0", "norm1"),
33 | ("in_layers.2", "conv1"),
34 | ("out_layers.0", "norm2"),
35 | ("out_layers.3", "conv2"),
36 | ("emb_layers.1", "time_emb_proj"),
37 | ("skip_connection", "conv_shortcut"),
38 | ]
39 |
40 | unet_conversion_map_layer = []
41 | # hardcoded number of downblocks and resnets/attentions...
42 | # would need smarter logic for other networks.
43 | for i in range(4):
44 | # loop over downblocks/upblocks
45 |
46 | for j in range(2):
47 | # loop over resnets/attentions for downblocks
48 | hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
49 | sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
50 | unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
51 |
52 | if i < 3:
53 | # no attention layers in down_blocks.3
54 | hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
55 | sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
56 | unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
57 |
58 | for j in range(3):
59 | # loop over resnets/attentions for upblocks
60 | hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
61 | sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
62 | unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
63 |
64 | if i > 0:
65 | # no attention layers in up_blocks.0
66 | hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
67 | sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
68 | unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
69 |
70 | if i < 3:
71 | # no downsample in down_blocks.3
72 | hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
73 | sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
74 | unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
75 |
76 | # no upsample in up_blocks.3
77 | hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
78 | sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
79 | unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
80 |
81 | hf_mid_atn_prefix = "mid_block.attentions.0."
82 | sd_mid_atn_prefix = "middle_block.1."
83 | unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
84 |
85 | for j in range(2):
86 | hf_mid_res_prefix = f"mid_block.resnets.{j}."
87 | sd_mid_res_prefix = f"middle_block.{2*j}."
88 | unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
89 |
90 |
91 | def convert_unet_state_dict(unet_state_dict):
92 | # buyer beware: this is a *brittle* function,
93 | # and correct output requires that all of these pieces interact in
94 | # the exact order in which I have arranged them.
95 | mapping = {k: k for k in unet_state_dict.keys()}
96 | for sd_name, hf_name in unet_conversion_map:
97 | mapping[hf_name] = sd_name
98 | for k, v in mapping.items():
99 | if "resnets" in k:
100 | for sd_part, hf_part in unet_conversion_map_resnet:
101 | v = v.replace(hf_part, sd_part)
102 | mapping[k] = v
103 | for k, v in mapping.items():
104 | for sd_part, hf_part in unet_conversion_map_layer:
105 | v = v.replace(hf_part, sd_part)
106 | mapping[k] = v
107 | new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
108 | return new_state_dict
109 |
110 |
111 | # ================#
112 | # VAE Conversion #
113 | # ================#
114 |
115 | vae_conversion_map = [
116 | # (stable-diffusion, HF Diffusers)
117 | ("nin_shortcut", "conv_shortcut"),
118 | ("norm_out", "conv_norm_out"),
119 | ("mid.attn_1.", "mid_block.attentions.0."),
120 | ]
121 |
122 | for i in range(4):
123 | # down_blocks have two resnets
124 | for j in range(2):
125 | hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
126 | sd_down_prefix = f"encoder.down.{i}.block.{j}."
127 | vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
128 |
129 | if i < 3:
130 | hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
131 | sd_downsample_prefix = f"down.{i}.downsample."
132 | vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
133 |
134 | hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
135 | sd_upsample_prefix = f"up.{3-i}.upsample."
136 | vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
137 |
138 | # up_blocks have three resnets
139 | # also, up blocks in hf are numbered in reverse from sd
140 | for j in range(3):
141 | hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
142 | sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
143 | vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
144 |
145 | # this part accounts for mid blocks in both the encoder and the decoder
146 | for i in range(2):
147 | hf_mid_res_prefix = f"mid_block.resnets.{i}."
148 | sd_mid_res_prefix = f"mid.block_{i+1}."
149 | vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
150 |
151 |
152 | vae_conversion_map_attn = [
153 | # (stable-diffusion, HF Diffusers)
154 | ("norm.", "group_norm."),
155 | ("q.", "query."),
156 | ("k.", "key."),
157 | ("v.", "value."),
158 | ("proj_out.", "proj_attn."),
159 | ]
160 |
161 |
162 | def reshape_weight_for_sd(w):
163 | # convert HF linear weights to SD conv2d weights
164 | return w.reshape(*w.shape, 1, 1)
165 |
166 |
167 | def convert_vae_state_dict(vae_state_dict):
168 | mapping = {k: k for k in vae_state_dict.keys()}
169 | for k, v in mapping.items():
170 | for sd_part, hf_part in vae_conversion_map:
171 | v = v.replace(hf_part, sd_part)
172 | mapping[k] = v
173 | for k, v in mapping.items():
174 | if "attentions" in k:
175 | for sd_part, hf_part in vae_conversion_map_attn:
176 | v = v.replace(hf_part, sd_part)
177 | mapping[k] = v
178 | new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
179 | weights_to_convert = ["q", "k", "v", "proj_out"]
180 | for k, v in new_state_dict.items():
181 | for weight_name in weights_to_convert:
182 | if f"mid.attn_1.{weight_name}.weight" in k:
183 | print(f"Reshaping {k} for SD format")
184 | new_state_dict[k] = reshape_weight_for_sd(v)
185 | return new_state_dict
186 |
187 |
188 | # =========================#
189 | # Text Encoder Conversion #
190 | # =========================#
191 | # pretty much a no-op
192 |
193 |
194 | def convert_text_enc_state_dict(text_enc_dict):
195 | return text_enc_dict
196 |
197 |
198 | if __name__ == "__main__":
199 | parser = argparse.ArgumentParser()
200 |
201 | parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
202 | parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
203 | parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
204 |
205 | args = parser.parse_args()
206 |
207 | assert args.model_path is not None, "Must provide a model path!"
208 |
209 | assert args.checkpoint_path is not None, "Must provide a checkpoint path!"
210 |
211 | unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
212 | vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
213 | text_enc_path = osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
214 |
215 | # Convert the UNet model
216 | unet_state_dict = torch.load(unet_path, map_location='cpu')
217 | unet_state_dict = convert_unet_state_dict(unet_state_dict)
218 | unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
219 |
220 | # Convert the VAE model
221 | vae_state_dict = torch.load(vae_path, map_location='cpu')
222 | vae_state_dict = convert_vae_state_dict(vae_state_dict)
223 | vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
224 |
225 | # Convert the text encoder model
226 | text_enc_dict = torch.load(text_enc_path, map_location='cpu')
227 | text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
228 | text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
229 |
230 | # Put together new checkpoint
231 | state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
232 | if args.half:
233 | state_dict = {k:v.half() for k,v in state_dict.items()}
234 | state_dict = {"state_dict": state_dict}
235 | torch.save(state_dict, args.checkpoint_path)
--------------------------------------------------------------------------------
/src/zeropainter/dreamshaper.py:
--------------------------------------------------------------------------------
1 | import src.smplfusion
2 | from src.smplfusion import scheduler
3 | from src.smplfusion.common import *
4 | from types import SimpleNamespace
5 |
6 | import importlib
7 | from omegaconf import OmegaConf
8 | import torch
9 | import safetensors
10 | import safetensors.torch
11 |
12 | from os.path import dirname
13 |
14 | print("Loading model: Dreamshaper Inpainting V8")
15 | PROJECT_DIR = dirname(dirname(dirname(dirname(__file__))))
16 | print(PROJECT_DIR)
17 | CONFIG_FOLDER = f"{PROJECT_DIR}/lib/smplfusion/config/"
18 | MODEL_FOLDER = f"{PROJECT_DIR}/assets/models/"
19 | ASSETS_FOLDER = f"{PROJECT_DIR}/assets/"
20 |
21 |
22 | def get_inpainting_condition(model, image, mask):
23 | latent_size = [x // 8 for x in image.size]
24 | condition_x0 = (
25 | model.vae.encode(image.torch().cuda() * ~mask.torch(0).bool().cuda()).mean
26 | * model.config.scale_factor
27 | )
28 | condition_mask = mask.resize(latent_size[::-1]).cuda().torch(0).bool().float()
29 | return torch.cat([condition_mask, condition_x0], 1)
30 |
31 |
32 | def get_obj_from_str(string):
33 | module, cls = string.rsplit(".", 1)
34 | return getattr(importlib.import_module(module, package=None), cls)
35 |
36 | def load_obj(path):
37 | objyaml = OmegaConf.load(path)
38 | return get_obj_from_str(objyaml["__class__"])(**objyaml.get("__init__", {}))
39 |
40 |
41 | def get_t2i_model():
42 | model_t2i = SimpleNamespace()
43 | state_dict = safetensors.torch.load_file(
44 | f"{MODEL_FOLDER}/dreamshaper/dreamshaper_8.safetensors", device="cuda"
45 | )
46 |
47 | model_t2i.config = OmegaConf.load(f"{CONFIG_FOLDER}/ddpm/v1.yaml")
48 | model_t2i.unet = load_obj(f"{CONFIG_FOLDER}/unet/v1.yaml").eval().cuda()
49 | model_t2i.vae = load_obj(f"{CONFIG_FOLDER}/vae.yaml").eval().cuda()
50 | model_t2i.encoder = load_obj(f"{CONFIG_FOLDER}/encoders/clip.yaml").eval().cuda()
51 |
52 | extract = lambda state_dict, model: {
53 | x[len(model) + 1 :]: y for x, y in state_dict.items() if model in x
54 | }
55 | unet_state = extract(state_dict, "model.diffusion_model")
56 | encoder_state = extract(state_dict, "cond_stage_model")
57 | vae_state = extract(state_dict, "first_stage_model")
58 |
59 | model_t2i.unet.load_state_dict(unet_state, strict=False)
60 | model_t2i.encoder.load_state_dict(encoder_state, strict=False)
61 | model_t2i.vae.load_state_dict(vae_state, strict=False)
62 | model_t2i.unet = model_t2i.unet.requires_grad_(False)
63 | model_t2i.encoder = model_t2i.encoder.requires_grad_(False)
64 | model_t2i.vae = model_t2i.vae.requires_grad_(False)
65 |
66 | model_t2i.schedule = scheduler.linear(
67 | model_t2i.config.timesteps,
68 | model_t2i.config.linear_start,
69 | model_t2i.config.linear_end,
70 | )
71 | return model_t2i, unet_state
72 |
73 |
74 | def get_inpainting_model():
75 | model_inp = SimpleNamespace()
76 | state_dict = safetensors.torch.load_file(
77 | f"{MODEL_FOLDER}/dreamshaper/dreamshaper_8Inpainting.safetensors", device="cuda"
78 | )
79 |
80 | model_inp.config = OmegaConf.load(f"{CONFIG_FOLDER}/ddpm/v1.yaml")
81 | model_inp.unet = load_obj(f"{CONFIG_FOLDER}/unet/inpainting/v1.yaml").eval().cuda()
82 | model_inp.vae = load_obj(f"{CONFIG_FOLDER}/vae.yaml").eval().cuda()
83 | model_inp.encoder = load_obj(f"{CONFIG_FOLDER}/encoders/clip.yaml").eval().cuda()
84 |
85 | extract = lambda state_dict, model: {
86 | x[len(model) + 1 :]: y for x, y in state_dict.items() if model in x
87 | }
88 | unet_state = extract(state_dict, "model.diffusion_model")
89 | encoder_state = extract(state_dict, "cond_stage_model")
90 | vae_state = extract(state_dict, "first_stage_model")
91 |
92 | model_inp.unet.load_state_dict(unet_state, strict=False)
93 | model_inp.encoder.load_state_dict(encoder_state, strict=False)
94 | model_inp.vae.load_state_dict(vae_state, strict=False)
95 | model_inp.unet = model_inp.unet.requires_grad_(False)
96 | model_inp.encoder = model_inp.encoder.requires_grad_(False)
97 | model_inp.vae = model_inp.vae.requires_grad_(False)
98 |
99 | model_inp.schedule = scheduler.linear(
100 | model_inp.config.timesteps,
101 | model_inp.config.linear_start,
102 | model_inp.config.linear_end,
103 | )
104 | return model_inp, unet_state
105 |
--------------------------------------------------------------------------------
/src/zeropainter/inpainting.py:
--------------------------------------------------------------------------------
1 | from src.smplfusion.common import *
2 | import torch
3 | from src.smplfusion import share
4 | from src.smplfusion.patches import attentionpatch
5 | from tqdm import tqdm
6 | from pytorch_lightning import seed_everything
7 | from src.smplfusion.patches import router
8 | from src.smplfusion import IImage
9 |
10 | # negative_prompt = "worst quality, ugly, gross, disfigured, deformed, dehydrated, extra limbs, fused body parts, mutilated, malformed, mutated, bad anatomy, bad proportions, low quality, cropped, low resolution, out of frame, poorly drawn, text, watermark, letters, jpeg artifacts"
11 | # positive_prompt = ", realistic, HD, Full HD, 4K, high quality, high resolution, masterpiece, trending on artstation, realistic lighting"
12 | negative_prompt = ''
13 | positive_prompt = ''
14 | VERBOSE = True
15 |
16 | class AttnForward:
17 | def __init__(self, masks, object_context, object_uc_context):
18 | self.masks = masks
19 | self.object_context = object_context
20 | self.object_uc_context = object_uc_context
21 |
22 | def __call__(data, self, x, context=None, mask=None):
23 | att_type = "self" if context is None else "cross"
24 | batch_size = x.shape[0]
25 |
26 | if att_type == 'cross' and x.shape[1] in [share.input_shape.res16]: # For cross attention
27 | out = torch.zeros_like(x)
28 | for i in range(len(data.masks)):
29 | if data.masks[i].sum()>0:
30 | if batch_size == 1:
31 | out[:,data.masks[i]] = attentionpatch.default.forward(
32 | self,
33 | x[:,data.masks[i] > 0],
34 | data.object_context[i][None]
35 | ).float()
36 | elif batch_size == 2:
37 | out[:,data.masks[i]] = attentionpatch.default.forward(
38 | self,
39 | x[:,data.masks[i] > 0],
40 | torch.stack([data.object_uc_context[i],data.object_context[i]])
41 | ).float()
42 | else:
43 | raise NotImplementedError("Batch Size > 1 not yet supported!")
44 | return out
45 | else:
46 | return attentionpatch.default.forward(self, x, context, mask)
47 |
48 | def gen_filled_image(model, prompt, image, mask, zp_masks, seed, T = 899, dt = 20, model_t2i = None, guidance_scale = 7.5, use_lcm_multistep = False):
49 | masks = [x.modified_mask.val16.flatten()>0 for x in zp_masks]
50 | masks.append(torch.stack(masks).sum(0) == 0)
51 |
52 | context = model.encoder.encode(['', f"realistic photo of a {prompt}" + positive_prompt])
53 | object_context = model.encoder.encode([x.local_prompt for x in zp_masks] + [f"realistic photo of a {prompt}" + positive_prompt])
54 | object_uc_context = model.encoder.encode([negative_prompt] * len(zp_masks) + [', '.join([x.local_prompt for x in zp_masks])])
55 |
56 | seed_everything(seed)
57 |
58 | eps = torch.randn((1,4,64,64)).cuda()
59 | # zT = schedule.sqrt_alphas[799] * condition_mask + schedule.sqrt_one_minus_alphas[799] * eps
60 | # zT = (condition_mask < 0) * zT + (condition_mask > 0) * eps
61 | # zT = schedule.sqrt_one_minus_alphas[899] * eps
62 | # zT = schedule.sqrt_alphas[899] * IImage(255*orig_masks_all).resize(64).alpha().torch().cuda() + schedule.sqrt_one_minus_alphas[899] * eps
63 | condition_x0 = model.vae.encode(image.torch().cuda() * ~mask.torch(0).bool().cuda()).mean * model.config.scale_factor
64 | condition_xT = model.schedule.sqrt_alphas[T] * condition_x0 + model.schedule.sqrt_one_minus_alphas[T] * eps
65 | condition_mask = mask.resize(64).cuda().torch(0).bool().float()
66 | zT = (condition_mask == 0) * condition_xT + (condition_mask > 0) * eps
67 | # zT = eps
68 |
69 | router.attention_forward = AttnForward(masks, object_context, object_uc_context)
70 | # router.attention_forward = attentionpatch.default.forward
71 |
72 | with torch.autocast('cuda'), torch.no_grad():
73 | zt = zT
74 | timesteps = list(range(T, 0, -dt))
75 | pbar = tqdm(timesteps) if VERBOSE else timesteps
76 | for index,t in enumerate(pbar):
77 | if index == 0:
78 | current_mask = ~(~mask).dilate(2)
79 | condition_x0 = model.vae.encode(image.torch().cuda() * ~mask.cuda().torch(0).bool().cuda()).mean * model.config.scale_factor
80 | condition_mask = current_mask.resize(64).cuda().torch(0).bool().float()
81 | condition = torch.cat([condition_mask, condition_x0], 1)
82 | if index == 5:
83 | current_mask = ~(~mask).dilate(0)
84 | condition_x0 = model.vae.encode(image.torch().cuda() * ~mask.cuda().torch(0).bool().cuda()).mean * model.config.scale_factor
85 | condition_mask = current_mask.resize(64).cuda().torch(0).bool().float()
86 | condition = torch.cat([condition_mask, condition_x0], 1)
87 | if index == len(timesteps) - 5:
88 | if model_t2i is not None:
89 | model = model_t2i
90 | condition = None
91 | else:
92 | current_mask = mask.dilate(512)
93 | condition_x0 = model.vae.encode(image.torch().cuda() * ~current_mask.cuda().torch(0).bool().cuda()).mean * model.config.scale_factor
94 | condition_mask = current_mask.resize(64).cuda().torch(0).bool().float()
95 | condition = torch.cat([condition_mask, condition_x0], 1)
96 |
97 | _zt = zt if condition is None else torch.cat([zt, condition], 1)
98 |
99 | if use_lcm_multistep:
100 | eps = model.unet(
101 | _zt,
102 | timesteps = torch.tensor([t]).cuda(),
103 | context = context[1][None]
104 | )
105 | else:
106 | eps_uncond, eps = model.unet(
107 | torch.cat([_zt, _zt]),
108 | timesteps = torch.tensor([t, t]).cuda(),
109 | context = context
110 | ).chunk(2)
111 | eps = (eps_uncond + guidance_scale * (eps - eps_uncond))
112 |
113 | z0 = (zt - model.schedule.sqrt_one_minus_alphas[t] * eps) / model.schedule.sqrt_alphas[t]
114 | zt = model.schedule.sqrt_alphas[t - dt] * z0 + model.schedule.sqrt_one_minus_alphas[t - dt] * eps
115 |
116 | out = IImage(model.vae.decode(z0 / model.config.scale_factor))
117 | return out
118 | # # out.save('../../assets/paper_data/visuals_for_paper/'+str(seed_num)+'_'+name)
119 |
120 | # np_out = np.array(out.data[0])
121 | # print(np_out.max(), orig_masks_all.max())
122 | # blending_result_helper = 0.4*(np_out/255)+0.6*orig_masks_all
123 | # res = np.hstack([np_out,blending_result_helper*255])
124 | # n = name.split('.')[0]
125 | # cv2.imwrite('../../assets/paper_data/visuals_for_paper/'+n+'_'+str(my_seed_gen)+'__'+str(seed_num)+'_.png',res[:,:,::-1])
126 | # (IImage(255 * blending_result_helper) | IImage(np_out))
--------------------------------------------------------------------------------
/src/zeropainter/models.py:
--------------------------------------------------------------------------------
1 | from types import SimpleNamespace
2 | from . import convert_diffusers
3 | import safetensors.torch
4 |
5 | import torch
6 | from src.smplfusion import scheduler
7 | from os.path import dirname
8 |
9 | from omegaconf import OmegaConf
10 | import importlib
11 |
12 | def load_obj(objyaml):
13 | if "__init__" in objyaml:
14 | return get_obj_from_str(objyaml["__class__"])(**objyaml["__init__"])
15 | else:
16 | return get_obj_from_str(objyaml["__class__"])()
17 |
18 |
19 | def get_obj_from_str(string):
20 | module, cls = string.rsplit(".", 1)
21 | try:
22 | return getattr(importlib.import_module(module, package=None), cls)
23 | except Exception as e:
24 | return getattr(importlib.import_module("src." + module, package=None), cls)
25 |
26 |
27 | def get_inpainting_condition(model, image, mask):
28 | latent_size = [x // 8 for x in image.size]
29 | condition_x0 = (
30 | model.vae.encode(image.torch().cuda() * ~mask.torch(0).bool().cuda()).mean
31 | * model.config.scale_factor
32 | )
33 | condition_mask = mask.resize(latent_size[::-1]).cuda().torch(0).bool().float()
34 | return torch.cat([condition_mask, condition_x0], 1)
35 |
36 |
37 | def get_t2i_model(config_folder, model_folder):
38 | model_t2i = SimpleNamespace()
39 |
40 | model_t2i.config = OmegaConf.load(
41 | f"{config_folder}/ddpm/v1.yaml"
42 | ) # smplfusion.options.ddpm.v1_yaml
43 | model_t2i.schedule = scheduler.linear(
44 | model_t2i.config.timesteps,
45 | model_t2i.config.linear_start,
46 | model_t2i.config.linear_end,
47 | )
48 |
49 | model_t2i.unet = load_obj(
50 | OmegaConf.load(f"{config_folder}/unet/v1.yaml")
51 | ).cuda() # smplfusion.options.unet.v1_yaml.cuda()
52 | model_t2i.vae = load_obj(
53 | OmegaConf.load(f"{config_folder}/vae.yaml")
54 | ).cuda() # smplfusion.options.vae_yaml.cuda()
55 | model_t2i.encoder = load_obj(
56 | OmegaConf.load(f"{config_folder}/encoders/clip.yaml")
57 | ).cuda() # smplfusion.options.encoders.clip_yaml.cuda()
58 |
59 | unet_state_dict = torch.load(
60 | f"{model_folder}/unet.ckpt"
61 | ) # assets.models.sd_1_5_inpainting.unet_ckpt
62 | vae_state_dict = torch.load(
63 | f"{model_folder}/vae.ckpt"
64 | ) # assets.models.sd_1_5_inpainting.vae_ckpt
65 | encoder_state_dict = torch.load(
66 | f"{model_folder}/encoder.ckpt"
67 | ) # assets.models.sd_1_5_inpainting.encoder_ckpt
68 |
69 | model_t2i.unet.load_state_dict(unet_state_dict)
70 | model_t2i.vae.load_state_dict(vae_state_dict)
71 | model_t2i.encoder.load_state_dict(encoder_state_dict, strict=False)
72 |
73 | model_t2i.unet = model_t2i.unet.requires_grad_(False).eval()
74 | model_t2i.vae = model_t2i.vae.requires_grad_(False).eval()
75 | model_t2i.encoder = model_t2i.encoder.requires_grad_(False).eval()
76 | return model_t2i, unet_state_dict
77 |
78 |
79 | def get_inpainting_model(config_folder, model_folder):
80 | model_inp = SimpleNamespace()
81 | model_inp.config = OmegaConf.load(
82 | f"{config_folder}/ddpm/v1.yaml"
83 | ) # smplfusion.options.ddpm.v1_yaml
84 | model_inp.schedule = scheduler.linear(
85 | model_inp.config.timesteps,
86 | model_inp.config.linear_start,
87 | model_inp.config.linear_end,
88 | )
89 |
90 | model_inp.unet = load_obj(
91 | OmegaConf.load(f"{config_folder}/unet/inpainting/v1.yaml")
92 | ).cuda() # smplfusion.options.unet.inpainting.v1_yaml.cuda()
93 | model_inp.vae = load_obj(
94 | OmegaConf.load(f"{config_folder}/vae.yaml")
95 | ).cuda() # smplfusion.options.vae_yaml.cuda()
96 | model_inp.encoder = load_obj(
97 | OmegaConf.load(f"{config_folder}/encoders/clip.yaml")
98 | ).cuda() # smplfusion.options.encoders.clip_yaml.cuda()
99 |
100 | unet_state_dict = torch.load(
101 | f"{model_folder}/unet.ckpt"
102 | ) # assets.models.sd_1_5_inpainting.unet_ckpt
103 | vae_state_dict = torch.load(
104 | f"{model_folder}/vae.ckpt"
105 | ) # assets.models.sd_1_5_inpainting.vae_ckpt
106 | encoder_state_dict = torch.load(
107 | f"{model_folder}/encoder.ckpt"
108 | ) # assets.models.sd_1_5_inpainting.encoder_ckpt
109 |
110 | model_inp.unet.load_state_dict(unet_state_dict)
111 | model_inp.vae.load_state_dict(vae_state_dict)
112 | model_inp.encoder.load_state_dict(encoder_state_dict, strict=False)
113 |
114 | model_inp.unet = model_inp.unet.requires_grad_(False).eval()
115 | model_inp.vae = model_inp.vae.requires_grad_(False).eval()
116 | model_inp.encoder = model_inp.encoder.requires_grad_(False).eval()
117 |
118 | return model_inp, unet_state_dict
119 |
120 |
121 | def get_lora(lora_path):
122 | _lora_state_dict = safetensors.torch.load_file(lora_path)
123 |
124 | # Dictionary for converting
125 | unet_conversion_dict = {}
126 | unet_conversion_dict.update(convert_diffusers.unet_conversion_map_layer)
127 | unet_conversion_dict.update(convert_diffusers.unet_conversion_map_resnet)
128 | unet_conversion_dict.update(convert_diffusers.unet_conversion_map)
129 | unet_conversion_dict = {
130 | y.replace(".", "_"): x for x, y in unet_conversion_dict.items()
131 | }
132 | unet_conversion_dict["lora_unet_"] = ""
133 |
134 | lora_state_dict = {}
135 | for key in _lora_state_dict:
136 | key_converted = key
137 | for x, y in unet_conversion_dict.items():
138 | key_converted = key_converted.replace(x, y)
139 | lora_state_dict[key_converted] = _lora_state_dict[key]
140 |
141 | return lora_state_dict
142 |
--------------------------------------------------------------------------------
/src/zeropainter/segmentation.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from os.path import dirname
3 |
4 |
5 | from segment_anything import SamPredictor, sam_model_registry
6 |
7 | def get_bounding_box(mask):
8 | all_y,all_x = np.where(mask == 1)
9 | all_y,all_x = sorted(all_y),sorted(all_x)
10 | x1,y1,x2,y2 = all_x[0],all_y[0],all_x[-1],all_y[-1]
11 | x,y,w,h = x1,y1,x2-x1,y2-y1
12 | return x,y,w,h
13 |
14 | def get_segmentation_model(path_to_check):
15 | # path_to_check = "sam_vit_h_4b8939.pth"
16 | sam = sam_model_registry["vit_h"](checkpoint=path_to_check)
17 | predictor = SamPredictor(sam)
18 | predictor.model.cuda();
19 | return predictor
20 |
21 | def get_segmentation(predictor, im, mask):
22 | im_np = im.data[0]
23 |
24 | # SAM prediction
25 | # mask = obj['data_512'][i]['mask']
26 | x,y,w,h = get_bounding_box(mask)
27 | predictor.set_image(im_np)
28 | input_box = np.array([x,y,x+w,y+h])
29 | masks, scores, other = predictor.predict(box=input_box)
30 |
31 | _masks = np.concatenate([masks, ~masks])
32 | _scores = np.concatenate([scores, scores])
33 | _scores = np.array([y for x,y in zip(_masks, _scores) if ((1 - mask) * x).sum() <= (mask * x).sum()])
34 | _masks = np.array([x for x in _masks if ((1 - mask) * x).sum() <= (mask * x).sum()])
35 | if len(_masks) > 0:
36 | masks,scores = _masks,_scores
37 |
38 | pred_seg_mask = masks[scores.argmax()]
39 | pred_seg_mask[pred_seg_mask > 0] = 1
40 | pred_seg_mask = np.stack([pred_seg_mask] * 3, axis=-1) * 1
41 |
42 | return pred_seg_mask * mask[...,None]
43 |
--------------------------------------------------------------------------------
/src/zeropainter/zero_painter_dataset.py:
--------------------------------------------------------------------------------
1 | import json
2 | import torch
3 | import cv2
4 | import numpy as np
5 | from src.smplfusion import share
6 |
7 | class ZeroPainterMask:
8 | def __init__(self, color, local_prompt, img_grey,rgb,image_w,image_h):
9 | self.img_grey = img_grey
10 | self.color = eval(color)
11 | self.local_prompt = local_prompt
12 | self.bbox = None
13 | self.bbox_64 = None
14 | self.area = None
15 | self.mask = None
16 | self.mask_64 = None
17 | self.token_idx = None
18 | self.modified_mask = None
19 | self.inverse_mask = None
20 | self.modified_indexses_of_prompt = None
21 | self.sot_index = [0]
22 | self.w_positive = 1
23 | self.w_negative = 1
24 | self.image_w = image_w
25 | self.image_h = image_h
26 | self.rgb = rgb
27 | self.modify_info()
28 |
29 | def get_bounding_box(self, mask):
30 | mask = mask*255.0
31 | all_y, all_x = np.where(mask == 255.0)
32 | all_y, all_x = sorted(all_y), sorted(all_x)
33 | x1, y1, x2, y2 = all_x[0], all_y[0], all_x[-1], all_y[-1]
34 | x, y, w, h = x1, y1, x2 - x1, y2 - y1
35 | return x, y, w, h
36 |
37 |
38 | def modify_info(self):
39 | # Define the color to be searched for
40 | if self.rgb:
41 | color_to_search = np.array([self.color[0],self.color[1],self.color[2]])
42 | mask_1d = np.all(self.img_grey == color_to_search, axis=2)*1.0
43 | else:
44 | mask_1d = (self.img_grey == self.color) * 1.0
45 | mask_1d_64 = cv2.resize(mask_1d, (64, 64), interpolation=cv2.INTER_NEAREST)
46 | self.bbox = self.get_bounding_box(mask_1d)
47 | self.bbox_64 = self.get_bounding_box(mask_1d_64)
48 | self.area = self.bbox[2] * self.bbox[3]
49 | self.mask = mask_1d.copy()
50 | self.mask_64 = mask_1d_64 // 255
51 |
52 | splited_prompt = self.local_prompt.split()
53 | self.token_idx = np.arange(1, len(splited_prompt) + 1, 1, dtype=int)
54 |
55 | #dif part
56 | mask_1d = mask_1d[None,None,:,:]
57 | mask_1d = torch.from_numpy(mask_1d)
58 |
59 | self.modified_indexses_of_prompt = np.arange(5, len(splited_prompt)+1)
60 | self.modified_mask = share.InputMask(mask_1d)
61 | self.inverse_mask = share.InputMask(1-mask_1d)
62 |
63 | return
64 |
65 | class ZeroPainterSample:
66 | #Example of sample
67 | #{'prompt': 'Brown gift box beside red candle.',
68 | # 'color_context_dict': {'1': 'Brown gift box', '2': 'red candle'}}
69 | def __init__(self, item, img_grey, image_h=512,image_w=512):
70 |
71 | self.item = item
72 | self.global_prompt = self.item['prompt']
73 | self.image_w = image_w
74 | self.image_h = image_h
75 | self.img_grey = img_grey
76 | self.masks = self.load_masks()
77 |
78 |
79 | def load_masks(self):
80 | data_samples = []
81 |
82 | self.img_grey = cv2.resize(self.img_grey, (self.image_w, self.image_h),interpolation=cv2.INTER_NEAREST)
83 | for color, local_prompt in self.item['color_context_dict'].items():
84 | data_samples.append(ZeroPainterMask(color,local_prompt,self.img_grey,self.item['rgb'],self.image_w,self.image_h))
85 |
86 | return data_samples
87 |
88 |
89 | class ZeroPainterDataset:
90 | def __init__(self, root_path_img, json_path,rgb=True):
91 | self.root_path_img = root_path_img
92 | self.json_path = json_path
93 | self.rgb = rgb
94 |
95 | if isinstance(json_path, dict) or isinstance(json_path, list):
96 | self.json_data = json_path
97 | else:
98 | with open(self.json_path, 'r') as file:
99 | self.json_data = json.load(file)
100 | def __len__(self):
101 | return len(self.json_data)
102 |
103 | def __getitem__(self, index):
104 | item = self.json_data[index]
105 | item['rgb'] = self.rgb
106 | if isinstance(self.root_path_img, str):
107 | self.img_path = self.root_path_img
108 | if self.rgb:
109 | self.img_grey = cv2.imread(self.img_path)
110 | else:
111 | self.img_grey = cv2.imread(self.img_path,0)
112 | else:
113 | self.img_grey = np.array(self.root_path_img)
114 |
115 | return ZeroPainterSample(item, self.img_grey)
116 |
117 |
118 |
--------------------------------------------------------------------------------
/src/zeropainter/zero_painter_pipline.py:
--------------------------------------------------------------------------------
1 | from . import segmentation, generation, inpainting
2 |
3 | import torch
4 | import numpy as np
5 | from src.smplfusion import libimage
6 | from src.smplfusion import IImage
7 |
8 |
9 | class ZeroPainter:
10 | def __init__(self, model_t2i, model_inp, model_sam):
11 | self.model_t2i = model_t2i
12 | self.model_inp = model_inp
13 | self.model_sam = model_sam
14 |
15 |
16 | def gen_sample(
17 | self,
18 | sample,
19 | object_seed,
20 | image_seed,
21 | num_ddim_steps_t2i,
22 | num_ddim_steps_inp,
23 | cfg_scale_t2i=7.5,
24 | cfg_scale_inp=7.5,
25 | use_lcm_multistep_t2i=False,
26 | use_lcm_multistep_inp=False,
27 | ):
28 |
29 |
30 | gen_obj_list = []
31 | gen_mask_list = []
32 | real_mask_list = []
33 |
34 | if isinstance(object_seed, int):
35 | object_seed = [object_seed] * len(sample.masks)
36 |
37 | for i in range(len(sample.masks)):
38 | eps_list, z0_list, zt_list = generation.gen_single_object(
39 | self.model_t2i,
40 | sample.masks[i],
41 | object_seed[i],
42 | dt=1000 // num_ddim_steps_t2i,
43 | guidance_scale=cfg_scale_t2i,
44 | use_lcm_multistep=use_lcm_multistep_t2i,
45 | )
46 | mask = sample.masks[i].mask
47 |
48 | gen_image = IImage(
49 | self.model_t2i.vae.decode(z0_list[-1] / self.model_t2i.config.scale_factor)
50 | )
51 | gen_mask = segmentation.get_segmentation(self.model_sam, gen_image, mask)
52 | gen_object = gen_image * IImage(255 * gen_mask)
53 |
54 | gen_obj_list.append(gen_object)
55 | gen_mask_list.append(gen_mask)
56 | real_mask_list.append(mask)
57 |
58 | gen_image = IImage(libimage.stack(gen_obj_list).data.sum(0))
59 | gen_mask = IImage(255 * np.sum(gen_mask_list, 0))
60 | real_mask = IImage(255 * np.sum(real_mask_list, 0))
61 |
62 | output = inpainting.gen_filled_image(
63 | self.model_inp,
64 | sample.global_prompt,
65 | gen_image,
66 | (~gen_mask.alpha()).dilate(3),
67 | sample.masks,
68 | image_seed,
69 | dt=1000 // num_ddim_steps_inp,
70 | guidance_scale=cfg_scale_inp,
71 | use_lcm_multistep=use_lcm_multistep_inp,
72 | )
73 |
74 | return output
75 |
76 |
77 |
78 |
--------------------------------------------------------------------------------
/zero_painter.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | from src.zeropainter.zero_painter_pipline import ZeroPainter
3 | from src.zeropainter import models, dreamshaper,segmentation
4 | from src.zeropainter import zero_painter_dataset
5 |
6 | def get_args():
7 | parser = argparse.ArgumentParser()
8 |
9 | parser.add_argument('--mask-path', default='data/masks/1_rgb.png',help='Mask path.')
10 | parser.add_argument('--metadata', default='data/metadata/1.json', type=str, help='Text prompt.')
11 | parser.add_argument('--output-dir', default='data/outputs/',help='Output dir.')
12 |
13 | #load models
14 | parser.add_argument('--config-folder-for-models', type=str, default='config', help='Path to configs')
15 | parser.add_argument('--model-folder-inpiting', type=str, default='models/sd-1-5-inpainting', help='Path to load inpainting model')
16 | parser.add_argument('--model-folder-generation', type=str, default='models/sd-1-4', help='Path to load generation model')
17 | parser.add_argument('--segment-anything-model', type=str, default='models/sam_vit_h_4b8939.pth', help='Path to load segmentation model')
18 |
19 | return parser.parse_args()
20 |
21 |
22 | def main():
23 | args = get_args()
24 |
25 | model_inp,_ = models.get_inpainting_model(args.config_folder_for_models,args.model_folder_inpiting)
26 | model_t2i,_ = models.get_t2i_model(args.config_folder_for_models,args.model_folder_generation)
27 | model_sam = segmentation.get_segmentation_model(args.segment_anything_model)
28 | zero_painter_model = ZeroPainter(model_t2i, model_inp, model_sam)
29 |
30 | data = zero_painter_dataset.ZeroPainterDataset(args.mask_path, args.metadata)
31 | name = args.mask_path.split('/')[-1]
32 | result = zero_painter_model.gen_sample(data[0], 42, 42,30,30)
33 | result.save(args.output_dir+name)
34 |
35 |
36 | if __name__ == '__main__':
37 | main()
38 |
--------------------------------------------------------------------------------
19 | 
77 |
78 | ---
79 |
80 | ## 🎓 Citation
81 | If you use our work in your research, please cite our publication:
82 | ```
83 | @article{Zeropainter,
84 | title={Zero-Painter: Training-Free Layout Control for Text-to-Image Synthesis},
85 | url={http://arxiv.org/abs/2406.04032},
86 | publisher={arXiv},
87 | author={Ohanyan, Marianna and Manukyan, Hayk and Wang, Zhangyang and Navasardyan, Shant and Shi, Humphrey},
88 | year={2024}}
89 |
90 | ```
--------------------------------------------------------------------------------
/__assets__/github/method_arch.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/__assets__/github/method_arch.png
--------------------------------------------------------------------------------
/__assets__/github/teaser.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/__assets__/github/teaser.png
--------------------------------------------------------------------------------
/config/ddpm/v1.yaml:
--------------------------------------------------------------------------------
1 | linear_start: 0.00085
2 | linear_end: 0.0120
3 | num_timesteps_cond: 1
4 | log_every_t: 200
5 | timesteps: 1000
6 | first_stage_key: "jpg"
7 | cond_stage_key: "txt"
8 | image_size: 64
9 | channels: 4
10 | cond_stage_trainable: false
11 | conditioning_key: crossattn
12 | monitor: val/loss_simple_ema
13 | scale_factor: 0.18215
14 | use_ema: False # we set this to false because this is an inference only config
--------------------------------------------------------------------------------
/config/ddpm/v2-upsample.yaml:
--------------------------------------------------------------------------------
1 | parameterization: "v"
2 | low_scale_key: "lr"
3 | linear_start: 0.0001
4 | linear_end: 0.02
5 | num_timesteps_cond: 1
6 | log_every_t: 200
7 | timesteps: 1000
8 | first_stage_key: "jpg"
9 | cond_stage_key: "txt"
10 | image_size: 128
11 | channels: 4
12 | cond_stage_trainable: false
13 | conditioning_key: "hybrid-adm"
14 | monitor: val/loss_simple_ema
15 | scale_factor: 0.08333
16 | use_ema: False
17 |
18 | low_scale_config:
19 | target: ldm.modules.diffusionmodules.upscaling.ImageConcatWithNoiseAugmentation
20 | params:
21 | noise_schedule_config: # image space
22 | linear_start: 0.0001
23 | linear_end: 0.02
24 | max_noise_level: 350
25 |
--------------------------------------------------------------------------------
/config/encoders/clip.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.encoders.clip_embedder.FrozenCLIPEmbedder
--------------------------------------------------------------------------------
/config/encoders/openclip.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.encoders.open_clip_embedder.FrozenOpenCLIPEmbedder
2 | __init__:
3 | freeze: True
4 | layer: "penultimate"
--------------------------------------------------------------------------------
/config/unet/inpainting/v1.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.unet.UNetModel
2 | __init__:
3 | image_size: 32 # unused
4 | in_channels: 9 # 4 data + 4 downscaled image + 1 mask
5 | out_channels: 4
6 | model_channels: 320
7 | attention_resolutions: [ 4, 2, 1 ]
8 | num_res_blocks: 2
9 | channel_mult: [ 1, 2, 4, 4 ]
10 | num_heads: 8
11 | use_spatial_transformer: True
12 | transformer_depth: 1
13 | context_dim: 768
14 | use_checkpoint: False
15 | legacy: False
--------------------------------------------------------------------------------
/config/unet/inpainting/v2.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.unet.UNetModel
2 | __init__:
3 | use_checkpoint: False
4 | image_size: 32 # unused
5 | in_channels: 9
6 | out_channels: 4
7 | model_channels: 320
8 | attention_resolutions: [ 4, 2, 1 ]
9 | num_res_blocks: 2
10 | channel_mult: [ 1, 2, 4, 4 ]
11 | num_head_channels: 64 # need to fix for flash-attn
12 | use_spatial_transformer: True
13 | use_linear_in_transformer: True
14 | transformer_depth: 1
15 | context_dim: 1024
16 | legacy: False
--------------------------------------------------------------------------------
/config/unet/upsample/v2.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.unet.UNetModel
2 | __init__:
3 | use_checkpoint: False
4 | num_classes: 1000 # timesteps for noise conditioning (here constant, just need one)
5 | image_size: 128
6 | in_channels: 7
7 | out_channels: 4
8 | model_channels: 256
9 | attention_resolutions: [ 2,4,8]
10 | num_res_blocks: 2
11 | channel_mult: [ 1, 2, 2, 4]
12 | disable_self_attentions: [True, True, True, False]
13 | disable_middle_self_attn: False
14 | num_heads: 8
15 | use_spatial_transformer: True
16 | transformer_depth: 1
17 | context_dim: 1024
18 | legacy: False
19 | use_linear_in_transformer: True
--------------------------------------------------------------------------------
/config/unet/v/v2.yaml:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/config/unet/v/v2.yaml
--------------------------------------------------------------------------------
/config/unet/v1.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.unet.UNetModel
2 | __init__:
3 | image_size: 32 # unused
4 | in_channels: 4
5 | out_channels: 4
6 | model_channels: 320
7 | attention_resolutions: [ 4, 2, 1 ]
8 | num_res_blocks: 2
9 | channel_mult: [ 1, 2, 4, 4 ]
10 | num_heads: 8
11 | use_spatial_transformer: True
12 | transformer_depth: 1
13 | context_dim: 768
14 | use_checkpoint: False
15 | legacy: False
--------------------------------------------------------------------------------
/config/unet/v2.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.unet.UNetModel
2 | __init__:
3 | use_checkpoint: False
4 | use_fp16: True
5 | image_size: 32 # unused
6 | in_channels: 4
7 | out_channels: 4
8 | model_channels: 320
9 | attention_resolutions: [ 4, 2, 1 ]
10 | num_res_blocks: 2
11 | channel_mult: [ 1, 2, 4, 4 ]
12 | num_head_channels: 64
13 | use_spatial_transformer: True
14 | use_linear_in_transformer: True
15 | transformer_depth: 1
16 | context_dim: 1024
17 | legacy: False
--------------------------------------------------------------------------------
/config/vae-upsample.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.vae.AutoencoderKL
2 | __init__:
3 | embed_dim: 4
4 | ddconfig:
5 | double_z: True
6 | z_channels: 4
7 | resolution: 256
8 | in_channels: 3
9 | out_ch: 3
10 | ch: 128
11 | ch_mult: [ 1,2,4 ]
12 | num_res_blocks: 2
13 | attn_resolutions: [ ]
14 | dropout: 0.0
15 | lossconfig:
16 | target: torch.nn.Identity
--------------------------------------------------------------------------------
/config/vae.yaml:
--------------------------------------------------------------------------------
1 | __class__: smplfusion.models.vae.AutoencoderKL
2 | __init__:
3 | embed_dim: 4
4 | monitor: val/rec_loss
5 | ddconfig:
6 | double_z: true
7 | z_channels: 4
8 | resolution: 256
9 | in_channels: 3
10 | out_ch: 3
11 | ch: 128
12 | ch_mult: [1,2,4,4]
13 | num_res_blocks: 2
14 | attn_resolutions: []
15 | dropout: 0.0
16 | lossconfig:
17 | target: torch.nn.Identity
--------------------------------------------------------------------------------
/data/masks/0_rgb.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/data/masks/0_rgb.png
--------------------------------------------------------------------------------
/data/masks/1_rgb.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/data/masks/1_rgb.png
--------------------------------------------------------------------------------
/data/masks/2_rgb.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/data/masks/2_rgb.png
--------------------------------------------------------------------------------
/data/metadata/0.json:
--------------------------------------------------------------------------------
1 | [{
2 | "prompt": "Brown gift box beside red candle.",
3 | "color_context_dict": {
4 | "(244, 54, 32)": "Brown gift box",
5 | "(54, 245, 32)": "red candle"
6 | }
7 | }]
8 |
--------------------------------------------------------------------------------
/data/metadata/1.json:
--------------------------------------------------------------------------------
1 | [ {
2 | "prompt": "Brown tabby cat on white stairs photo",
3 | "color_context_dict": {
4 | "(244, 54, 32)": "Brown tabby cat",
5 | "(54, 245, 32)": "blue vase",
6 | "(4,54,200)": "red apple"
7 | }
8 | }]
--------------------------------------------------------------------------------
/data/metadata/2.json:
--------------------------------------------------------------------------------
1 | [{
2 | "prompt": "Green succulent on white and pink pot photo",
3 | "color_context_dict": {
4 | "(213, 24, 207)": "white and pink pot",
5 | "(78, 213, 24)": "Green succulent "
6 | }
7 | }]
--------------------------------------------------------------------------------
/data/outputs/0_rgb.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/data/outputs/0_rgb.png
--------------------------------------------------------------------------------
/data/outputs/1_rgb.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/data/outputs/1_rgb.png
--------------------------------------------------------------------------------
/data/outputs/2_rgb.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Picsart-AI-Research/Zero-Painter/9d391845e633ef69d85ed41ab07ff026d2ce722a/data/outputs/2_rgb.png
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | torch==2.0.1
2 | torchvision
3 | einops
4 | pytorch_lightning
5 | xformers
6 | open_clip_torch
7 | ipywidgets
8 | transformers==4.40.2
9 | segment-anything
10 | imageio
11 | scipy
12 | opencv-python
13 | matplotlib
14 | omegaconf
--------------------------------------------------------------------------------
/src/smplfusion/__init__.py:
--------------------------------------------------------------------------------
1 | # TODO: Remove farancia.Path depency!
2 | from .libpath import Path
3 | from .libimage import IImage
4 |
5 | print (Path.cwd(__file__))
6 | config = Path.cwd(__file__).config
7 | models = Path.cwd(__file__).config
8 | options = Path.cwd(__file__).config
--------------------------------------------------------------------------------
/src/smplfusion/animation.py:
--------------------------------------------------------------------------------
1 | import matplotlib.pyplot as plt
2 | from matplotlib import animation
3 | from IPython.display import HTML, Image, display
4 |
5 | class Animation:
6 | JS = 0
7 | HTML = 1
8 | ANIMATION_MODE = HTML
9 | def __init__(self, frames, fps = 30):
10 | """_summary_
11 |
12 | Args:
13 | frames (np.ndarray): _description_
14 | """
15 | self.frames = frames
16 | self.fps = fps
17 | self.anim_obj = None
18 | self.anim_str = None
19 | def render(self):
20 | size = (self.frames.shape[2],self.frames.shape[1])
21 | self.fig = plt.figure(figsize = size, dpi = 1)
22 | plt.axis('off')
23 | img = plt.imshow(self.frames[0], cmap = 'gray')
24 | self.fig.subplots_adjust(0,0,1,1)
25 | self.anim_obj = animation.FuncAnimation(
26 | self.fig,
27 | lambda i: img.set_data(self.frames[i,:,:,:]),
28 | frames=self.frames.shape[0],
29 | interval = 1000 / self.fps
30 | )
31 | plt.close()
32 | if Animation.ANIMATION_MODE == Animation.HTML:
33 | self.anim_str = self.anim_obj.to_html5_video()
34 | elif Animation.ANIMATION_MODE == Animation.JS:
35 | self.anim_str = self.anim_obj.to_jshtml()
36 | return self.anim_obj
37 | def _repr_html_(self):
38 | if self.anim_obj is None: self.render()
39 | return self.anim_str
--------------------------------------------------------------------------------
/src/smplfusion/api/__init__.py:
--------------------------------------------------------------------------------
1 | # API goes here!
--------------------------------------------------------------------------------
/src/smplfusion/common.py:
--------------------------------------------------------------------------------
1 | from . import share, scheduler
2 | from .ddim import DDIM
3 | from .patches import router, attentionpatch, transformerpatch
4 | from . import options
5 | from pytorch_lightning import seed_everything
6 | import open_clip
7 |
8 | def count_tokens(prompt):
9 | tokens = open_clip.tokenize(prompt)[0]
10 | return (tokens > 0).sum()
11 |
12 | def tokenize(prompt):
13 | tokens = open_clip.tokenize(prompt)[0]
14 | return [open_clip.tokenizer._tokenizer.decoder[x.item()] for x in tokens]
15 |
16 | def get_token_idx(prompt, prefix, positive_prompt):
17 | prompt = prefix.format(prompt)
18 | return list(range(1 + prefix.split(' ').index('{}'), tokenize(prompt).index('