Spaces:
Runtime error
Runtime error
liuyuan-pal
commited on
Commit
•
8bb8404
1
Parent(s):
2dbbd42
init
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- .gitignore +5 -0
- app.py +106 -0
- assets/crop_size.jpg +0 -0
- assets/elevation.jpg +0 -0
- assets/teaser.jpg +0 -0
- blender_script.py +282 -0
- configs/nerf.yaml +25 -0
- configs/neus.yaml +26 -0
- configs/syncdreamer-train.yaml +63 -0
- configs/syncdreamer.yaml +45 -0
- examples/monkey.png +0 -0
- foreground_segment.py +50 -0
- generate.py +62 -0
- hf_demo/examples/monkey.png +0 -0
- hf_demo/style.css +33 -0
- ldm/base_utils.py +158 -0
- ldm/data/__init__.py +0 -0
- ldm/data/base.py +40 -0
- ldm/data/coco.py +253 -0
- ldm/data/dummy.py +34 -0
- ldm/data/imagenet.py +394 -0
- ldm/data/inpainting/__init__.py +0 -0
- ldm/data/inpainting/synthetic_mask.py +166 -0
- ldm/data/laion.py +537 -0
- ldm/data/lsun.py +92 -0
- ldm/data/nerf_like.py +165 -0
- ldm/data/simple.py +526 -0
- ldm/data/sync_dreamer.py +132 -0
- ldm/lr_scheduler.py +98 -0
- ldm/models/autoencoder.py +443 -0
- ldm/models/diffusion/__init__.py +0 -0
- ldm/models/diffusion/sync_dreamer.py +665 -0
- ldm/models/diffusion/sync_dreamer_attention.py +142 -0
- ldm/models/diffusion/sync_dreamer_network.py +186 -0
- ldm/models/diffusion/sync_dreamer_utils.py +103 -0
- ldm/modules/attention.py +336 -0
- ldm/modules/diffusionmodules/__init__.py +0 -0
- ldm/modules/diffusionmodules/model.py +835 -0
- ldm/modules/diffusionmodules/openaimodel.py +996 -0
- ldm/modules/diffusionmodules/util.py +267 -0
- ldm/modules/distributions/__init__.py +0 -0
- ldm/modules/distributions/distributions.py +92 -0
- ldm/modules/encoders/__init__.py +0 -0
- ldm/modules/encoders/modules.py +550 -0
- ldm/modules/x_transformer.py +641 -0
- ldm/thirdp/psp/helpers.py +121 -0
- ldm/thirdp/psp/id_loss.py +23 -0
- ldm/thirdp/psp/model_irse.py +86 -0
- ldm/util.py +302 -0
- meta_info/camera-16.pkl +3 -0
.gitignore
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
.idea
|
2 |
+
training_examples
|
3 |
+
objaverse_examples
|
4 |
+
ldm/__pycache__/
|
5 |
+
|
app.py
ADDED
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from functools import partial
|
2 |
+
|
3 |
+
from PIL import Image
|
4 |
+
import numpy as np
|
5 |
+
import gradio as gr
|
6 |
+
import torch
|
7 |
+
import os
|
8 |
+
import fire
|
9 |
+
|
10 |
+
from ldm.util import add_margin
|
11 |
+
|
12 |
+
_TITLE = '''SyncDreamer: Generating Multiview-consistent Images from a Single-view Image'''
|
13 |
+
_DESCRIPTION = '''
|
14 |
+
<div>
|
15 |
+
<a style="display:inline-block" href="https://liuyuan-pal.github.io/SyncDreamer/"><img src="https://img.shields.io/badge/SyncDremer-Homepage-blue"></a>
|
16 |
+
<a style="display:inline-block; margin-left: .5em" href="https://arxiv.org/abs/2309.03453"><img src="https://img.shields.io/badge/2309.03453-f9f7f7?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADcAAABMCAYAAADJPi9EAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAuIwAALiMBeKU/dgAAABl0RVh0U29mdHdhcmUAd3d3Lmlua3NjYXBlLm9yZ5vuPBoAAAa2SURBVHja3Zt7bBRFGMAXUCDGF4rY7m7bAwuhlggKStFgLBgFEkCIIRJEEoOBYHwRFYKilUgEReVNJEGCJJpehHI3M9vZvd3bUP1DjNhEIRQQsQgSHiJgQZ5dv7krWEvvdmZ7d7vHJN+ft/f99pv5XvOtJMFCqvoCUpTdIEeRLC+L9Ox5i3Q9LACaCeK0kXoSChVcD3C/tQPHpAEsquQ73IkUcEz2kcLCknyGW5MGjkljRFVL8xJOKyi4CwCOuQAeAkfTP1+tNxLkogvgEbDgffkJqKqvuMA5ifOpqg/5qWecRstNg7xoUTI1Fovdxg8oy2s5AP8CGeYHmGngeZaOL4I4LXLcpHg4149/GDz4xqgsb+UAbMKKUpkrqHA43MUyyJpWUK0EHeG2YKRXr7tB+QMcgGewLD+ebTDbtrtbBt7UPlhS4rV4IvcDI7J8P1OeA/AcAI7LHljN7aB8XTowJmZt9EFRD/o0SDMH4HlwMhMyDWZZSAHFf3YDs3RS49WDLuaAY3IJq+qzmQKLxXAZKN7oDoYbdV3v5elPqiSpMyiOuAEVZVqHXb1OhloUH+MA+ztO0cAO/RkrfyBE7OAEbAZvO8vzVtTRWFD6DAfY5biBM3PWiaL0a4lvXICwnV8WjmE6ntYmhqX2jjp5LbMZjCw/wbYeN6CizOa2GMVzQOlmHjB4Ceuyk6LJ8huccEmR5Xddg7OOV/NAtchW+E3XbOag60QA4Qwuarca0bRuEJyr+cFQwzcY98huxhAKdQelt4kAQpj4qJ3gvFXAYn+aJumXk1yPlpQUgtIHhbYoFMUstNRRWgjnpl4A7IKlayNymqFHFaWCpV9CFry3LGxR1CgA5kB5M8OX2goApwpaz6mdOMGxtAgXWJySxb4WuQD4qTDgU+N5AAnzpr7ChSWpCyisiQJqY0Y7FtmSKpbV23b45kC0KHBxcQ9QeI8w4KgnHRPVtIU7rOtbioLVg5Hl/qDwSVFAMqLSMSObroCdZYlzIJtMRFVHCaRo/wFWPgaAXzdbBpkc2A4aKzCNd97+URQuESYGDDhIVfWOQIKZJu4D2+oXlgDTV1865gUQZDts756BArMNMoR1oa46BYqbyPixZz1ZUFV3sgwoGBajuBKATl3btIn8QYYMuezRgrsiRUWyr2BxA40EkPMpA/Hm6gbUu7fjEXA3azP6AsbKD9bxdUuhjM9W7fII52BF+daRpE4+WA3P501+jbfmHvQKyFqMuXf7Ot4mkN2fr50y+bRH61X7AXdUpHSxaPQ4GVbR5AGw3g+434XgQGKfr72I+vQRhfsu92dOx7WicInzt3CBg1RVpMm0NveWo2SqFzgmdNZMbriILD+S+zoueWf2vSdAipzacWN5nMl6XxNlUHa/J8DoJodUDE0HR8Ll5V0lPxcrLEHZPV4AzS83OLis7FowVa3RSku7BSNxJqQAlN3hBTC2apmDSkpaw22wJemGQFUG7J4MlP3JC6A+f96V7vRyX9It3nzT/GrjIU8edM7rMSnIi10f476lzbE1K7yEiEuWro0OJBguLCwDuFOJc1Na6sRWL/cCeMIwUN9ggSVbe3v/5/EgzTKWLvEAiBrYRUkgwNI2ZaFQNT75UDxEUEx97zYnzpmiLEmbaYCbNxYtFAb0/Z4AztgUrhyxuNgxPnhfHFDHz/vTgFWUQZxTRkkJhQ6YNdVUEPAfO6ZV5BRss6LcCVb7VaAma9giy0XJZBt9IQh42NY0NSdgbLIPlLUF6rEdrdt0CUCK1wsCbkcI3ZSLc7ZSwGLbmJXbPsNxnE5xilYKAobZ77LpGZ8TAIun+/iCKQoF71IxQDI3K2CCd+ARNvXg9sykBcnHAoCZG4u66hlDoQLe6QV4CRtFSxZQ+D0BwNO2jgdkzoGoah1nj3FVlSR19taTSYxI8QLut23U8dsgzqHulJNCQpcqBnpTALCuQ6NSYLHpmR5i42gZzuIdcrMMvMJbQlxe3jXxyZnLACl7ARm/FjPIDOY8ODtpM71sxwfcZpvBeUzKWmfNINM5AS+wO0Khh7dMqKccu4+qatarZjYAwDlgetzStHtEt+XedsBOQtU9XMrRgjg4KTnc5nr+dmqadit/4C4uLm8DuA9koJTj1TL7fI5nDL+qqoo/FLGAzL7dYT17PzvAcQONYSUQRxW/QMrHZVIyik0ZuQA2mzp+Ji8BW4YM3Mbzm9inaHkJCGfrUZZjujiYailfFwA8DHIy3acwUj4v9vUVa+SmgNsl5fuyDTKovW9/IAmfLV0Pi2UncA515kjYdrwC9i9rpuHiq3JwtAAAAABJRU5ErkJggg=="></a>
|
17 |
+
<a style="display:inline-block; margin-left: .5em" href='https://github.com/liuyuan-pal/SyncDreamer'><img src='https://img.shields.io/github/stars/liuyuan-pal/SyncDreamer?style=social' /></a>
|
18 |
+
</div>
|
19 |
+
Given a single-view image, SyncDreamer is able to generate multiview-consistent images, which enables direct 3D reconstruction with NeuS or NeRF without SDS loss'''
|
20 |
+
_USER_GUIDE0 = "Step0: Please upload an image in the block above (or choose an example above). We use alpha values as object masks if given."
|
21 |
+
_USER_GUIDE1 = "Step1: Please select a crop size using the glider."
|
22 |
+
_USER_GUIDE2 = "Step2: Please choose a suitable elevation angle and then click the Generate button."
|
23 |
+
|
24 |
+
|
25 |
+
def mask_prediction(mask_predictor, image_in: Image.Image):
|
26 |
+
if image_in.mode=='RGBA':
|
27 |
+
return image_in
|
28 |
+
else:
|
29 |
+
raise NotImplementedError
|
30 |
+
|
31 |
+
def resize_inputs(image_input, crop_size):
|
32 |
+
alpha_np = np.asarray(image_input)[:, :, 3]
|
33 |
+
coords = np.stack(np.nonzero(alpha_np), 1)[:, (1, 0)]
|
34 |
+
min_x, min_y = np.min(coords, 0)
|
35 |
+
max_x, max_y = np.max(coords, 0)
|
36 |
+
ref_img_ = image_input.crop((min_x, min_y, max_x, max_y))
|
37 |
+
h, w = ref_img_.height, ref_img_.width
|
38 |
+
scale = crop_size / max(h, w)
|
39 |
+
h_, w_ = int(scale * h), int(scale * w)
|
40 |
+
ref_img_ = ref_img_.resize((w_, h_), resample=Image.BICUBIC)
|
41 |
+
results = add_margin(ref_img_, size=256)
|
42 |
+
return results
|
43 |
+
|
44 |
+
def generate(syncdreamer_model, image, elevation):
|
45 |
+
|
46 |
+
|
47 |
+
def run_demo():
|
48 |
+
device = f"cuda:0" if torch.cuda.is_available() else "cpu"
|
49 |
+
models = None # init_model(device, os.path.join(code_dir, ckpt))
|
50 |
+
|
51 |
+
# init sam model
|
52 |
+
mask_predictor = None # sam_init(device_idx)
|
53 |
+
|
54 |
+
# with open('instructions_12345.md', 'r') as f:
|
55 |
+
# article = f.read()
|
56 |
+
|
57 |
+
# NOTE: Examples must match inputs
|
58 |
+
example_folder = os.path.join(os.path.dirname(__file__), 'hf_demo', 'examples')
|
59 |
+
example_fns = os.listdir(example_folder)
|
60 |
+
example_fns.sort()
|
61 |
+
examples_full = [os.path.join(example_folder, x) for x in example_fns if x.endswith('.png')]
|
62 |
+
|
63 |
+
# Compose demo layout & data flow.
|
64 |
+
with gr.Blocks(title=_TITLE, css="hf_demo/style.css") as demo:
|
65 |
+
with gr.Row():
|
66 |
+
with gr.Column(scale=1):
|
67 |
+
gr.Markdown('# ' + _TITLE)
|
68 |
+
# with gr.Column(scale=0):
|
69 |
+
# gr.DuplicateButton(value='Duplicate Space for private use', elem_id='duplicate-button')
|
70 |
+
gr.Markdown(_DESCRIPTION)
|
71 |
+
|
72 |
+
with gr.Row(variant='panel'):
|
73 |
+
with gr.Column(scale=1):
|
74 |
+
image_block = gr.Image(type='pil', image_mode='RGBA', height=256, label='Input image', tool=None, interactive=True)
|
75 |
+
guide_text = gr.Markdown(_USER_GUIDE0, visible=True)
|
76 |
+
gr.Examples(
|
77 |
+
examples=examples_full, # NOTE: elements must match inputs list!
|
78 |
+
inputs=[image_block],
|
79 |
+
outputs=[image_block],
|
80 |
+
cache_examples=False,
|
81 |
+
label='Examples (click one of the images below to start)',
|
82 |
+
examples_per_page=40
|
83 |
+
)
|
84 |
+
|
85 |
+
with gr.Column(scale=1):
|
86 |
+
sam_block = gr.Image(type='pil', image_mode='RGBA', label="SAM output", height=256, interactive=False)
|
87 |
+
crop_size_slider = gr.Slider(120, 240, 200, step=10, label='Crop size', interactive=True)
|
88 |
+
|
89 |
+
with gr.Column(scale=1):
|
90 |
+
input_block = gr.Image(type='pil', image_mode='RGB', label="Input to SyncDreamer", height=256, interactive=False)
|
91 |
+
elevation_slider = gr.Slider(-10, 40, 30, step=5, label='Elevation angle', interactive=True)
|
92 |
+
run_btn = gr.Button('Run Generation', variant='primary', interactive=False)
|
93 |
+
|
94 |
+
update_guide = lambda GUIDE_TEXT: gr.update(value=GUIDE_TEXT)
|
95 |
+
image_block.change(fn=partial(mask_prediction, mask_predictor), inputs=[image_block], outputs=[sam_block], queue=False)\
|
96 |
+
.success(fn=partial(update_guide, _USER_GUIDE1), outputs=[guide_text], queue=False)
|
97 |
+
|
98 |
+
crop_size_slider.change(fn=resize_inputs, inputs=[sam_block, crop_size_slider], outputs=[input_block], queue=False)\
|
99 |
+
.success(fn=partial(update_guide, _USER_GUIDE2), outputs=[guide_text], queue=False)
|
100 |
+
|
101 |
+
run_btn.click
|
102 |
+
|
103 |
+
demo.queue().launch(share=False, max_threads=80) # auth=("admin", os.environ['PASSWD'])
|
104 |
+
|
105 |
+
if __name__=="__main__":
|
106 |
+
fire.Fire(run_demo)
|
assets/crop_size.jpg
ADDED
assets/elevation.jpg
ADDED
assets/teaser.jpg
ADDED
blender_script.py
ADDED
@@ -0,0 +1,282 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""Blender script to render images of 3D models.
|
2 |
+
|
3 |
+
This script is used to render images of 3D models. It takes in a list of paths
|
4 |
+
to .glb files and renders images of each model. The images are from rotating the
|
5 |
+
object around the origin. The images are saved to the output directory.
|
6 |
+
|
7 |
+
Example usage:
|
8 |
+
blender -b -P blender_script.py -- \
|
9 |
+
--object_path my_object.glb \
|
10 |
+
--output_dir ./views \
|
11 |
+
--engine CYCLES \
|
12 |
+
--scale 0.8 \
|
13 |
+
--num_images 12 \
|
14 |
+
--camera_dist 1.2
|
15 |
+
|
16 |
+
Here, input_model_paths.json is a json file containing a list of paths to .glb.
|
17 |
+
"""
|
18 |
+
|
19 |
+
import argparse
|
20 |
+
import json
|
21 |
+
import math
|
22 |
+
import os
|
23 |
+
import random
|
24 |
+
import sys
|
25 |
+
import time
|
26 |
+
import urllib.request
|
27 |
+
from pathlib import Path
|
28 |
+
|
29 |
+
from mathutils import Vector, Matrix
|
30 |
+
import numpy as np
|
31 |
+
|
32 |
+
import bpy
|
33 |
+
from mathutils import Vector
|
34 |
+
import pickle
|
35 |
+
|
36 |
+
def read_pickle(pkl_path):
|
37 |
+
with open(pkl_path, 'rb') as f:
|
38 |
+
return pickle.load(f)
|
39 |
+
|
40 |
+
def save_pickle(data, pkl_path):
|
41 |
+
# os.system('mkdir -p {}'.format(os.path.dirname(pkl_path)))
|
42 |
+
with open(pkl_path, 'wb') as f:
|
43 |
+
pickle.dump(data, f)
|
44 |
+
|
45 |
+
parser = argparse.ArgumentParser()
|
46 |
+
parser.add_argument("--object_path", type=str, required=True)
|
47 |
+
parser.add_argument("--output_dir", type=str, required=True)
|
48 |
+
parser.add_argument("--engine", type=str, default="CYCLES", choices=["CYCLES", "BLENDER_EEVEE"])
|
49 |
+
parser.add_argument("--camera_type", type=str, default='even')
|
50 |
+
parser.add_argument("--num_images", type=int, default=16)
|
51 |
+
parser.add_argument("--elevation", type=float, default=30)
|
52 |
+
parser.add_argument("--elevation_start", type=float, default=-10)
|
53 |
+
parser.add_argument("--elevation_end", type=float, default=40)
|
54 |
+
parser.add_argument("--device", type=str, default='CUDA')
|
55 |
+
|
56 |
+
argv = sys.argv[sys.argv.index("--") + 1 :]
|
57 |
+
args = parser.parse_args(argv)
|
58 |
+
|
59 |
+
print('===================', args.engine, '===================')
|
60 |
+
|
61 |
+
context = bpy.context
|
62 |
+
scene = context.scene
|
63 |
+
render = scene.render
|
64 |
+
|
65 |
+
cam = scene.objects["Camera"]
|
66 |
+
cam.location = (0, 1.2, 0)
|
67 |
+
cam.data.lens = 35
|
68 |
+
cam.data.sensor_width = 32
|
69 |
+
|
70 |
+
cam_constraint = cam.constraints.new(type="TRACK_TO")
|
71 |
+
cam_constraint.track_axis = "TRACK_NEGATIVE_Z"
|
72 |
+
cam_constraint.up_axis = "UP_Y"
|
73 |
+
|
74 |
+
render.engine = args.engine
|
75 |
+
render.image_settings.file_format = "PNG"
|
76 |
+
render.image_settings.color_mode = "RGBA"
|
77 |
+
render.resolution_x = 256
|
78 |
+
render.resolution_y = 256
|
79 |
+
render.resolution_percentage = 100
|
80 |
+
|
81 |
+
scene.cycles.device = "GPU"
|
82 |
+
scene.cycles.samples = 128
|
83 |
+
scene.cycles.diffuse_bounces = 1
|
84 |
+
scene.cycles.glossy_bounces = 1
|
85 |
+
scene.cycles.transparent_max_bounces = 3
|
86 |
+
scene.cycles.transmission_bounces = 3
|
87 |
+
scene.cycles.filter_width = 0.01
|
88 |
+
scene.cycles.use_denoising = True
|
89 |
+
scene.render.film_transparent = True
|
90 |
+
|
91 |
+
bpy.context.preferences.addons["cycles"].preferences.get_devices()
|
92 |
+
# Set the device_type
|
93 |
+
bpy.context.preferences.addons["cycles"].preferences.compute_device_type = args.device # or "OPENCL"
|
94 |
+
bpy.context.scene.cycles.tile_size = 8192
|
95 |
+
|
96 |
+
|
97 |
+
def az_el_to_points(azimuths, elevations):
|
98 |
+
x = np.cos(azimuths)*np.cos(elevations)
|
99 |
+
y = np.sin(azimuths)*np.cos(elevations)
|
100 |
+
z = np.sin(elevations)
|
101 |
+
return np.stack([x,y,z],-1) #
|
102 |
+
|
103 |
+
def set_camera_location(cam_pt):
|
104 |
+
# from https://blender.stackexchange.com/questions/18530/
|
105 |
+
x, y, z = cam_pt # sample_spherical(radius_min=1.5, radius_max=2.2, maxz=2.2, minz=-2.2)
|
106 |
+
camera = bpy.data.objects["Camera"]
|
107 |
+
camera.location = x, y, z
|
108 |
+
|
109 |
+
return camera
|
110 |
+
|
111 |
+
def get_calibration_matrix_K_from_blender(camera):
|
112 |
+
f_in_mm = camera.data.lens
|
113 |
+
scene = bpy.context.scene
|
114 |
+
resolution_x_in_px = scene.render.resolution_x
|
115 |
+
resolution_y_in_px = scene.render.resolution_y
|
116 |
+
scale = scene.render.resolution_percentage / 100
|
117 |
+
sensor_width_in_mm = camera.data.sensor_width
|
118 |
+
sensor_height_in_mm = camera.data.sensor_height
|
119 |
+
pixel_aspect_ratio = scene.render.pixel_aspect_x / scene.render.pixel_aspect_y
|
120 |
+
|
121 |
+
if camera.data.sensor_fit == 'VERTICAL':
|
122 |
+
# the sensor height is fixed (sensor fit is horizontal),
|
123 |
+
# the sensor width is effectively changed with the pixel aspect ratio
|
124 |
+
s_u = resolution_x_in_px * scale / sensor_width_in_mm / pixel_aspect_ratio
|
125 |
+
s_v = resolution_y_in_px * scale / sensor_height_in_mm
|
126 |
+
else: # 'HORIZONTAL' and 'AUTO'
|
127 |
+
# the sensor width is fixed (sensor fit is horizontal),
|
128 |
+
# the sensor height is effectively changed with the pixel aspect ratio
|
129 |
+
s_u = resolution_x_in_px * scale / sensor_width_in_mm
|
130 |
+
s_v = resolution_y_in_px * scale * pixel_aspect_ratio / sensor_height_in_mm
|
131 |
+
|
132 |
+
# Parameters of intrinsic calibration matrix K
|
133 |
+
alpha_u = f_in_mm * s_u
|
134 |
+
alpha_v = f_in_mm * s_u
|
135 |
+
u_0 = resolution_x_in_px * scale / 2
|
136 |
+
v_0 = resolution_y_in_px * scale / 2
|
137 |
+
skew = 0 # only use rectangular pixels
|
138 |
+
|
139 |
+
K = np.asarray(((alpha_u, skew, u_0),
|
140 |
+
(0, alpha_v, v_0),
|
141 |
+
(0, 0, 1)),np.float32)
|
142 |
+
return K
|
143 |
+
|
144 |
+
|
145 |
+
def reset_scene() -> None:
|
146 |
+
"""Resets the scene to a clean state."""
|
147 |
+
# delete everything that isn't part of a camera or a light
|
148 |
+
for obj in bpy.data.objects:
|
149 |
+
if obj.type not in {"CAMERA", "LIGHT"}:
|
150 |
+
bpy.data.objects.remove(obj, do_unlink=True)
|
151 |
+
# delete all the materials
|
152 |
+
for material in bpy.data.materials:
|
153 |
+
bpy.data.materials.remove(material, do_unlink=True)
|
154 |
+
# delete all the textures
|
155 |
+
for texture in bpy.data.textures:
|
156 |
+
bpy.data.textures.remove(texture, do_unlink=True)
|
157 |
+
# delete all the images
|
158 |
+
for image in bpy.data.images:
|
159 |
+
bpy.data.images.remove(image, do_unlink=True)
|
160 |
+
|
161 |
+
|
162 |
+
# load the glb model
|
163 |
+
def load_object(object_path: str) -> None:
|
164 |
+
"""Loads a glb model into the scene."""
|
165 |
+
if object_path.endswith(".glb"):
|
166 |
+
bpy.ops.import_scene.gltf(filepath=object_path, merge_vertices=True)
|
167 |
+
elif object_path.endswith(".fbx"):
|
168 |
+
bpy.ops.import_scene.fbx(filepath=object_path)
|
169 |
+
else:
|
170 |
+
raise ValueError(f"Unsupported file type: {object_path}")
|
171 |
+
|
172 |
+
|
173 |
+
def scene_bbox(single_obj=None, ignore_matrix=False):
|
174 |
+
bbox_min = (math.inf,) * 3
|
175 |
+
bbox_max = (-math.inf,) * 3
|
176 |
+
found = False
|
177 |
+
for obj in scene_meshes() if single_obj is None else [single_obj]:
|
178 |
+
found = True
|
179 |
+
for coord in obj.bound_box:
|
180 |
+
coord = Vector(coord)
|
181 |
+
if not ignore_matrix:
|
182 |
+
coord = obj.matrix_world @ coord
|
183 |
+
bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
|
184 |
+
bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
|
185 |
+
if not found:
|
186 |
+
raise RuntimeError("no objects in scene to compute bounding box for")
|
187 |
+
return Vector(bbox_min), Vector(bbox_max)
|
188 |
+
|
189 |
+
|
190 |
+
def scene_root_objects():
|
191 |
+
for obj in bpy.context.scene.objects.values():
|
192 |
+
if not obj.parent:
|
193 |
+
yield obj
|
194 |
+
|
195 |
+
|
196 |
+
def scene_meshes():
|
197 |
+
for obj in bpy.context.scene.objects.values():
|
198 |
+
if isinstance(obj.data, (bpy.types.Mesh)):
|
199 |
+
yield obj
|
200 |
+
|
201 |
+
# function from https://github.com/panmari/stanford-shapenet-renderer/blob/master/render_blender.py
|
202 |
+
def get_3x4_RT_matrix_from_blender(cam):
|
203 |
+
bpy.context.view_layer.update()
|
204 |
+
location, rotation = cam.matrix_world.decompose()[0:2]
|
205 |
+
R = np.asarray(rotation.to_matrix())
|
206 |
+
t = np.asarray(location)
|
207 |
+
|
208 |
+
cam_rec = np.asarray([[1, 0, 0], [0, -1, 0], [0, 0, -1]], np.float32)
|
209 |
+
R = R.T
|
210 |
+
t = -R @ t
|
211 |
+
R_world2cv = cam_rec @ R
|
212 |
+
t_world2cv = cam_rec @ t
|
213 |
+
|
214 |
+
RT = np.concatenate([R_world2cv,t_world2cv[:,None]],1)
|
215 |
+
return RT
|
216 |
+
|
217 |
+
def normalize_scene():
|
218 |
+
bbox_min, bbox_max = scene_bbox()
|
219 |
+
scale = 1 / max(bbox_max - bbox_min)
|
220 |
+
for obj in scene_root_objects():
|
221 |
+
obj.scale = obj.scale * scale
|
222 |
+
# Apply scale to matrix_world.
|
223 |
+
bpy.context.view_layer.update()
|
224 |
+
bbox_min, bbox_max = scene_bbox()
|
225 |
+
offset = -(bbox_min + bbox_max) / 2
|
226 |
+
for obj in scene_root_objects():
|
227 |
+
obj.matrix_world.translation += offset
|
228 |
+
bpy.ops.object.select_all(action="DESELECT")
|
229 |
+
|
230 |
+
def save_images(object_file: str) -> None:
|
231 |
+
object_uid = os.path.basename(object_file).split(".")[0]
|
232 |
+
os.makedirs(args.output_dir, exist_ok=True)
|
233 |
+
|
234 |
+
reset_scene()
|
235 |
+
# load the object
|
236 |
+
load_object(object_file)
|
237 |
+
# object_uid = os.path.basename(object_file).split(".")[0]
|
238 |
+
normalize_scene()
|
239 |
+
|
240 |
+
# create an empty object to track
|
241 |
+
empty = bpy.data.objects.new("Empty", None)
|
242 |
+
scene.collection.objects.link(empty)
|
243 |
+
cam_constraint.target = empty
|
244 |
+
|
245 |
+
world_tree = bpy.context.scene.world.node_tree
|
246 |
+
back_node = world_tree.nodes['Background']
|
247 |
+
env_light = 0.5
|
248 |
+
back_node.inputs['Color'].default_value = Vector([env_light, env_light, env_light, 1.0])
|
249 |
+
back_node.inputs['Strength'].default_value = 1.0
|
250 |
+
|
251 |
+
distances = np.asarray([1.5 for _ in range(args.num_images)])
|
252 |
+
if args.camera_type=='fixed':
|
253 |
+
azimuths = (np.arange(args.num_images)/args.num_images*np.pi*2).astype(np.float32)
|
254 |
+
elevations = np.deg2rad(np.asarray([args.elevation] * args.num_images).astype(np.float32))
|
255 |
+
elif args.camera_type=='random':
|
256 |
+
azimuths = (np.arange(args.num_images) / args.num_images * np.pi * 2).astype(np.float32)
|
257 |
+
elevations = np.random.uniform(args.elevation_start, args.elevation_end, args.num_images)
|
258 |
+
elevations = np.deg2rad(elevations)
|
259 |
+
else:
|
260 |
+
raise NotImplementedError
|
261 |
+
|
262 |
+
cam_pts = az_el_to_points(azimuths, elevations) * distances[:,None]
|
263 |
+
cam_poses = []
|
264 |
+
(Path(args.output_dir) / object_uid).mkdir(exist_ok=True, parents=True)
|
265 |
+
for i in range(args.num_images):
|
266 |
+
# set camera
|
267 |
+
camera = set_camera_location(cam_pts[i])
|
268 |
+
RT = get_3x4_RT_matrix_from_blender(camera)
|
269 |
+
cam_poses.append(RT)
|
270 |
+
|
271 |
+
render_path = os.path.join(args.output_dir, object_uid, f"{i:03d}.png")
|
272 |
+
if os.path.exists(render_path): continue
|
273 |
+
scene.render.filepath = os.path.abspath(render_path)
|
274 |
+
bpy.ops.render.render(write_still=True)
|
275 |
+
|
276 |
+
if args.camera_type=='random':
|
277 |
+
K = get_calibration_matrix_K_from_blender(camera)
|
278 |
+
cam_poses = np.stack(cam_poses, 0)
|
279 |
+
save_pickle([K, azimuths, elevations, distances, cam_poses], os.path.join(args.output_dir, object_uid, "meta.pkl"))
|
280 |
+
|
281 |
+
if __name__ == "__main__":
|
282 |
+
save_images(args.object_path)
|
configs/nerf.yaml
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
model:
|
2 |
+
base_lr: 1.0e-2
|
3 |
+
target: renderer.renderer.RendererTrainer
|
4 |
+
params:
|
5 |
+
total_steps: 2000
|
6 |
+
warm_up_steps: 100
|
7 |
+
train_batch_num: 40960
|
8 |
+
test_batch_num: 40960
|
9 |
+
renderer: ngp
|
10 |
+
cube_bound: 0.6
|
11 |
+
use_mask: true
|
12 |
+
lambda_rgb_loss: 0.5
|
13 |
+
lambda_mask_loss: 10.0
|
14 |
+
|
15 |
+
data:
|
16 |
+
target: renderer.dummy_dataset.DummyDataset
|
17 |
+
params: {}
|
18 |
+
|
19 |
+
callbacks:
|
20 |
+
save_interval: 5000
|
21 |
+
|
22 |
+
trainer:
|
23 |
+
val_check_interval: 500
|
24 |
+
max_steps: 2000
|
25 |
+
|
configs/neus.yaml
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
model:
|
2 |
+
base_lr: 5.0e-4
|
3 |
+
target: renderer.renderer.RendererTrainer
|
4 |
+
params:
|
5 |
+
total_steps: 2000
|
6 |
+
warm_up_steps: 100
|
7 |
+
train_batch_num: 3584
|
8 |
+
train_batch_fg_num: 512
|
9 |
+
test_batch_num: 4096
|
10 |
+
use_mask: true
|
11 |
+
lambda_rgb_loss: 0.5
|
12 |
+
lambda_mask_loss: 1.0
|
13 |
+
lambda_eikonal_loss: 0.1
|
14 |
+
use_warm_up: true
|
15 |
+
|
16 |
+
data:
|
17 |
+
target: renderer.dummy_dataset.DummyDataset
|
18 |
+
params: {}
|
19 |
+
|
20 |
+
callbacks:
|
21 |
+
save_interval: 500
|
22 |
+
|
23 |
+
trainer:
|
24 |
+
val_check_interval: 500
|
25 |
+
max_steps: 2000
|
26 |
+
|
configs/syncdreamer-train.yaml
ADDED
@@ -0,0 +1,63 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
model:
|
2 |
+
base_learning_rate: 5.0e-05
|
3 |
+
target: ldm.models.diffusion.sync_dreamer.SyncMultiviewDiffusion
|
4 |
+
params:
|
5 |
+
view_num: 16
|
6 |
+
image_size: 256
|
7 |
+
cfg_scale: 2.0
|
8 |
+
output_num: 8
|
9 |
+
batch_view_num: 4
|
10 |
+
finetune_unet: false
|
11 |
+
finetune_projection: false
|
12 |
+
drop_conditions: false
|
13 |
+
clip_image_encoder_path: ckpt/ViT-L-14.pt
|
14 |
+
|
15 |
+
scheduler_config: # 10000 warmup steps
|
16 |
+
target: ldm.lr_scheduler.LambdaLinearScheduler
|
17 |
+
params:
|
18 |
+
warm_up_steps: [ 100 ]
|
19 |
+
cycle_lengths: [ 100000 ]
|
20 |
+
f_start: [ 0.02 ]
|
21 |
+
f_max: [ 1.0 ]
|
22 |
+
f_min: [ 1.0 ]
|
23 |
+
|
24 |
+
unet_config:
|
25 |
+
target: ldm.models.diffusion.sync_dreamer_attention.DepthWiseAttention
|
26 |
+
params:
|
27 |
+
volume_dims: [64, 128, 256, 512]
|
28 |
+
image_size: 32
|
29 |
+
in_channels: 8
|
30 |
+
out_channels: 4
|
31 |
+
model_channels: 320
|
32 |
+
attention_resolutions: [ 4, 2, 1 ]
|
33 |
+
num_res_blocks: 2
|
34 |
+
channel_mult: [ 1, 2, 4, 4 ]
|
35 |
+
num_heads: 8
|
36 |
+
use_spatial_transformer: True
|
37 |
+
transformer_depth: 1
|
38 |
+
context_dim: 768
|
39 |
+
use_checkpoint: True
|
40 |
+
legacy: False
|
41 |
+
|
42 |
+
data:
|
43 |
+
target: ldm.data.sync_dreamer.SyncDreamerDataset
|
44 |
+
params:
|
45 |
+
target_dir: training_examples/target # renderings of target views
|
46 |
+
input_dir: training_examples/input # renderings of input views
|
47 |
+
uid_set_pkl: training_examples/uid_set.pkl # a list of uids
|
48 |
+
validation_dir: validation_set # directory of validation data
|
49 |
+
batch_size: 24 # batch size for a single gpu
|
50 |
+
num_workers: 8
|
51 |
+
|
52 |
+
lightning:
|
53 |
+
modelcheckpoint:
|
54 |
+
params:
|
55 |
+
every_n_train_steps: 1000 # we will save models every 1k steps
|
56 |
+
callbacks:
|
57 |
+
{}
|
58 |
+
|
59 |
+
trainer:
|
60 |
+
benchmark: True
|
61 |
+
val_check_interval: 1000 # we will run validation every 1k steps, the validation will output images to <log_dir>/<images>/val
|
62 |
+
num_sanity_val_steps: 0
|
63 |
+
check_val_every_n_epoch: null
|
configs/syncdreamer.yaml
ADDED
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
model:
|
2 |
+
base_learning_rate: 5.0e-05
|
3 |
+
target: ldm.models.diffusion.sync_dreamer.SyncMultiviewDiffusion
|
4 |
+
params:
|
5 |
+
view_num: 16
|
6 |
+
image_size: 256
|
7 |
+
cfg_scale: 2.0
|
8 |
+
output_num: 8
|
9 |
+
batch_view_num: 4
|
10 |
+
finetune_unet: false
|
11 |
+
finetune_projection: false
|
12 |
+
drop_conditions: false
|
13 |
+
clip_image_encoder_path: ckpt/ViT-L-14.pt
|
14 |
+
|
15 |
+
scheduler_config: # 10000 warmup steps
|
16 |
+
target: ldm.lr_scheduler.LambdaLinearScheduler
|
17 |
+
params:
|
18 |
+
warm_up_steps: [ 100 ]
|
19 |
+
cycle_lengths: [ 100000 ]
|
20 |
+
f_start: [ 0.02 ]
|
21 |
+
f_max: [ 1.0 ]
|
22 |
+
f_min: [ 1.0 ]
|
23 |
+
|
24 |
+
unet_config:
|
25 |
+
target: ldm.models.diffusion.sync_dreamer_attention.DepthWiseAttention
|
26 |
+
params:
|
27 |
+
volume_dims: [64, 128, 256, 512]
|
28 |
+
image_size: 32
|
29 |
+
in_channels: 8
|
30 |
+
out_channels: 4
|
31 |
+
model_channels: 320
|
32 |
+
attention_resolutions: [ 4, 2, 1 ]
|
33 |
+
num_res_blocks: 2
|
34 |
+
channel_mult: [ 1, 2, 4, 4 ]
|
35 |
+
num_heads: 8
|
36 |
+
use_spatial_transformer: True
|
37 |
+
transformer_depth: 1
|
38 |
+
context_dim: 768
|
39 |
+
use_checkpoint: True
|
40 |
+
legacy: False
|
41 |
+
|
42 |
+
data: {}
|
43 |
+
|
44 |
+
lightning:
|
45 |
+
trainer: {}
|
examples/monkey.png
ADDED
foreground_segment.py
ADDED
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import argparse
|
3 |
+
import numpy as np
|
4 |
+
|
5 |
+
import torch
|
6 |
+
from PIL import Image
|
7 |
+
|
8 |
+
|
9 |
+
class BackgroundRemoval:
|
10 |
+
def __init__(self, device='cuda'):
|
11 |
+
from carvekit.api.high import HiInterface
|
12 |
+
self.interface = HiInterface(
|
13 |
+
object_type="object", # Can be "object" or "hairs-like".
|
14 |
+
batch_size_seg=5,
|
15 |
+
batch_size_matting=1,
|
16 |
+
device=device,
|
17 |
+
seg_mask_size=640, # Use 640 for Tracer B7 and 320 for U2Net
|
18 |
+
matting_mask_size=2048,
|
19 |
+
trimap_prob_threshold=231,
|
20 |
+
trimap_dilation=30,
|
21 |
+
trimap_erosion_iters=5,
|
22 |
+
fp16=True,
|
23 |
+
)
|
24 |
+
|
25 |
+
@torch.no_grad()
|
26 |
+
def __call__(self, image):
|
27 |
+
# image: [H, W, 3] array in [0, 255].
|
28 |
+
image = Image.fromarray(image)
|
29 |
+
image = self.interface([image])[0]
|
30 |
+
image = np.array(image)
|
31 |
+
return image
|
32 |
+
|
33 |
+
def process(image_path, mask_path):
|
34 |
+
mask_predictor = BackgroundRemoval()
|
35 |
+
image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
|
36 |
+
if image.shape[-1] == 4:
|
37 |
+
image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
|
38 |
+
else:
|
39 |
+
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
40 |
+
rgba = mask_predictor(image) # [H, W, 4]
|
41 |
+
cv2.imwrite(mask_path, cv2.cvtColor(rgba, cv2.COLOR_RGBA2BGRA))
|
42 |
+
|
43 |
+
|
44 |
+
if __name__ == '__main__':
|
45 |
+
parser = argparse.ArgumentParser()
|
46 |
+
parser.add_argument('--input', required=True, type=str)
|
47 |
+
parser.add_argument('--output', required=True, type=str)
|
48 |
+
opt = parser.parse_args()
|
49 |
+
|
50 |
+
process(opt.input, opt.output)
|
generate.py
ADDED
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import argparse
|
2 |
+
from pathlib import Path
|
3 |
+
|
4 |
+
import numpy as np
|
5 |
+
import torch
|
6 |
+
from omegaconf import OmegaConf
|
7 |
+
from skimage.io import imsave
|
8 |
+
|
9 |
+
from ldm.models.diffusion.sync_dreamer import SyncMultiviewDiffusion
|
10 |
+
from ldm.util import instantiate_from_config, prepare_inputs
|
11 |
+
|
12 |
+
|
13 |
+
def load_model(cfg,ckpt,strict=True):
|
14 |
+
config = OmegaConf.load(cfg)
|
15 |
+
model = instantiate_from_config(config.model)
|
16 |
+
print(f'loading model from {ckpt} ...')
|
17 |
+
ckpt = torch.load(ckpt,map_location='cpu')
|
18 |
+
model.load_state_dict(ckpt['state_dict'],strict=strict)
|
19 |
+
model = model.cuda().eval()
|
20 |
+
return model
|
21 |
+
|
22 |
+
def main():
|
23 |
+
parser = argparse.ArgumentParser()
|
24 |
+
parser.add_argument('--cfg',type=str, default='configs/syncdreamer.yaml')
|
25 |
+
parser.add_argument('--ckpt',type=str, default='ckpt/syncdreamer-step80k.ckpt')
|
26 |
+
parser.add_argument('--output', type=str, required=True)
|
27 |
+
parser.add_argument('--input', type=str, required=True)
|
28 |
+
parser.add_argument('--elevation', type=float, required=True)
|
29 |
+
|
30 |
+
parser.add_argument('--sample_num', type=int, default=4)
|
31 |
+
parser.add_argument('--crop_size', type=int, default=-1)
|
32 |
+
parser.add_argument('--cfg_scale', type=float, default=2.0)
|
33 |
+
parser.add_argument('--batch_view_num', type=int, default=8)
|
34 |
+
parser.add_argument('--seed', type=int, default=6033)
|
35 |
+
flags = parser.parse_args()
|
36 |
+
|
37 |
+
torch.random.manual_seed(flags.seed)
|
38 |
+
np.random.seed(flags.seed)
|
39 |
+
|
40 |
+
model = load_model(flags.cfg, flags.ckpt, strict=True)
|
41 |
+
assert isinstance(model, SyncMultiviewDiffusion)
|
42 |
+
Path(f'{flags.output}').mkdir(exist_ok=True, parents=True)
|
43 |
+
|
44 |
+
# prepare data
|
45 |
+
data = prepare_inputs(flags.input, flags.elevation, flags.crop_size)
|
46 |
+
for k, v in data.items():
|
47 |
+
data[k] = v.unsqueeze(0).cuda()
|
48 |
+
data[k] = torch.repeat_interleave(data[k], flags.sample_num, dim=0)
|
49 |
+
x_sample = model.sample(data, flags.cfg_scale, flags.batch_view_num)
|
50 |
+
|
51 |
+
B, N, _, H, W = x_sample.shape
|
52 |
+
x_sample = (torch.clamp(x_sample,max=1.0,min=-1.0) + 1) * 0.5
|
53 |
+
x_sample = x_sample.permute(0,1,3,4,2).cpu().numpy() * 255
|
54 |
+
x_sample = x_sample.astype(np.uint8)
|
55 |
+
|
56 |
+
for bi in range(B):
|
57 |
+
output_fn = Path(flags.output)/ f'{bi}.png'
|
58 |
+
imsave(output_fn, np.concatenate([x_sample[bi,ni] for ni in range(N)], 1))
|
59 |
+
|
60 |
+
if __name__=="__main__":
|
61 |
+
main()
|
62 |
+
|
hf_demo/examples/monkey.png
ADDED
hf_demo/style.css
ADDED
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#model-3d-out {
|
2 |
+
height: 400px;
|
3 |
+
}
|
4 |
+
|
5 |
+
#plot-out {
|
6 |
+
height: 450px;
|
7 |
+
}
|
8 |
+
|
9 |
+
#duplicate-button {
|
10 |
+
margin-left: auto;
|
11 |
+
color: #fff;
|
12 |
+
background: #1565c0;
|
13 |
+
}
|
14 |
+
|
15 |
+
.footer {
|
16 |
+
margin-bottom: 45px;
|
17 |
+
margin-top: 10px;
|
18 |
+
text-align: center;
|
19 |
+
border-bottom: 1px solid #e5e5e5;
|
20 |
+
}
|
21 |
+
.footer>p {
|
22 |
+
font-size: .8rem;
|
23 |
+
display: inline-block;
|
24 |
+
padding: 0 10px;
|
25 |
+
transform: translateY(15px);
|
26 |
+
background: white;
|
27 |
+
}
|
28 |
+
.dark .footer {
|
29 |
+
border-color: #303030;
|
30 |
+
}
|
31 |
+
.dark .footer>p {
|
32 |
+
background: #0b0f19;
|
33 |
+
}
|
ldm/base_utils.py
ADDED
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import pickle
|
2 |
+
import numpy as np
|
3 |
+
import cv2
|
4 |
+
from skimage.io import imread
|
5 |
+
|
6 |
+
|
7 |
+
def save_pickle(data, pkl_path):
|
8 |
+
# os.system('mkdir -p {}'.format(os.path.dirname(pkl_path)))
|
9 |
+
with open(pkl_path, 'wb') as f:
|
10 |
+
pickle.dump(data, f)
|
11 |
+
|
12 |
+
def read_pickle(pkl_path):
|
13 |
+
with open(pkl_path, 'rb') as f:
|
14 |
+
return pickle.load(f)
|
15 |
+
|
16 |
+
def draw_epipolar_line(F, img0, img1, pt0, color):
|
17 |
+
h1,w1=img1.shape[:2]
|
18 |
+
hpt = np.asarray([pt0[0], pt0[1], 1], dtype=np.float32)[:, None]
|
19 |
+
l = F @ hpt
|
20 |
+
l = l[:, 0]
|
21 |
+
a, b, c = l[0], l[1], l[2]
|
22 |
+
pt1 = np.asarray([0, -c / b]).astype(np.int32)
|
23 |
+
pt2 = np.asarray([w1, (-a * w1 - c) / b]).astype(np.int32)
|
24 |
+
|
25 |
+
img0 = cv2.circle(img0, tuple(pt0.astype(np.int32)), 5, color, 2)
|
26 |
+
img1 = cv2.line(img1, tuple(pt1), tuple(pt2), color, 2)
|
27 |
+
return img0, img1
|
28 |
+
|
29 |
+
def draw_epipolar_lines(F, img0, img1,num=20):
|
30 |
+
img0,img1=img0.copy(),img1.copy()
|
31 |
+
h0, w0, _ = img0.shape
|
32 |
+
h1, w1, _ = img1.shape
|
33 |
+
|
34 |
+
for k in range(num):
|
35 |
+
color = np.random.randint(0, 255, [3], dtype=np.int32)
|
36 |
+
color = [int(c) for c in color]
|
37 |
+
pt = np.random.uniform(0, 1, 2)
|
38 |
+
pt[0] *= w0
|
39 |
+
pt[1] *= h0
|
40 |
+
pt = pt.astype(np.int32)
|
41 |
+
img0, img1 = draw_epipolar_line(F, img0, img1, pt, color)
|
42 |
+
|
43 |
+
return img0, img1
|
44 |
+
|
45 |
+
def compute_F(K1, K2, Rt0, Rt1=None):
|
46 |
+
if Rt1 is None:
|
47 |
+
R, t = Rt0[:,:3], Rt0[:,3:]
|
48 |
+
else:
|
49 |
+
Rt = compute_dR_dt(Rt0,Rt1)
|
50 |
+
R, t = Rt[:,:3], Rt[:,3:]
|
51 |
+
A = K1 @ R.T @ t # [3,1]
|
52 |
+
C = np.asarray([[0,-A[2,0],A[1,0]],
|
53 |
+
[A[2,0],0,-A[0,0]],
|
54 |
+
[-A[1,0],A[0,0],0]])
|
55 |
+
F = (np.linalg.inv(K2)).T @ R @ K1.T @ C
|
56 |
+
return F
|
57 |
+
|
58 |
+
def compute_dR_dt(Rt0, Rt1):
|
59 |
+
R0, t0 = Rt0[:,:3], Rt0[:,3:]
|
60 |
+
R1, t1 = Rt1[:,:3], Rt1[:,3:]
|
61 |
+
dR = np.dot(R1, R0.T)
|
62 |
+
dt = t1 - np.dot(dR, t0)
|
63 |
+
return np.concatenate([dR, dt], -1)
|
64 |
+
|
65 |
+
def concat_images(img0,img1,vert=False):
|
66 |
+
if not vert:
|
67 |
+
h0,h1=img0.shape[0],img1.shape[0],
|
68 |
+
if h0<h1: img0=cv2.copyMakeBorder(img0,0,h1-h0,0,0,borderType=cv2.BORDER_CONSTANT,value=0)
|
69 |
+
if h1<h0: img1=cv2.copyMakeBorder(img1,0,h0-h1,0,0,borderType=cv2.BORDER_CONSTANT,value=0)
|
70 |
+
img = np.concatenate([img0, img1], axis=1)
|
71 |
+
else:
|
72 |
+
w0,w1=img0.shape[1],img1.shape[1]
|
73 |
+
if w0<w1: img0=cv2.copyMakeBorder(img0,0,0,0,w1-w0,borderType=cv2.BORDER_CONSTANT,value=0)
|
74 |
+
if w1<w0: img1=cv2.copyMakeBorder(img1,0,0,0,w0-w1,borderType=cv2.BORDER_CONSTANT,value=0)
|
75 |
+
img = np.concatenate([img0, img1], axis=0)
|
76 |
+
|
77 |
+
return img
|
78 |
+
|
79 |
+
def concat_images_list(*args,vert=False):
|
80 |
+
if len(args)==1: return args[0]
|
81 |
+
img_out=args[0]
|
82 |
+
for img in args[1:]:
|
83 |
+
img_out=concat_images(img_out,img,vert)
|
84 |
+
return img_out
|
85 |
+
|
86 |
+
|
87 |
+
def pose_inverse(pose):
|
88 |
+
R = pose[:,:3].T
|
89 |
+
t = - R @ pose[:,3:]
|
90 |
+
return np.concatenate([R,t],-1)
|
91 |
+
|
92 |
+
def project_points(pts,RT,K):
|
93 |
+
pts = np.matmul(pts,RT[:,:3].transpose())+RT[:,3:].transpose()
|
94 |
+
pts = np.matmul(pts,K.transpose())
|
95 |
+
dpt = pts[:,2]
|
96 |
+
mask0 = (np.abs(dpt)<1e-4) & (np.abs(dpt)>0)
|
97 |
+
if np.sum(mask0)>0: dpt[mask0]=1e-4
|
98 |
+
mask1=(np.abs(dpt) > -1e-4) & (np.abs(dpt) < 0)
|
99 |
+
if np.sum(mask1)>0: dpt[mask1]=-1e-4
|
100 |
+
pts2d = pts[:,:2]/dpt[:,None]
|
101 |
+
return pts2d, dpt
|
102 |
+
|
103 |
+
|
104 |
+
def draw_keypoints(img, kps, colors=None, radius=2):
|
105 |
+
out_img=img.copy()
|
106 |
+
for pi, pt in enumerate(kps):
|
107 |
+
pt = np.round(pt).astype(np.int32)
|
108 |
+
if colors is not None:
|
109 |
+
color=[int(c) for c in colors[pi]]
|
110 |
+
cv2.circle(out_img, tuple(pt), radius, color, -1)
|
111 |
+
else:
|
112 |
+
cv2.circle(out_img, tuple(pt), radius, (0,255,0), -1)
|
113 |
+
return out_img
|
114 |
+
|
115 |
+
|
116 |
+
def output_points(fn,pts,colors=None):
|
117 |
+
with open(fn, 'w') as f:
|
118 |
+
for pi, pt in enumerate(pts):
|
119 |
+
f.write(f'{pt[0]:.6f} {pt[1]:.6f} {pt[2]:.6f} ')
|
120 |
+
if colors is not None:
|
121 |
+
f.write(f'{int(colors[pi,0])} {int(colors[pi,1])} {int(colors[pi,2])}')
|
122 |
+
f.write('\n')
|
123 |
+
|
124 |
+
DEPTH_MAX, DEPTH_MIN = 2.4, 0.6
|
125 |
+
DEPTH_VALID_MAX, DEPTH_VALID_MIN = 2.37, 0.63
|
126 |
+
def read_depth_objaverse(depth_fn):
|
127 |
+
depth = imread(depth_fn)
|
128 |
+
depth = depth.astype(np.float32) / 65535 * (DEPTH_MAX-DEPTH_MIN) + DEPTH_MIN
|
129 |
+
mask = (depth > DEPTH_VALID_MIN) & (depth < DEPTH_VALID_MAX)
|
130 |
+
return depth, mask
|
131 |
+
|
132 |
+
|
133 |
+
def mask_depth_to_pts(mask,depth,K,rgb=None):
|
134 |
+
hs,ws=np.nonzero(mask)
|
135 |
+
depth=depth[hs,ws]
|
136 |
+
pts=np.asarray([ws,hs,depth],np.float32).transpose()
|
137 |
+
pts[:,:2]*=pts[:,2:]
|
138 |
+
if rgb is not None:
|
139 |
+
return np.dot(pts, np.linalg.inv(K).transpose()), rgb[hs,ws]
|
140 |
+
else:
|
141 |
+
return np.dot(pts, np.linalg.inv(K).transpose())
|
142 |
+
|
143 |
+
def transform_points_pose(pts, pose):
|
144 |
+
R, t = pose[:, :3], pose[:, 3]
|
145 |
+
if len(pts.shape)==1:
|
146 |
+
return (R @ pts[:,None] + t[:,None])[:,0]
|
147 |
+
return pts @ R.T + t[None,:]
|
148 |
+
|
149 |
+
def pose_apply(pose,pts):
|
150 |
+
return transform_points_pose(pts, pose)
|
151 |
+
|
152 |
+
def downsample_gaussian_blur(img, ratio):
|
153 |
+
sigma = (1 / ratio) / 3
|
154 |
+
# ksize=np.ceil(2*sigma)
|
155 |
+
ksize = int(np.ceil(((sigma - 0.8) / 0.3 + 1) * 2 + 1))
|
156 |
+
ksize = ksize + 1 if ksize % 2 == 0 else ksize
|
157 |
+
img = cv2.GaussianBlur(img, (ksize, ksize), sigma, borderType=cv2.BORDER_REFLECT101)
|
158 |
+
return img
|
ldm/data/__init__.py
ADDED
File without changes
|
ldm/data/base.py
ADDED
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import numpy as np
|
3 |
+
from abc import abstractmethod
|
4 |
+
from torch.utils.data import Dataset, ConcatDataset, ChainDataset, IterableDataset
|
5 |
+
|
6 |
+
|
7 |
+
class Txt2ImgIterableBaseDataset(IterableDataset):
|
8 |
+
'''
|
9 |
+
Define an interface to make the IterableDatasets for text2img data chainable
|
10 |
+
'''
|
11 |
+
def __init__(self, num_records=0, valid_ids=None, size=256):
|
12 |
+
super().__init__()
|
13 |
+
self.num_records = num_records
|
14 |
+
self.valid_ids = valid_ids
|
15 |
+
self.sample_ids = valid_ids
|
16 |
+
self.size = size
|
17 |
+
|
18 |
+
print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')
|
19 |
+
|
20 |
+
def __len__(self):
|
21 |
+
return self.num_records
|
22 |
+
|
23 |
+
@abstractmethod
|
24 |
+
def __iter__(self):
|
25 |
+
pass
|
26 |
+
|
27 |
+
|
28 |
+
class PRNGMixin(object):
|
29 |
+
"""
|
30 |
+
Adds a prng property which is a numpy RandomState which gets
|
31 |
+
reinitialized whenever the pid changes to avoid synchronized sampling
|
32 |
+
behavior when used in conjunction with multiprocessing.
|
33 |
+
"""
|
34 |
+
@property
|
35 |
+
def prng(self):
|
36 |
+
currentpid = os.getpid()
|
37 |
+
if getattr(self, "_initpid", None) != currentpid:
|
38 |
+
self._initpid = currentpid
|
39 |
+
self._prng = np.random.RandomState()
|
40 |
+
return self._prng
|
ldm/data/coco.py
ADDED
@@ -0,0 +1,253 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import json
|
3 |
+
import albumentations
|
4 |
+
import numpy as np
|
5 |
+
from PIL import Image
|
6 |
+
from tqdm import tqdm
|
7 |
+
from torch.utils.data import Dataset
|
8 |
+
from abc import abstractmethod
|
9 |
+
|
10 |
+
|
11 |
+
class CocoBase(Dataset):
|
12 |
+
"""needed for (image, caption, segmentation) pairs"""
|
13 |
+
def __init__(self, size=None, dataroot="", datajson="", onehot_segmentation=False, use_stuffthing=False,
|
14 |
+
crop_size=None, force_no_crop=False, given_files=None, use_segmentation=True,crop_type=None):
|
15 |
+
self.split = self.get_split()
|
16 |
+
self.size = size
|
17 |
+
if crop_size is None:
|
18 |
+
self.crop_size = size
|
19 |
+
else:
|
20 |
+
self.crop_size = crop_size
|
21 |
+
|
22 |
+
assert crop_type in [None, 'random', 'center']
|
23 |
+
self.crop_type = crop_type
|
24 |
+
self.use_segmenation = use_segmentation
|
25 |
+
self.onehot = onehot_segmentation # return segmentation as rgb or one hot
|
26 |
+
self.stuffthing = use_stuffthing # include thing in segmentation
|
27 |
+
if self.onehot and not self.stuffthing:
|
28 |
+
raise NotImplemented("One hot mode is only supported for the "
|
29 |
+
"stuffthings version because labels are stored "
|
30 |
+
"a bit different.")
|
31 |
+
|
32 |
+
data_json = datajson
|
33 |
+
with open(data_json) as json_file:
|
34 |
+
self.json_data = json.load(json_file)
|
35 |
+
self.img_id_to_captions = dict()
|
36 |
+
self.img_id_to_filepath = dict()
|
37 |
+
self.img_id_to_segmentation_filepath = dict()
|
38 |
+
|
39 |
+
assert data_json.split("/")[-1] in [f"captions_train{self.year()}.json",
|
40 |
+
f"captions_val{self.year()}.json"]
|
41 |
+
# TODO currently hardcoded paths, would be better to follow logic in
|
42 |
+
# cocstuff pixelmaps
|
43 |
+
if self.use_segmenation:
|
44 |
+
if self.stuffthing:
|
45 |
+
self.segmentation_prefix = (
|
46 |
+
f"data/cocostuffthings/val{self.year()}" if
|
47 |
+
data_json.endswith(f"captions_val{self.year()}.json") else
|
48 |
+
f"data/cocostuffthings/train{self.year()}")
|
49 |
+
else:
|
50 |
+
self.segmentation_prefix = (
|
51 |
+
f"data/coco/annotations/stuff_val{self.year()}_pixelmaps" if
|
52 |
+
data_json.endswith(f"captions_val{self.year()}.json") else
|
53 |
+
f"data/coco/annotations/stuff_train{self.year()}_pixelmaps")
|
54 |
+
|
55 |
+
imagedirs = self.json_data["images"]
|
56 |
+
self.labels = {"image_ids": list()}
|
57 |
+
for imgdir in tqdm(imagedirs, desc="ImgToPath"):
|
58 |
+
self.img_id_to_filepath[imgdir["id"]] = os.path.join(dataroot, imgdir["file_name"])
|
59 |
+
self.img_id_to_captions[imgdir["id"]] = list()
|
60 |
+
pngfilename = imgdir["file_name"].replace("jpg", "png")
|
61 |
+
if self.use_segmenation:
|
62 |
+
self.img_id_to_segmentation_filepath[imgdir["id"]] = os.path.join(
|
63 |
+
self.segmentation_prefix, pngfilename)
|
64 |
+
if given_files is not None:
|
65 |
+
if pngfilename in given_files:
|
66 |
+
self.labels["image_ids"].append(imgdir["id"])
|
67 |
+
else:
|
68 |
+
self.labels["image_ids"].append(imgdir["id"])
|
69 |
+
|
70 |
+
capdirs = self.json_data["annotations"]
|
71 |
+
for capdir in tqdm(capdirs, desc="ImgToCaptions"):
|
72 |
+
# there are in average 5 captions per image
|
73 |
+
#self.img_id_to_captions[capdir["image_id"]].append(np.array([capdir["caption"]]))
|
74 |
+
self.img_id_to_captions[capdir["image_id"]].append(capdir["caption"])
|
75 |
+
|
76 |
+
self.rescaler = albumentations.SmallestMaxSize(max_size=self.size)
|
77 |
+
if self.split=="validation":
|
78 |
+
self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
|
79 |
+
else:
|
80 |
+
# default option for train is random crop
|
81 |
+
if self.crop_type in [None, 'random']:
|
82 |
+
self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size)
|
83 |
+
else:
|
84 |
+
self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
|
85 |
+
self.preprocessor = albumentations.Compose(
|
86 |
+
[self.rescaler, self.cropper],
|
87 |
+
additional_targets={"segmentation": "image"})
|
88 |
+
if force_no_crop:
|
89 |
+
self.rescaler = albumentations.Resize(height=self.size, width=self.size)
|
90 |
+
self.preprocessor = albumentations.Compose(
|
91 |
+
[self.rescaler],
|
92 |
+
additional_targets={"segmentation": "image"})
|
93 |
+
|
94 |
+
@abstractmethod
|
95 |
+
def year(self):
|
96 |
+
raise NotImplementedError()
|
97 |
+
|
98 |
+
def __len__(self):
|
99 |
+
return len(self.labels["image_ids"])
|
100 |
+
|
101 |
+
def preprocess_image(self, image_path, segmentation_path=None):
|
102 |
+
image = Image.open(image_path)
|
103 |
+
if not image.mode == "RGB":
|
104 |
+
image = image.convert("RGB")
|
105 |
+
image = np.array(image).astype(np.uint8)
|
106 |
+
if segmentation_path:
|
107 |
+
segmentation = Image.open(segmentation_path)
|
108 |
+
if not self.onehot and not segmentation.mode == "RGB":
|
109 |
+
segmentation = segmentation.convert("RGB")
|
110 |
+
segmentation = np.array(segmentation).astype(np.uint8)
|
111 |
+
if self.onehot:
|
112 |
+
assert self.stuffthing
|
113 |
+
# stored in caffe format: unlabeled==255. stuff and thing from
|
114 |
+
# 0-181. to be compatible with the labels in
|
115 |
+
# https://github.com/nightrome/cocostuff/blob/master/labels.txt
|
116 |
+
# we shift stuffthing one to the right and put unlabeled in zero
|
117 |
+
# as long as segmentation is uint8 shifting to right handles the
|
118 |
+
# latter too
|
119 |
+
assert segmentation.dtype == np.uint8
|
120 |
+
segmentation = segmentation + 1
|
121 |
+
|
122 |
+
processed = self.preprocessor(image=image, segmentation=segmentation)
|
123 |
+
|
124 |
+
image, segmentation = processed["image"], processed["segmentation"]
|
125 |
+
else:
|
126 |
+
image = self.preprocessor(image=image,)['image']
|
127 |
+
|
128 |
+
image = (image / 127.5 - 1.0).astype(np.float32)
|
129 |
+
if segmentation_path:
|
130 |
+
if self.onehot:
|
131 |
+
assert segmentation.dtype == np.uint8
|
132 |
+
# make it one hot
|
133 |
+
n_labels = 183
|
134 |
+
flatseg = np.ravel(segmentation)
|
135 |
+
onehot = np.zeros((flatseg.size, n_labels), dtype=np.bool)
|
136 |
+
onehot[np.arange(flatseg.size), flatseg] = True
|
137 |
+
onehot = onehot.reshape(segmentation.shape + (n_labels,)).astype(int)
|
138 |
+
segmentation = onehot
|
139 |
+
else:
|
140 |
+
segmentation = (segmentation / 127.5 - 1.0).astype(np.float32)
|
141 |
+
return image, segmentation
|
142 |
+
else:
|
143 |
+
return image
|
144 |
+
|
145 |
+
def __getitem__(self, i):
|
146 |
+
img_path = self.img_id_to_filepath[self.labels["image_ids"][i]]
|
147 |
+
if self.use_segmenation:
|
148 |
+
seg_path = self.img_id_to_segmentation_filepath[self.labels["image_ids"][i]]
|
149 |
+
image, segmentation = self.preprocess_image(img_path, seg_path)
|
150 |
+
else:
|
151 |
+
image = self.preprocess_image(img_path)
|
152 |
+
captions = self.img_id_to_captions[self.labels["image_ids"][i]]
|
153 |
+
# randomly draw one of all available captions per image
|
154 |
+
caption = captions[np.random.randint(0, len(captions))]
|
155 |
+
example = {"image": image,
|
156 |
+
#"caption": [str(caption[0])],
|
157 |
+
"caption": caption,
|
158 |
+
"img_path": img_path,
|
159 |
+
"filename_": img_path.split(os.sep)[-1]
|
160 |
+
}
|
161 |
+
if self.use_segmenation:
|
162 |
+
example.update({"seg_path": seg_path, 'segmentation': segmentation})
|
163 |
+
return example
|
164 |
+
|
165 |
+
|
166 |
+
class CocoImagesAndCaptionsTrain2017(CocoBase):
|
167 |
+
"""returns a pair of (image, caption)"""
|
168 |
+
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,):
|
169 |
+
super().__init__(size=size,
|
170 |
+
dataroot="data/coco/train2017",
|
171 |
+
datajson="data/coco/annotations/captions_train2017.json",
|
172 |
+
onehot_segmentation=onehot_segmentation,
|
173 |
+
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop)
|
174 |
+
|
175 |
+
def get_split(self):
|
176 |
+
return "train"
|
177 |
+
|
178 |
+
def year(self):
|
179 |
+
return '2017'
|
180 |
+
|
181 |
+
|
182 |
+
class CocoImagesAndCaptionsValidation2017(CocoBase):
|
183 |
+
"""returns a pair of (image, caption)"""
|
184 |
+
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,
|
185 |
+
given_files=None):
|
186 |
+
super().__init__(size=size,
|
187 |
+
dataroot="data/coco/val2017",
|
188 |
+
datajson="data/coco/annotations/captions_val2017.json",
|
189 |
+
onehot_segmentation=onehot_segmentation,
|
190 |
+
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
|
191 |
+
given_files=given_files)
|
192 |
+
|
193 |
+
def get_split(self):
|
194 |
+
return "validation"
|
195 |
+
|
196 |
+
def year(self):
|
197 |
+
return '2017'
|
198 |
+
|
199 |
+
|
200 |
+
|
201 |
+
class CocoImagesAndCaptionsTrain2014(CocoBase):
|
202 |
+
"""returns a pair of (image, caption)"""
|
203 |
+
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,crop_type='random'):
|
204 |
+
super().__init__(size=size,
|
205 |
+
dataroot="data/coco/train2014",
|
206 |
+
datajson="data/coco/annotations2014/annotations/captions_train2014.json",
|
207 |
+
onehot_segmentation=onehot_segmentation,
|
208 |
+
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
|
209 |
+
use_segmentation=False,
|
210 |
+
crop_type=crop_type)
|
211 |
+
|
212 |
+
def get_split(self):
|
213 |
+
return "train"
|
214 |
+
|
215 |
+
def year(self):
|
216 |
+
return '2014'
|
217 |
+
|
218 |
+
class CocoImagesAndCaptionsValidation2014(CocoBase):
|
219 |
+
"""returns a pair of (image, caption)"""
|
220 |
+
def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,
|
221 |
+
given_files=None,crop_type='center',**kwargs):
|
222 |
+
super().__init__(size=size,
|
223 |
+
dataroot="data/coco/val2014",
|
224 |
+
datajson="data/coco/annotations2014/annotations/captions_val2014.json",
|
225 |
+
onehot_segmentation=onehot_segmentation,
|
226 |
+
use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
|
227 |
+
given_files=given_files,
|
228 |
+
use_segmentation=False,
|
229 |
+
crop_type=crop_type)
|
230 |
+
|
231 |
+
def get_split(self):
|
232 |
+
return "validation"
|
233 |
+
|
234 |
+
def year(self):
|
235 |
+
return '2014'
|
236 |
+
|
237 |
+
if __name__ == '__main__':
|
238 |
+
with open("data/coco/annotations2014/annotations/captions_val2014.json", "r") as json_file:
|
239 |
+
json_data = json.load(json_file)
|
240 |
+
capdirs = json_data["annotations"]
|
241 |
+
import pudb; pudb.set_trace()
|
242 |
+
#d2 = CocoImagesAndCaptionsTrain2014(size=256)
|
243 |
+
d2 = CocoImagesAndCaptionsValidation2014(size=256)
|
244 |
+
print("constructed dataset.")
|
245 |
+
print(f"length of {d2.__class__.__name__}: {len(d2)}")
|
246 |
+
|
247 |
+
ex2 = d2[0]
|
248 |
+
# ex3 = d3[0]
|
249 |
+
# print(ex1["image"].shape)
|
250 |
+
print(ex2["image"].shape)
|
251 |
+
# print(ex3["image"].shape)
|
252 |
+
# print(ex1["segmentation"].shape)
|
253 |
+
print(ex2["caption"].__class__.__name__)
|
ldm/data/dummy.py
ADDED
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import random
|
3 |
+
import string
|
4 |
+
from torch.utils.data import Dataset, Subset
|
5 |
+
|
6 |
+
class DummyData(Dataset):
|
7 |
+
def __init__(self, length, size):
|
8 |
+
self.length = length
|
9 |
+
self.size = size
|
10 |
+
|
11 |
+
def __len__(self):
|
12 |
+
return self.length
|
13 |
+
|
14 |
+
def __getitem__(self, i):
|
15 |
+
x = np.random.randn(*self.size)
|
16 |
+
letters = string.ascii_lowercase
|
17 |
+
y = ''.join(random.choice(string.ascii_lowercase) for i in range(10))
|
18 |
+
return {"jpg": x, "txt": y}
|
19 |
+
|
20 |
+
|
21 |
+
class DummyDataWithEmbeddings(Dataset):
|
22 |
+
def __init__(self, length, size, emb_size):
|
23 |
+
self.length = length
|
24 |
+
self.size = size
|
25 |
+
self.emb_size = emb_size
|
26 |
+
|
27 |
+
def __len__(self):
|
28 |
+
return self.length
|
29 |
+
|
30 |
+
def __getitem__(self, i):
|
31 |
+
x = np.random.randn(*self.size)
|
32 |
+
y = np.random.randn(*self.emb_size).astype(np.float32)
|
33 |
+
return {"jpg": x, "txt": y}
|
34 |
+
|
ldm/data/imagenet.py
ADDED
@@ -0,0 +1,394 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os, yaml, pickle, shutil, tarfile, glob
|
2 |
+
import cv2
|
3 |
+
import albumentations
|
4 |
+
import PIL
|
5 |
+
import numpy as np
|
6 |
+
import torchvision.transforms.functional as TF
|
7 |
+
from omegaconf import OmegaConf
|
8 |
+
from functools import partial
|
9 |
+
from PIL import Image
|
10 |
+
from tqdm import tqdm
|
11 |
+
from torch.utils.data import Dataset, Subset
|
12 |
+
|
13 |
+
import taming.data.utils as tdu
|
14 |
+
from taming.data.imagenet import str_to_indices, give_synsets_from_indices, download, retrieve
|
15 |
+
from taming.data.imagenet import ImagePaths
|
16 |
+
|
17 |
+
from ldm.modules.image_degradation import degradation_fn_bsr, degradation_fn_bsr_light
|
18 |
+
|
19 |
+
|
20 |
+
def synset2idx(path_to_yaml="data/index_synset.yaml"):
|
21 |
+
with open(path_to_yaml) as f:
|
22 |
+
di2s = yaml.load(f)
|
23 |
+
return dict((v,k) for k,v in di2s.items())
|
24 |
+
|
25 |
+
|
26 |
+
class ImageNetBase(Dataset):
|
27 |
+
def __init__(self, config=None):
|
28 |
+
self.config = config or OmegaConf.create()
|
29 |
+
if not type(self.config)==dict:
|
30 |
+
self.config = OmegaConf.to_container(self.config)
|
31 |
+
self.keep_orig_class_label = self.config.get("keep_orig_class_label", False)
|
32 |
+
self.process_images = True # if False we skip loading & processing images and self.data contains filepaths
|
33 |
+
self._prepare()
|
34 |
+
self._prepare_synset_to_human()
|
35 |
+
self._prepare_idx_to_synset()
|
36 |
+
self._prepare_human_to_integer_label()
|
37 |
+
self._load()
|
38 |
+
|
39 |
+
def __len__(self):
|
40 |
+
return len(self.data)
|
41 |
+
|
42 |
+
def __getitem__(self, i):
|
43 |
+
return self.data[i]
|
44 |
+
|
45 |
+
def _prepare(self):
|
46 |
+
raise NotImplementedError()
|
47 |
+
|
48 |
+
def _filter_relpaths(self, relpaths):
|
49 |
+
ignore = set([
|
50 |
+
"n06596364_9591.JPEG",
|
51 |
+
])
|
52 |
+
relpaths = [rpath for rpath in relpaths if not rpath.split("/")[-1] in ignore]
|
53 |
+
if "sub_indices" in self.config:
|
54 |
+
indices = str_to_indices(self.config["sub_indices"])
|
55 |
+
synsets = give_synsets_from_indices(indices, path_to_yaml=self.idx2syn) # returns a list of strings
|
56 |
+
self.synset2idx = synset2idx(path_to_yaml=self.idx2syn)
|
57 |
+
files = []
|
58 |
+
for rpath in relpaths:
|
59 |
+
syn = rpath.split("/")[0]
|
60 |
+
if syn in synsets:
|
61 |
+
files.append(rpath)
|
62 |
+
return files
|
63 |
+
else:
|
64 |
+
return relpaths
|
65 |
+
|
66 |
+
def _prepare_synset_to_human(self):
|
67 |
+
SIZE = 2655750
|
68 |
+
URL = "https://heibox.uni-heidelberg.de/f/9f28e956cd304264bb82/?dl=1"
|
69 |
+
self.human_dict = os.path.join(self.root, "synset_human.txt")
|
70 |
+
if (not os.path.exists(self.human_dict) or
|
71 |
+
not os.path.getsize(self.human_dict)==SIZE):
|
72 |
+
download(URL, self.human_dict)
|
73 |
+
|
74 |
+
def _prepare_idx_to_synset(self):
|
75 |
+
URL = "https://heibox.uni-heidelberg.de/f/d835d5b6ceda4d3aa910/?dl=1"
|
76 |
+
self.idx2syn = os.path.join(self.root, "index_synset.yaml")
|
77 |
+
if (not os.path.exists(self.idx2syn)):
|
78 |
+
download(URL, self.idx2syn)
|
79 |
+
|
80 |
+
def _prepare_human_to_integer_label(self):
|
81 |
+
URL = "https://heibox.uni-heidelberg.de/f/2362b797d5be43b883f6/?dl=1"
|
82 |
+
self.human2integer = os.path.join(self.root, "imagenet1000_clsidx_to_labels.txt")
|
83 |
+
if (not os.path.exists(self.human2integer)):
|
84 |
+
download(URL, self.human2integer)
|
85 |
+
with open(self.human2integer, "r") as f:
|
86 |
+
lines = f.read().splitlines()
|
87 |
+
assert len(lines) == 1000
|
88 |
+
self.human2integer_dict = dict()
|
89 |
+
for line in lines:
|
90 |
+
value, key = line.split(":")
|
91 |
+
self.human2integer_dict[key] = int(value)
|
92 |
+
|
93 |
+
def _load(self):
|
94 |
+
with open(self.txt_filelist, "r") as f:
|
95 |
+
self.relpaths = f.read().splitlines()
|
96 |
+
l1 = len(self.relpaths)
|
97 |
+
self.relpaths = self._filter_relpaths(self.relpaths)
|
98 |
+
print("Removed {} files from filelist during filtering.".format(l1 - len(self.relpaths)))
|
99 |
+
|
100 |
+
self.synsets = [p.split("/")[0] for p in self.relpaths]
|
101 |
+
self.abspaths = [os.path.join(self.datadir, p) for p in self.relpaths]
|
102 |
+
|
103 |
+
unique_synsets = np.unique(self.synsets)
|
104 |
+
class_dict = dict((synset, i) for i, synset in enumerate(unique_synsets))
|
105 |
+
if not self.keep_orig_class_label:
|
106 |
+
self.class_labels = [class_dict[s] for s in self.synsets]
|
107 |
+
else:
|
108 |
+
self.class_labels = [self.synset2idx[s] for s in self.synsets]
|
109 |
+
|
110 |
+
with open(self.human_dict, "r") as f:
|
111 |
+
human_dict = f.read().splitlines()
|
112 |
+
human_dict = dict(line.split(maxsplit=1) for line in human_dict)
|
113 |
+
|
114 |
+
self.human_labels = [human_dict[s] for s in self.synsets]
|
115 |
+
|
116 |
+
labels = {
|
117 |
+
"relpath": np.array(self.relpaths),
|
118 |
+
"synsets": np.array(self.synsets),
|
119 |
+
"class_label": np.array(self.class_labels),
|
120 |
+
"human_label": np.array(self.human_labels),
|
121 |
+
}
|
122 |
+
|
123 |
+
if self.process_images:
|
124 |
+
self.size = retrieve(self.config, "size", default=256)
|
125 |
+
self.data = ImagePaths(self.abspaths,
|
126 |
+
labels=labels,
|
127 |
+
size=self.size,
|
128 |
+
random_crop=self.random_crop,
|
129 |
+
)
|
130 |
+
else:
|
131 |
+
self.data = self.abspaths
|
132 |
+
|
133 |
+
|
134 |
+
class ImageNetTrain(ImageNetBase):
|
135 |
+
NAME = "ILSVRC2012_train"
|
136 |
+
URL = "http://www.image-net.org/challenges/LSVRC/2012/"
|
137 |
+
AT_HASH = "a306397ccf9c2ead27155983c254227c0fd938e2"
|
138 |
+
FILES = [
|
139 |
+
"ILSVRC2012_img_train.tar",
|
140 |
+
]
|
141 |
+
SIZES = [
|
142 |
+
147897477120,
|
143 |
+
]
|
144 |
+
|
145 |
+
def __init__(self, process_images=True, data_root=None, **kwargs):
|
146 |
+
self.process_images = process_images
|
147 |
+
self.data_root = data_root
|
148 |
+
super().__init__(**kwargs)
|
149 |
+
|
150 |
+
def _prepare(self):
|
151 |
+
if self.data_root:
|
152 |
+
self.root = os.path.join(self.data_root, self.NAME)
|
153 |
+
else:
|
154 |
+
cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
155 |
+
self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
|
156 |
+
|
157 |
+
self.datadir = os.path.join(self.root, "data")
|
158 |
+
self.txt_filelist = os.path.join(self.root, "filelist.txt")
|
159 |
+
self.expected_length = 1281167
|
160 |
+
self.random_crop = retrieve(self.config, "ImageNetTrain/random_crop",
|
161 |
+
default=True)
|
162 |
+
if not tdu.is_prepared(self.root):
|
163 |
+
# prep
|
164 |
+
print("Preparing dataset {} in {}".format(self.NAME, self.root))
|
165 |
+
|
166 |
+
datadir = self.datadir
|
167 |
+
if not os.path.exists(datadir):
|
168 |
+
path = os.path.join(self.root, self.FILES[0])
|
169 |
+
if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
|
170 |
+
import academictorrents as at
|
171 |
+
atpath = at.get(self.AT_HASH, datastore=self.root)
|
172 |
+
assert atpath == path
|
173 |
+
|
174 |
+
print("Extracting {} to {}".format(path, datadir))
|
175 |
+
os.makedirs(datadir, exist_ok=True)
|
176 |
+
with tarfile.open(path, "r:") as tar:
|
177 |
+
tar.extractall(path=datadir)
|
178 |
+
|
179 |
+
print("Extracting sub-tars.")
|
180 |
+
subpaths = sorted(glob.glob(os.path.join(datadir, "*.tar")))
|
181 |
+
for subpath in tqdm(subpaths):
|
182 |
+
subdir = subpath[:-len(".tar")]
|
183 |
+
os.makedirs(subdir, exist_ok=True)
|
184 |
+
with tarfile.open(subpath, "r:") as tar:
|
185 |
+
tar.extractall(path=subdir)
|
186 |
+
|
187 |
+
filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
|
188 |
+
filelist = [os.path.relpath(p, start=datadir) for p in filelist]
|
189 |
+
filelist = sorted(filelist)
|
190 |
+
filelist = "\n".join(filelist)+"\n"
|
191 |
+
with open(self.txt_filelist, "w") as f:
|
192 |
+
f.write(filelist)
|
193 |
+
|
194 |
+
tdu.mark_prepared(self.root)
|
195 |
+
|
196 |
+
|
197 |
+
class ImageNetValidation(ImageNetBase):
|
198 |
+
NAME = "ILSVRC2012_validation"
|
199 |
+
URL = "http://www.image-net.org/challenges/LSVRC/2012/"
|
200 |
+
AT_HASH = "5d6d0df7ed81efd49ca99ea4737e0ae5e3a5f2e5"
|
201 |
+
VS_URL = "https://heibox.uni-heidelberg.de/f/3e0f6e9c624e45f2bd73/?dl=1"
|
202 |
+
FILES = [
|
203 |
+
"ILSVRC2012_img_val.tar",
|
204 |
+
"validation_synset.txt",
|
205 |
+
]
|
206 |
+
SIZES = [
|
207 |
+
6744924160,
|
208 |
+
1950000,
|
209 |
+
]
|
210 |
+
|
211 |
+
def __init__(self, process_images=True, data_root=None, **kwargs):
|
212 |
+
self.data_root = data_root
|
213 |
+
self.process_images = process_images
|
214 |
+
super().__init__(**kwargs)
|
215 |
+
|
216 |
+
def _prepare(self):
|
217 |
+
if self.data_root:
|
218 |
+
self.root = os.path.join(self.data_root, self.NAME)
|
219 |
+
else:
|
220 |
+
cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
221 |
+
self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
|
222 |
+
self.datadir = os.path.join(self.root, "data")
|
223 |
+
self.txt_filelist = os.path.join(self.root, "filelist.txt")
|
224 |
+
self.expected_length = 50000
|
225 |
+
self.random_crop = retrieve(self.config, "ImageNetValidation/random_crop",
|
226 |
+
default=False)
|
227 |
+
if not tdu.is_prepared(self.root):
|
228 |
+
# prep
|
229 |
+
print("Preparing dataset {} in {}".format(self.NAME, self.root))
|
230 |
+
|
231 |
+
datadir = self.datadir
|
232 |
+
if not os.path.exists(datadir):
|
233 |
+
path = os.path.join(self.root, self.FILES[0])
|
234 |
+
if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
|
235 |
+
import academictorrents as at
|
236 |
+
atpath = at.get(self.AT_HASH, datastore=self.root)
|
237 |
+
assert atpath == path
|
238 |
+
|
239 |
+
print("Extracting {} to {}".format(path, datadir))
|
240 |
+
os.makedirs(datadir, exist_ok=True)
|
241 |
+
with tarfile.open(path, "r:") as tar:
|
242 |
+
tar.extractall(path=datadir)
|
243 |
+
|
244 |
+
vspath = os.path.join(self.root, self.FILES[1])
|
245 |
+
if not os.path.exists(vspath) or not os.path.getsize(vspath)==self.SIZES[1]:
|
246 |
+
download(self.VS_URL, vspath)
|
247 |
+
|
248 |
+
with open(vspath, "r") as f:
|
249 |
+
synset_dict = f.read().splitlines()
|
250 |
+
synset_dict = dict(line.split() for line in synset_dict)
|
251 |
+
|
252 |
+
print("Reorganizing into synset folders")
|
253 |
+
synsets = np.unique(list(synset_dict.values()))
|
254 |
+
for s in synsets:
|
255 |
+
os.makedirs(os.path.join(datadir, s), exist_ok=True)
|
256 |
+
for k, v in synset_dict.items():
|
257 |
+
src = os.path.join(datadir, k)
|
258 |
+
dst = os.path.join(datadir, v)
|
259 |
+
shutil.move(src, dst)
|
260 |
+
|
261 |
+
filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
|
262 |
+
filelist = [os.path.relpath(p, start=datadir) for p in filelist]
|
263 |
+
filelist = sorted(filelist)
|
264 |
+
filelist = "\n".join(filelist)+"\n"
|
265 |
+
with open(self.txt_filelist, "w") as f:
|
266 |
+
f.write(filelist)
|
267 |
+
|
268 |
+
tdu.mark_prepared(self.root)
|
269 |
+
|
270 |
+
|
271 |
+
|
272 |
+
class ImageNetSR(Dataset):
|
273 |
+
def __init__(self, size=None,
|
274 |
+
degradation=None, downscale_f=4, min_crop_f=0.5, max_crop_f=1.,
|
275 |
+
random_crop=True):
|
276 |
+
"""
|
277 |
+
Imagenet Superresolution Dataloader
|
278 |
+
Performs following ops in order:
|
279 |
+
1. crops a crop of size s from image either as random or center crop
|
280 |
+
2. resizes crop to size with cv2.area_interpolation
|
281 |
+
3. degrades resized crop with degradation_fn
|
282 |
+
|
283 |
+
:param size: resizing to size after cropping
|
284 |
+
:param degradation: degradation_fn, e.g. cv_bicubic or bsrgan_light
|
285 |
+
:param downscale_f: Low Resolution Downsample factor
|
286 |
+
:param min_crop_f: determines crop size s,
|
287 |
+
where s = c * min_img_side_len with c sampled from interval (min_crop_f, max_crop_f)
|
288 |
+
:param max_crop_f: ""
|
289 |
+
:param data_root:
|
290 |
+
:param random_crop:
|
291 |
+
"""
|
292 |
+
self.base = self.get_base()
|
293 |
+
assert size
|
294 |
+
assert (size / downscale_f).is_integer()
|
295 |
+
self.size = size
|
296 |
+
self.LR_size = int(size / downscale_f)
|
297 |
+
self.min_crop_f = min_crop_f
|
298 |
+
self.max_crop_f = max_crop_f
|
299 |
+
assert(max_crop_f <= 1.)
|
300 |
+
self.center_crop = not random_crop
|
301 |
+
|
302 |
+
self.image_rescaler = albumentations.SmallestMaxSize(max_size=size, interpolation=cv2.INTER_AREA)
|
303 |
+
|
304 |
+
self.pil_interpolation = False # gets reset later if incase interp_op is from pillow
|
305 |
+
|
306 |
+
if degradation == "bsrgan":
|
307 |
+
self.degradation_process = partial(degradation_fn_bsr, sf=downscale_f)
|
308 |
+
|
309 |
+
elif degradation == "bsrgan_light":
|
310 |
+
self.degradation_process = partial(degradation_fn_bsr_light, sf=downscale_f)
|
311 |
+
|
312 |
+
else:
|
313 |
+
interpolation_fn = {
|
314 |
+
"cv_nearest": cv2.INTER_NEAREST,
|
315 |
+
"cv_bilinear": cv2.INTER_LINEAR,
|
316 |
+
"cv_bicubic": cv2.INTER_CUBIC,
|
317 |
+
"cv_area": cv2.INTER_AREA,
|
318 |
+
"cv_lanczos": cv2.INTER_LANCZOS4,
|
319 |
+
"pil_nearest": PIL.Image.NEAREST,
|
320 |
+
"pil_bilinear": PIL.Image.BILINEAR,
|
321 |
+
"pil_bicubic": PIL.Image.BICUBIC,
|
322 |
+
"pil_box": PIL.Image.BOX,
|
323 |
+
"pil_hamming": PIL.Image.HAMMING,
|
324 |
+
"pil_lanczos": PIL.Image.LANCZOS,
|
325 |
+
}[degradation]
|
326 |
+
|
327 |
+
self.pil_interpolation = degradation.startswith("pil_")
|
328 |
+
|
329 |
+
if self.pil_interpolation:
|
330 |
+
self.degradation_process = partial(TF.resize, size=self.LR_size, interpolation=interpolation_fn)
|
331 |
+
|
332 |
+
else:
|
333 |
+
self.degradation_process = albumentations.SmallestMaxSize(max_size=self.LR_size,
|
334 |
+
interpolation=interpolation_fn)
|
335 |
+
|
336 |
+
def __len__(self):
|
337 |
+
return len(self.base)
|
338 |
+
|
339 |
+
def __getitem__(self, i):
|
340 |
+
example = self.base[i]
|
341 |
+
image = Image.open(example["file_path_"])
|
342 |
+
|
343 |
+
if not image.mode == "RGB":
|
344 |
+
image = image.convert("RGB")
|
345 |
+
|
346 |
+
image = np.array(image).astype(np.uint8)
|
347 |
+
|
348 |
+
min_side_len = min(image.shape[:2])
|
349 |
+
crop_side_len = min_side_len * np.random.uniform(self.min_crop_f, self.max_crop_f, size=None)
|
350 |
+
crop_side_len = int(crop_side_len)
|
351 |
+
|
352 |
+
if self.center_crop:
|
353 |
+
self.cropper = albumentations.CenterCrop(height=crop_side_len, width=crop_side_len)
|
354 |
+
|
355 |
+
else:
|
356 |
+
self.cropper = albumentations.RandomCrop(height=crop_side_len, width=crop_side_len)
|
357 |
+
|
358 |
+
image = self.cropper(image=image)["image"]
|
359 |
+
image = self.image_rescaler(image=image)["image"]
|
360 |
+
|
361 |
+
if self.pil_interpolation:
|
362 |
+
image_pil = PIL.Image.fromarray(image)
|
363 |
+
LR_image = self.degradation_process(image_pil)
|
364 |
+
LR_image = np.array(LR_image).astype(np.uint8)
|
365 |
+
|
366 |
+
else:
|
367 |
+
LR_image = self.degradation_process(image=image)["image"]
|
368 |
+
|
369 |
+
example["image"] = (image/127.5 - 1.0).astype(np.float32)
|
370 |
+
example["LR_image"] = (LR_image/127.5 - 1.0).astype(np.float32)
|
371 |
+
example["caption"] = example["human_label"] # dummy caption
|
372 |
+
return example
|
373 |
+
|
374 |
+
|
375 |
+
class ImageNetSRTrain(ImageNetSR):
|
376 |
+
def __init__(self, **kwargs):
|
377 |
+
super().__init__(**kwargs)
|
378 |
+
|
379 |
+
def get_base(self):
|
380 |
+
with open("data/imagenet_train_hr_indices.p", "rb") as f:
|
381 |
+
indices = pickle.load(f)
|
382 |
+
dset = ImageNetTrain(process_images=False,)
|
383 |
+
return Subset(dset, indices)
|
384 |
+
|
385 |
+
|
386 |
+
class ImageNetSRValidation(ImageNetSR):
|
387 |
+
def __init__(self, **kwargs):
|
388 |
+
super().__init__(**kwargs)
|
389 |
+
|
390 |
+
def get_base(self):
|
391 |
+
with open("data/imagenet_val_hr_indices.p", "rb") as f:
|
392 |
+
indices = pickle.load(f)
|
393 |
+
dset = ImageNetValidation(process_images=False,)
|
394 |
+
return Subset(dset, indices)
|
ldm/data/inpainting/__init__.py
ADDED
File without changes
|
ldm/data/inpainting/synthetic_mask.py
ADDED
@@ -0,0 +1,166 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from PIL import Image, ImageDraw
|
2 |
+
import numpy as np
|
3 |
+
|
4 |
+
settings = {
|
5 |
+
"256narrow": {
|
6 |
+
"p_irr": 1,
|
7 |
+
"min_n_irr": 4,
|
8 |
+
"max_n_irr": 50,
|
9 |
+
"max_l_irr": 40,
|
10 |
+
"max_w_irr": 10,
|
11 |
+
"min_n_box": None,
|
12 |
+
"max_n_box": None,
|
13 |
+
"min_s_box": None,
|
14 |
+
"max_s_box": None,
|
15 |
+
"marg": None,
|
16 |
+
},
|
17 |
+
"256train": {
|
18 |
+
"p_irr": 0.5,
|
19 |
+
"min_n_irr": 1,
|
20 |
+
"max_n_irr": 5,
|
21 |
+
"max_l_irr": 200,
|
22 |
+
"max_w_irr": 100,
|
23 |
+
"min_n_box": 1,
|
24 |
+
"max_n_box": 4,
|
25 |
+
"min_s_box": 30,
|
26 |
+
"max_s_box": 150,
|
27 |
+
"marg": 10,
|
28 |
+
},
|
29 |
+
"512train": { # TODO: experimental
|
30 |
+
"p_irr": 0.5,
|
31 |
+
"min_n_irr": 1,
|
32 |
+
"max_n_irr": 5,
|
33 |
+
"max_l_irr": 450,
|
34 |
+
"max_w_irr": 250,
|
35 |
+
"min_n_box": 1,
|
36 |
+
"max_n_box": 4,
|
37 |
+
"min_s_box": 30,
|
38 |
+
"max_s_box": 300,
|
39 |
+
"marg": 10,
|
40 |
+
},
|
41 |
+
"512train-large": { # TODO: experimental
|
42 |
+
"p_irr": 0.5,
|
43 |
+
"min_n_irr": 1,
|
44 |
+
"max_n_irr": 5,
|
45 |
+
"max_l_irr": 450,
|
46 |
+
"max_w_irr": 400,
|
47 |
+
"min_n_box": 1,
|
48 |
+
"max_n_box": 4,
|
49 |
+
"min_s_box": 75,
|
50 |
+
"max_s_box": 450,
|
51 |
+
"marg": 10,
|
52 |
+
},
|
53 |
+
}
|
54 |
+
|
55 |
+
|
56 |
+
def gen_segment_mask(mask, start, end, brush_width):
|
57 |
+
mask = mask > 0
|
58 |
+
mask = (255 * mask).astype(np.uint8)
|
59 |
+
mask = Image.fromarray(mask)
|
60 |
+
draw = ImageDraw.Draw(mask)
|
61 |
+
draw.line([start, end], fill=255, width=brush_width, joint="curve")
|
62 |
+
mask = np.array(mask) / 255
|
63 |
+
return mask
|
64 |
+
|
65 |
+
|
66 |
+
def gen_box_mask(mask, masked):
|
67 |
+
x_0, y_0, w, h = masked
|
68 |
+
mask[y_0:y_0 + h, x_0:x_0 + w] = 1
|
69 |
+
return mask
|
70 |
+
|
71 |
+
|
72 |
+
def gen_round_mask(mask, masked, radius):
|
73 |
+
x_0, y_0, w, h = masked
|
74 |
+
xy = [(x_0, y_0), (x_0 + w, y_0 + w)]
|
75 |
+
|
76 |
+
mask = mask > 0
|
77 |
+
mask = (255 * mask).astype(np.uint8)
|
78 |
+
mask = Image.fromarray(mask)
|
79 |
+
draw = ImageDraw.Draw(mask)
|
80 |
+
draw.rounded_rectangle(xy, radius=radius, fill=255)
|
81 |
+
mask = np.array(mask) / 255
|
82 |
+
return mask
|
83 |
+
|
84 |
+
|
85 |
+
def gen_large_mask(prng, img_h, img_w,
|
86 |
+
marg, p_irr, min_n_irr, max_n_irr, max_l_irr, max_w_irr,
|
87 |
+
min_n_box, max_n_box, min_s_box, max_s_box):
|
88 |
+
"""
|
89 |
+
img_h: int, an image height
|
90 |
+
img_w: int, an image width
|
91 |
+
marg: int, a margin for a box starting coordinate
|
92 |
+
p_irr: float, 0 <= p_irr <= 1, a probability of a polygonal chain mask
|
93 |
+
|
94 |
+
min_n_irr: int, min number of segments
|
95 |
+
max_n_irr: int, max number of segments
|
96 |
+
max_l_irr: max length of a segment in polygonal chain
|
97 |
+
max_w_irr: max width of a segment in polygonal chain
|
98 |
+
|
99 |
+
min_n_box: int, min bound for the number of box primitives
|
100 |
+
max_n_box: int, max bound for the number of box primitives
|
101 |
+
min_s_box: int, min length of a box side
|
102 |
+
max_s_box: int, max length of a box side
|
103 |
+
"""
|
104 |
+
|
105 |
+
mask = np.zeros((img_h, img_w))
|
106 |
+
uniform = prng.randint
|
107 |
+
|
108 |
+
if np.random.uniform(0, 1) < p_irr: # generate polygonal chain
|
109 |
+
n = uniform(min_n_irr, max_n_irr) # sample number of segments
|
110 |
+
|
111 |
+
for _ in range(n):
|
112 |
+
y = uniform(0, img_h) # sample a starting point
|
113 |
+
x = uniform(0, img_w)
|
114 |
+
|
115 |
+
a = uniform(0, 360) # sample angle
|
116 |
+
l = uniform(10, max_l_irr) # sample segment length
|
117 |
+
w = uniform(5, max_w_irr) # sample a segment width
|
118 |
+
|
119 |
+
# draw segment starting from (x,y) to (x_,y_) using brush of width w
|
120 |
+
x_ = x + l * np.sin(a)
|
121 |
+
y_ = y + l * np.cos(a)
|
122 |
+
|
123 |
+
mask = gen_segment_mask(mask, start=(x, y), end=(x_, y_), brush_width=w)
|
124 |
+
x, y = x_, y_
|
125 |
+
else: # generate Box masks
|
126 |
+
n = uniform(min_n_box, max_n_box) # sample number of rectangles
|
127 |
+
|
128 |
+
for _ in range(n):
|
129 |
+
h = uniform(min_s_box, max_s_box) # sample box shape
|
130 |
+
w = uniform(min_s_box, max_s_box)
|
131 |
+
|
132 |
+
x_0 = uniform(marg, img_w - marg - w) # sample upper-left coordinates of box
|
133 |
+
y_0 = uniform(marg, img_h - marg - h)
|
134 |
+
|
135 |
+
if np.random.uniform(0, 1) < 0.5:
|
136 |
+
mask = gen_box_mask(mask, masked=(x_0, y_0, w, h))
|
137 |
+
else:
|
138 |
+
r = uniform(0, 60) # sample radius
|
139 |
+
mask = gen_round_mask(mask, masked=(x_0, y_0, w, h), radius=r)
|
140 |
+
return mask
|
141 |
+
|
142 |
+
|
143 |
+
make_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256train"])
|
144 |
+
make_narrow_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256narrow"])
|
145 |
+
make_512_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train"])
|
146 |
+
make_512_lama_mask_large = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train-large"])
|
147 |
+
|
148 |
+
|
149 |
+
MASK_MODES = {
|
150 |
+
"256train": make_lama_mask,
|
151 |
+
"256narrow": make_narrow_lama_mask,
|
152 |
+
"512train": make_512_lama_mask,
|
153 |
+
"512train-large": make_512_lama_mask_large
|
154 |
+
}
|
155 |
+
|
156 |
+
if __name__ == "__main__":
|
157 |
+
import sys
|
158 |
+
|
159 |
+
out = sys.argv[1]
|
160 |
+
|
161 |
+
prng = np.random.RandomState(1)
|
162 |
+
kwargs = settings["256train"]
|
163 |
+
mask = gen_large_mask(prng, 256, 256, **kwargs)
|
164 |
+
mask = (255 * mask).astype(np.uint8)
|
165 |
+
mask = Image.fromarray(mask)
|
166 |
+
mask.save(out)
|
ldm/data/laion.py
ADDED
@@ -0,0 +1,537 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import webdataset as wds
|
2 |
+
import kornia
|
3 |
+
from PIL import Image
|
4 |
+
import io
|
5 |
+
import os
|
6 |
+
import torchvision
|
7 |
+
from PIL import Image
|
8 |
+
import glob
|
9 |
+
import random
|
10 |
+
import numpy as np
|
11 |
+
import pytorch_lightning as pl
|
12 |
+
from tqdm import tqdm
|
13 |
+
from omegaconf import OmegaConf
|
14 |
+
from einops import rearrange
|
15 |
+
import torch
|
16 |
+
from webdataset.handlers import warn_and_continue
|
17 |
+
|
18 |
+
|
19 |
+
from ldm.util import instantiate_from_config
|
20 |
+
from ldm.data.inpainting.synthetic_mask import gen_large_mask, MASK_MODES
|
21 |
+
from ldm.data.base import PRNGMixin
|
22 |
+
|
23 |
+
|
24 |
+
class DataWithWings(torch.utils.data.IterableDataset):
|
25 |
+
def __init__(self, min_size, transform=None, target_transform=None):
|
26 |
+
self.min_size = min_size
|
27 |
+
self.transform = transform if transform is not None else nn.Identity()
|
28 |
+
self.target_transform = target_transform if target_transform is not None else nn.Identity()
|
29 |
+
self.kv = OnDiskKV(file='/home/ubuntu/laion5B-watermark-safety-ordered', key_format='q', value_format='ee')
|
30 |
+
self.kv_aesthetic = OnDiskKV(file='/home/ubuntu/laion5B-aesthetic-tags-kv', key_format='q', value_format='e')
|
31 |
+
self.pwatermark_threshold = 0.8
|
32 |
+
self.punsafe_threshold = 0.5
|
33 |
+
self.aesthetic_threshold = 5.
|
34 |
+
self.total_samples = 0
|
35 |
+
self.samples = 0
|
36 |
+
location = 'pipe:aws s3 cp --quiet s3://s-datasets/laion5b/laion2B-data/{000000..231349}.tar -'
|
37 |
+
|
38 |
+
self.inner_dataset = wds.DataPipeline(
|
39 |
+
wds.ResampledShards(location),
|
40 |
+
wds.tarfile_to_samples(handler=wds.warn_and_continue),
|
41 |
+
wds.shuffle(1000, handler=wds.warn_and_continue),
|
42 |
+
wds.decode('pilrgb', handler=wds.warn_and_continue),
|
43 |
+
wds.map(self._add_tags, handler=wds.ignore_and_continue),
|
44 |
+
wds.select(self._filter_predicate),
|
45 |
+
wds.map_dict(jpg=self.transform, txt=self.target_transform, punsafe=self._punsafe_to_class, handler=wds.warn_and_continue),
|
46 |
+
wds.to_tuple('jpg', 'txt', 'punsafe', handler=wds.warn_and_continue),
|
47 |
+
)
|
48 |
+
|
49 |
+
@staticmethod
|
50 |
+
def _compute_hash(url, text):
|
51 |
+
if url is None:
|
52 |
+
url = ''
|
53 |
+
if text is None:
|
54 |
+
text = ''
|
55 |
+
total = (url + text).encode('utf-8')
|
56 |
+
return mmh3.hash64(total)[0]
|
57 |
+
|
58 |
+
def _add_tags(self, x):
|
59 |
+
hsh = self._compute_hash(x['json']['url'], x['txt'])
|
60 |
+
pwatermark, punsafe = self.kv[hsh]
|
61 |
+
aesthetic = self.kv_aesthetic[hsh][0]
|
62 |
+
return {**x, 'pwatermark': pwatermark, 'punsafe': punsafe, 'aesthetic': aesthetic}
|
63 |
+
|
64 |
+
def _punsafe_to_class(self, punsafe):
|
65 |
+
return torch.tensor(punsafe >= self.punsafe_threshold).long()
|
66 |
+
|
67 |
+
def _filter_predicate(self, x):
|
68 |
+
try:
|
69 |
+
return x['pwatermark'] < self.pwatermark_threshold and x['aesthetic'] >= self.aesthetic_threshold and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size
|
70 |
+
except:
|
71 |
+
return False
|
72 |
+
|
73 |
+
def __iter__(self):
|
74 |
+
return iter(self.inner_dataset)
|
75 |
+
|
76 |
+
|
77 |
+
def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True):
|
78 |
+
"""Take a list of samples (as dictionary) and create a batch, preserving the keys.
|
79 |
+
If `tensors` is True, `ndarray` objects are combined into
|
80 |
+
tensor batches.
|
81 |
+
:param dict samples: list of samples
|
82 |
+
:param bool tensors: whether to turn lists of ndarrays into a single ndarray
|
83 |
+
:returns: single sample consisting of a batch
|
84 |
+
:rtype: dict
|
85 |
+
"""
|
86 |
+
keys = set.intersection(*[set(sample.keys()) for sample in samples])
|
87 |
+
batched = {key: [] for key in keys}
|
88 |
+
|
89 |
+
for s in samples:
|
90 |
+
[batched[key].append(s[key]) for key in batched]
|
91 |
+
|
92 |
+
result = {}
|
93 |
+
for key in batched:
|
94 |
+
if isinstance(batched[key][0], (int, float)):
|
95 |
+
if combine_scalars:
|
96 |
+
result[key] = np.array(list(batched[key]))
|
97 |
+
elif isinstance(batched[key][0], torch.Tensor):
|
98 |
+
if combine_tensors:
|
99 |
+
result[key] = torch.stack(list(batched[key]))
|
100 |
+
elif isinstance(batched[key][0], np.ndarray):
|
101 |
+
if combine_tensors:
|
102 |
+
result[key] = np.array(list(batched[key]))
|
103 |
+
else:
|
104 |
+
result[key] = list(batched[key])
|
105 |
+
return result
|
106 |
+
|
107 |
+
|
108 |
+
class WebDataModuleFromConfig(pl.LightningDataModule):
|
109 |
+
def __init__(self, tar_base, batch_size, train=None, validation=None,
|
110 |
+
test=None, num_workers=4, multinode=True, min_size=None,
|
111 |
+
max_pwatermark=1.0,
|
112 |
+
**kwargs):
|
113 |
+
super().__init__(self)
|
114 |
+
print(f'Setting tar base to {tar_base}')
|
115 |
+
self.tar_base = tar_base
|
116 |
+
self.batch_size = batch_size
|
117 |
+
self.num_workers = num_workers
|
118 |
+
self.train = train
|
119 |
+
self.validation = validation
|
120 |
+
self.test = test
|
121 |
+
self.multinode = multinode
|
122 |
+
self.min_size = min_size # filter out very small images
|
123 |
+
self.max_pwatermark = max_pwatermark # filter out watermarked images
|
124 |
+
|
125 |
+
def make_loader(self, dataset_config, train=True):
|
126 |
+
if 'image_transforms' in dataset_config:
|
127 |
+
image_transforms = [instantiate_from_config(tt) for tt in dataset_config.image_transforms]
|
128 |
+
else:
|
129 |
+
image_transforms = []
|
130 |
+
|
131 |
+
image_transforms.extend([torchvision.transforms.ToTensor(),
|
132 |
+
torchvision.transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
|
133 |
+
image_transforms = torchvision.transforms.Compose(image_transforms)
|
134 |
+
|
135 |
+
if 'transforms' in dataset_config:
|
136 |
+
transforms_config = OmegaConf.to_container(dataset_config.transforms)
|
137 |
+
else:
|
138 |
+
transforms_config = dict()
|
139 |
+
|
140 |
+
transform_dict = {dkey: load_partial_from_config(transforms_config[dkey])
|
141 |
+
if transforms_config[dkey] != 'identity' else identity
|
142 |
+
for dkey in transforms_config}
|
143 |
+
img_key = dataset_config.get('image_key', 'jpeg')
|
144 |
+
transform_dict.update({img_key: image_transforms})
|
145 |
+
|
146 |
+
if 'postprocess' in dataset_config:
|
147 |
+
postprocess = instantiate_from_config(dataset_config['postprocess'])
|
148 |
+
else:
|
149 |
+
postprocess = None
|
150 |
+
|
151 |
+
shuffle = dataset_config.get('shuffle', 0)
|
152 |
+
shardshuffle = shuffle > 0
|
153 |
+
|
154 |
+
nodesplitter = wds.shardlists.split_by_node if self.multinode else wds.shardlists.single_node_only
|
155 |
+
|
156 |
+
if self.tar_base == "__improvedaesthetic__":
|
157 |
+
print("## Warning, loading the same improved aesthetic dataset "
|
158 |
+
"for all splits and ignoring shards parameter.")
|
159 |
+
tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
|
160 |
+
else:
|
161 |
+
tars = os.path.join(self.tar_base, dataset_config.shards)
|
162 |
+
|
163 |
+
dset = wds.WebDataset(
|
164 |
+
tars,
|
165 |
+
nodesplitter=nodesplitter,
|
166 |
+
shardshuffle=shardshuffle,
|
167 |
+
handler=wds.warn_and_continue).repeat().shuffle(shuffle)
|
168 |
+
print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.')
|
169 |
+
|
170 |
+
dset = (dset
|
171 |
+
.select(self.filter_keys)
|
172 |
+
.decode('pil', handler=wds.warn_and_continue)
|
173 |
+
.select(self.filter_size)
|
174 |
+
.map_dict(**transform_dict, handler=wds.warn_and_continue)
|
175 |
+
)
|
176 |
+
if postprocess is not None:
|
177 |
+
dset = dset.map(postprocess)
|
178 |
+
dset = (dset
|
179 |
+
.batched(self.batch_size, partial=False,
|
180 |
+
collation_fn=dict_collation_fn)
|
181 |
+
)
|
182 |
+
|
183 |
+
loader = wds.WebLoader(dset, batch_size=None, shuffle=False,
|
184 |
+
num_workers=self.num_workers)
|
185 |
+
|
186 |
+
return loader
|
187 |
+
|
188 |
+
def filter_size(self, x):
|
189 |
+
try:
|
190 |
+
valid = True
|
191 |
+
if self.min_size is not None and self.min_size > 1:
|
192 |
+
try:
|
193 |
+
valid = valid and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size
|
194 |
+
except Exception:
|
195 |
+
valid = False
|
196 |
+
if self.max_pwatermark is not None and self.max_pwatermark < 1.0:
|
197 |
+
try:
|
198 |
+
valid = valid and x['json']['pwatermark'] <= self.max_pwatermark
|
199 |
+
except Exception:
|
200 |
+
valid = False
|
201 |
+
return valid
|
202 |
+
except Exception:
|
203 |
+
return False
|
204 |
+
|
205 |
+
def filter_keys(self, x):
|
206 |
+
try:
|
207 |
+
return ("jpg" in x) and ("txt" in x)
|
208 |
+
except Exception:
|
209 |
+
return False
|
210 |
+
|
211 |
+
def train_dataloader(self):
|
212 |
+
return self.make_loader(self.train)
|
213 |
+
|
214 |
+
def val_dataloader(self):
|
215 |
+
return self.make_loader(self.validation, train=False)
|
216 |
+
|
217 |
+
def test_dataloader(self):
|
218 |
+
return self.make_loader(self.test, train=False)
|
219 |
+
|
220 |
+
|
221 |
+
from ldm.modules.image_degradation import degradation_fn_bsr_light
|
222 |
+
import cv2
|
223 |
+
|
224 |
+
class AddLR(object):
|
225 |
+
def __init__(self, factor, output_size, initial_size=None, image_key="jpg"):
|
226 |
+
self.factor = factor
|
227 |
+
self.output_size = output_size
|
228 |
+
self.image_key = image_key
|
229 |
+
self.initial_size = initial_size
|
230 |
+
|
231 |
+
def pt2np(self, x):
|
232 |
+
x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy()
|
233 |
+
return x
|
234 |
+
|
235 |
+
def np2pt(self, x):
|
236 |
+
x = torch.from_numpy(x)/127.5-1.0
|
237 |
+
return x
|
238 |
+
|
239 |
+
def __call__(self, sample):
|
240 |
+
# sample['jpg'] is tensor hwc in [-1, 1] at this point
|
241 |
+
x = self.pt2np(sample[self.image_key])
|
242 |
+
if self.initial_size is not None:
|
243 |
+
x = cv2.resize(x, (self.initial_size, self.initial_size), interpolation=2)
|
244 |
+
x = degradation_fn_bsr_light(x, sf=self.factor)['image']
|
245 |
+
x = cv2.resize(x, (self.output_size, self.output_size), interpolation=2)
|
246 |
+
x = self.np2pt(x)
|
247 |
+
sample['lr'] = x
|
248 |
+
return sample
|
249 |
+
|
250 |
+
class AddBW(object):
|
251 |
+
def __init__(self, image_key="jpg"):
|
252 |
+
self.image_key = image_key
|
253 |
+
|
254 |
+
def pt2np(self, x):
|
255 |
+
x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy()
|
256 |
+
return x
|
257 |
+
|
258 |
+
def np2pt(self, x):
|
259 |
+
x = torch.from_numpy(x)/127.5-1.0
|
260 |
+
return x
|
261 |
+
|
262 |
+
def __call__(self, sample):
|
263 |
+
# sample['jpg'] is tensor hwc in [-1, 1] at this point
|
264 |
+
x = sample[self.image_key]
|
265 |
+
w = torch.rand(3, device=x.device)
|
266 |
+
w /= w.sum()
|
267 |
+
out = torch.einsum('hwc,c->hw', x, w)
|
268 |
+
|
269 |
+
# Keep as 3ch so we can pass to encoder, also we might want to add hints
|
270 |
+
sample['lr'] = out.unsqueeze(-1).tile(1,1,3)
|
271 |
+
return sample
|
272 |
+
|
273 |
+
class AddMask(PRNGMixin):
|
274 |
+
def __init__(self, mode="512train", p_drop=0.):
|
275 |
+
super().__init__()
|
276 |
+
assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"'
|
277 |
+
self.make_mask = MASK_MODES[mode]
|
278 |
+
self.p_drop = p_drop
|
279 |
+
|
280 |
+
def __call__(self, sample):
|
281 |
+
# sample['jpg'] is tensor hwc in [-1, 1] at this point
|
282 |
+
x = sample['jpg']
|
283 |
+
mask = self.make_mask(self.prng, x.shape[0], x.shape[1])
|
284 |
+
if self.prng.choice(2, p=[1 - self.p_drop, self.p_drop]):
|
285 |
+
mask = np.ones_like(mask)
|
286 |
+
mask[mask < 0.5] = 0
|
287 |
+
mask[mask > 0.5] = 1
|
288 |
+
mask = torch.from_numpy(mask[..., None])
|
289 |
+
sample['mask'] = mask
|
290 |
+
sample['masked_image'] = x * (mask < 0.5)
|
291 |
+
return sample
|
292 |
+
|
293 |
+
|
294 |
+
class AddEdge(PRNGMixin):
|
295 |
+
def __init__(self, mode="512train", mask_edges=True):
|
296 |
+
super().__init__()
|
297 |
+
assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"'
|
298 |
+
self.make_mask = MASK_MODES[mode]
|
299 |
+
self.n_down_choices = [0]
|
300 |
+
self.sigma_choices = [1, 2]
|
301 |
+
self.mask_edges = mask_edges
|
302 |
+
|
303 |
+
@torch.no_grad()
|
304 |
+
def __call__(self, sample):
|
305 |
+
# sample['jpg'] is tensor hwc in [-1, 1] at this point
|
306 |
+
x = sample['jpg']
|
307 |
+
|
308 |
+
mask = self.make_mask(self.prng, x.shape[0], x.shape[1])
|
309 |
+
mask[mask < 0.5] = 0
|
310 |
+
mask[mask > 0.5] = 1
|
311 |
+
mask = torch.from_numpy(mask[..., None])
|
312 |
+
sample['mask'] = mask
|
313 |
+
|
314 |
+
n_down_idx = self.prng.choice(len(self.n_down_choices))
|
315 |
+
sigma_idx = self.prng.choice(len(self.sigma_choices))
|
316 |
+
|
317 |
+
n_choices = len(self.n_down_choices)*len(self.sigma_choices)
|
318 |
+
raveled_idx = np.ravel_multi_index((n_down_idx, sigma_idx),
|
319 |
+
(len(self.n_down_choices), len(self.sigma_choices)))
|
320 |
+
normalized_idx = raveled_idx/max(1, n_choices-1)
|
321 |
+
|
322 |
+
n_down = self.n_down_choices[n_down_idx]
|
323 |
+
sigma = self.sigma_choices[sigma_idx]
|
324 |
+
|
325 |
+
kernel_size = 4*sigma+1
|
326 |
+
kernel_size = (kernel_size, kernel_size)
|
327 |
+
sigma = (sigma, sigma)
|
328 |
+
canny = kornia.filters.Canny(
|
329 |
+
low_threshold=0.1,
|
330 |
+
high_threshold=0.2,
|
331 |
+
kernel_size=kernel_size,
|
332 |
+
sigma=sigma,
|
333 |
+
hysteresis=True,
|
334 |
+
)
|
335 |
+
y = (x+1.0)/2.0 # in 01
|
336 |
+
y = y.unsqueeze(0).permute(0, 3, 1, 2).contiguous()
|
337 |
+
|
338 |
+
# down
|
339 |
+
for i_down in range(n_down):
|
340 |
+
size = min(y.shape[-2], y.shape[-1])//2
|
341 |
+
y = kornia.geometry.transform.resize(y, size, antialias=True)
|
342 |
+
|
343 |
+
# edge
|
344 |
+
_, y = canny(y)
|
345 |
+
|
346 |
+
if n_down > 0:
|
347 |
+
size = x.shape[0], x.shape[1]
|
348 |
+
y = kornia.geometry.transform.resize(y, size, interpolation="nearest")
|
349 |
+
|
350 |
+
y = y.permute(0, 2, 3, 1)[0].expand(-1, -1, 3).contiguous()
|
351 |
+
y = y*2.0-1.0
|
352 |
+
|
353 |
+
if self.mask_edges:
|
354 |
+
sample['masked_image'] = y * (mask < 0.5)
|
355 |
+
else:
|
356 |
+
sample['masked_image'] = y
|
357 |
+
sample['mask'] = torch.zeros_like(sample['mask'])
|
358 |
+
|
359 |
+
# concat normalized idx
|
360 |
+
sample['smoothing_strength'] = torch.ones_like(sample['mask'])*normalized_idx
|
361 |
+
|
362 |
+
return sample
|
363 |
+
|
364 |
+
|
365 |
+
def example00():
|
366 |
+
url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/000000.tar -"
|
367 |
+
dataset = wds.WebDataset(url)
|
368 |
+
example = next(iter(dataset))
|
369 |
+
for k in example:
|
370 |
+
print(k, type(example[k]))
|
371 |
+
|
372 |
+
print(example["__key__"])
|
373 |
+
for k in ["json", "txt"]:
|
374 |
+
print(example[k].decode())
|
375 |
+
|
376 |
+
image = Image.open(io.BytesIO(example["jpg"]))
|
377 |
+
outdir = "tmp"
|
378 |
+
os.makedirs(outdir, exist_ok=True)
|
379 |
+
image.save(os.path.join(outdir, example["__key__"] + ".png"))
|
380 |
+
|
381 |
+
|
382 |
+
def load_example(example):
|
383 |
+
return {
|
384 |
+
"key": example["__key__"],
|
385 |
+
"image": Image.open(io.BytesIO(example["jpg"])),
|
386 |
+
"text": example["txt"].decode(),
|
387 |
+
}
|
388 |
+
|
389 |
+
|
390 |
+
for i, example in tqdm(enumerate(dataset)):
|
391 |
+
ex = load_example(example)
|
392 |
+
print(ex["image"].size, ex["text"])
|
393 |
+
if i >= 100:
|
394 |
+
break
|
395 |
+
|
396 |
+
|
397 |
+
def example01():
|
398 |
+
# the first laion shards contain ~10k examples each
|
399 |
+
url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/{000000..000002}.tar -"
|
400 |
+
|
401 |
+
batch_size = 3
|
402 |
+
shuffle_buffer = 10000
|
403 |
+
dset = wds.WebDataset(
|
404 |
+
url,
|
405 |
+
nodesplitter=wds.shardlists.split_by_node,
|
406 |
+
shardshuffle=True,
|
407 |
+
)
|
408 |
+
dset = (dset
|
409 |
+
.shuffle(shuffle_buffer, initial=shuffle_buffer)
|
410 |
+
.decode('pil', handler=warn_and_continue)
|
411 |
+
.batched(batch_size, partial=False,
|
412 |
+
collation_fn=dict_collation_fn)
|
413 |
+
)
|
414 |
+
|
415 |
+
num_workers = 2
|
416 |
+
loader = wds.WebLoader(dset, batch_size=None, shuffle=False, num_workers=num_workers)
|
417 |
+
|
418 |
+
batch_sizes = list()
|
419 |
+
keys_per_epoch = list()
|
420 |
+
for epoch in range(5):
|
421 |
+
keys = list()
|
422 |
+
for batch in tqdm(loader):
|
423 |
+
batch_sizes.append(len(batch["__key__"]))
|
424 |
+
keys.append(batch["__key__"])
|
425 |
+
|
426 |
+
for bs in batch_sizes:
|
427 |
+
assert bs==batch_size
|
428 |
+
print(f"{len(batch_sizes)} batches of size {batch_size}.")
|
429 |
+
batch_sizes = list()
|
430 |
+
|
431 |
+
keys_per_epoch.append(keys)
|
432 |
+
for i_batch in [0, 1, -1]:
|
433 |
+
print(f"Batch {i_batch} of epoch {epoch}:")
|
434 |
+
print(keys[i_batch])
|
435 |
+
print("next epoch.")
|
436 |
+
|
437 |
+
|
438 |
+
def example02():
|
439 |
+
from omegaconf import OmegaConf
|
440 |
+
from torch.utils.data.distributed import DistributedSampler
|
441 |
+
from torch.utils.data import IterableDataset
|
442 |
+
from torch.utils.data import DataLoader, RandomSampler, Sampler, SequentialSampler
|
443 |
+
from pytorch_lightning.trainer.supporters import CombinedLoader, CycleIterator
|
444 |
+
|
445 |
+
#config = OmegaConf.load("configs/stable-diffusion/txt2img-1p4B-multinode-clip-encoder-high-res-512.yaml")
|
446 |
+
#config = OmegaConf.load("configs/stable-diffusion/txt2img-upscale-clip-encoder-f16-1024.yaml")
|
447 |
+
config = OmegaConf.load("configs/stable-diffusion/txt2img-v2-clip-encoder-improved_aesthetics-256.yaml")
|
448 |
+
datamod = WebDataModuleFromConfig(**config["data"]["params"])
|
449 |
+
dataloader = datamod.train_dataloader()
|
450 |
+
|
451 |
+
for batch in dataloader:
|
452 |
+
print(batch.keys())
|
453 |
+
print(batch["jpg"].shape)
|
454 |
+
break
|
455 |
+
|
456 |
+
|
457 |
+
def example03():
|
458 |
+
# improved aesthetics
|
459 |
+
tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
|
460 |
+
dataset = wds.WebDataset(tars)
|
461 |
+
|
462 |
+
def filter_keys(x):
|
463 |
+
try:
|
464 |
+
return ("jpg" in x) and ("txt" in x)
|
465 |
+
except Exception:
|
466 |
+
return False
|
467 |
+
|
468 |
+
def filter_size(x):
|
469 |
+
try:
|
470 |
+
return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
|
471 |
+
except Exception:
|
472 |
+
return False
|
473 |
+
|
474 |
+
def filter_watermark(x):
|
475 |
+
try:
|
476 |
+
return x['json']['pwatermark'] < 0.5
|
477 |
+
except Exception:
|
478 |
+
return False
|
479 |
+
|
480 |
+
dataset = (dataset
|
481 |
+
.select(filter_keys)
|
482 |
+
.decode('pil', handler=wds.warn_and_continue))
|
483 |
+
n_save = 20
|
484 |
+
n_total = 0
|
485 |
+
n_large = 0
|
486 |
+
n_large_nowm = 0
|
487 |
+
for i, example in enumerate(dataset):
|
488 |
+
n_total += 1
|
489 |
+
if filter_size(example):
|
490 |
+
n_large += 1
|
491 |
+
if filter_watermark(example):
|
492 |
+
n_large_nowm += 1
|
493 |
+
if n_large_nowm < n_save+1:
|
494 |
+
image = example["jpg"]
|
495 |
+
image.save(os.path.join("tmp", f"{n_large_nowm-1:06}.png"))
|
496 |
+
|
497 |
+
if i%500 == 0:
|
498 |
+
print(i)
|
499 |
+
print(f"Large: {n_large}/{n_total} | {n_large/n_total*100:.2f}%")
|
500 |
+
if n_large > 0:
|
501 |
+
print(f"No Watermark: {n_large_nowm}/{n_large} | {n_large_nowm/n_large*100:.2f}%")
|
502 |
+
|
503 |
+
|
504 |
+
|
505 |
+
def example04():
|
506 |
+
# improved aesthetics
|
507 |
+
for i_shard in range(60208)[::-1]:
|
508 |
+
print(i_shard)
|
509 |
+
tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{:06}.tar -".format(i_shard)
|
510 |
+
dataset = wds.WebDataset(tars)
|
511 |
+
|
512 |
+
def filter_keys(x):
|
513 |
+
try:
|
514 |
+
return ("jpg" in x) and ("txt" in x)
|
515 |
+
except Exception:
|
516 |
+
return False
|
517 |
+
|
518 |
+
def filter_size(x):
|
519 |
+
try:
|
520 |
+
return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
|
521 |
+
except Exception:
|
522 |
+
return False
|
523 |
+
|
524 |
+
dataset = (dataset
|
525 |
+
.select(filter_keys)
|
526 |
+
.decode('pil', handler=wds.warn_and_continue))
|
527 |
+
try:
|
528 |
+
example = next(iter(dataset))
|
529 |
+
except Exception:
|
530 |
+
print(f"Error @ {i_shard}")
|
531 |
+
|
532 |
+
|
533 |
+
if __name__ == "__main__":
|
534 |
+
#example01()
|
535 |
+
#example02()
|
536 |
+
example03()
|
537 |
+
#example04()
|
ldm/data/lsun.py
ADDED
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import numpy as np
|
3 |
+
import PIL
|
4 |
+
from PIL import Image
|
5 |
+
from torch.utils.data import Dataset
|
6 |
+
from torchvision import transforms
|
7 |
+
|
8 |
+
|
9 |
+
class LSUNBase(Dataset):
|
10 |
+
def __init__(self,
|
11 |
+
txt_file,
|
12 |
+
data_root,
|
13 |
+
size=None,
|
14 |
+
interpolation="bicubic",
|
15 |
+
flip_p=0.5
|
16 |
+
):
|
17 |
+
self.data_paths = txt_file
|
18 |
+
self.data_root = data_root
|
19 |
+
with open(self.data_paths, "r") as f:
|
20 |
+
self.image_paths = f.read().splitlines()
|
21 |
+
self._length = len(self.image_paths)
|
22 |
+
self.labels = {
|
23 |
+
"relative_file_path_": [l for l in self.image_paths],
|
24 |
+
"file_path_": [os.path.join(self.data_root, l)
|
25 |
+
for l in self.image_paths],
|
26 |
+
}
|
27 |
+
|
28 |
+
self.size = size
|
29 |
+
self.interpolation = {"linear": PIL.Image.LINEAR,
|
30 |
+
"bilinear": PIL.Image.BILINEAR,
|
31 |
+
"bicubic": PIL.Image.BICUBIC,
|
32 |
+
"lanczos": PIL.Image.LANCZOS,
|
33 |
+
}[interpolation]
|
34 |
+
self.flip = transforms.RandomHorizontalFlip(p=flip_p)
|
35 |
+
|
36 |
+
def __len__(self):
|
37 |
+
return self._length
|
38 |
+
|
39 |
+
def __getitem__(self, i):
|
40 |
+
example = dict((k, self.labels[k][i]) for k in self.labels)
|
41 |
+
image = Image.open(example["file_path_"])
|
42 |
+
if not image.mode == "RGB":
|
43 |
+
image = image.convert("RGB")
|
44 |
+
|
45 |
+
# default to score-sde preprocessing
|
46 |
+
img = np.array(image).astype(np.uint8)
|
47 |
+
crop = min(img.shape[0], img.shape[1])
|
48 |
+
h, w, = img.shape[0], img.shape[1]
|
49 |
+
img = img[(h - crop) // 2:(h + crop) // 2,
|
50 |
+
(w - crop) // 2:(w + crop) // 2]
|
51 |
+
|
52 |
+
image = Image.fromarray(img)
|
53 |
+
if self.size is not None:
|
54 |
+
image = image.resize((self.size, self.size), resample=self.interpolation)
|
55 |
+
|
56 |
+
image = self.flip(image)
|
57 |
+
image = np.array(image).astype(np.uint8)
|
58 |
+
example["image"] = (image / 127.5 - 1.0).astype(np.float32)
|
59 |
+
return example
|
60 |
+
|
61 |
+
|
62 |
+
class LSUNChurchesTrain(LSUNBase):
|
63 |
+
def __init__(self, **kwargs):
|
64 |
+
super().__init__(txt_file="data/lsun/church_outdoor_train.txt", data_root="data/lsun/churches", **kwargs)
|
65 |
+
|
66 |
+
|
67 |
+
class LSUNChurchesValidation(LSUNBase):
|
68 |
+
def __init__(self, flip_p=0., **kwargs):
|
69 |
+
super().__init__(txt_file="data/lsun/church_outdoor_val.txt", data_root="data/lsun/churches",
|
70 |
+
flip_p=flip_p, **kwargs)
|
71 |
+
|
72 |
+
|
73 |
+
class LSUNBedroomsTrain(LSUNBase):
|
74 |
+
def __init__(self, **kwargs):
|
75 |
+
super().__init__(txt_file="data/lsun/bedrooms_train.txt", data_root="data/lsun/bedrooms", **kwargs)
|
76 |
+
|
77 |
+
|
78 |
+
class LSUNBedroomsValidation(LSUNBase):
|
79 |
+
def __init__(self, flip_p=0.0, **kwargs):
|
80 |
+
super().__init__(txt_file="data/lsun/bedrooms_val.txt", data_root="data/lsun/bedrooms",
|
81 |
+
flip_p=flip_p, **kwargs)
|
82 |
+
|
83 |
+
|
84 |
+
class LSUNCatsTrain(LSUNBase):
|
85 |
+
def __init__(self, **kwargs):
|
86 |
+
super().__init__(txt_file="data/lsun/cat_train.txt", data_root="data/lsun/cats", **kwargs)
|
87 |
+
|
88 |
+
|
89 |
+
class LSUNCatsValidation(LSUNBase):
|
90 |
+
def __init__(self, flip_p=0., **kwargs):
|
91 |
+
super().__init__(txt_file="data/lsun/cat_val.txt", data_root="data/lsun/cats",
|
92 |
+
flip_p=flip_p, **kwargs)
|
ldm/data/nerf_like.py
ADDED
@@ -0,0 +1,165 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from torch.utils.data import Dataset
|
2 |
+
import os
|
3 |
+
import json
|
4 |
+
import numpy as np
|
5 |
+
import torch
|
6 |
+
import imageio
|
7 |
+
import math
|
8 |
+
import cv2
|
9 |
+
from torchvision import transforms
|
10 |
+
|
11 |
+
def cartesian_to_spherical(xyz):
|
12 |
+
ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
|
13 |
+
xy = xyz[:,0]**2 + xyz[:,1]**2
|
14 |
+
z = np.sqrt(xy + xyz[:,2]**2)
|
15 |
+
theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
|
16 |
+
#ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
|
17 |
+
azimuth = np.arctan2(xyz[:,1], xyz[:,0])
|
18 |
+
return np.array([theta, azimuth, z])
|
19 |
+
|
20 |
+
|
21 |
+
def get_T(T_target, T_cond):
|
22 |
+
theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :])
|
23 |
+
theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
|
24 |
+
|
25 |
+
d_theta = theta_target - theta_cond
|
26 |
+
d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
|
27 |
+
d_z = z_target - z_cond
|
28 |
+
|
29 |
+
d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
|
30 |
+
return d_T
|
31 |
+
|
32 |
+
def get_spherical(T_target, T_cond):
|
33 |
+
theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :])
|
34 |
+
theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
|
35 |
+
|
36 |
+
d_theta = theta_target - theta_cond
|
37 |
+
d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
|
38 |
+
d_z = z_target - z_cond
|
39 |
+
|
40 |
+
d_T = torch.tensor([math.degrees(d_theta.item()), math.degrees(d_azimuth.item()), d_z.item()])
|
41 |
+
return d_T
|
42 |
+
|
43 |
+
class RTMV(Dataset):
|
44 |
+
def __init__(self, root_dir='datasets/RTMV/google_scanned',\
|
45 |
+
first_K=64, resolution=256, load_target=False):
|
46 |
+
self.root_dir = root_dir
|
47 |
+
self.scene_list = sorted(next(os.walk(root_dir))[1])
|
48 |
+
self.resolution = resolution
|
49 |
+
self.first_K = first_K
|
50 |
+
self.load_target = load_target
|
51 |
+
|
52 |
+
def __len__(self):
|
53 |
+
return len(self.scene_list)
|
54 |
+
|
55 |
+
def __getitem__(self, idx):
|
56 |
+
scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
|
57 |
+
with open(os.path.join(scene_dir, 'transforms.json'), "r") as f:
|
58 |
+
meta = json.load(f)
|
59 |
+
imgs = []
|
60 |
+
poses = []
|
61 |
+
for i_img in range(self.first_K):
|
62 |
+
meta_img = meta['frames'][i_img]
|
63 |
+
|
64 |
+
if i_img == 0 or self.load_target:
|
65 |
+
img_path = os.path.join(scene_dir, meta_img['file_path'])
|
66 |
+
img = imageio.imread(img_path)
|
67 |
+
img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
|
68 |
+
imgs.append(img)
|
69 |
+
|
70 |
+
c2w = meta_img['transform_matrix']
|
71 |
+
poses.append(c2w)
|
72 |
+
|
73 |
+
imgs = (np.array(imgs) / 255.).astype(np.float32) # (RGBA) imgs
|
74 |
+
imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
|
75 |
+
imgs = imgs * 2 - 1. # convert to stable diffusion range
|
76 |
+
poses = torch.tensor(np.array(poses).astype(np.float32))
|
77 |
+
return imgs, poses
|
78 |
+
|
79 |
+
def blend_rgba(self, img):
|
80 |
+
img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:]) # blend A to RGB
|
81 |
+
return img
|
82 |
+
|
83 |
+
|
84 |
+
class GSO(Dataset):
|
85 |
+
def __init__(self, root_dir='datasets/GoogleScannedObjects',\
|
86 |
+
split='val', first_K=5, resolution=256, load_target=False, name='render_mvs'):
|
87 |
+
self.root_dir = root_dir
|
88 |
+
with open(os.path.join(root_dir, '%s.json' % split), "r") as f:
|
89 |
+
self.scene_list = json.load(f)
|
90 |
+
self.resolution = resolution
|
91 |
+
self.first_K = first_K
|
92 |
+
self.load_target = load_target
|
93 |
+
self.name = name
|
94 |
+
|
95 |
+
def __len__(self):
|
96 |
+
return len(self.scene_list)
|
97 |
+
|
98 |
+
def __getitem__(self, idx):
|
99 |
+
scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
|
100 |
+
with open(os.path.join(scene_dir, 'transforms_%s.json' % self.name), "r") as f:
|
101 |
+
meta = json.load(f)
|
102 |
+
imgs = []
|
103 |
+
poses = []
|
104 |
+
for i_img in range(self.first_K):
|
105 |
+
meta_img = meta['frames'][i_img]
|
106 |
+
|
107 |
+
if i_img == 0 or self.load_target:
|
108 |
+
img_path = os.path.join(scene_dir, meta_img['file_path'])
|
109 |
+
img = imageio.imread(img_path)
|
110 |
+
img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
|
111 |
+
imgs.append(img)
|
112 |
+
|
113 |
+
c2w = meta_img['transform_matrix']
|
114 |
+
poses.append(c2w)
|
115 |
+
|
116 |
+
imgs = (np.array(imgs) / 255.).astype(np.float32) # (RGBA) imgs
|
117 |
+
mask = imgs[:, :, :, -1]
|
118 |
+
imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
|
119 |
+
imgs = imgs * 2 - 1. # convert to stable diffusion range
|
120 |
+
poses = torch.tensor(np.array(poses).astype(np.float32))
|
121 |
+
return imgs, poses
|
122 |
+
|
123 |
+
def blend_rgba(self, img):
|
124 |
+
img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:]) # blend A to RGB
|
125 |
+
return img
|
126 |
+
|
127 |
+
class WILD(Dataset):
|
128 |
+
def __init__(self, root_dir='data/nerf_wild',\
|
129 |
+
first_K=33, resolution=256, load_target=False):
|
130 |
+
self.root_dir = root_dir
|
131 |
+
self.scene_list = sorted(next(os.walk(root_dir))[1])
|
132 |
+
self.resolution = resolution
|
133 |
+
self.first_K = first_K
|
134 |
+
self.load_target = load_target
|
135 |
+
|
136 |
+
def __len__(self):
|
137 |
+
return len(self.scene_list)
|
138 |
+
|
139 |
+
def __getitem__(self, idx):
|
140 |
+
scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
|
141 |
+
with open(os.path.join(scene_dir, 'transforms_train.json'), "r") as f:
|
142 |
+
meta = json.load(f)
|
143 |
+
imgs = []
|
144 |
+
poses = []
|
145 |
+
for i_img in range(self.first_K):
|
146 |
+
meta_img = meta['frames'][i_img]
|
147 |
+
|
148 |
+
if i_img == 0 or self.load_target:
|
149 |
+
img_path = os.path.join(scene_dir, meta_img['file_path'])
|
150 |
+
img = imageio.imread(img_path + '.png')
|
151 |
+
img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
|
152 |
+
imgs.append(img)
|
153 |
+
|
154 |
+
c2w = meta_img['transform_matrix']
|
155 |
+
poses.append(c2w)
|
156 |
+
|
157 |
+
imgs = (np.array(imgs) / 255.).astype(np.float32) # (RGBA) imgs
|
158 |
+
imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
|
159 |
+
imgs = imgs * 2 - 1. # convert to stable diffusion range
|
160 |
+
poses = torch.tensor(np.array(poses).astype(np.float32))
|
161 |
+
return imgs, poses
|
162 |
+
|
163 |
+
def blend_rgba(self, img):
|
164 |
+
img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:]) # blend A to RGB
|
165 |
+
return img
|
ldm/data/simple.py
ADDED
@@ -0,0 +1,526 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Dict
|
2 |
+
import webdataset as wds
|
3 |
+
import numpy as np
|
4 |
+
from omegaconf import DictConfig, ListConfig
|
5 |
+
import torch
|
6 |
+
from torch.utils.data import Dataset
|
7 |
+
from pathlib import Path
|
8 |
+
import json
|
9 |
+
from PIL import Image
|
10 |
+
from torchvision import transforms
|
11 |
+
import torchvision
|
12 |
+
from einops import rearrange
|
13 |
+
from ldm.util import instantiate_from_config
|
14 |
+
from datasets import load_dataset
|
15 |
+
import pytorch_lightning as pl
|
16 |
+
import copy
|
17 |
+
import csv
|
18 |
+
import cv2
|
19 |
+
import random
|
20 |
+
import matplotlib.pyplot as plt
|
21 |
+
from torch.utils.data import DataLoader
|
22 |
+
import json
|
23 |
+
import os
|
24 |
+
import webdataset as wds
|
25 |
+
import math
|
26 |
+
from torch.utils.data.distributed import DistributedSampler
|
27 |
+
|
28 |
+
# Some hacky things to make experimentation easier
|
29 |
+
def make_transform_multi_folder_data(paths, caption_files=None, **kwargs):
|
30 |
+
ds = make_multi_folder_data(paths, caption_files, **kwargs)
|
31 |
+
return TransformDataset(ds)
|
32 |
+
|
33 |
+
def make_nfp_data(base_path):
|
34 |
+
dirs = list(Path(base_path).glob("*/"))
|
35 |
+
print(f"Found {len(dirs)} folders")
|
36 |
+
print(dirs)
|
37 |
+
tforms = [transforms.Resize(512), transforms.CenterCrop(512)]
|
38 |
+
datasets = [NfpDataset(x, image_transforms=copy.copy(tforms), default_caption="A view from a train window") for x in dirs]
|
39 |
+
return torch.utils.data.ConcatDataset(datasets)
|
40 |
+
|
41 |
+
|
42 |
+
class VideoDataset(Dataset):
|
43 |
+
def __init__(self, root_dir, image_transforms, caption_file, offset=8, n=2):
|
44 |
+
self.root_dir = Path(root_dir)
|
45 |
+
self.caption_file = caption_file
|
46 |
+
self.n = n
|
47 |
+
ext = "mp4"
|
48 |
+
self.paths = sorted(list(self.root_dir.rglob(f"*.{ext}")))
|
49 |
+
self.offset = offset
|
50 |
+
|
51 |
+
if isinstance(image_transforms, ListConfig):
|
52 |
+
image_transforms = [instantiate_from_config(tt) for tt in image_transforms]
|
53 |
+
image_transforms.extend([transforms.ToTensor(),
|
54 |
+
transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
|
55 |
+
image_transforms = transforms.Compose(image_transforms)
|
56 |
+
self.tform = image_transforms
|
57 |
+
with open(self.caption_file) as f:
|
58 |
+
reader = csv.reader(f)
|
59 |
+
rows = [row for row in reader]
|
60 |
+
self.captions = dict(rows)
|
61 |
+
|
62 |
+
def __len__(self):
|
63 |
+
return len(self.paths)
|
64 |
+
|
65 |
+
def __getitem__(self, index):
|
66 |
+
for i in range(10):
|
67 |
+
try:
|
68 |
+
return self._load_sample(index)
|
69 |
+
except Exception:
|
70 |
+
# Not really good enough but...
|
71 |
+
print("uh oh")
|
72 |
+
|
73 |
+
def _load_sample(self, index):
|
74 |
+
n = self.n
|
75 |
+
filename = self.paths[index]
|
76 |
+
min_frame = 2*self.offset + 2
|
77 |
+
vid = cv2.VideoCapture(str(filename))
|
78 |
+
max_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
|
79 |
+
curr_frame_n = random.randint(min_frame, max_frames)
|
80 |
+
vid.set(cv2.CAP_PROP_POS_FRAMES,curr_frame_n)
|
81 |
+
_, curr_frame = vid.read()
|
82 |
+
|
83 |
+
prev_frames = []
|
84 |
+
for i in range(n):
|
85 |
+
prev_frame_n = curr_frame_n - (i+1)*self.offset
|
86 |
+
vid.set(cv2.CAP_PROP_POS_FRAMES,prev_frame_n)
|
87 |
+
_, prev_frame = vid.read()
|
88 |
+
prev_frame = self.tform(Image.fromarray(prev_frame[...,::-1]))
|
89 |
+
prev_frames.append(prev_frame)
|
90 |
+
|
91 |
+
vid.release()
|
92 |
+
caption = self.captions[filename.name]
|
93 |
+
data = {
|
94 |
+
"image": self.tform(Image.fromarray(curr_frame[...,::-1])),
|
95 |
+
"prev": torch.cat(prev_frames, dim=-1),
|
96 |
+
"txt": caption
|
97 |
+
}
|
98 |
+
return data
|
99 |
+
|
100 |
+
# end hacky things
|
101 |
+
|
102 |
+
|
103 |
+
def make_tranforms(image_transforms):
|
104 |
+
# if isinstance(image_transforms, ListConfig):
|
105 |
+
# image_transforms = [instantiate_from_config(tt) for tt in image_transforms]
|
106 |
+
image_transforms = []
|
107 |
+
image_transforms.extend([transforms.ToTensor(),
|
108 |
+
transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
|
109 |
+
image_transforms = transforms.Compose(image_transforms)
|
110 |
+
return image_transforms
|
111 |
+
|
112 |
+
|
113 |
+
def make_multi_folder_data(paths, caption_files=None, **kwargs):
|
114 |
+
"""Make a concat dataset from multiple folders
|
115 |
+
Don't suport captions yet
|
116 |
+
|
117 |
+
If paths is a list, that's ok, if it's a Dict interpret it as:
|
118 |
+
k=folder v=n_times to repeat that
|
119 |
+
"""
|
120 |
+
list_of_paths = []
|
121 |
+
if isinstance(paths, (Dict, DictConfig)):
|
122 |
+
assert caption_files is None, \
|
123 |
+
"Caption files not yet supported for repeats"
|
124 |
+
for folder_path, repeats in paths.items():
|
125 |
+
list_of_paths.extend([folder_path]*repeats)
|
126 |
+
paths = list_of_paths
|
127 |
+
|
128 |
+
if caption_files is not None:
|
129 |
+
datasets = [FolderData(p, caption_file=c, **kwargs) for (p, c) in zip(paths, caption_files)]
|
130 |
+
else:
|
131 |
+
datasets = [FolderData(p, **kwargs) for p in paths]
|
132 |
+
return torch.utils.data.ConcatDataset(datasets)
|
133 |
+
|
134 |
+
|
135 |
+
|
136 |
+
class NfpDataset(Dataset):
|
137 |
+
def __init__(self,
|
138 |
+
root_dir,
|
139 |
+
image_transforms=[],
|
140 |
+
ext="jpg",
|
141 |
+
default_caption="",
|
142 |
+
) -> None:
|
143 |
+
"""assume sequential frames and a deterministic transform"""
|
144 |
+
|
145 |
+
self.root_dir = Path(root_dir)
|
146 |
+
self.default_caption = default_caption
|
147 |
+
|
148 |
+
self.paths = sorted(list(self.root_dir.rglob(f"*.{ext}")))
|
149 |
+
self.tform = make_tranforms(image_transforms)
|
150 |
+
|
151 |
+
def __len__(self):
|
152 |
+
return len(self.paths) - 1
|
153 |
+
|
154 |
+
|
155 |
+
def __getitem__(self, index):
|
156 |
+
prev = self.paths[index]
|
157 |
+
curr = self.paths[index+1]
|
158 |
+
data = {}
|
159 |
+
data["image"] = self._load_im(curr)
|
160 |
+
data["prev"] = self._load_im(prev)
|
161 |
+
data["txt"] = self.default_caption
|
162 |
+
return data
|
163 |
+
|
164 |
+
def _load_im(self, filename):
|
165 |
+
im = Image.open(filename).convert("RGB")
|
166 |
+
return self.tform(im)
|
167 |
+
|
168 |
+
class ObjaverseDataModuleFromConfig(pl.LightningDataModule):
|
169 |
+
def __init__(self, root_dir, batch_size, total_view, train=None, validation=None,
|
170 |
+
test=None, num_workers=4, **kwargs):
|
171 |
+
super().__init__(self)
|
172 |
+
self.root_dir = root_dir
|
173 |
+
self.batch_size = batch_size
|
174 |
+
self.num_workers = num_workers
|
175 |
+
self.total_view = total_view
|
176 |
+
|
177 |
+
if train is not None:
|
178 |
+
dataset_config = train
|
179 |
+
if validation is not None:
|
180 |
+
dataset_config = validation
|
181 |
+
|
182 |
+
if 'image_transforms' in dataset_config:
|
183 |
+
image_transforms = [torchvision.transforms.Resize(dataset_config.image_transforms.size)]
|
184 |
+
else:
|
185 |
+
image_transforms = []
|
186 |
+
image_transforms.extend([transforms.ToTensor(),
|
187 |
+
transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
|
188 |
+
self.image_transforms = torchvision.transforms.Compose(image_transforms)
|
189 |
+
|
190 |
+
|
191 |
+
def train_dataloader(self):
|
192 |
+
dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=False, \
|
193 |
+
image_transforms=self.image_transforms)
|
194 |
+
sampler = DistributedSampler(dataset)
|
195 |
+
return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, sampler=sampler)
|
196 |
+
|
197 |
+
def val_dataloader(self):
|
198 |
+
dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=True, \
|
199 |
+
image_transforms=self.image_transforms)
|
200 |
+
sampler = DistributedSampler(dataset)
|
201 |
+
return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
|
202 |
+
|
203 |
+
def test_dataloader(self):
|
204 |
+
return wds.WebLoader(ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=self.validation),\
|
205 |
+
batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
|
206 |
+
|
207 |
+
|
208 |
+
class ObjaverseData(Dataset):
|
209 |
+
def __init__(self,
|
210 |
+
root_dir='.objaverse/hf-objaverse-v1/views',
|
211 |
+
image_transforms=[],
|
212 |
+
ext="png",
|
213 |
+
default_trans=torch.zeros(3),
|
214 |
+
postprocess=None,
|
215 |
+
return_paths=False,
|
216 |
+
total_view=4,
|
217 |
+
validation=False
|
218 |
+
) -> None:
|
219 |
+
"""Create a dataset from a folder of images.
|
220 |
+
If you pass in a root directory it will be searched for images
|
221 |
+
ending in ext (ext can be a list)
|
222 |
+
"""
|
223 |
+
self.root_dir = Path(root_dir)
|
224 |
+
self.default_trans = default_trans
|
225 |
+
self.return_paths = return_paths
|
226 |
+
if isinstance(postprocess, DictConfig):
|
227 |
+
postprocess = instantiate_from_config(postprocess)
|
228 |
+
self.postprocess = postprocess
|
229 |
+
self.total_view = total_view
|
230 |
+
|
231 |
+
if not isinstance(ext, (tuple, list, ListConfig)):
|
232 |
+
ext = [ext]
|
233 |
+
|
234 |
+
with open(os.path.join(root_dir, 'valid_paths.json')) as f:
|
235 |
+
self.paths = json.load(f)
|
236 |
+
|
237 |
+
total_objects = len(self.paths)
|
238 |
+
if validation:
|
239 |
+
self.paths = self.paths[math.floor(total_objects / 100. * 99.):] # used last 1% as validation
|
240 |
+
else:
|
241 |
+
self.paths = self.paths[:math.floor(total_objects / 100. * 99.)] # used first 99% as training
|
242 |
+
print('============= length of dataset %d =============' % len(self.paths))
|
243 |
+
self.tform = image_transforms
|
244 |
+
|
245 |
+
def __len__(self):
|
246 |
+
return len(self.paths)
|
247 |
+
|
248 |
+
def cartesian_to_spherical(self, xyz):
|
249 |
+
ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
|
250 |
+
xy = xyz[:,0]**2 + xyz[:,1]**2
|
251 |
+
z = np.sqrt(xy + xyz[:,2]**2)
|
252 |
+
theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
|
253 |
+
#ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
|
254 |
+
azimuth = np.arctan2(xyz[:,1], xyz[:,0])
|
255 |
+
return np.array([theta, azimuth, z])
|
256 |
+
|
257 |
+
def get_T(self, target_RT, cond_RT):
|
258 |
+
R, T = target_RT[:3, :3], target_RT[:, -1]
|
259 |
+
T_target = -R.T @ T
|
260 |
+
|
261 |
+
R, T = cond_RT[:3, :3], cond_RT[:, -1]
|
262 |
+
T_cond = -R.T @ T
|
263 |
+
|
264 |
+
theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
|
265 |
+
theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
|
266 |
+
|
267 |
+
d_theta = theta_target - theta_cond
|
268 |
+
d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
|
269 |
+
d_z = z_target - z_cond
|
270 |
+
|
271 |
+
d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
|
272 |
+
return d_T
|
273 |
+
|
274 |
+
def load_im(self, path, color):
|
275 |
+
'''
|
276 |
+
replace background pixel with random color in rendering
|
277 |
+
'''
|
278 |
+
try:
|
279 |
+
img = plt.imread(path)
|
280 |
+
except:
|
281 |
+
print(path)
|
282 |
+
sys.exit()
|
283 |
+
img[img[:, :, -1] == 0.] = color
|
284 |
+
img = Image.fromarray(np.uint8(img[:, :, :3] * 255.))
|
285 |
+
return img
|
286 |
+
|
287 |
+
def __getitem__(self, index):
|
288 |
+
|
289 |
+
data = {}
|
290 |
+
if self.paths[index][-2:] == '_1': # dirty fix for rendering dataset twice
|
291 |
+
total_view = 8
|
292 |
+
else:
|
293 |
+
total_view = 4
|
294 |
+
index_target, index_cond = random.sample(range(total_view), 2) # without replacement
|
295 |
+
filename = os.path.join(self.root_dir, self.paths[index])
|
296 |
+
|
297 |
+
# print(self.paths[index])
|
298 |
+
|
299 |
+
if self.return_paths:
|
300 |
+
data["path"] = str(filename)
|
301 |
+
|
302 |
+
color = [1., 1., 1., 1.]
|
303 |
+
|
304 |
+
try:
|
305 |
+
target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
|
306 |
+
cond_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_cond), color))
|
307 |
+
target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
|
308 |
+
cond_RT = np.load(os.path.join(filename, '%03d.npy' % index_cond))
|
309 |
+
except:
|
310 |
+
# very hacky solution, sorry about this
|
311 |
+
filename = os.path.join(self.root_dir, '692db5f2d3a04bb286cb977a7dba903e_1') # this one we know is valid
|
312 |
+
target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
|
313 |
+
cond_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_cond), color))
|
314 |
+
target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
|
315 |
+
cond_RT = np.load(os.path.join(filename, '%03d.npy' % index_cond))
|
316 |
+
target_im = torch.zeros_like(target_im)
|
317 |
+
cond_im = torch.zeros_like(cond_im)
|
318 |
+
|
319 |
+
data["image_target"] = target_im
|
320 |
+
data["image_cond"] = cond_im
|
321 |
+
data["T"] = self.get_T(target_RT, cond_RT)
|
322 |
+
|
323 |
+
if self.postprocess is not None:
|
324 |
+
data = self.postprocess(data)
|
325 |
+
|
326 |
+
return data
|
327 |
+
|
328 |
+
def process_im(self, im):
|
329 |
+
im = im.convert("RGB")
|
330 |
+
return self.tform(im)
|
331 |
+
|
332 |
+
class FolderData(Dataset):
|
333 |
+
def __init__(self,
|
334 |
+
root_dir,
|
335 |
+
caption_file=None,
|
336 |
+
image_transforms=[],
|
337 |
+
ext="jpg",
|
338 |
+
default_caption="",
|
339 |
+
postprocess=None,
|
340 |
+
return_paths=False,
|
341 |
+
) -> None:
|
342 |
+
"""Create a dataset from a folder of images.
|
343 |
+
If you pass in a root directory it will be searched for images
|
344 |
+
ending in ext (ext can be a list)
|
345 |
+
"""
|
346 |
+
self.root_dir = Path(root_dir)
|
347 |
+
self.default_caption = default_caption
|
348 |
+
self.return_paths = return_paths
|
349 |
+
if isinstance(postprocess, DictConfig):
|
350 |
+
postprocess = instantiate_from_config(postprocess)
|
351 |
+
self.postprocess = postprocess
|
352 |
+
if caption_file is not None:
|
353 |
+
with open(caption_file, "rt") as f:
|
354 |
+
ext = Path(caption_file).suffix.lower()
|
355 |
+
if ext == ".json":
|
356 |
+
captions = json.load(f)
|
357 |
+
elif ext == ".jsonl":
|
358 |
+
lines = f.readlines()
|
359 |
+
lines = [json.loads(x) for x in lines]
|
360 |
+
captions = {x["file_name"]: x["text"].strip("\n") for x in lines}
|
361 |
+
else:
|
362 |
+
raise ValueError(f"Unrecognised format: {ext}")
|
363 |
+
self.captions = captions
|
364 |
+
else:
|
365 |
+
self.captions = None
|
366 |
+
|
367 |
+
if not isinstance(ext, (tuple, list, ListConfig)):
|
368 |
+
ext = [ext]
|
369 |
+
|
370 |
+
# Only used if there is no caption file
|
371 |
+
self.paths = []
|
372 |
+
for e in ext:
|
373 |
+
self.paths.extend(sorted(list(self.root_dir.rglob(f"*.{e}"))))
|
374 |
+
self.tform = make_tranforms(image_transforms)
|
375 |
+
|
376 |
+
def __len__(self):
|
377 |
+
if self.captions is not None:
|
378 |
+
return len(self.captions.keys())
|
379 |
+
else:
|
380 |
+
return len(self.paths)
|
381 |
+
|
382 |
+
def __getitem__(self, index):
|
383 |
+
data = {}
|
384 |
+
if self.captions is not None:
|
385 |
+
chosen = list(self.captions.keys())[index]
|
386 |
+
caption = self.captions.get(chosen, None)
|
387 |
+
if caption is None:
|
388 |
+
caption = self.default_caption
|
389 |
+
filename = self.root_dir/chosen
|
390 |
+
else:
|
391 |
+
filename = self.paths[index]
|
392 |
+
|
393 |
+
if self.return_paths:
|
394 |
+
data["path"] = str(filename)
|
395 |
+
|
396 |
+
im = Image.open(filename).convert("RGB")
|
397 |
+
im = self.process_im(im)
|
398 |
+
data["image"] = im
|
399 |
+
|
400 |
+
if self.captions is not None:
|
401 |
+
data["txt"] = caption
|
402 |
+
else:
|
403 |
+
data["txt"] = self.default_caption
|
404 |
+
|
405 |
+
if self.postprocess is not None:
|
406 |
+
data = self.postprocess(data)
|
407 |
+
|
408 |
+
return data
|
409 |
+
|
410 |
+
def process_im(self, im):
|
411 |
+
im = im.convert("RGB")
|
412 |
+
return self.tform(im)
|
413 |
+
import random
|
414 |
+
|
415 |
+
class TransformDataset():
|
416 |
+
def __init__(self, ds, extra_label="sksbspic"):
|
417 |
+
self.ds = ds
|
418 |
+
self.extra_label = extra_label
|
419 |
+
self.transforms = {
|
420 |
+
"align": transforms.Resize(768),
|
421 |
+
"centerzoom": transforms.CenterCrop(768),
|
422 |
+
"randzoom": transforms.RandomCrop(768),
|
423 |
+
}
|
424 |
+
|
425 |
+
|
426 |
+
def __getitem__(self, index):
|
427 |
+
data = self.ds[index]
|
428 |
+
|
429 |
+
im = data['image']
|
430 |
+
im = im.permute(2,0,1)
|
431 |
+
# In case data is smaller than expected
|
432 |
+
im = transforms.Resize(1024)(im)
|
433 |
+
|
434 |
+
tform_name = random.choice(list(self.transforms.keys()))
|
435 |
+
im = self.transforms[tform_name](im)
|
436 |
+
|
437 |
+
im = im.permute(1,2,0)
|
438 |
+
|
439 |
+
data['image'] = im
|
440 |
+
data['txt'] = data['txt'] + f" {self.extra_label} {tform_name}"
|
441 |
+
|
442 |
+
return data
|
443 |
+
|
444 |
+
def __len__(self):
|
445 |
+
return len(self.ds)
|
446 |
+
|
447 |
+
def hf_dataset(
|
448 |
+
name,
|
449 |
+
image_transforms=[],
|
450 |
+
image_column="image",
|
451 |
+
text_column="text",
|
452 |
+
split='train',
|
453 |
+
image_key='image',
|
454 |
+
caption_key='txt',
|
455 |
+
):
|
456 |
+
"""Make huggingface dataset with appropriate list of transforms applied
|
457 |
+
"""
|
458 |
+
ds = load_dataset(name, split=split)
|
459 |
+
tform = make_tranforms(image_transforms)
|
460 |
+
|
461 |
+
assert image_column in ds.column_names, f"Didn't find column {image_column} in {ds.column_names}"
|
462 |
+
assert text_column in ds.column_names, f"Didn't find column {text_column} in {ds.column_names}"
|
463 |
+
|
464 |
+
def pre_process(examples):
|
465 |
+
processed = {}
|
466 |
+
processed[image_key] = [tform(im) for im in examples[image_column]]
|
467 |
+
processed[caption_key] = examples[text_column]
|
468 |
+
return processed
|
469 |
+
|
470 |
+
ds.set_transform(pre_process)
|
471 |
+
return ds
|
472 |
+
|
473 |
+
class TextOnly(Dataset):
|
474 |
+
def __init__(self, captions, output_size, image_key="image", caption_key="txt", n_gpus=1):
|
475 |
+
"""Returns only captions with dummy images"""
|
476 |
+
self.output_size = output_size
|
477 |
+
self.image_key = image_key
|
478 |
+
self.caption_key = caption_key
|
479 |
+
if isinstance(captions, Path):
|
480 |
+
self.captions = self._load_caption_file(captions)
|
481 |
+
else:
|
482 |
+
self.captions = captions
|
483 |
+
|
484 |
+
if n_gpus > 1:
|
485 |
+
# hack to make sure that all the captions appear on each gpu
|
486 |
+
repeated = [n_gpus*[x] for x in self.captions]
|
487 |
+
self.captions = []
|
488 |
+
[self.captions.extend(x) for x in repeated]
|
489 |
+
|
490 |
+
def __len__(self):
|
491 |
+
return len(self.captions)
|
492 |
+
|
493 |
+
def __getitem__(self, index):
|
494 |
+
dummy_im = torch.zeros(3, self.output_size, self.output_size)
|
495 |
+
dummy_im = rearrange(dummy_im * 2. - 1., 'c h w -> h w c')
|
496 |
+
return {self.image_key: dummy_im, self.caption_key: self.captions[index]}
|
497 |
+
|
498 |
+
def _load_caption_file(self, filename):
|
499 |
+
with open(filename, 'rt') as f:
|
500 |
+
captions = f.readlines()
|
501 |
+
return [x.strip('\n') for x in captions]
|
502 |
+
|
503 |
+
|
504 |
+
|
505 |
+
import random
|
506 |
+
import json
|
507 |
+
class IdRetreivalDataset(FolderData):
|
508 |
+
def __init__(self, ret_file, *args, **kwargs):
|
509 |
+
super().__init__(*args, **kwargs)
|
510 |
+
with open(ret_file, "rt") as f:
|
511 |
+
self.ret = json.load(f)
|
512 |
+
|
513 |
+
def __getitem__(self, index):
|
514 |
+
data = super().__getitem__(index)
|
515 |
+
key = self.paths[index].name
|
516 |
+
matches = self.ret[key]
|
517 |
+
if len(matches) > 0:
|
518 |
+
retreived = random.choice(matches)
|
519 |
+
else:
|
520 |
+
retreived = key
|
521 |
+
filename = self.root_dir/retreived
|
522 |
+
im = Image.open(filename).convert("RGB")
|
523 |
+
im = self.process_im(im)
|
524 |
+
# data["match"] = im
|
525 |
+
data["match"] = torch.cat((data["image"], im), dim=-1)
|
526 |
+
return data
|
ldm/data/sync_dreamer.py
ADDED
@@ -0,0 +1,132 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import pytorch_lightning as pl
|
2 |
+
import numpy as np
|
3 |
+
import torch
|
4 |
+
import PIL
|
5 |
+
import os
|
6 |
+
from skimage.io import imread
|
7 |
+
import webdataset as wds
|
8 |
+
import PIL.Image as Image
|
9 |
+
from torch.utils.data import Dataset
|
10 |
+
from torch.utils.data.distributed import DistributedSampler
|
11 |
+
from pathlib import Path
|
12 |
+
|
13 |
+
from ldm.base_utils import read_pickle, pose_inverse
|
14 |
+
import torchvision.transforms as transforms
|
15 |
+
import torchvision
|
16 |
+
from einops import rearrange
|
17 |
+
|
18 |
+
from ldm.util import prepare_inputs
|
19 |
+
|
20 |
+
|
21 |
+
class SyncDreamerTrainData(Dataset):
|
22 |
+
def __init__(self, target_dir, input_dir, uid_set_pkl, image_size=256):
|
23 |
+
self.default_image_size = 256
|
24 |
+
self.image_size = image_size
|
25 |
+
self.target_dir = Path(target_dir)
|
26 |
+
self.input_dir = Path(input_dir)
|
27 |
+
|
28 |
+
self.uids = read_pickle(uid_set_pkl)
|
29 |
+
print('============= length of dataset %d =============' % len(self.uids))
|
30 |
+
|
31 |
+
image_transforms = []
|
32 |
+
image_transforms.extend([transforms.ToTensor(), transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
|
33 |
+
self.image_transforms = torchvision.transforms.Compose(image_transforms)
|
34 |
+
self.num_images = 16
|
35 |
+
|
36 |
+
def __len__(self):
|
37 |
+
return len(self.uids)
|
38 |
+
|
39 |
+
def load_im(self, path):
|
40 |
+
img = imread(path)
|
41 |
+
img = img.astype(np.float32) / 255.0
|
42 |
+
mask = img[:,:,3:]
|
43 |
+
img[:,:,:3] = img[:,:,:3] * mask + 1 - mask # white background
|
44 |
+
img = Image.fromarray(np.uint8(img[:, :, :3] * 255.))
|
45 |
+
return img, mask
|
46 |
+
|
47 |
+
def process_im(self, im):
|
48 |
+
im = im.convert("RGB")
|
49 |
+
im = im.resize((self.image_size, self.image_size), resample=PIL.Image.BICUBIC)
|
50 |
+
return self.image_transforms(im)
|
51 |
+
|
52 |
+
def load_index(self, filename, index):
|
53 |
+
img, _ = self.load_im(os.path.join(filename, '%03d.png' % index))
|
54 |
+
img = self.process_im(img)
|
55 |
+
return img
|
56 |
+
|
57 |
+
def get_data_for_index(self, index):
|
58 |
+
target_dir = os.path.join(self.target_dir, self.uids[index])
|
59 |
+
input_dir = os.path.join(self.input_dir, self.uids[index])
|
60 |
+
|
61 |
+
views = np.arange(0, self.num_images)
|
62 |
+
start_view_index = np.random.randint(0, self.num_images)
|
63 |
+
views = (views + start_view_index) % self.num_images
|
64 |
+
|
65 |
+
target_images = []
|
66 |
+
for si, target_index in enumerate(views):
|
67 |
+
img = self.load_index(target_dir, target_index)
|
68 |
+
target_images.append(img)
|
69 |
+
target_images = torch.stack(target_images, 0)
|
70 |
+
input_img = self.load_index(input_dir, start_view_index)
|
71 |
+
|
72 |
+
K, azimuths, elevations, distances, cam_poses = read_pickle(os.path.join(input_dir, f'meta.pkl'))
|
73 |
+
input_elevation = torch.from_numpy(elevations[start_view_index:start_view_index+1].astype(np.float32))
|
74 |
+
return {"target_image": target_images, "input_image": input_img, "input_elevation": input_elevation}
|
75 |
+
|
76 |
+
def __getitem__(self, index):
|
77 |
+
data = self.get_data_for_index(index)
|
78 |
+
return data
|
79 |
+
|
80 |
+
class SyncDreamerEvalData(Dataset):
|
81 |
+
def __init__(self, image_dir):
|
82 |
+
self.image_size = 256
|
83 |
+
self.image_dir = Path(image_dir)
|
84 |
+
self.crop_size = 20
|
85 |
+
|
86 |
+
self.fns = []
|
87 |
+
for fn in Path(image_dir).iterdir():
|
88 |
+
if fn.suffix=='.png':
|
89 |
+
self.fns.append(fn)
|
90 |
+
print('============= length of dataset %d =============' % len(self.fns))
|
91 |
+
|
92 |
+
def __len__(self):
|
93 |
+
return len(self.fns)
|
94 |
+
|
95 |
+
def get_data_for_index(self, index):
|
96 |
+
input_img_fn = self.fns[index]
|
97 |
+
elevation = int(Path(input_img_fn).stem.split('-')[-1])
|
98 |
+
return prepare_inputs(input_img_fn, elevation, 200)
|
99 |
+
|
100 |
+
def __getitem__(self, index):
|
101 |
+
return self.get_data_for_index(index)
|
102 |
+
|
103 |
+
class SyncDreamerDataset(pl.LightningDataModule):
|
104 |
+
def __init__(self, target_dir, input_dir, validation_dir, batch_size, uid_set_pkl, image_size=256, num_workers=4, seed=0, **kwargs):
|
105 |
+
super().__init__()
|
106 |
+
self.target_dir = target_dir
|
107 |
+
self.input_dir = input_dir
|
108 |
+
self.validation_dir = validation_dir
|
109 |
+
self.batch_size = batch_size
|
110 |
+
self.num_workers = num_workers
|
111 |
+
self.uid_set_pkl = uid_set_pkl
|
112 |
+
self.seed = seed
|
113 |
+
self.additional_args = kwargs
|
114 |
+
self.image_size = image_size
|
115 |
+
|
116 |
+
def setup(self, stage):
|
117 |
+
if stage in ['fit']:
|
118 |
+
self.train_dataset = SyncDreamerTrainData(self.target_dir, self.input_dir, uid_set_pkl=self.uid_set_pkl, image_size=256)
|
119 |
+
self.val_dataset = SyncDreamerEvalData(image_dir=self.validation_dir)
|
120 |
+
else:
|
121 |
+
raise NotImplementedError
|
122 |
+
|
123 |
+
def train_dataloader(self):
|
124 |
+
sampler = DistributedSampler(self.train_dataset, seed=self.seed)
|
125 |
+
return wds.WebLoader(self.train_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, sampler=sampler)
|
126 |
+
|
127 |
+
def val_dataloader(self):
|
128 |
+
loader = wds.WebLoader(self.val_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
|
129 |
+
return loader
|
130 |
+
|
131 |
+
def test_dataloader(self):
|
132 |
+
return wds.WebLoader(self.val_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
|
ldm/lr_scheduler.py
ADDED
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
|
3 |
+
|
4 |
+
class LambdaWarmUpCosineScheduler:
|
5 |
+
"""
|
6 |
+
note: use with a base_lr of 1.0
|
7 |
+
"""
|
8 |
+
def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
|
9 |
+
self.lr_warm_up_steps = warm_up_steps
|
10 |
+
self.lr_start = lr_start
|
11 |
+
self.lr_min = lr_min
|
12 |
+
self.lr_max = lr_max
|
13 |
+
self.lr_max_decay_steps = max_decay_steps
|
14 |
+
self.last_lr = 0.
|
15 |
+
self.verbosity_interval = verbosity_interval
|
16 |
+
|
17 |
+
def schedule(self, n, **kwargs):
|
18 |
+
if self.verbosity_interval > 0:
|
19 |
+
if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
|
20 |
+
if n < self.lr_warm_up_steps:
|
21 |
+
lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
|
22 |
+
self.last_lr = lr
|
23 |
+
return lr
|
24 |
+
else:
|
25 |
+
t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
|
26 |
+
t = min(t, 1.0)
|
27 |
+
lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
|
28 |
+
1 + np.cos(t * np.pi))
|
29 |
+
self.last_lr = lr
|
30 |
+
return lr
|
31 |
+
|
32 |
+
def __call__(self, n, **kwargs):
|
33 |
+
return self.schedule(n,**kwargs)
|
34 |
+
|
35 |
+
|
36 |
+
class LambdaWarmUpCosineScheduler2:
|
37 |
+
"""
|
38 |
+
supports repeated iterations, configurable via lists
|
39 |
+
note: use with a base_lr of 1.0.
|
40 |
+
"""
|
41 |
+
def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
|
42 |
+
assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
|
43 |
+
self.lr_warm_up_steps = warm_up_steps
|
44 |
+
self.f_start = f_start
|
45 |
+
self.f_min = f_min
|
46 |
+
self.f_max = f_max
|
47 |
+
self.cycle_lengths = cycle_lengths
|
48 |
+
self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
|
49 |
+
self.last_f = 0.
|
50 |
+
self.verbosity_interval = verbosity_interval
|
51 |
+
|
52 |
+
def find_in_interval(self, n):
|
53 |
+
interval = 0
|
54 |
+
for cl in self.cum_cycles[1:]:
|
55 |
+
if n <= cl:
|
56 |
+
return interval
|
57 |
+
interval += 1
|
58 |
+
|
59 |
+
def schedule(self, n, **kwargs):
|
60 |
+
cycle = self.find_in_interval(n)
|
61 |
+
n = n - self.cum_cycles[cycle]
|
62 |
+
if self.verbosity_interval > 0:
|
63 |
+
if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
|
64 |
+
f"current cycle {cycle}")
|
65 |
+
if n < self.lr_warm_up_steps[cycle]:
|
66 |
+
f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
|
67 |
+
self.last_f = f
|
68 |
+
return f
|
69 |
+
else:
|
70 |
+
t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
|
71 |
+
t = min(t, 1.0)
|
72 |
+
f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
|
73 |
+
1 + np.cos(t * np.pi))
|
74 |
+
self.last_f = f
|
75 |
+
return f
|
76 |
+
|
77 |
+
def __call__(self, n, **kwargs):
|
78 |
+
return self.schedule(n, **kwargs)
|
79 |
+
|
80 |
+
|
81 |
+
class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
|
82 |
+
|
83 |
+
def schedule(self, n, **kwargs):
|
84 |
+
cycle = self.find_in_interval(n)
|
85 |
+
n = n - self.cum_cycles[cycle]
|
86 |
+
if self.verbosity_interval > 0:
|
87 |
+
if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
|
88 |
+
f"current cycle {cycle}")
|
89 |
+
|
90 |
+
if n < self.lr_warm_up_steps[cycle]:
|
91 |
+
f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
|
92 |
+
self.last_f = f
|
93 |
+
return f
|
94 |
+
else:
|
95 |
+
f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle])
|
96 |
+
self.last_f = f
|
97 |
+
return f
|
98 |
+
|
ldm/models/autoencoder.py
ADDED
@@ -0,0 +1,443 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import pytorch_lightning as pl
|
3 |
+
import torch.nn.functional as F
|
4 |
+
from contextlib import contextmanager
|
5 |
+
|
6 |
+
from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
|
7 |
+
|
8 |
+
from ldm.modules.diffusionmodules.model import Encoder, Decoder
|
9 |
+
from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
|
10 |
+
|
11 |
+
from ldm.util import instantiate_from_config
|
12 |
+
|
13 |
+
|
14 |
+
class VQModel(pl.LightningModule):
|
15 |
+
def __init__(self,
|
16 |
+
ddconfig,
|
17 |
+
lossconfig,
|
18 |
+
n_embed,
|
19 |
+
embed_dim,
|
20 |
+
ckpt_path=None,
|
21 |
+
ignore_keys=[],
|
22 |
+
image_key="image",
|
23 |
+
colorize_nlabels=None,
|
24 |
+
monitor=None,
|
25 |
+
batch_resize_range=None,
|
26 |
+
scheduler_config=None,
|
27 |
+
lr_g_factor=1.0,
|
28 |
+
remap=None,
|
29 |
+
sane_index_shape=False, # tell vector quantizer to return indices as bhw
|
30 |
+
use_ema=False
|
31 |
+
):
|
32 |
+
super().__init__()
|
33 |
+
self.embed_dim = embed_dim
|
34 |
+
self.n_embed = n_embed
|
35 |
+
self.image_key = image_key
|
36 |
+
self.encoder = Encoder(**ddconfig)
|
37 |
+
self.decoder = Decoder(**ddconfig)
|
38 |
+
self.loss = instantiate_from_config(lossconfig)
|
39 |
+
self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
|
40 |
+
remap=remap,
|
41 |
+
sane_index_shape=sane_index_shape)
|
42 |
+
self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
|
43 |
+
self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
|
44 |
+
if colorize_nlabels is not None:
|
45 |
+
assert type(colorize_nlabels)==int
|
46 |
+
self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
|
47 |
+
if monitor is not None:
|
48 |
+
self.monitor = monitor
|
49 |
+
self.batch_resize_range = batch_resize_range
|
50 |
+
if self.batch_resize_range is not None:
|
51 |
+
print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.")
|
52 |
+
|
53 |
+
self.use_ema = use_ema
|
54 |
+
if self.use_ema:
|
55 |
+
self.model_ema = LitEma(self)
|
56 |
+
print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
|
57 |
+
|
58 |
+
if ckpt_path is not None:
|
59 |
+
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
|
60 |
+
self.scheduler_config = scheduler_config
|
61 |
+
self.lr_g_factor = lr_g_factor
|
62 |
+
|
63 |
+
@contextmanager
|
64 |
+
def ema_scope(self, context=None):
|
65 |
+
if self.use_ema:
|
66 |
+
self.model_ema.store(self.parameters())
|
67 |
+
self.model_ema.copy_to(self)
|
68 |
+
if context is not None:
|
69 |
+
print(f"{context}: Switched to EMA weights")
|
70 |
+
try:
|
71 |
+
yield None
|
72 |
+
finally:
|
73 |
+
if self.use_ema:
|
74 |
+
self.model_ema.restore(self.parameters())
|
75 |
+
if context is not None:
|
76 |
+
print(f"{context}: Restored training weights")
|
77 |
+
|
78 |
+
def init_from_ckpt(self, path, ignore_keys=list()):
|
79 |
+
sd = torch.load(path, map_location="cpu")["state_dict"]
|
80 |
+
keys = list(sd.keys())
|
81 |
+
for k in keys:
|
82 |
+
for ik in ignore_keys:
|
83 |
+
if k.startswith(ik):
|
84 |
+
print("Deleting key {} from state_dict.".format(k))
|
85 |
+
del sd[k]
|
86 |
+
missing, unexpected = self.load_state_dict(sd, strict=False)
|
87 |
+
print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
|
88 |
+
if len(missing) > 0:
|
89 |
+
print(f"Missing Keys: {missing}")
|
90 |
+
print(f"Unexpected Keys: {unexpected}")
|
91 |
+
|
92 |
+
def on_train_batch_end(self, *args, **kwargs):
|
93 |
+
if self.use_ema:
|
94 |
+
self.model_ema(self)
|
95 |
+
|
96 |
+
def encode(self, x):
|
97 |
+
h = self.encoder(x)
|
98 |
+
h = self.quant_conv(h)
|
99 |
+
quant, emb_loss, info = self.quantize(h)
|
100 |
+
return quant, emb_loss, info
|
101 |
+
|
102 |
+
def encode_to_prequant(self, x):
|
103 |
+
h = self.encoder(x)
|
104 |
+
h = self.quant_conv(h)
|
105 |
+
return h
|
106 |
+
|
107 |
+
def decode(self, quant):
|
108 |
+
quant = self.post_quant_conv(quant)
|
109 |
+
dec = self.decoder(quant)
|
110 |
+
return dec
|
111 |
+
|
112 |
+
def decode_code(self, code_b):
|
113 |
+
quant_b = self.quantize.embed_code(code_b)
|
114 |
+
dec = self.decode(quant_b)
|
115 |
+
return dec
|
116 |
+
|
117 |
+
def forward(self, input, return_pred_indices=False):
|
118 |
+
quant, diff, (_,_,ind) = self.encode(input)
|
119 |
+
dec = self.decode(quant)
|
120 |
+
if return_pred_indices:
|
121 |
+
return dec, diff, ind
|
122 |
+
return dec, diff
|
123 |
+
|
124 |
+
def get_input(self, batch, k):
|
125 |
+
x = batch[k]
|
126 |
+
if len(x.shape) == 3:
|
127 |
+
x = x[..., None]
|
128 |
+
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
|
129 |
+
if self.batch_resize_range is not None:
|
130 |
+
lower_size = self.batch_resize_range[0]
|
131 |
+
upper_size = self.batch_resize_range[1]
|
132 |
+
if self.global_step <= 4:
|
133 |
+
# do the first few batches with max size to avoid later oom
|
134 |
+
new_resize = upper_size
|
135 |
+
else:
|
136 |
+
new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16))
|
137 |
+
if new_resize != x.shape[2]:
|
138 |
+
x = F.interpolate(x, size=new_resize, mode="bicubic")
|
139 |
+
x = x.detach()
|
140 |
+
return x
|
141 |
+
|
142 |
+
def training_step(self, batch, batch_idx, optimizer_idx):
|
143 |
+
# https://github.com/pytorch/pytorch/issues/37142
|
144 |
+
# try not to fool the heuristics
|
145 |
+
x = self.get_input(batch, self.image_key)
|
146 |
+
xrec, qloss, ind = self(x, return_pred_indices=True)
|
147 |
+
|
148 |
+
if optimizer_idx == 0:
|
149 |
+
# autoencode
|
150 |
+
aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
|
151 |
+
last_layer=self.get_last_layer(), split="train",
|
152 |
+
predicted_indices=ind)
|
153 |
+
|
154 |
+
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
155 |
+
return aeloss
|
156 |
+
|
157 |
+
if optimizer_idx == 1:
|
158 |
+
# discriminator
|
159 |
+
discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
|
160 |
+
last_layer=self.get_last_layer(), split="train")
|
161 |
+
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
162 |
+
return discloss
|
163 |
+
|
164 |
+
def validation_step(self, batch, batch_idx):
|
165 |
+
log_dict = self._validation_step(batch, batch_idx)
|
166 |
+
with self.ema_scope():
|
167 |
+
log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema")
|
168 |
+
return log_dict
|
169 |
+
|
170 |
+
def _validation_step(self, batch, batch_idx, suffix=""):
|
171 |
+
x = self.get_input(batch, self.image_key)
|
172 |
+
xrec, qloss, ind = self(x, return_pred_indices=True)
|
173 |
+
aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0,
|
174 |
+
self.global_step,
|
175 |
+
last_layer=self.get_last_layer(),
|
176 |
+
split="val"+suffix,
|
177 |
+
predicted_indices=ind
|
178 |
+
)
|
179 |
+
|
180 |
+
discloss, log_dict_disc = self.loss(qloss, x, xrec, 1,
|
181 |
+
self.global_step,
|
182 |
+
last_layer=self.get_last_layer(),
|
183 |
+
split="val"+suffix,
|
184 |
+
predicted_indices=ind
|
185 |
+
)
|
186 |
+
rec_loss = log_dict_ae[f"val{suffix}/rec_loss"]
|
187 |
+
self.log(f"val{suffix}/rec_loss", rec_loss,
|
188 |
+
prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
|
189 |
+
self.log(f"val{suffix}/aeloss", aeloss,
|
190 |
+
prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
|
191 |
+
if version.parse(pl.__version__) >= version.parse('1.4.0'):
|
192 |
+
del log_dict_ae[f"val{suffix}/rec_loss"]
|
193 |
+
self.log_dict(log_dict_ae)
|
194 |
+
self.log_dict(log_dict_disc)
|
195 |
+
return self.log_dict
|
196 |
+
|
197 |
+
def configure_optimizers(self):
|
198 |
+
lr_d = self.learning_rate
|
199 |
+
lr_g = self.lr_g_factor*self.learning_rate
|
200 |
+
print("lr_d", lr_d)
|
201 |
+
print("lr_g", lr_g)
|
202 |
+
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
|
203 |
+
list(self.decoder.parameters())+
|
204 |
+
list(self.quantize.parameters())+
|
205 |
+
list(self.quant_conv.parameters())+
|
206 |
+
list(self.post_quant_conv.parameters()),
|
207 |
+
lr=lr_g, betas=(0.5, 0.9))
|
208 |
+
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
|
209 |
+
lr=lr_d, betas=(0.5, 0.9))
|
210 |
+
|
211 |
+
if self.scheduler_config is not None:
|
212 |
+
scheduler = instantiate_from_config(self.scheduler_config)
|
213 |
+
|
214 |
+
print("Setting up LambdaLR scheduler...")
|
215 |
+
scheduler = [
|
216 |
+
{
|
217 |
+
'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule),
|
218 |
+
'interval': 'step',
|
219 |
+
'frequency': 1
|
220 |
+
},
|
221 |
+
{
|
222 |
+
'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule),
|
223 |
+
'interval': 'step',
|
224 |
+
'frequency': 1
|
225 |
+
},
|
226 |
+
]
|
227 |
+
return [opt_ae, opt_disc], scheduler
|
228 |
+
return [opt_ae, opt_disc], []
|
229 |
+
|
230 |
+
def get_last_layer(self):
|
231 |
+
return self.decoder.conv_out.weight
|
232 |
+
|
233 |
+
def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs):
|
234 |
+
log = dict()
|
235 |
+
x = self.get_input(batch, self.image_key)
|
236 |
+
x = x.to(self.device)
|
237 |
+
if only_inputs:
|
238 |
+
log["inputs"] = x
|
239 |
+
return log
|
240 |
+
xrec, _ = self(x)
|
241 |
+
if x.shape[1] > 3:
|
242 |
+
# colorize with random projection
|
243 |
+
assert xrec.shape[1] > 3
|
244 |
+
x = self.to_rgb(x)
|
245 |
+
xrec = self.to_rgb(xrec)
|
246 |
+
log["inputs"] = x
|
247 |
+
log["reconstructions"] = xrec
|
248 |
+
if plot_ema:
|
249 |
+
with self.ema_scope():
|
250 |
+
xrec_ema, _ = self(x)
|
251 |
+
if x.shape[1] > 3: xrec_ema = self.to_rgb(xrec_ema)
|
252 |
+
log["reconstructions_ema"] = xrec_ema
|
253 |
+
return log
|
254 |
+
|
255 |
+
def to_rgb(self, x):
|
256 |
+
assert self.image_key == "segmentation"
|
257 |
+
if not hasattr(self, "colorize"):
|
258 |
+
self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
|
259 |
+
x = F.conv2d(x, weight=self.colorize)
|
260 |
+
x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
|
261 |
+
return x
|
262 |
+
|
263 |
+
|
264 |
+
class VQModelInterface(VQModel):
|
265 |
+
def __init__(self, embed_dim, *args, **kwargs):
|
266 |
+
super().__init__(embed_dim=embed_dim, *args, **kwargs)
|
267 |
+
self.embed_dim = embed_dim
|
268 |
+
|
269 |
+
def encode(self, x):
|
270 |
+
h = self.encoder(x)
|
271 |
+
h = self.quant_conv(h)
|
272 |
+
return h
|
273 |
+
|
274 |
+
def decode(self, h, force_not_quantize=False):
|
275 |
+
# also go through quantization layer
|
276 |
+
if not force_not_quantize:
|
277 |
+
quant, emb_loss, info = self.quantize(h)
|
278 |
+
else:
|
279 |
+
quant = h
|
280 |
+
quant = self.post_quant_conv(quant)
|
281 |
+
dec = self.decoder(quant)
|
282 |
+
return dec
|
283 |
+
|
284 |
+
|
285 |
+
class AutoencoderKL(pl.LightningModule):
|
286 |
+
def __init__(self,
|
287 |
+
ddconfig,
|
288 |
+
lossconfig,
|
289 |
+
embed_dim,
|
290 |
+
ckpt_path=None,
|
291 |
+
ignore_keys=[],
|
292 |
+
image_key="image",
|
293 |
+
colorize_nlabels=None,
|
294 |
+
monitor=None,
|
295 |
+
):
|
296 |
+
super().__init__()
|
297 |
+
self.image_key = image_key
|
298 |
+
self.encoder = Encoder(**ddconfig)
|
299 |
+
self.decoder = Decoder(**ddconfig)
|
300 |
+
self.loss = instantiate_from_config(lossconfig)
|
301 |
+
assert ddconfig["double_z"]
|
302 |
+
self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
|
303 |
+
self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
|
304 |
+
self.embed_dim = embed_dim
|
305 |
+
if colorize_nlabels is not None:
|
306 |
+
assert type(colorize_nlabels)==int
|
307 |
+
self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
|
308 |
+
if monitor is not None:
|
309 |
+
self.monitor = monitor
|
310 |
+
if ckpt_path is not None:
|
311 |
+
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
|
312 |
+
|
313 |
+
def init_from_ckpt(self, path, ignore_keys=list()):
|
314 |
+
sd = torch.load(path, map_location="cpu")["state_dict"]
|
315 |
+
keys = list(sd.keys())
|
316 |
+
for k in keys:
|
317 |
+
for ik in ignore_keys:
|
318 |
+
if k.startswith(ik):
|
319 |
+
print("Deleting key {} from state_dict.".format(k))
|
320 |
+
del sd[k]
|
321 |
+
self.load_state_dict(sd, strict=False)
|
322 |
+
print(f"Restored from {path}")
|
323 |
+
|
324 |
+
def encode(self, x):
|
325 |
+
h = self.encoder(x)
|
326 |
+
moments = self.quant_conv(h)
|
327 |
+
posterior = DiagonalGaussianDistribution(moments)
|
328 |
+
return posterior
|
329 |
+
|
330 |
+
def decode(self, z):
|
331 |
+
z = self.post_quant_conv(z)
|
332 |
+
dec = self.decoder(z)
|
333 |
+
return dec
|
334 |
+
|
335 |
+
def forward(self, input, sample_posterior=True):
|
336 |
+
posterior = self.encode(input)
|
337 |
+
if sample_posterior:
|
338 |
+
z = posterior.sample()
|
339 |
+
else:
|
340 |
+
z = posterior.mode()
|
341 |
+
dec = self.decode(z)
|
342 |
+
return dec, posterior
|
343 |
+
|
344 |
+
def get_input(self, batch, k):
|
345 |
+
x = batch[k]
|
346 |
+
if len(x.shape) == 3:
|
347 |
+
x = x[..., None]
|
348 |
+
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
|
349 |
+
return x
|
350 |
+
|
351 |
+
def training_step(self, batch, batch_idx, optimizer_idx):
|
352 |
+
inputs = self.get_input(batch, self.image_key)
|
353 |
+
reconstructions, posterior = self(inputs)
|
354 |
+
|
355 |
+
if optimizer_idx == 0:
|
356 |
+
# train encoder+decoder+logvar
|
357 |
+
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
|
358 |
+
last_layer=self.get_last_layer(), split="train")
|
359 |
+
self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
360 |
+
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
|
361 |
+
return aeloss
|
362 |
+
|
363 |
+
if optimizer_idx == 1:
|
364 |
+
# train the discriminator
|
365 |
+
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
|
366 |
+
last_layer=self.get_last_layer(), split="train")
|
367 |
+
|
368 |
+
self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
369 |
+
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
|
370 |
+
return discloss
|
371 |
+
|
372 |
+
def validation_step(self, batch, batch_idx):
|
373 |
+
inputs = self.get_input(batch, self.image_key)
|
374 |
+
reconstructions, posterior = self(inputs)
|
375 |
+
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
|
376 |
+
last_layer=self.get_last_layer(), split="val")
|
377 |
+
|
378 |
+
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
|
379 |
+
last_layer=self.get_last_layer(), split="val")
|
380 |
+
|
381 |
+
self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
|
382 |
+
self.log_dict(log_dict_ae)
|
383 |
+
self.log_dict(log_dict_disc)
|
384 |
+
return self.log_dict
|
385 |
+
|
386 |
+
def configure_optimizers(self):
|
387 |
+
lr = self.learning_rate
|
388 |
+
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
|
389 |
+
list(self.decoder.parameters())+
|
390 |
+
list(self.quant_conv.parameters())+
|
391 |
+
list(self.post_quant_conv.parameters()),
|
392 |
+
lr=lr, betas=(0.5, 0.9))
|
393 |
+
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
|
394 |
+
lr=lr, betas=(0.5, 0.9))
|
395 |
+
return [opt_ae, opt_disc], []
|
396 |
+
|
397 |
+
def get_last_layer(self):
|
398 |
+
return self.decoder.conv_out.weight
|
399 |
+
|
400 |
+
@torch.no_grad()
|
401 |
+
def log_images(self, batch, only_inputs=False, **kwargs):
|
402 |
+
log = dict()
|
403 |
+
x = self.get_input(batch, self.image_key)
|
404 |
+
x = x.to(self.device)
|
405 |
+
if not only_inputs:
|
406 |
+
xrec, posterior = self(x)
|
407 |
+
if x.shape[1] > 3:
|
408 |
+
# colorize with random projection
|
409 |
+
assert xrec.shape[1] > 3
|
410 |
+
x = self.to_rgb(x)
|
411 |
+
xrec = self.to_rgb(xrec)
|
412 |
+
log["samples"] = self.decode(torch.randn_like(posterior.sample()))
|
413 |
+
log["reconstructions"] = xrec
|
414 |
+
log["inputs"] = x
|
415 |
+
return log
|
416 |
+
|
417 |
+
def to_rgb(self, x):
|
418 |
+
assert self.image_key == "segmentation"
|
419 |
+
if not hasattr(self, "colorize"):
|
420 |
+
self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
|
421 |
+
x = F.conv2d(x, weight=self.colorize)
|
422 |
+
x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
|
423 |
+
return x
|
424 |
+
|
425 |
+
|
426 |
+
class IdentityFirstStage(torch.nn.Module):
|
427 |
+
def __init__(self, *args, vq_interface=False, **kwargs):
|
428 |
+
self.vq_interface = vq_interface # TODO: Should be true by default but check to not break older stuff
|
429 |
+
super().__init__()
|
430 |
+
|
431 |
+
def encode(self, x, *args, **kwargs):
|
432 |
+
return x
|
433 |
+
|
434 |
+
def decode(self, x, *args, **kwargs):
|
435 |
+
return x
|
436 |
+
|
437 |
+
def quantize(self, x, *args, **kwargs):
|
438 |
+
if self.vq_interface:
|
439 |
+
return x, None, [None, None, None]
|
440 |
+
return x
|
441 |
+
|
442 |
+
def forward(self, x, *args, **kwargs):
|
443 |
+
return x
|
ldm/models/diffusion/__init__.py
ADDED
File without changes
|
ldm/models/diffusion/sync_dreamer.py
ADDED
@@ -0,0 +1,665 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from pathlib import Path
|
2 |
+
|
3 |
+
import pytorch_lightning as pl
|
4 |
+
import torch
|
5 |
+
import torch.nn as nn
|
6 |
+
import torch.nn.functional as F
|
7 |
+
import numpy as np
|
8 |
+
from skimage.io import imsave
|
9 |
+
from torch.optim.lr_scheduler import LambdaLR
|
10 |
+
from tqdm import tqdm
|
11 |
+
|
12 |
+
from ldm.base_utils import read_pickle, concat_images_list
|
13 |
+
from ldm.models.diffusion.sync_dreamer_utils import get_warp_coordinates, create_target_volume
|
14 |
+
from ldm.models.diffusion.sync_dreamer_network import NoisyTargetViewEncoder, SpatialTime3DNet, FrustumTV3DNet
|
15 |
+
from ldm.modules.diffusionmodules.util import make_ddim_timesteps, timestep_embedding
|
16 |
+
from ldm.modules.encoders.modules import FrozenCLIPImageEmbedder
|
17 |
+
from ldm.util import instantiate_from_config
|
18 |
+
|
19 |
+
def disabled_train(self, mode=True):
|
20 |
+
"""Overwrite model.train with this function to make sure train/eval mode
|
21 |
+
does not change anymore."""
|
22 |
+
return self
|
23 |
+
|
24 |
+
def disable_training_module(module: nn.Module):
|
25 |
+
module = module.eval()
|
26 |
+
module.train = disabled_train
|
27 |
+
for para in module.parameters():
|
28 |
+
para.requires_grad = False
|
29 |
+
return module
|
30 |
+
|
31 |
+
def repeat_to_batch(tensor, B, VN):
|
32 |
+
t_shape = tensor.shape
|
33 |
+
ones = [1 for _ in range(len(t_shape)-1)]
|
34 |
+
tensor_new = tensor.view(B,1,*t_shape[1:]).repeat(1,VN,*ones).view(B*VN,*t_shape[1:])
|
35 |
+
return tensor_new
|
36 |
+
|
37 |
+
class UNetWrapper(nn.Module):
|
38 |
+
def __init__(self, diff_model_config, drop_conditions=False, drop_scheme='default', use_zero_123=True):
|
39 |
+
super().__init__()
|
40 |
+
self.diffusion_model = instantiate_from_config(diff_model_config)
|
41 |
+
self.drop_conditions = drop_conditions
|
42 |
+
self.drop_scheme=drop_scheme
|
43 |
+
self.use_zero_123 = use_zero_123
|
44 |
+
|
45 |
+
|
46 |
+
def drop(self, cond, mask):
|
47 |
+
shape = cond.shape
|
48 |
+
B = shape[0]
|
49 |
+
cond = mask.view(B,*[1 for _ in range(len(shape)-1)]) * cond
|
50 |
+
return cond
|
51 |
+
|
52 |
+
def get_trainable_parameters(self):
|
53 |
+
return self.diffusion_model.get_trainable_parameters()
|
54 |
+
|
55 |
+
def get_drop_scheme(self, B, device):
|
56 |
+
if self.drop_scheme=='default':
|
57 |
+
random = torch.rand(B, dtype=torch.float32, device=device)
|
58 |
+
drop_clip = (random > 0.15) & (random <= 0.2)
|
59 |
+
drop_volume = (random > 0.1) & (random <= 0.15)
|
60 |
+
drop_concat = (random > 0.05) & (random <= 0.1)
|
61 |
+
drop_all = random <= 0.05
|
62 |
+
else:
|
63 |
+
raise NotImplementedError
|
64 |
+
return drop_clip, drop_volume, drop_concat, drop_all
|
65 |
+
|
66 |
+
def forward(self, x, t, clip_embed, volume_feats, x_concat, is_train=False):
|
67 |
+
"""
|
68 |
+
|
69 |
+
@param x: B,4,H,W
|
70 |
+
@param t: B,
|
71 |
+
@param clip_embed: B,M,768
|
72 |
+
@param volume_feats: B,C,D,H,W
|
73 |
+
@param x_concat: B,C,H,W
|
74 |
+
@param is_train:
|
75 |
+
@return:
|
76 |
+
"""
|
77 |
+
if self.drop_conditions and is_train:
|
78 |
+
B = x.shape[0]
|
79 |
+
drop_clip, drop_volume, drop_concat, drop_all = self.get_drop_scheme(B, x.device)
|
80 |
+
|
81 |
+
clip_mask = 1.0 - (drop_clip | drop_all).float()
|
82 |
+
clip_embed = self.drop(clip_embed, clip_mask)
|
83 |
+
|
84 |
+
volume_mask = 1.0 - (drop_volume | drop_all).float()
|
85 |
+
for k, v in volume_feats.items():
|
86 |
+
volume_feats[k] = self.drop(v, mask=volume_mask)
|
87 |
+
|
88 |
+
concat_mask = 1.0 - (drop_concat | drop_all).float()
|
89 |
+
x_concat = self.drop(x_concat, concat_mask)
|
90 |
+
|
91 |
+
if self.use_zero_123:
|
92 |
+
# zero123 does not multiply this when encoding, maybe a bug for zero123
|
93 |
+
first_stage_scale_factor = 0.18215
|
94 |
+
x_concat_ = x_concat * 1.0
|
95 |
+
x_concat_[:, :4] = x_concat_[:, :4] / first_stage_scale_factor
|
96 |
+
else:
|
97 |
+
x_concat_ = x_concat
|
98 |
+
|
99 |
+
x = torch.cat([x, x_concat_], 1)
|
100 |
+
pred = self.diffusion_model(x, t, clip_embed, source_dict=volume_feats)
|
101 |
+
return pred
|
102 |
+
|
103 |
+
def predict_with_unconditional_scale(self, x, t, clip_embed, volume_feats, x_concat, unconditional_scale):
|
104 |
+
x_ = torch.cat([x] * 2, 0)
|
105 |
+
t_ = torch.cat([t] * 2, 0)
|
106 |
+
clip_embed_ = torch.cat([clip_embed, torch.zeros_like(clip_embed)], 0)
|
107 |
+
|
108 |
+
v_ = {}
|
109 |
+
for k, v in volume_feats.items():
|
110 |
+
v_[k] = torch.cat([v, torch.zeros_like(v)], 0)
|
111 |
+
|
112 |
+
x_concat_ = torch.cat([x_concat, torch.zeros_like(x_concat)], 0)
|
113 |
+
if self.use_zero_123:
|
114 |
+
# zero123 does not multiply this when encoding, maybe a bug for zero123
|
115 |
+
first_stage_scale_factor = 0.18215
|
116 |
+
x_concat_[:, :4] = x_concat_[:, :4] / first_stage_scale_factor
|
117 |
+
x_ = torch.cat([x_, x_concat_], 1)
|
118 |
+
s, s_uc = self.diffusion_model(x_, t_, clip_embed_, source_dict=v_).chunk(2)
|
119 |
+
s = s_uc + unconditional_scale * (s - s_uc)
|
120 |
+
return s
|
121 |
+
|
122 |
+
|
123 |
+
class SpatialVolumeNet(nn.Module):
|
124 |
+
def __init__(self, time_dim, view_dim, view_num,
|
125 |
+
input_image_size=256, frustum_volume_depth=48,
|
126 |
+
spatial_volume_size=32, spatial_volume_length=0.5,
|
127 |
+
frustum_volume_length=0.86603 # sqrt(3)/2
|
128 |
+
):
|
129 |
+
super().__init__()
|
130 |
+
self.target_encoder = NoisyTargetViewEncoder(time_dim, view_dim, output_dim=16)
|
131 |
+
self.spatial_volume_feats = SpatialTime3DNet(input_dim=16 * view_num, time_dim=time_dim, dims=(64, 128, 256, 512))
|
132 |
+
self.frustum_volume_feats = FrustumTV3DNet(64, time_dim, view_dim, dims=(64, 128, 256, 512))
|
133 |
+
|
134 |
+
self.frustum_volume_length = frustum_volume_length
|
135 |
+
self.input_image_size = input_image_size
|
136 |
+
self.spatial_volume_size = spatial_volume_size
|
137 |
+
self.spatial_volume_length = spatial_volume_length
|
138 |
+
|
139 |
+
self.frustum_volume_size = self.input_image_size // 8
|
140 |
+
self.frustum_volume_depth = frustum_volume_depth
|
141 |
+
self.time_dim = time_dim
|
142 |
+
self.view_dim = view_dim
|
143 |
+
self.default_origin_depth = 1.5 # our rendered images are 1.5 away from the origin, we assume camera is 1.5 away from the origin
|
144 |
+
|
145 |
+
def construct_spatial_volume(self, x, t_embed, v_embed, target_poses, target_Ks):
|
146 |
+
"""
|
147 |
+
@param x: B,N,4,H,W
|
148 |
+
@param t_embed: B,t_dim
|
149 |
+
@param v_embed: B,N,v_dim
|
150 |
+
@param target_poses: N,3,4
|
151 |
+
@param target_Ks: N,3,3
|
152 |
+
@return:
|
153 |
+
"""
|
154 |
+
B, N, _, H, W = x.shape
|
155 |
+
V = self.spatial_volume_size
|
156 |
+
device = x.device
|
157 |
+
|
158 |
+
spatial_volume_verts = torch.linspace(-self.spatial_volume_length, self.spatial_volume_length, V, dtype=torch.float32, device=device)
|
159 |
+
spatial_volume_verts = torch.stack(torch.meshgrid(spatial_volume_verts, spatial_volume_verts, spatial_volume_verts), -1)
|
160 |
+
spatial_volume_verts = spatial_volume_verts.reshape(1, V ** 3, 3)[:, :, (2, 1, 0)]
|
161 |
+
spatial_volume_verts = spatial_volume_verts.view(1, V, V, V, 3).permute(0, 4, 1, 2, 3).repeat(B, 1, 1, 1, 1)
|
162 |
+
|
163 |
+
# encode source features
|
164 |
+
t_embed_ = t_embed.view(B, 1, self.time_dim).repeat(1, N, 1).view(B, N, self.time_dim)
|
165 |
+
# v_embed_ = v_embed.view(1, N, self.view_dim).repeat(B, 1, 1).view(B, N, self.view_dim)
|
166 |
+
v_embed_ = v_embed
|
167 |
+
target_Ks = target_Ks.unsqueeze(0).repeat(B, 1, 1, 1)
|
168 |
+
target_poses = target_poses.unsqueeze(0).repeat(B, 1, 1, 1)
|
169 |
+
|
170 |
+
# extract 2D image features
|
171 |
+
spatial_volume_feats = []
|
172 |
+
# project source features
|
173 |
+
for ni in range(0, N):
|
174 |
+
pose_source_ = target_poses[:, ni]
|
175 |
+
K_source_ = target_Ks[:, ni]
|
176 |
+
x_ = self.target_encoder(x[:, ni], t_embed_[:, ni], v_embed_[:, ni])
|
177 |
+
C = x_.shape[1]
|
178 |
+
|
179 |
+
coords_source = get_warp_coordinates(spatial_volume_verts, x_.shape[-1], self.input_image_size, K_source_, pose_source_).view(B, V, V * V, 2)
|
180 |
+
unproj_feats_ = F.grid_sample(x_, coords_source, mode='bilinear', padding_mode='zeros', align_corners=True)
|
181 |
+
unproj_feats_ = unproj_feats_.view(B, C, V, V, V)
|
182 |
+
spatial_volume_feats.append(unproj_feats_)
|
183 |
+
|
184 |
+
spatial_volume_feats = torch.stack(spatial_volume_feats, 1) # B,N,C,V,V,V
|
185 |
+
N = spatial_volume_feats.shape[1]
|
186 |
+
spatial_volume_feats = spatial_volume_feats.view(B, N*C, V, V, V)
|
187 |
+
|
188 |
+
spatial_volume_feats = self.spatial_volume_feats(spatial_volume_feats, t_embed) # b,64,32,32,32
|
189 |
+
return spatial_volume_feats
|
190 |
+
|
191 |
+
def construct_view_frustum_volume(self, spatial_volume, t_embed, v_embed, poses, Ks, target_indices):
|
192 |
+
"""
|
193 |
+
@param spatial_volume: B,C,V,V,V
|
194 |
+
@param t_embed: B,t_dim
|
195 |
+
@param v_embed: B,N,v_dim
|
196 |
+
@param poses: N,3,4
|
197 |
+
@param Ks: N,3,3
|
198 |
+
@param target_indices: B,TN
|
199 |
+
@return: B*TN,C,H,W
|
200 |
+
"""
|
201 |
+
B, TN = target_indices.shape
|
202 |
+
H, W = self.frustum_volume_size, self.frustum_volume_size
|
203 |
+
D = self.frustum_volume_depth
|
204 |
+
V = self.spatial_volume_size
|
205 |
+
|
206 |
+
near = torch.ones(B * TN, 1, H, W, dtype=spatial_volume.dtype, device=spatial_volume.device) * self.default_origin_depth - self.frustum_volume_length
|
207 |
+
far = torch.ones(B * TN, 1, H, W, dtype=spatial_volume.dtype, device=spatial_volume.device) * self.default_origin_depth + self.frustum_volume_length
|
208 |
+
|
209 |
+
target_indices = target_indices.view(B*TN) # B*TN
|
210 |
+
poses_ = poses[target_indices] # B*TN,3,4
|
211 |
+
Ks_ = Ks[target_indices] # B*TN,3,4
|
212 |
+
volume_xyz, volume_depth = create_target_volume(D, self.frustum_volume_size, self.input_image_size, poses_, Ks_, near, far) # B*TN,3 or 1,D,H,W
|
213 |
+
|
214 |
+
volume_xyz_ = volume_xyz / self.spatial_volume_length # since the spatial volume is constructed in [-spatial_volume_length,spatial_volume_length]
|
215 |
+
volume_xyz_ = volume_xyz_.permute(0, 2, 3, 4, 1) # B*TN,D,H,W,3
|
216 |
+
spatial_volume_ = spatial_volume.unsqueeze(1).repeat(1, TN, 1, 1, 1, 1).view(B * TN, -1, V, V, V)
|
217 |
+
volume_feats = F.grid_sample(spatial_volume_, volume_xyz_, mode='bilinear', padding_mode='zeros', align_corners=True) # B*TN,C,D,H,W
|
218 |
+
|
219 |
+
v_embed_ = v_embed[torch.arange(B)[:,None], target_indices.view(B,TN)].view(B*TN, -1) # B*TN
|
220 |
+
t_embed_ = t_embed.unsqueeze(1).repeat(1,TN,1).view(B*TN,-1)
|
221 |
+
volume_feats_dict = self.frustum_volume_feats(volume_feats, t_embed_, v_embed_)
|
222 |
+
return volume_feats_dict, volume_depth
|
223 |
+
|
224 |
+
class SyncMultiviewDiffusion(pl.LightningModule):
|
225 |
+
def __init__(self, unet_config, scheduler_config,
|
226 |
+
finetune_unet=False, finetune_projection=True,
|
227 |
+
view_num=16, image_size=256,
|
228 |
+
cfg_scale=3.0, output_num=8, batch_view_num=4,
|
229 |
+
drop_conditions=False, drop_scheme='default',
|
230 |
+
clip_image_encoder_path="/apdcephfs/private_rondyliu/projects/clip/ViT-L-14.pt"):
|
231 |
+
super().__init__()
|
232 |
+
|
233 |
+
self.finetune_unet = finetune_unet
|
234 |
+
self.finetune_projection = finetune_projection
|
235 |
+
|
236 |
+
self.view_num = view_num
|
237 |
+
self.viewpoint_dim = 4
|
238 |
+
self.output_num = output_num
|
239 |
+
self.image_size = image_size
|
240 |
+
|
241 |
+
self.batch_view_num = batch_view_num
|
242 |
+
self.cfg_scale = cfg_scale
|
243 |
+
|
244 |
+
self.clip_image_encoder_path = clip_image_encoder_path
|
245 |
+
|
246 |
+
self._init_time_step_embedding()
|
247 |
+
self._init_first_stage()
|
248 |
+
self._init_schedule()
|
249 |
+
self._init_multiview()
|
250 |
+
self._init_clip_image_encoder()
|
251 |
+
self._init_clip_projection()
|
252 |
+
|
253 |
+
self.spatial_volume = SpatialVolumeNet(self.time_embed_dim, self.viewpoint_dim, self.view_num)
|
254 |
+
self.model = UNetWrapper(unet_config, drop_conditions=drop_conditions, drop_scheme=drop_scheme)
|
255 |
+
self.scheduler_config = scheduler_config
|
256 |
+
|
257 |
+
latent_size = image_size//8
|
258 |
+
self.ddim = SyncDDIMSampler(self, 200, "uniform", 1.0, latent_size=latent_size)
|
259 |
+
|
260 |
+
def _init_clip_projection(self):
|
261 |
+
self.cc_projection = nn.Linear(772, 768)
|
262 |
+
nn.init.eye_(list(self.cc_projection.parameters())[0][:768, :768])
|
263 |
+
nn.init.zeros_(list(self.cc_projection.parameters())[1])
|
264 |
+
self.cc_projection.requires_grad_(True)
|
265 |
+
|
266 |
+
if not self.finetune_projection:
|
267 |
+
disable_training_module(self.cc_projection)
|
268 |
+
|
269 |
+
def _init_multiview(self):
|
270 |
+
K, azs, _, _, poses = read_pickle(f'meta_info/camera-{self.view_num}.pkl')
|
271 |
+
default_image_size = 256
|
272 |
+
ratio = self.image_size/default_image_size
|
273 |
+
K = np.diag([ratio,ratio,1]) @ K
|
274 |
+
K = torch.from_numpy(K.astype(np.float32)) # [3,3]
|
275 |
+
K = K.unsqueeze(0).repeat(self.view_num,1,1) # N,3,3
|
276 |
+
poses = torch.from_numpy(poses.astype(np.float32)) # N,3,4
|
277 |
+
self.register_buffer('poses', poses)
|
278 |
+
self.register_buffer('Ks', K)
|
279 |
+
azs = (azs + np.pi) % (np.pi * 2) - np.pi # scale to [-pi,pi] and the index=0 has az=0
|
280 |
+
self.register_buffer('azimuth', torch.from_numpy(azs.astype(np.float32)))
|
281 |
+
|
282 |
+
def get_viewpoint_embedding(self, batch_size, elevation_ref):
|
283 |
+
"""
|
284 |
+
@param batch_size:
|
285 |
+
@param elevation_ref: B
|
286 |
+
@return:
|
287 |
+
"""
|
288 |
+
azimuth_input = self.azimuth[0].unsqueeze(0) # 1
|
289 |
+
azimuth_target = self.azimuth # N
|
290 |
+
elevation_input = -elevation_ref # note that zero123 use a negative elevation here!!!
|
291 |
+
elevation_target = -np.deg2rad(30)
|
292 |
+
d_e = elevation_target - elevation_input # B
|
293 |
+
N = self.azimuth.shape[0]
|
294 |
+
B = batch_size
|
295 |
+
d_e = d_e.unsqueeze(1).repeat(1, N)
|
296 |
+
d_a = azimuth_target - azimuth_input # N
|
297 |
+
d_a = d_a.unsqueeze(0).repeat(B, 1)
|
298 |
+
d_z = torch.zeros_like(d_a)
|
299 |
+
embedding = torch.stack([d_e, torch.sin(d_a), torch.cos(d_a), d_z], -1) # B,N,4
|
300 |
+
return embedding
|
301 |
+
|
302 |
+
def _init_first_stage(self):
|
303 |
+
first_stage_config={
|
304 |
+
"target": "ldm.models.autoencoder.AutoencoderKL",
|
305 |
+
"params": {
|
306 |
+
"embed_dim": 4,
|
307 |
+
"monitor": "val/rec_loss",
|
308 |
+
"ddconfig":{
|
309 |
+
"double_z": True,
|
310 |
+
"z_channels": 4,
|
311 |
+
"resolution": self.image_size,
|
312 |
+
"in_channels": 3,
|
313 |
+
"out_ch": 3,
|
314 |
+
"ch": 128,
|
315 |
+
"ch_mult": [1,2,4,4],
|
316 |
+
"num_res_blocks": 2,
|
317 |
+
"attn_resolutions": [],
|
318 |
+
"dropout": 0.0
|
319 |
+
},
|
320 |
+
"lossconfig": {"target": "torch.nn.Identity"},
|
321 |
+
}
|
322 |
+
}
|
323 |
+
self.first_stage_scale_factor = 0.18215
|
324 |
+
self.first_stage_model = instantiate_from_config(first_stage_config)
|
325 |
+
self.first_stage_model = disable_training_module(self.first_stage_model)
|
326 |
+
|
327 |
+
def _init_clip_image_encoder(self):
|
328 |
+
self.clip_image_encoder = FrozenCLIPImageEmbedder(model=self.clip_image_encoder_path)
|
329 |
+
self.clip_image_encoder = disable_training_module(self.clip_image_encoder)
|
330 |
+
|
331 |
+
def _init_schedule(self):
|
332 |
+
self.num_timesteps = 1000
|
333 |
+
linear_start = 0.00085
|
334 |
+
linear_end = 0.0120
|
335 |
+
num_timesteps = 1000
|
336 |
+
betas = torch.linspace(linear_start ** 0.5, linear_end ** 0.5, num_timesteps, dtype=torch.float32) ** 2 # T
|
337 |
+
assert betas.shape[0] == self.num_timesteps
|
338 |
+
|
339 |
+
# all in float64 first
|
340 |
+
alphas = 1. - betas
|
341 |
+
alphas_cumprod = torch.cumprod(alphas, dim=0) # T
|
342 |
+
alphas_cumprod_prev = torch.cat([torch.ones(1, dtype=torch.float64), alphas_cumprod[:-1]], 0)
|
343 |
+
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) # T
|
344 |
+
posterior_log_variance_clipped = torch.log(torch.clamp(posterior_variance, min=1e-20))
|
345 |
+
posterior_log_variance_clipped = torch.clamp(posterior_log_variance_clipped, min=-10)
|
346 |
+
|
347 |
+
self.register_buffer("betas", betas.float())
|
348 |
+
self.register_buffer("alphas", alphas.float())
|
349 |
+
self.register_buffer("alphas_cumprod", alphas_cumprod.float())
|
350 |
+
self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod).float())
|
351 |
+
self.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1 - alphas_cumprod).float())
|
352 |
+
self.register_buffer("posterior_variance", posterior_variance.float())
|
353 |
+
self.register_buffer('posterior_log_variance_clipped', posterior_log_variance_clipped.float())
|
354 |
+
|
355 |
+
def _init_time_step_embedding(self):
|
356 |
+
self.time_embed_dim = 256
|
357 |
+
self.time_embed = nn.Sequential(
|
358 |
+
nn.Linear(self.time_embed_dim, self.time_embed_dim),
|
359 |
+
nn.SiLU(True),
|
360 |
+
nn.Linear(self.time_embed_dim, self.time_embed_dim),
|
361 |
+
)
|
362 |
+
|
363 |
+
def encode_first_stage(self, x, sample=True):
|
364 |
+
with torch.no_grad():
|
365 |
+
posterior = self.first_stage_model.encode(x) # b,4,h//8,w//8
|
366 |
+
if sample:
|
367 |
+
return posterior.sample().detach() * self.first_stage_scale_factor
|
368 |
+
else:
|
369 |
+
return posterior.mode().detach() * self.first_stage_scale_factor
|
370 |
+
|
371 |
+
def decode_first_stage(self, z):
|
372 |
+
with torch.no_grad():
|
373 |
+
z = 1. / self.first_stage_scale_factor * z
|
374 |
+
return self.first_stage_model.decode(z)
|
375 |
+
|
376 |
+
def prepare(self, batch):
|
377 |
+
# encode target
|
378 |
+
if 'target_image' in batch:
|
379 |
+
image_target = batch['target_image'].permute(0, 1, 4, 2, 3) # b,n,3,h,w
|
380 |
+
N = image_target.shape[1]
|
381 |
+
x = [self.encode_first_stage(image_target[:,ni], True) for ni in range(N)]
|
382 |
+
x = torch.stack(x, 1) # b,n,4,h//8,w//8
|
383 |
+
else:
|
384 |
+
x = None
|
385 |
+
|
386 |
+
image_input = batch['input_image'].permute(0, 3, 1, 2)
|
387 |
+
elevation_input = batch['input_elevation'][:, 0] # b
|
388 |
+
x_input = self.encode_first_stage(image_input)
|
389 |
+
input_info = {'image': image_input, 'elevation': elevation_input, 'x': x_input}
|
390 |
+
with torch.no_grad():
|
391 |
+
clip_embed = self.clip_image_encoder.encode(image_input)
|
392 |
+
return x, clip_embed, input_info
|
393 |
+
|
394 |
+
def embed_time(self, t):
|
395 |
+
t_embed = timestep_embedding(t, self.time_embed_dim, repeat_only=False) # B,TED
|
396 |
+
t_embed = self.time_embed(t_embed) # B,TED
|
397 |
+
return t_embed
|
398 |
+
|
399 |
+
def get_target_view_feats(self, x_input, spatial_volume, clip_embed, t_embed, v_embed, target_index):
|
400 |
+
"""
|
401 |
+
@param x_input: B,4,H,W
|
402 |
+
@param spatial_volume: B,C,V,V,V
|
403 |
+
@param clip_embed: B,1,768
|
404 |
+
@param t_embed: B,t_dim
|
405 |
+
@param v_embed: B,N,v_dim
|
406 |
+
@param target_index: B,TN
|
407 |
+
@return:
|
408 |
+
tensors of size B*TN,*
|
409 |
+
"""
|
410 |
+
B, _, H, W = x_input.shape
|
411 |
+
frustum_volume_feats, frustum_volume_depth = self.spatial_volume.construct_view_frustum_volume(spatial_volume, t_embed, v_embed, self.poses, self.Ks, target_index)
|
412 |
+
|
413 |
+
# clip
|
414 |
+
TN = target_index.shape[1]
|
415 |
+
v_embed_ = v_embed[torch.arange(B)[:,None], target_index].view(B*TN, self.viewpoint_dim) # B*TN,v_dim
|
416 |
+
clip_embed_ = clip_embed.unsqueeze(1).repeat(1,TN,1,1).view(B*TN,1,768)
|
417 |
+
clip_embed_ = self.cc_projection(torch.cat([clip_embed_, v_embed_.unsqueeze(1)], -1)) # B*TN,1,768
|
418 |
+
|
419 |
+
x_input_ = x_input.unsqueeze(1).repeat(1, TN, 1, 1, 1).view(B * TN, 4, H, W)
|
420 |
+
|
421 |
+
x_concat = x_input_
|
422 |
+
return clip_embed_, frustum_volume_feats, x_concat
|
423 |
+
|
424 |
+
def training_step(self, batch):
|
425 |
+
B = batch['target_image'].shape[0]
|
426 |
+
time_steps = torch.randint(0, self.num_timesteps, (B,), device=self.device).long()
|
427 |
+
|
428 |
+
x, clip_embed, input_info = self.prepare(batch)
|
429 |
+
x_noisy, noise = self.add_noise(x, time_steps) # B,N,4,H,W
|
430 |
+
|
431 |
+
N = self.view_num
|
432 |
+
target_index = torch.randint(0, N, (B, 1), device=self.device).long() # B, 1
|
433 |
+
v_embed = self.get_viewpoint_embedding(B, input_info['elevation']) # N,v_dim
|
434 |
+
|
435 |
+
t_embed = self.embed_time(time_steps)
|
436 |
+
spatial_volume = self.spatial_volume.construct_spatial_volume(x_noisy, t_embed, v_embed, self.poses, self.Ks)
|
437 |
+
|
438 |
+
clip_embed, volume_feats, x_concat = self.get_target_view_feats(input_info['x'], spatial_volume, clip_embed, t_embed, v_embed, target_index)
|
439 |
+
|
440 |
+
x_noisy_ = x_noisy[torch.arange(B)[:,None],target_index][:,0] # B,4,H,W
|
441 |
+
noise_predict = self.model(x_noisy_, time_steps, clip_embed, volume_feats, x_concat, is_train=True) # B,4,H,W
|
442 |
+
|
443 |
+
noise_target = noise[torch.arange(B)[:,None],target_index][:,0] # B,4,H,W
|
444 |
+
# loss simple for diffusion
|
445 |
+
loss_simple = torch.nn.functional.mse_loss(noise_target, noise_predict, reduction='none')
|
446 |
+
loss = loss_simple.mean()
|
447 |
+
self.log('sim', loss_simple.mean(), prog_bar=True, logger=True, on_step=True, on_epoch=True, rank_zero_only=True)
|
448 |
+
|
449 |
+
# log others
|
450 |
+
lr = self.optimizers().param_groups[0]['lr']
|
451 |
+
self.log('lr', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False, rank_zero_only=True)
|
452 |
+
self.log("step", self.global_step, prog_bar=True, logger=True, on_step=True, on_epoch=False, rank_zero_only=True)
|
453 |
+
return loss
|
454 |
+
|
455 |
+
def add_noise(self, x_start, t):
|
456 |
+
"""
|
457 |
+
@param x_start: B,*
|
458 |
+
@param t: B,
|
459 |
+
@return:
|
460 |
+
"""
|
461 |
+
B = x_start.shape[0]
|
462 |
+
noise = torch.randn_like(x_start) # B,*
|
463 |
+
|
464 |
+
sqrt_alphas_cumprod_ = self.sqrt_alphas_cumprod[t] # B,
|
465 |
+
sqrt_one_minus_alphas_cumprod_ = self.sqrt_one_minus_alphas_cumprod[t] # B
|
466 |
+
sqrt_alphas_cumprod_ = sqrt_alphas_cumprod_.view(B, *[1 for _ in range(len(x_start.shape)-1)])
|
467 |
+
sqrt_one_minus_alphas_cumprod_ = sqrt_one_minus_alphas_cumprod_.view(B, *[1 for _ in range(len(x_start.shape)-1)])
|
468 |
+
x_noisy = sqrt_alphas_cumprod_ * x_start + sqrt_one_minus_alphas_cumprod_ * noise
|
469 |
+
return x_noisy, noise
|
470 |
+
|
471 |
+
def sample(self, batch, cfg_scale, batch_view_num, use_ddim=True,
|
472 |
+
return_inter_results=False, inter_interval=50, inter_view_interval=2):
|
473 |
+
_, clip_embed, input_info = self.prepare(batch)
|
474 |
+
if use_ddim:
|
475 |
+
x_sample, inter = self.ddim.sample(input_info, clip_embed, unconditional_scale=cfg_scale, log_every_t=inter_interval, batch_view_num=batch_view_num)
|
476 |
+
else:
|
477 |
+
raise NotImplementedError
|
478 |
+
|
479 |
+
N = x_sample.shape[1]
|
480 |
+
x_sample = torch.stack([self.decode_first_stage(x_sample[:, ni]) for ni in range(N)], 1)
|
481 |
+
if return_inter_results:
|
482 |
+
torch.cuda.synchronize()
|
483 |
+
torch.cuda.empty_cache()
|
484 |
+
inter = torch.stack(inter['x_inter'], 2) # # B,N,T,C,H,W
|
485 |
+
B,N,T,C,H,W = inter.shape
|
486 |
+
inter_results = []
|
487 |
+
for ni in tqdm(range(0, N, inter_view_interval)):
|
488 |
+
inter_results_ = []
|
489 |
+
for ti in range(T):
|
490 |
+
inter_results_.append(self.decode_first_stage(inter[:, ni, ti]))
|
491 |
+
inter_results.append(torch.stack(inter_results_, 1)) # B,T,3,H,W
|
492 |
+
inter_results = torch.stack(inter_results,1) # B,N,T,3,H,W
|
493 |
+
return x_sample, inter_results
|
494 |
+
else:
|
495 |
+
return x_sample
|
496 |
+
|
497 |
+
def log_image(self, x_sample, batch, step, output_dir):
|
498 |
+
process = lambda x: ((torch.clip(x, min=-1, max=1).cpu().numpy() * 0.5 + 0.5) * 255).astype(np.uint8)
|
499 |
+
B = x_sample.shape[0]
|
500 |
+
N = x_sample.shape[1]
|
501 |
+
image_cond = []
|
502 |
+
for bi in range(B):
|
503 |
+
img_pr_ = concat_images_list(process(batch['input_image'][bi]),*[process(x_sample[bi, ni].permute(1, 2, 0)) for ni in range(N)])
|
504 |
+
image_cond.append(img_pr_)
|
505 |
+
|
506 |
+
output_dir = Path(output_dir)
|
507 |
+
imsave(str(output_dir/f'{step}.jpg'), concat_images_list(*image_cond, vert=True))
|
508 |
+
|
509 |
+
@torch.no_grad()
|
510 |
+
def validation_step(self, batch, batch_idx):
|
511 |
+
if batch_idx==0 and self.global_rank==0:
|
512 |
+
self.eval()
|
513 |
+
step = self.global_step
|
514 |
+
batch_ = {}
|
515 |
+
for k, v in batch.items(): batch_[k] = v[:self.output_num]
|
516 |
+
x_sample = self.sample(batch_, self.cfg_scale, self.batch_view_num)
|
517 |
+
output_dir = Path(self.image_dir) / 'images' / 'val'
|
518 |
+
output_dir.mkdir(exist_ok=True, parents=True)
|
519 |
+
self.log_image(x_sample, batch, step, output_dir=output_dir)
|
520 |
+
|
521 |
+
def configure_optimizers(self):
|
522 |
+
lr = self.learning_rate
|
523 |
+
print(f'setting learning rate to {lr:.4f} ...')
|
524 |
+
paras = []
|
525 |
+
if self.finetune_projection:
|
526 |
+
paras.append({"params": self.cc_projection.parameters(), "lr": lr},)
|
527 |
+
if self.finetune_unet:
|
528 |
+
paras.append({"params": self.model.parameters(), "lr": lr},)
|
529 |
+
else:
|
530 |
+
paras.append({"params": self.model.get_trainable_parameters(), "lr": lr},)
|
531 |
+
|
532 |
+
paras.append({"params": self.time_embed.parameters(), "lr": lr*10.0},)
|
533 |
+
paras.append({"params": self.spatial_volume.parameters(), "lr": lr*10.0},)
|
534 |
+
|
535 |
+
opt = torch.optim.AdamW(paras, lr=lr)
|
536 |
+
|
537 |
+
scheduler = instantiate_from_config(self.scheduler_config)
|
538 |
+
print("Setting up LambdaLR scheduler...")
|
539 |
+
scheduler = [{'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule), 'interval': 'step', 'frequency': 1}]
|
540 |
+
return [opt], scheduler
|
541 |
+
|
542 |
+
class SyncDDIMSampler:
|
543 |
+
def __init__(self, model: SyncMultiviewDiffusion, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., latent_size=32):
|
544 |
+
super().__init__()
|
545 |
+
self.model = model
|
546 |
+
self.ddpm_num_timesteps = model.num_timesteps
|
547 |
+
self.latent_size = latent_size
|
548 |
+
self._make_schedule(ddim_num_steps, ddim_discretize, ddim_eta)
|
549 |
+
self.eta = ddim_eta
|
550 |
+
|
551 |
+
def _make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
|
552 |
+
self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps, num_ddpm_timesteps=self.ddpm_num_timesteps, verbose=verbose) # DT
|
553 |
+
ddim_timesteps_ = torch.from_numpy(self.ddim_timesteps.astype(np.int64)) # DT
|
554 |
+
|
555 |
+
alphas_cumprod = self.model.alphas_cumprod # T
|
556 |
+
assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
|
557 |
+
self.ddim_alphas = alphas_cumprod[ddim_timesteps_].double() # DT
|
558 |
+
self.ddim_alphas_prev = torch.cat([alphas_cumprod[0:1], alphas_cumprod[ddim_timesteps_[:-1]]], 0) # DT
|
559 |
+
self.ddim_sigmas = ddim_eta * torch.sqrt((1 - self.ddim_alphas_prev) / (1 - self.ddim_alphas) * (1 - self.ddim_alphas / self.ddim_alphas_prev))
|
560 |
+
|
561 |
+
self.ddim_alphas_raw = self.model.alphas[ddim_timesteps_].float() # DT
|
562 |
+
self.ddim_sigmas = self.ddim_sigmas.float()
|
563 |
+
self.ddim_alphas = self.ddim_alphas.float()
|
564 |
+
self.ddim_alphas_prev = self.ddim_alphas_prev.float()
|
565 |
+
self.ddim_sqrt_one_minus_alphas = torch.sqrt(1. - self.ddim_alphas).float()
|
566 |
+
|
567 |
+
|
568 |
+
@torch.no_grad()
|
569 |
+
def denoise_apply_impl(self, x_target_noisy, index, noise_pred, is_step0=False):
|
570 |
+
"""
|
571 |
+
@param x_target_noisy: B,N,4,H,W
|
572 |
+
@param index: index
|
573 |
+
@param noise_pred: B,N,4,H,W
|
574 |
+
@param is_step0: bool
|
575 |
+
@return:
|
576 |
+
"""
|
577 |
+
device = x_target_noisy.device
|
578 |
+
B,N,_,H,W = x_target_noisy.shape
|
579 |
+
|
580 |
+
# apply noise
|
581 |
+
a_t = self.ddim_alphas[index].to(device).float().view(1,1,1,1,1)
|
582 |
+
a_prev = self.ddim_alphas_prev[index].to(device).float().view(1,1,1,1,1)
|
583 |
+
sqrt_one_minus_at = self.ddim_sqrt_one_minus_alphas[index].to(device).float().view(1,1,1,1,1)
|
584 |
+
sigma_t = self.ddim_sigmas[index].to(device).float().view(1,1,1,1,1)
|
585 |
+
|
586 |
+
pred_x0 = (x_target_noisy - sqrt_one_minus_at * noise_pred) / a_t.sqrt()
|
587 |
+
dir_xt = torch.clamp(1. - a_prev - sigma_t**2, min=1e-7).sqrt() * noise_pred
|
588 |
+
x_prev = a_prev.sqrt() * pred_x0 + dir_xt
|
589 |
+
if not is_step0:
|
590 |
+
noise = sigma_t * torch.randn_like(x_target_noisy)
|
591 |
+
x_prev = x_prev + noise
|
592 |
+
return x_prev
|
593 |
+
|
594 |
+
@torch.no_grad()
|
595 |
+
def denoise_apply(self, x_target_noisy, input_info, clip_embed, time_steps, index, unconditional_scale, batch_view_num=1, is_step0=False):
|
596 |
+
"""
|
597 |
+
@param x_target_noisy: B,N,4,H,W
|
598 |
+
@param input_info:
|
599 |
+
@param clip_embed: B,M,768
|
600 |
+
@param time_steps: B,
|
601 |
+
@param index: int
|
602 |
+
@param unconditional_scale:
|
603 |
+
@param batch_view_num: int
|
604 |
+
@param is_step0: bool
|
605 |
+
@return:
|
606 |
+
"""
|
607 |
+
x_input, elevation_input = input_info['x'], input_info['elevation']
|
608 |
+
B, N, C, H, W = x_target_noisy.shape
|
609 |
+
|
610 |
+
# construct source data
|
611 |
+
v_embed = self.model.get_viewpoint_embedding(B, elevation_input) # B,N,v_dim
|
612 |
+
t_embed = self.model.embed_time(time_steps) # B,t_dim
|
613 |
+
spatial_volume = self.model.spatial_volume.construct_spatial_volume(x_target_noisy, t_embed, v_embed, self.model.poses, self.model.Ks)
|
614 |
+
|
615 |
+
e_t = []
|
616 |
+
target_indices = torch.arange(N) # N
|
617 |
+
for ni in range(0, N, batch_view_num):
|
618 |
+
x_target_noisy_ = x_target_noisy[:, ni:ni + batch_view_num]
|
619 |
+
VN = x_target_noisy_.shape[1]
|
620 |
+
x_target_noisy_ = x_target_noisy_.reshape(B*VN,C,H,W)
|
621 |
+
|
622 |
+
time_steps_ = repeat_to_batch(time_steps, B, VN)
|
623 |
+
target_indices_ = target_indices[ni:ni+batch_view_num].unsqueeze(0).repeat(B,1)
|
624 |
+
clip_embed_, volume_feats_, x_concat_ = self.model.get_target_view_feats(x_input, spatial_volume, clip_embed, t_embed, v_embed, target_indices_)
|
625 |
+
if unconditional_scale!=1.0:
|
626 |
+
noise = self.model.model.predict_with_unconditional_scale(x_target_noisy_, time_steps_, clip_embed_, volume_feats_, x_concat_, unconditional_scale)
|
627 |
+
else:
|
628 |
+
noise = self.model.model(x_target_noisy_, time_steps_, clip_embed_, volume_feats_, x_concat_, is_train=False)
|
629 |
+
e_t.append(noise.view(B,VN,4,H,W))
|
630 |
+
|
631 |
+
e_t = torch.cat(e_t, 1)
|
632 |
+
x_prev = self.denoise_apply_impl(x_target_noisy, index, e_t, is_step0)
|
633 |
+
return x_prev
|
634 |
+
|
635 |
+
@torch.no_grad()
|
636 |
+
def sample(self, input_info, clip_embed, unconditional_scale=1.0, log_every_t=50, batch_view_num=1):
|
637 |
+
"""
|
638 |
+
@param input_info: x, elevation
|
639 |
+
@param clip_embed: B,M,768
|
640 |
+
@param unconditional_scale:
|
641 |
+
@param log_every_t:
|
642 |
+
@param batch_view_num:
|
643 |
+
@return:
|
644 |
+
"""
|
645 |
+
print(f"unconditional scale {unconditional_scale:.1f}")
|
646 |
+
C, H, W = 4, self.latent_size, self.latent_size
|
647 |
+
B = clip_embed.shape[0]
|
648 |
+
N = self.model.view_num
|
649 |
+
device = self.model.device
|
650 |
+
x_target_noisy = torch.randn([B, N, C, H, W], device=device)
|
651 |
+
|
652 |
+
timesteps = self.ddim_timesteps
|
653 |
+
intermediates = {'x_inter': []}
|
654 |
+
time_range = np.flip(timesteps)
|
655 |
+
total_steps = timesteps.shape[0]
|
656 |
+
|
657 |
+
iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
|
658 |
+
for i, step in enumerate(iterator):
|
659 |
+
index = total_steps - i - 1 # index in ddim state
|
660 |
+
time_steps = torch.full((B,), step, device=device, dtype=torch.long)
|
661 |
+
x_target_noisy = self.denoise_apply(x_target_noisy, input_info, clip_embed, time_steps, index, unconditional_scale, batch_view_num=batch_view_num, is_step0=index==0)
|
662 |
+
if index % log_every_t == 0 or index == total_steps - 1:
|
663 |
+
intermediates['x_inter'].append(x_target_noisy)
|
664 |
+
|
665 |
+
return x_target_noisy, intermediates
|
ldm/models/diffusion/sync_dreamer_attention.py
ADDED
@@ -0,0 +1,142 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
from ldm.modules.attention import default, zero_module, checkpoint
|
5 |
+
from ldm.modules.diffusionmodules.openaimodel import UNetModel
|
6 |
+
from ldm.modules.diffusionmodules.util import timestep_embedding
|
7 |
+
|
8 |
+
class DepthAttention(nn.Module):
|
9 |
+
def __init__(self, query_dim, context_dim, heads, dim_head, output_bias=True):
|
10 |
+
super().__init__()
|
11 |
+
inner_dim = dim_head * heads
|
12 |
+
context_dim = default(context_dim, query_dim)
|
13 |
+
|
14 |
+
self.scale = dim_head ** -0.5
|
15 |
+
self.heads = heads
|
16 |
+
self.dim_head = dim_head
|
17 |
+
|
18 |
+
self.to_q = nn.Conv2d(query_dim, inner_dim, 1, 1, bias=False)
|
19 |
+
self.to_k = nn.Conv3d(context_dim, inner_dim, 1, 1, bias=False)
|
20 |
+
self.to_v = nn.Conv3d(context_dim, inner_dim, 1, 1, bias=False)
|
21 |
+
if output_bias:
|
22 |
+
self.to_out = nn.Conv2d(inner_dim, query_dim, 1, 1)
|
23 |
+
else:
|
24 |
+
self.to_out = nn.Conv2d(inner_dim, query_dim, 1, 1, bias=False)
|
25 |
+
|
26 |
+
def forward(self, x, context):
|
27 |
+
"""
|
28 |
+
|
29 |
+
@param x: b,f0,h,w
|
30 |
+
@param context: b,f1,d,h,w
|
31 |
+
@return:
|
32 |
+
"""
|
33 |
+
hn, hd = self.heads, self.dim_head
|
34 |
+
b, _, h, w = x.shape
|
35 |
+
b, _, d, h, w = context.shape
|
36 |
+
|
37 |
+
q = self.to_q(x).reshape(b,hn,hd,h,w) # b,t,h,w
|
38 |
+
k = self.to_k(context).reshape(b,hn,hd,d,h,w) # b,t,d,h,w
|
39 |
+
v = self.to_v(context).reshape(b,hn,hd,d,h,w) # b,t,d,h,w
|
40 |
+
|
41 |
+
sim = torch.sum(q.unsqueeze(3) * k, 2) * self.scale # b,hn,d,h,w
|
42 |
+
attn = sim.softmax(dim=2)
|
43 |
+
|
44 |
+
# b,hn,hd,d,h,w * b,hn,1,d,h,w
|
45 |
+
out = torch.sum(v * attn.unsqueeze(2), 3) # b,hn,hd,h,w
|
46 |
+
out = out.reshape(b,hn*hd,h,w)
|
47 |
+
return self.to_out(out)
|
48 |
+
|
49 |
+
|
50 |
+
class DepthTransformer(nn.Module):
|
51 |
+
def __init__(self, dim, n_heads, d_head, context_dim=None, checkpoint=True):
|
52 |
+
super().__init__()
|
53 |
+
inner_dim = n_heads * d_head
|
54 |
+
self.proj_in = nn.Sequential(
|
55 |
+
nn.Conv2d(dim, inner_dim, 1, 1),
|
56 |
+
nn.GroupNorm(8, inner_dim),
|
57 |
+
nn.SiLU(True),
|
58 |
+
)
|
59 |
+
self.proj_context = nn.Sequential(
|
60 |
+
nn.Conv3d(context_dim, context_dim, 1, 1, bias=False), # no bias
|
61 |
+
nn.GroupNorm(8, context_dim),
|
62 |
+
nn.ReLU(True), # only relu, because we want input is 0, output is 0
|
63 |
+
)
|
64 |
+
self.depth_attn = DepthAttention(query_dim=inner_dim, heads=n_heads, dim_head=d_head, context_dim=context_dim, output_bias=False) # is a self-attention if not self.disable_self_attn
|
65 |
+
self.proj_out = nn.Sequential(
|
66 |
+
nn.GroupNorm(8, inner_dim),
|
67 |
+
nn.ReLU(True),
|
68 |
+
nn.Conv2d(inner_dim, inner_dim, 3, 1, 1, bias=False),
|
69 |
+
nn.GroupNorm(8, inner_dim),
|
70 |
+
nn.ReLU(True),
|
71 |
+
zero_module(nn.Conv2d(inner_dim, dim, 3, 1, 1, bias=False)),
|
72 |
+
)
|
73 |
+
self.checkpoint = checkpoint
|
74 |
+
|
75 |
+
def forward(self, x, context=None):
|
76 |
+
return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
|
77 |
+
|
78 |
+
def _forward(self, x, context):
|
79 |
+
x_in = x
|
80 |
+
x = self.proj_in(x)
|
81 |
+
context = self.proj_context(context)
|
82 |
+
x = self.depth_attn(x, context)
|
83 |
+
x = self.proj_out(x) + x_in
|
84 |
+
return x
|
85 |
+
|
86 |
+
|
87 |
+
class DepthWiseAttention(UNetModel):
|
88 |
+
def __init__(self, volume_dims=(5,16,32,64), *args, **kwargs):
|
89 |
+
super().__init__(*args, **kwargs)
|
90 |
+
# num_heads = 4
|
91 |
+
model_channels = kwargs['model_channels']
|
92 |
+
channel_mult = kwargs['channel_mult']
|
93 |
+
d0,d1,d2,d3 = volume_dims
|
94 |
+
|
95 |
+
# 4
|
96 |
+
ch = model_channels*channel_mult[2]
|
97 |
+
self.middle_conditions = DepthTransformer(ch, 4, d3 // 2, context_dim=d3)
|
98 |
+
|
99 |
+
self.output_conditions=nn.ModuleList()
|
100 |
+
self.output_b2c = {3:0,4:1,5:2,6:3,7:4,8:5,9:6,10:7,11:8}
|
101 |
+
# 8
|
102 |
+
ch = model_channels*channel_mult[2]
|
103 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d2 // 2, context_dim=d2)) # 0
|
104 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d2 // 2, context_dim=d2)) # 1
|
105 |
+
# 16
|
106 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d1 // 2, context_dim=d1)) # 2
|
107 |
+
ch = model_channels*channel_mult[1]
|
108 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d1 // 2, context_dim=d1)) # 3
|
109 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d1 // 2, context_dim=d1)) # 4
|
110 |
+
# 32
|
111 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d0 // 2, context_dim=d0)) # 5
|
112 |
+
ch = model_channels*channel_mult[0]
|
113 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d0 // 2, context_dim=d0)) # 6
|
114 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d0 // 2, context_dim=d0)) # 7
|
115 |
+
self.output_conditions.append(DepthTransformer(ch, 4, d0 // 2, context_dim=d0)) # 8
|
116 |
+
|
117 |
+
def forward(self, x, timesteps=None, context=None, source_dict=None, **kwargs):
|
118 |
+
hs = []
|
119 |
+
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
|
120 |
+
emb = self.time_embed(t_emb)
|
121 |
+
|
122 |
+
h = x.type(self.dtype)
|
123 |
+
for index, module in enumerate(self.input_blocks):
|
124 |
+
h = module(h, emb, context)
|
125 |
+
hs.append(h)
|
126 |
+
|
127 |
+
h = self.middle_block(h, emb, context)
|
128 |
+
h = self.middle_conditions(h, context=source_dict[h.shape[-1]])
|
129 |
+
|
130 |
+
for index, module in enumerate(self.output_blocks):
|
131 |
+
h = torch.cat([h, hs.pop()], dim=1)
|
132 |
+
h = module(h, emb, context)
|
133 |
+
if index in self.output_b2c:
|
134 |
+
layer = self.output_conditions[self.output_b2c[index]]
|
135 |
+
h = layer(h, context=source_dict[h.shape[-1]])
|
136 |
+
|
137 |
+
h = h.type(x.dtype)
|
138 |
+
return self.out(h)
|
139 |
+
|
140 |
+
def get_trainable_parameters(self):
|
141 |
+
paras = [para for para in self.middle_conditions.parameters()] + [para for para in self.output_conditions.parameters()]
|
142 |
+
return paras
|
ldm/models/diffusion/sync_dreamer_network.py
ADDED
@@ -0,0 +1,186 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
class Image2DResBlockWithTV(nn.Module):
|
5 |
+
def __init__(self, dim, tdim, vdim):
|
6 |
+
super().__init__()
|
7 |
+
norm = lambda c: nn.GroupNorm(8, c)
|
8 |
+
self.time_embed = nn.Conv2d(tdim, dim, 1, 1)
|
9 |
+
self.view_embed = nn.Conv2d(vdim, dim, 1, 1)
|
10 |
+
self.conv = nn.Sequential(
|
11 |
+
norm(dim),
|
12 |
+
nn.SiLU(True),
|
13 |
+
nn.Conv2d(dim, dim, 3, 1, 1),
|
14 |
+
norm(dim),
|
15 |
+
nn.SiLU(True),
|
16 |
+
nn.Conv2d(dim, dim, 3, 1, 1),
|
17 |
+
)
|
18 |
+
|
19 |
+
def forward(self, x, t, v):
|
20 |
+
return x+self.conv(x+self.time_embed(t)+self.view_embed(v))
|
21 |
+
|
22 |
+
|
23 |
+
class NoisyTargetViewEncoder(nn.Module):
|
24 |
+
def __init__(self, time_embed_dim, viewpoint_dim, run_dim=16, output_dim=8):
|
25 |
+
super().__init__()
|
26 |
+
|
27 |
+
self.init_conv = nn.Conv2d(4, run_dim, 3, 1, 1)
|
28 |
+
self.out_conv0 = Image2DResBlockWithTV(run_dim, time_embed_dim, viewpoint_dim)
|
29 |
+
self.out_conv1 = Image2DResBlockWithTV(run_dim, time_embed_dim, viewpoint_dim)
|
30 |
+
self.out_conv2 = Image2DResBlockWithTV(run_dim, time_embed_dim, viewpoint_dim)
|
31 |
+
self.final_out = nn.Sequential(
|
32 |
+
nn.GroupNorm(8, run_dim),
|
33 |
+
nn.SiLU(True),
|
34 |
+
nn.Conv2d(run_dim, output_dim, 3, 1, 1)
|
35 |
+
)
|
36 |
+
|
37 |
+
def forward(self, x, t, v):
|
38 |
+
B, DT = t.shape
|
39 |
+
t = t.view(B, DT, 1, 1)
|
40 |
+
B, DV = v.shape
|
41 |
+
v = v.view(B, DV, 1, 1)
|
42 |
+
|
43 |
+
x = self.init_conv(x)
|
44 |
+
x = self.out_conv0(x, t, v)
|
45 |
+
x = self.out_conv1(x, t, v)
|
46 |
+
x = self.out_conv2(x, t, v)
|
47 |
+
x = self.final_out(x)
|
48 |
+
return x
|
49 |
+
|
50 |
+
class SpatialUpTimeBlock(nn.Module):
|
51 |
+
def __init__(self, x_in_dim, t_in_dim, out_dim):
|
52 |
+
super().__init__()
|
53 |
+
norm_act = lambda c: nn.GroupNorm(8, c)
|
54 |
+
self.t_conv = nn.Conv3d(t_in_dim, x_in_dim, 1, 1) # 16
|
55 |
+
self.norm = norm_act(x_in_dim)
|
56 |
+
self.silu = nn.SiLU(True)
|
57 |
+
self.conv = nn.ConvTranspose3d(x_in_dim, out_dim, kernel_size=3, padding=1, output_padding=1, stride=2)
|
58 |
+
|
59 |
+
def forward(self, x, t):
|
60 |
+
x = x + self.t_conv(t)
|
61 |
+
return self.conv(self.silu(self.norm(x)))
|
62 |
+
|
63 |
+
class SpatialTimeBlock(nn.Module):
|
64 |
+
def __init__(self, x_in_dim, t_in_dim, out_dim, stride):
|
65 |
+
super().__init__()
|
66 |
+
norm_act = lambda c: nn.GroupNorm(8, c)
|
67 |
+
self.t_conv = nn.Conv3d(t_in_dim, x_in_dim, 1, 1) # 16
|
68 |
+
self.bn = norm_act(x_in_dim)
|
69 |
+
self.silu = nn.SiLU(True)
|
70 |
+
self.conv = nn.Conv3d(x_in_dim, out_dim, 3, stride=stride, padding=1)
|
71 |
+
|
72 |
+
def forward(self, x, t):
|
73 |
+
x = x + self.t_conv(t)
|
74 |
+
return self.conv(self.silu(self.bn(x)))
|
75 |
+
|
76 |
+
class SpatialTime3DNet(nn.Module):
|
77 |
+
def __init__(self, time_dim=256, input_dim=128, dims=(32, 64, 128, 256)):
|
78 |
+
super().__init__()
|
79 |
+
d0, d1, d2, d3 = dims
|
80 |
+
dt = time_dim
|
81 |
+
|
82 |
+
self.init_conv = nn.Conv3d(input_dim, d0, 3, 1, 1) # 32
|
83 |
+
self.conv0 = SpatialTimeBlock(d0, dt, d0, stride=1)
|
84 |
+
|
85 |
+
self.conv1 = SpatialTimeBlock(d0, dt, d1, stride=2)
|
86 |
+
self.conv2_0 = SpatialTimeBlock(d1, dt, d1, stride=1)
|
87 |
+
self.conv2_1 = SpatialTimeBlock(d1, dt, d1, stride=1)
|
88 |
+
|
89 |
+
self.conv3 = SpatialTimeBlock(d1, dt, d2, stride=2)
|
90 |
+
self.conv4_0 = SpatialTimeBlock(d2, dt, d2, stride=1)
|
91 |
+
self.conv4_1 = SpatialTimeBlock(d2, dt, d2, stride=1)
|
92 |
+
|
93 |
+
self.conv5 = SpatialTimeBlock(d2, dt, d3, stride=2)
|
94 |
+
self.conv6_0 = SpatialTimeBlock(d3, dt, d3, stride=1)
|
95 |
+
self.conv6_1 = SpatialTimeBlock(d3, dt, d3, stride=1)
|
96 |
+
|
97 |
+
self.conv7 = SpatialUpTimeBlock(d3, dt, d2)
|
98 |
+
self.conv8 = SpatialUpTimeBlock(d2, dt, d1)
|
99 |
+
self.conv9 = SpatialUpTimeBlock(d1, dt, d0)
|
100 |
+
|
101 |
+
def forward(self, x, t):
|
102 |
+
B, C = t.shape
|
103 |
+
t = t.view(B, C, 1, 1, 1)
|
104 |
+
|
105 |
+
x = self.init_conv(x)
|
106 |
+
conv0 = self.conv0(x, t)
|
107 |
+
|
108 |
+
x = self.conv1(conv0, t)
|
109 |
+
x = self.conv2_0(x, t)
|
110 |
+
conv2 = self.conv2_1(x, t)
|
111 |
+
|
112 |
+
x = self.conv3(conv2, t)
|
113 |
+
x = self.conv4_0(x, t)
|
114 |
+
conv4 = self.conv4_1(x, t)
|
115 |
+
|
116 |
+
x = self.conv5(conv4, t)
|
117 |
+
x = self.conv6_0(x, t)
|
118 |
+
x = self.conv6_1(x, t)
|
119 |
+
|
120 |
+
x = conv4 + self.conv7(x, t)
|
121 |
+
x = conv2 + self.conv8(x, t)
|
122 |
+
x = conv0 + self.conv9(x, t)
|
123 |
+
return x
|
124 |
+
|
125 |
+
class FrustumTVBlock(nn.Module):
|
126 |
+
def __init__(self, x_dim, t_dim, v_dim, out_dim, stride):
|
127 |
+
super().__init__()
|
128 |
+
norm_act = lambda c: nn.GroupNorm(8, c)
|
129 |
+
self.t_conv = nn.Conv3d(t_dim, x_dim, 1, 1) # 16
|
130 |
+
self.v_conv = nn.Conv3d(v_dim, x_dim, 1, 1) # 16
|
131 |
+
self.bn = norm_act(x_dim)
|
132 |
+
self.silu = nn.SiLU(True)
|
133 |
+
self.conv = nn.Conv3d(x_dim, out_dim, 3, stride=stride, padding=1)
|
134 |
+
|
135 |
+
def forward(self, x, t, v):
|
136 |
+
x = x + self.t_conv(t) + self.v_conv(v)
|
137 |
+
return self.conv(self.silu(self.bn(x)))
|
138 |
+
|
139 |
+
class FrustumTVUpBlock(nn.Module):
|
140 |
+
def __init__(self, x_dim, t_dim, v_dim, out_dim):
|
141 |
+
super().__init__()
|
142 |
+
norm_act = lambda c: nn.GroupNorm(8, c)
|
143 |
+
self.t_conv = nn.Conv3d(t_dim, x_dim, 1, 1) # 16
|
144 |
+
self.v_conv = nn.Conv3d(v_dim, x_dim, 1, 1) # 16
|
145 |
+
self.norm = norm_act(x_dim)
|
146 |
+
self.silu = nn.SiLU(True)
|
147 |
+
self.conv = nn.ConvTranspose3d(x_dim, out_dim, kernel_size=3, padding=1, output_padding=1, stride=2)
|
148 |
+
|
149 |
+
def forward(self, x, t, v):
|
150 |
+
x = x + self.t_conv(t) + self.v_conv(v)
|
151 |
+
return self.conv(self.silu(self.norm(x)))
|
152 |
+
|
153 |
+
class FrustumTV3DNet(nn.Module):
|
154 |
+
def __init__(self, in_dim, t_dim, v_dim, dims=(32, 64, 128, 256)):
|
155 |
+
super().__init__()
|
156 |
+
self.conv0 = nn.Conv3d(in_dim, dims[0], 3, 1, 1) # 32
|
157 |
+
|
158 |
+
self.conv1 = FrustumTVBlock(dims[0], t_dim, v_dim, dims[1], 2)
|
159 |
+
self.conv2 = FrustumTVBlock(dims[1], t_dim, v_dim, dims[1], 1)
|
160 |
+
|
161 |
+
self.conv3 = FrustumTVBlock(dims[1], t_dim, v_dim, dims[2], 2)
|
162 |
+
self.conv4 = FrustumTVBlock(dims[2], t_dim, v_dim, dims[2], 1)
|
163 |
+
|
164 |
+
self.conv5 = FrustumTVBlock(dims[2], t_dim, v_dim, dims[3], 2)
|
165 |
+
self.conv6 = FrustumTVBlock(dims[3], t_dim, v_dim, dims[3], 1)
|
166 |
+
|
167 |
+
self.up0 = FrustumTVUpBlock(dims[3], t_dim, v_dim, dims[2])
|
168 |
+
self.up1 = FrustumTVUpBlock(dims[2], t_dim, v_dim, dims[1])
|
169 |
+
self.up2 = FrustumTVUpBlock(dims[1], t_dim, v_dim, dims[0])
|
170 |
+
|
171 |
+
def forward(self, x, t, v):
|
172 |
+
B,DT = t.shape
|
173 |
+
t = t.view(B,DT,1,1,1)
|
174 |
+
B,DV = v.shape
|
175 |
+
v = v.view(B,DV,1,1,1)
|
176 |
+
|
177 |
+
b, _, d, h, w = x.shape
|
178 |
+
x0 = self.conv0(x)
|
179 |
+
x1 = self.conv2(self.conv1(x0, t, v), t, v)
|
180 |
+
x2 = self.conv4(self.conv3(x1, t, v), t, v)
|
181 |
+
x3 = self.conv6(self.conv5(x2, t, v), t, v)
|
182 |
+
|
183 |
+
x2 = self.up0(x3, t, v) + x2
|
184 |
+
x1 = self.up1(x2, t, v) + x1
|
185 |
+
x0 = self.up2(x1, t, v) + x0
|
186 |
+
return {w: x0, w//2: x1, w//4: x2, w//8: x3}
|
ldm/models/diffusion/sync_dreamer_utils.py
ADDED
@@ -0,0 +1,103 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from kornia import create_meshgrid
|
3 |
+
|
4 |
+
|
5 |
+
def project_and_normalize(ref_grid, src_proj, length):
|
6 |
+
"""
|
7 |
+
|
8 |
+
@param ref_grid: b 3 n
|
9 |
+
@param src_proj: b 4 4
|
10 |
+
@param length: int
|
11 |
+
@return: b, n, 2
|
12 |
+
"""
|
13 |
+
src_grid = src_proj[:, :3, :3] @ ref_grid + src_proj[:, :3, 3:] # b 3 n
|
14 |
+
div_val = src_grid[:, -1:]
|
15 |
+
div_val[div_val<1e-4] = 1e-4
|
16 |
+
src_grid = src_grid[:, :2] / div_val # divide by depth (b, 2, n)
|
17 |
+
src_grid[:, 0] = src_grid[:, 0]/((length - 1) / 2) - 1 # scale to -1~1
|
18 |
+
src_grid[:, 1] = src_grid[:, 1]/((length - 1) / 2) - 1 # scale to -1~1
|
19 |
+
src_grid = src_grid.permute(0, 2, 1) # (b, n, 2)
|
20 |
+
return src_grid
|
21 |
+
|
22 |
+
|
23 |
+
def construct_project_matrix(x_ratio, y_ratio, Ks, poses):
|
24 |
+
"""
|
25 |
+
@param x_ratio: float
|
26 |
+
@param y_ratio: float
|
27 |
+
@param Ks: b,3,3
|
28 |
+
@param poses: b,3,4
|
29 |
+
@return:
|
30 |
+
"""
|
31 |
+
rfn = Ks.shape[0]
|
32 |
+
scale_m = torch.tensor([x_ratio, y_ratio, 1.0], dtype=torch.float32, device=Ks.device)
|
33 |
+
scale_m = torch.diag(scale_m)
|
34 |
+
ref_prj = scale_m[None, :, :] @ Ks @ poses # rfn,3,4
|
35 |
+
pad_vals = torch.zeros([rfn, 1, 4], dtype=torch.float32, device=ref_prj.device)
|
36 |
+
pad_vals[:, :, 3] = 1.0
|
37 |
+
ref_prj = torch.cat([ref_prj, pad_vals], 1) # rfn,4,4
|
38 |
+
return ref_prj
|
39 |
+
|
40 |
+
def get_warp_coordinates(volume_xyz, warp_size, input_size, Ks, warp_pose):
|
41 |
+
B, _, D, H, W = volume_xyz.shape
|
42 |
+
ratio = warp_size / input_size
|
43 |
+
warp_proj = construct_project_matrix(ratio, ratio, Ks, warp_pose) # B,4,4
|
44 |
+
warp_coords = project_and_normalize(volume_xyz.view(B,3,D*H*W), warp_proj, warp_size).view(B, D, H, W, 2)
|
45 |
+
return warp_coords
|
46 |
+
|
47 |
+
|
48 |
+
def create_target_volume(depth_size, volume_size, input_image_size, pose_target, K, near=None, far=None):
|
49 |
+
device, dtype = pose_target.device, pose_target.dtype
|
50 |
+
|
51 |
+
# compute a depth range on the unit sphere
|
52 |
+
H, W, D, B = volume_size, volume_size, depth_size, pose_target.shape[0]
|
53 |
+
if near is not None and far is not None :
|
54 |
+
# near, far b,1,h,w
|
55 |
+
depth_values = torch.linspace(0, 1, steps=depth_size).to(near.device).to(near.dtype) # d
|
56 |
+
depth_values = depth_values.view(1, D, 1, 1) # 1,d,1,1
|
57 |
+
depth_values = depth_values * (far - near) + near # b d h w
|
58 |
+
depth_values = depth_values.view(B, 1, D, H * W)
|
59 |
+
else:
|
60 |
+
near, far = near_far_from_unit_sphere_using_camera_poses(pose_target) # b 1
|
61 |
+
depth_values = torch.linspace(0, 1, steps=depth_size).to(near.device).to(near.dtype) # d
|
62 |
+
depth_values = depth_values[None,:,None] * (far[:,None,:] - near[:,None,:]) + near[:,None,:] # b d 1
|
63 |
+
depth_values = depth_values.view(B, 1, D, 1).expand(B, 1, D, H*W)
|
64 |
+
|
65 |
+
ratio = volume_size / input_image_size
|
66 |
+
|
67 |
+
# creat a grid on the target (reference) view
|
68 |
+
# H, W, D, B = volume_size, volume_size, depth_values.shape[1], depth_values.shape[0]
|
69 |
+
|
70 |
+
# creat mesh grid: note reference also means target
|
71 |
+
ref_grid = create_meshgrid(H, W, normalized_coordinates=False) # (1, H, W, 2)
|
72 |
+
ref_grid = ref_grid.to(device).to(dtype)
|
73 |
+
ref_grid = ref_grid.permute(0, 3, 1, 2) # (1, 2, H, W)
|
74 |
+
ref_grid = ref_grid.reshape(1, 2, H*W) # (1, 2, H*W)
|
75 |
+
ref_grid = ref_grid.expand(B, -1, -1) # (B, 2, H*W)
|
76 |
+
ref_grid = torch.cat((ref_grid, torch.ones(B, 1, H*W, dtype=ref_grid.dtype, device=ref_grid.device)), dim=1) # (B, 3, H*W)
|
77 |
+
ref_grid = ref_grid.unsqueeze(2) * depth_values # (B, 3, D, H*W)
|
78 |
+
|
79 |
+
# unproject to space and transfer to world coordinates.
|
80 |
+
Ks = K
|
81 |
+
ref_proj = construct_project_matrix(ratio, ratio, Ks, pose_target) # B,4,4
|
82 |
+
ref_proj_inv = torch.inverse(ref_proj) # B,4,4
|
83 |
+
ref_grid = ref_proj_inv[:,:3,:3] @ ref_grid.view(B,3,D*H*W) + ref_proj_inv[:,:3,3:] # B,3,3 @ B,3,DHW + B,3,1 => B,3,DHW
|
84 |
+
return ref_grid.reshape(B,3,D,H,W), depth_values.view(B,1,D,H,W)
|
85 |
+
|
86 |
+
def near_far_from_unit_sphere_using_camera_poses(camera_poses):
|
87 |
+
"""
|
88 |
+
@param camera_poses: b 3 4
|
89 |
+
@return:
|
90 |
+
near: b,1
|
91 |
+
far: b,1
|
92 |
+
"""
|
93 |
+
R_w2c = camera_poses[..., :3, :3] # b 3 3
|
94 |
+
t_w2c = camera_poses[..., :3, 3:] # b 3 1
|
95 |
+
camera_origin = -R_w2c.permute(0,2,1) @ t_w2c # b 3 1
|
96 |
+
# R_w2c.T @ (0,0,1) = z_dir
|
97 |
+
camera_orient = R_w2c.permute(0,2,1)[...,:3,2:3] # b 3 1
|
98 |
+
camera_origin, camera_orient = camera_origin[...,0], camera_orient[..., 0] # b 3
|
99 |
+
a = torch.sum(camera_orient ** 2, dim=-1, keepdim=True) # b 1
|
100 |
+
b = -torch.sum(camera_orient * camera_origin, dim=-1, keepdim=True) # b 1
|
101 |
+
mid = b / a # b 1
|
102 |
+
near, far = mid - 1.0, mid + 1.0
|
103 |
+
return near, far
|
ldm/modules/attention.py
ADDED
@@ -0,0 +1,336 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from inspect import isfunction
|
2 |
+
import math
|
3 |
+
import torch
|
4 |
+
import torch.nn.functional as F
|
5 |
+
from torch import nn, einsum
|
6 |
+
from einops import rearrange, repeat
|
7 |
+
|
8 |
+
from ldm.modules.diffusionmodules.util import checkpoint
|
9 |
+
|
10 |
+
|
11 |
+
def exists(val):
|
12 |
+
return val is not None
|
13 |
+
|
14 |
+
|
15 |
+
def uniq(arr):
|
16 |
+
return{el: True for el in arr}.keys()
|
17 |
+
|
18 |
+
|
19 |
+
def default(val, d):
|
20 |
+
if exists(val):
|
21 |
+
return val
|
22 |
+
return d() if isfunction(d) else d
|
23 |
+
|
24 |
+
|
25 |
+
def max_neg_value(t):
|
26 |
+
return -torch.finfo(t.dtype).max
|
27 |
+
|
28 |
+
|
29 |
+
def init_(tensor):
|
30 |
+
dim = tensor.shape[-1]
|
31 |
+
std = 1 / math.sqrt(dim)
|
32 |
+
tensor.uniform_(-std, std)
|
33 |
+
return tensor
|
34 |
+
|
35 |
+
|
36 |
+
# feedforward
|
37 |
+
class GEGLU(nn.Module):
|
38 |
+
def __init__(self, dim_in, dim_out):
|
39 |
+
super().__init__()
|
40 |
+
self.proj = nn.Linear(dim_in, dim_out * 2)
|
41 |
+
|
42 |
+
def forward(self, x):
|
43 |
+
x, gate = self.proj(x).chunk(2, dim=-1)
|
44 |
+
return x * F.gelu(gate)
|
45 |
+
# feedforward
|
46 |
+
class ConvGEGLU(nn.Module):
|
47 |
+
def __init__(self, dim_in, dim_out):
|
48 |
+
super().__init__()
|
49 |
+
self.proj = nn.Conv2d(dim_in, dim_out * 2, 1, 1, 0)
|
50 |
+
|
51 |
+
def forward(self, x):
|
52 |
+
x, gate = self.proj(x).chunk(2, dim=1)
|
53 |
+
return x * F.gelu(gate)
|
54 |
+
|
55 |
+
|
56 |
+
class FeedForward(nn.Module):
|
57 |
+
def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
|
58 |
+
super().__init__()
|
59 |
+
inner_dim = int(dim * mult)
|
60 |
+
dim_out = default(dim_out, dim)
|
61 |
+
project_in = nn.Sequential(
|
62 |
+
nn.Linear(dim, inner_dim),
|
63 |
+
nn.GELU()
|
64 |
+
) if not glu else GEGLU(dim, inner_dim)
|
65 |
+
|
66 |
+
self.net = nn.Sequential(
|
67 |
+
project_in,
|
68 |
+
nn.Dropout(dropout),
|
69 |
+
nn.Linear(inner_dim, dim_out)
|
70 |
+
)
|
71 |
+
|
72 |
+
def forward(self, x):
|
73 |
+
return self.net(x)
|
74 |
+
|
75 |
+
|
76 |
+
def zero_module(module):
|
77 |
+
"""
|
78 |
+
Zero out the parameters of a module and return it.
|
79 |
+
"""
|
80 |
+
for p in module.parameters():
|
81 |
+
p.detach().zero_()
|
82 |
+
return module
|
83 |
+
|
84 |
+
|
85 |
+
def Normalize(in_channels):
|
86 |
+
return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
87 |
+
|
88 |
+
|
89 |
+
class LinearAttention(nn.Module):
|
90 |
+
def __init__(self, dim, heads=4, dim_head=32):
|
91 |
+
super().__init__()
|
92 |
+
self.heads = heads
|
93 |
+
hidden_dim = dim_head * heads
|
94 |
+
self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
|
95 |
+
self.to_out = nn.Conv2d(hidden_dim, dim, 1)
|
96 |
+
|
97 |
+
def forward(self, x):
|
98 |
+
b, c, h, w = x.shape
|
99 |
+
qkv = self.to_qkv(x)
|
100 |
+
q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
|
101 |
+
k = k.softmax(dim=-1)
|
102 |
+
context = torch.einsum('bhdn,bhen->bhde', k, v)
|
103 |
+
out = torch.einsum('bhde,bhdn->bhen', context, q)
|
104 |
+
out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
|
105 |
+
return self.to_out(out)
|
106 |
+
|
107 |
+
|
108 |
+
class SpatialSelfAttention(nn.Module):
|
109 |
+
def __init__(self, in_channels):
|
110 |
+
super().__init__()
|
111 |
+
self.in_channels = in_channels
|
112 |
+
|
113 |
+
self.norm = Normalize(in_channels)
|
114 |
+
self.q = torch.nn.Conv2d(in_channels,
|
115 |
+
in_channels,
|
116 |
+
kernel_size=1,
|
117 |
+
stride=1,
|
118 |
+
padding=0)
|
119 |
+
self.k = torch.nn.Conv2d(in_channels,
|
120 |
+
in_channels,
|
121 |
+
kernel_size=1,
|
122 |
+
stride=1,
|
123 |
+
padding=0)
|
124 |
+
self.v = torch.nn.Conv2d(in_channels,
|
125 |
+
in_channels,
|
126 |
+
kernel_size=1,
|
127 |
+
stride=1,
|
128 |
+
padding=0)
|
129 |
+
self.proj_out = torch.nn.Conv2d(in_channels,
|
130 |
+
in_channels,
|
131 |
+
kernel_size=1,
|
132 |
+
stride=1,
|
133 |
+
padding=0)
|
134 |
+
|
135 |
+
def forward(self, x):
|
136 |
+
h_ = x
|
137 |
+
h_ = self.norm(h_)
|
138 |
+
q = self.q(h_)
|
139 |
+
k = self.k(h_)
|
140 |
+
v = self.v(h_)
|
141 |
+
|
142 |
+
# compute attention
|
143 |
+
b,c,h,w = q.shape
|
144 |
+
q = rearrange(q, 'b c h w -> b (h w) c')
|
145 |
+
k = rearrange(k, 'b c h w -> b c (h w)')
|
146 |
+
w_ = torch.einsum('bij,bjk->bik', q, k)
|
147 |
+
|
148 |
+
w_ = w_ * (int(c)**(-0.5))
|
149 |
+
w_ = torch.nn.functional.softmax(w_, dim=2)
|
150 |
+
|
151 |
+
# attend to values
|
152 |
+
v = rearrange(v, 'b c h w -> b c (h w)')
|
153 |
+
w_ = rearrange(w_, 'b i j -> b j i')
|
154 |
+
h_ = torch.einsum('bij,bjk->bik', v, w_)
|
155 |
+
h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
|
156 |
+
h_ = self.proj_out(h_)
|
157 |
+
|
158 |
+
return x+h_
|
159 |
+
|
160 |
+
|
161 |
+
class CrossAttention(nn.Module):
|
162 |
+
def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
|
163 |
+
super().__init__()
|
164 |
+
inner_dim = dim_head * heads
|
165 |
+
context_dim = default(context_dim, query_dim)
|
166 |
+
|
167 |
+
self.scale = dim_head ** -0.5
|
168 |
+
self.heads = heads
|
169 |
+
|
170 |
+
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
|
171 |
+
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
|
172 |
+
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
|
173 |
+
|
174 |
+
self.to_out = nn.Sequential(
|
175 |
+
nn.Linear(inner_dim, query_dim),
|
176 |
+
nn.Dropout(dropout)
|
177 |
+
)
|
178 |
+
|
179 |
+
def forward(self, x, context=None, mask=None):
|
180 |
+
h = self.heads
|
181 |
+
|
182 |
+
q = self.to_q(x)
|
183 |
+
context = default(context, x)
|
184 |
+
k = self.to_k(context)
|
185 |
+
v = self.to_v(context)
|
186 |
+
|
187 |
+
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
|
188 |
+
|
189 |
+
sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
|
190 |
+
|
191 |
+
if exists(mask):
|
192 |
+
mask = mask>0
|
193 |
+
mask = rearrange(mask, 'b ... -> b (...)')
|
194 |
+
max_neg_value = -torch.finfo(sim.dtype).max
|
195 |
+
mask = repeat(mask, 'b j -> (b h) () j', h=h)
|
196 |
+
sim.masked_fill_(~mask, max_neg_value)
|
197 |
+
|
198 |
+
# attention, what we cannot get enough of
|
199 |
+
attn = sim.softmax(dim=-1)
|
200 |
+
|
201 |
+
out = einsum('b i j, b j d -> b i d', attn, v)
|
202 |
+
out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
|
203 |
+
return self.to_out(out)
|
204 |
+
|
205 |
+
class BasicSpatialTransformer(nn.Module):
|
206 |
+
def __init__(self, dim, n_heads, d_head, context_dim=None, checkpoint=True):
|
207 |
+
super().__init__()
|
208 |
+
inner_dim = n_heads * d_head
|
209 |
+
self.proj_in = nn.Sequential(
|
210 |
+
nn.GroupNorm(8, dim),
|
211 |
+
nn.Conv2d(dim, inner_dim, kernel_size=1, stride=1, padding=0),
|
212 |
+
nn.GroupNorm(8, inner_dim),
|
213 |
+
nn.ReLU(True),
|
214 |
+
)
|
215 |
+
self.attn = CrossAttention(query_dim=inner_dim, heads=n_heads, dim_head=d_head, context_dim=context_dim) # is a self-attention if not self.disable_self_attn
|
216 |
+
self.out_conv = nn.Sequential(
|
217 |
+
nn.GroupNorm(8, inner_dim),
|
218 |
+
nn.ReLU(True),
|
219 |
+
nn.Conv2d(inner_dim, inner_dim, 1, 1),
|
220 |
+
)
|
221 |
+
self.proj_out = nn.Sequential(
|
222 |
+
nn.GroupNorm(8, inner_dim),
|
223 |
+
nn.ReLU(True),
|
224 |
+
zero_module(nn.Conv2d(inner_dim, dim, kernel_size=1, stride=1, padding=0)),
|
225 |
+
)
|
226 |
+
self.checkpoint = checkpoint
|
227 |
+
|
228 |
+
def forward(self, x, context=None):
|
229 |
+
return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
|
230 |
+
|
231 |
+
def _forward(self, x, context):
|
232 |
+
# input
|
233 |
+
b,_,h,w = x.shape
|
234 |
+
x_in = x
|
235 |
+
x = self.proj_in(x)
|
236 |
+
|
237 |
+
# attention
|
238 |
+
x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
|
239 |
+
context = rearrange(context, 'b c h w -> b (h w) c').contiguous()
|
240 |
+
x = self.attn(x, context) + x
|
241 |
+
x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
|
242 |
+
|
243 |
+
# output
|
244 |
+
x = self.out_conv(x) + x
|
245 |
+
x = self.proj_out(x) + x_in
|
246 |
+
return x
|
247 |
+
|
248 |
+
class BasicTransformerBlock(nn.Module):
|
249 |
+
def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True, disable_self_attn=False):
|
250 |
+
super().__init__()
|
251 |
+
self.disable_self_attn = disable_self_attn
|
252 |
+
self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
|
253 |
+
context_dim=context_dim if self.disable_self_attn else None) # is a self-attention if not self.disable_self_attn
|
254 |
+
self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
|
255 |
+
self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
|
256 |
+
heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none
|
257 |
+
self.norm1 = nn.LayerNorm(dim)
|
258 |
+
self.norm2 = nn.LayerNorm(dim)
|
259 |
+
self.norm3 = nn.LayerNorm(dim)
|
260 |
+
self.checkpoint = checkpoint
|
261 |
+
|
262 |
+
def forward(self, x, context=None):
|
263 |
+
return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
|
264 |
+
|
265 |
+
def _forward(self, x, context=None):
|
266 |
+
x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
|
267 |
+
x = self.attn2(self.norm2(x), context=context) + x
|
268 |
+
x = self.ff(self.norm3(x)) + x
|
269 |
+
return x
|
270 |
+
|
271 |
+
class ConvFeedForward(nn.Module):
|
272 |
+
def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
|
273 |
+
super().__init__()
|
274 |
+
inner_dim = int(dim * mult)
|
275 |
+
dim_out = default(dim_out, dim)
|
276 |
+
project_in = nn.Sequential(
|
277 |
+
nn.Conv2d(dim, inner_dim, 1, 1, 0),
|
278 |
+
nn.GELU()
|
279 |
+
) if not glu else ConvGEGLU(dim, inner_dim)
|
280 |
+
|
281 |
+
self.net = nn.Sequential(
|
282 |
+
project_in,
|
283 |
+
nn.Dropout(dropout),
|
284 |
+
nn.Conv2d(inner_dim, dim_out, 1, 1, 0)
|
285 |
+
)
|
286 |
+
|
287 |
+
def forward(self, x):
|
288 |
+
return self.net(x)
|
289 |
+
|
290 |
+
|
291 |
+
class SpatialTransformer(nn.Module):
|
292 |
+
"""
|
293 |
+
Transformer block for image-like data.
|
294 |
+
First, project the input (aka embedding)
|
295 |
+
and reshape to b, t, d.
|
296 |
+
Then apply standard transformer action.
|
297 |
+
Finally, reshape to image
|
298 |
+
"""
|
299 |
+
def __init__(self, in_channels, n_heads, d_head,
|
300 |
+
depth=1, dropout=0., context_dim=None,
|
301 |
+
disable_self_attn=False):
|
302 |
+
super().__init__()
|
303 |
+
self.in_channels = in_channels
|
304 |
+
inner_dim = n_heads * d_head
|
305 |
+
self.norm = Normalize(in_channels)
|
306 |
+
|
307 |
+
self.proj_in = nn.Conv2d(in_channels,
|
308 |
+
inner_dim,
|
309 |
+
kernel_size=1,
|
310 |
+
stride=1,
|
311 |
+
padding=0)
|
312 |
+
|
313 |
+
self.transformer_blocks = nn.ModuleList(
|
314 |
+
[BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim,
|
315 |
+
disable_self_attn=disable_self_attn)
|
316 |
+
for d in range(depth)]
|
317 |
+
)
|
318 |
+
|
319 |
+
self.proj_out = zero_module(nn.Conv2d(inner_dim,
|
320 |
+
in_channels,
|
321 |
+
kernel_size=1,
|
322 |
+
stride=1,
|
323 |
+
padding=0))
|
324 |
+
|
325 |
+
def forward(self, x, context=None):
|
326 |
+
# note: if no context is given, cross-attention defaults to self-attention
|
327 |
+
b, c, h, w = x.shape
|
328 |
+
x_in = x
|
329 |
+
x = self.norm(x)
|
330 |
+
x = self.proj_in(x)
|
331 |
+
x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
|
332 |
+
for block in self.transformer_blocks:
|
333 |
+
x = block(x, context=context)
|
334 |
+
x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
|
335 |
+
x = self.proj_out(x)
|
336 |
+
return x + x_in
|
ldm/modules/diffusionmodules/__init__.py
ADDED
File without changes
|
ldm/modules/diffusionmodules/model.py
ADDED
@@ -0,0 +1,835 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# pytorch_diffusion + derived encoder decoder
|
2 |
+
import math
|
3 |
+
import torch
|
4 |
+
import torch.nn as nn
|
5 |
+
import numpy as np
|
6 |
+
from einops import rearrange
|
7 |
+
|
8 |
+
from ldm.util import instantiate_from_config
|
9 |
+
from ldm.modules.attention import LinearAttention
|
10 |
+
|
11 |
+
|
12 |
+
def get_timestep_embedding(timesteps, embedding_dim):
|
13 |
+
"""
|
14 |
+
This matches the implementation in Denoising Diffusion Probabilistic Models:
|
15 |
+
From Fairseq.
|
16 |
+
Build sinusoidal embeddings.
|
17 |
+
This matches the implementation in tensor2tensor, but differs slightly
|
18 |
+
from the description in Section 3.5 of "Attention Is All You Need".
|
19 |
+
"""
|
20 |
+
assert len(timesteps.shape) == 1
|
21 |
+
|
22 |
+
half_dim = embedding_dim // 2
|
23 |
+
emb = math.log(10000) / (half_dim - 1)
|
24 |
+
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
|
25 |
+
emb = emb.to(device=timesteps.device)
|
26 |
+
emb = timesteps.float()[:, None] * emb[None, :]
|
27 |
+
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
|
28 |
+
if embedding_dim % 2 == 1: # zero pad
|
29 |
+
emb = torch.nn.functional.pad(emb, (0,1,0,0))
|
30 |
+
return emb
|
31 |
+
|
32 |
+
|
33 |
+
def nonlinearity(x):
|
34 |
+
# swish
|
35 |
+
return x*torch.sigmoid(x)
|
36 |
+
|
37 |
+
|
38 |
+
def Normalize(in_channels, num_groups=32):
|
39 |
+
return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
40 |
+
|
41 |
+
|
42 |
+
class Upsample(nn.Module):
|
43 |
+
def __init__(self, in_channels, with_conv):
|
44 |
+
super().__init__()
|
45 |
+
self.with_conv = with_conv
|
46 |
+
if self.with_conv:
|
47 |
+
self.conv = torch.nn.Conv2d(in_channels,
|
48 |
+
in_channels,
|
49 |
+
kernel_size=3,
|
50 |
+
stride=1,
|
51 |
+
padding=1)
|
52 |
+
|
53 |
+
def forward(self, x):
|
54 |
+
x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
|
55 |
+
if self.with_conv:
|
56 |
+
x = self.conv(x)
|
57 |
+
return x
|
58 |
+
|
59 |
+
|
60 |
+
class Downsample(nn.Module):
|
61 |
+
def __init__(self, in_channels, with_conv):
|
62 |
+
super().__init__()
|
63 |
+
self.with_conv = with_conv
|
64 |
+
if self.with_conv:
|
65 |
+
# no asymmetric padding in torch conv, must do it ourselves
|
66 |
+
self.conv = torch.nn.Conv2d(in_channels,
|
67 |
+
in_channels,
|
68 |
+
kernel_size=3,
|
69 |
+
stride=2,
|
70 |
+
padding=0)
|
71 |
+
|
72 |
+
def forward(self, x):
|
73 |
+
if self.with_conv:
|
74 |
+
pad = (0,1,0,1)
|
75 |
+
x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
|
76 |
+
x = self.conv(x)
|
77 |
+
else:
|
78 |
+
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
|
79 |
+
return x
|
80 |
+
|
81 |
+
|
82 |
+
class ResnetBlock(nn.Module):
|
83 |
+
def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
|
84 |
+
dropout, temb_channels=512):
|
85 |
+
super().__init__()
|
86 |
+
self.in_channels = in_channels
|
87 |
+
out_channels = in_channels if out_channels is None else out_channels
|
88 |
+
self.out_channels = out_channels
|
89 |
+
self.use_conv_shortcut = conv_shortcut
|
90 |
+
|
91 |
+
self.norm1 = Normalize(in_channels)
|
92 |
+
self.conv1 = torch.nn.Conv2d(in_channels,
|
93 |
+
out_channels,
|
94 |
+
kernel_size=3,
|
95 |
+
stride=1,
|
96 |
+
padding=1)
|
97 |
+
if temb_channels > 0:
|
98 |
+
self.temb_proj = torch.nn.Linear(temb_channels,
|
99 |
+
out_channels)
|
100 |
+
self.norm2 = Normalize(out_channels)
|
101 |
+
self.dropout = torch.nn.Dropout(dropout)
|
102 |
+
self.conv2 = torch.nn.Conv2d(out_channels,
|
103 |
+
out_channels,
|
104 |
+
kernel_size=3,
|
105 |
+
stride=1,
|
106 |
+
padding=1)
|
107 |
+
if self.in_channels != self.out_channels:
|
108 |
+
if self.use_conv_shortcut:
|
109 |
+
self.conv_shortcut = torch.nn.Conv2d(in_channels,
|
110 |
+
out_channels,
|
111 |
+
kernel_size=3,
|
112 |
+
stride=1,
|
113 |
+
padding=1)
|
114 |
+
else:
|
115 |
+
self.nin_shortcut = torch.nn.Conv2d(in_channels,
|
116 |
+
out_channels,
|
117 |
+
kernel_size=1,
|
118 |
+
stride=1,
|
119 |
+
padding=0)
|
120 |
+
|
121 |
+
def forward(self, x, temb):
|
122 |
+
h = x
|
123 |
+
h = self.norm1(h)
|
124 |
+
h = nonlinearity(h)
|
125 |
+
h = self.conv1(h)
|
126 |
+
|
127 |
+
if temb is not None:
|
128 |
+
h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
|
129 |
+
|
130 |
+
h = self.norm2(h)
|
131 |
+
h = nonlinearity(h)
|
132 |
+
h = self.dropout(h)
|
133 |
+
h = self.conv2(h)
|
134 |
+
|
135 |
+
if self.in_channels != self.out_channels:
|
136 |
+
if self.use_conv_shortcut:
|
137 |
+
x = self.conv_shortcut(x)
|
138 |
+
else:
|
139 |
+
x = self.nin_shortcut(x)
|
140 |
+
|
141 |
+
return x+h
|
142 |
+
|
143 |
+
|
144 |
+
class LinAttnBlock(LinearAttention):
|
145 |
+
"""to match AttnBlock usage"""
|
146 |
+
def __init__(self, in_channels):
|
147 |
+
super().__init__(dim=in_channels, heads=1, dim_head=in_channels)
|
148 |
+
|
149 |
+
|
150 |
+
class AttnBlock(nn.Module):
|
151 |
+
def __init__(self, in_channels):
|
152 |
+
super().__init__()
|
153 |
+
self.in_channels = in_channels
|
154 |
+
|
155 |
+
self.norm = Normalize(in_channels)
|
156 |
+
self.q = torch.nn.Conv2d(in_channels,
|
157 |
+
in_channels,
|
158 |
+
kernel_size=1,
|
159 |
+
stride=1,
|
160 |
+
padding=0)
|
161 |
+
self.k = torch.nn.Conv2d(in_channels,
|
162 |
+
in_channels,
|
163 |
+
kernel_size=1,
|
164 |
+
stride=1,
|
165 |
+
padding=0)
|
166 |
+
self.v = torch.nn.Conv2d(in_channels,
|
167 |
+
in_channels,
|
168 |
+
kernel_size=1,
|
169 |
+
stride=1,
|
170 |
+
padding=0)
|
171 |
+
self.proj_out = torch.nn.Conv2d(in_channels,
|
172 |
+
in_channels,
|
173 |
+
kernel_size=1,
|
174 |
+
stride=1,
|
175 |
+
padding=0)
|
176 |
+
|
177 |
+
|
178 |
+
def forward(self, x):
|
179 |
+
h_ = x
|
180 |
+
h_ = self.norm(h_)
|
181 |
+
q = self.q(h_)
|
182 |
+
k = self.k(h_)
|
183 |
+
v = self.v(h_)
|
184 |
+
|
185 |
+
# compute attention
|
186 |
+
b,c,h,w = q.shape
|
187 |
+
q = q.reshape(b,c,h*w)
|
188 |
+
q = q.permute(0,2,1) # b,hw,c
|
189 |
+
k = k.reshape(b,c,h*w) # b,c,hw
|
190 |
+
w_ = torch.bmm(q,k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
|
191 |
+
w_ = w_ * (int(c)**(-0.5))
|
192 |
+
w_ = torch.nn.functional.softmax(w_, dim=2)
|
193 |
+
|
194 |
+
# attend to values
|
195 |
+
v = v.reshape(b,c,h*w)
|
196 |
+
w_ = w_.permute(0,2,1) # b,hw,hw (first hw of k, second of q)
|
197 |
+
h_ = torch.bmm(v,w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
|
198 |
+
h_ = h_.reshape(b,c,h,w)
|
199 |
+
|
200 |
+
h_ = self.proj_out(h_)
|
201 |
+
|
202 |
+
return x+h_
|
203 |
+
|
204 |
+
|
205 |
+
def make_attn(in_channels, attn_type="vanilla"):
|
206 |
+
assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
|
207 |
+
print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
|
208 |
+
if attn_type == "vanilla":
|
209 |
+
return AttnBlock(in_channels)
|
210 |
+
elif attn_type == "none":
|
211 |
+
return nn.Identity(in_channels)
|
212 |
+
else:
|
213 |
+
return LinAttnBlock(in_channels)
|
214 |
+
|
215 |
+
|
216 |
+
class Model(nn.Module):
|
217 |
+
def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
|
218 |
+
attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
|
219 |
+
resolution, use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
|
220 |
+
super().__init__()
|
221 |
+
if use_linear_attn: attn_type = "linear"
|
222 |
+
self.ch = ch
|
223 |
+
self.temb_ch = self.ch*4
|
224 |
+
self.num_resolutions = len(ch_mult)
|
225 |
+
self.num_res_blocks = num_res_blocks
|
226 |
+
self.resolution = resolution
|
227 |
+
self.in_channels = in_channels
|
228 |
+
|
229 |
+
self.use_timestep = use_timestep
|
230 |
+
if self.use_timestep:
|
231 |
+
# timestep embedding
|
232 |
+
self.temb = nn.Module()
|
233 |
+
self.temb.dense = nn.ModuleList([
|
234 |
+
torch.nn.Linear(self.ch,
|
235 |
+
self.temb_ch),
|
236 |
+
torch.nn.Linear(self.temb_ch,
|
237 |
+
self.temb_ch),
|
238 |
+
])
|
239 |
+
|
240 |
+
# downsampling
|
241 |
+
self.conv_in = torch.nn.Conv2d(in_channels,
|
242 |
+
self.ch,
|
243 |
+
kernel_size=3,
|
244 |
+
stride=1,
|
245 |
+
padding=1)
|
246 |
+
|
247 |
+
curr_res = resolution
|
248 |
+
in_ch_mult = (1,)+tuple(ch_mult)
|
249 |
+
self.down = nn.ModuleList()
|
250 |
+
for i_level in range(self.num_resolutions):
|
251 |
+
block = nn.ModuleList()
|
252 |
+
attn = nn.ModuleList()
|
253 |
+
block_in = ch*in_ch_mult[i_level]
|
254 |
+
block_out = ch*ch_mult[i_level]
|
255 |
+
for i_block in range(self.num_res_blocks):
|
256 |
+
block.append(ResnetBlock(in_channels=block_in,
|
257 |
+
out_channels=block_out,
|
258 |
+
temb_channels=self.temb_ch,
|
259 |
+
dropout=dropout))
|
260 |
+
block_in = block_out
|
261 |
+
if curr_res in attn_resolutions:
|
262 |
+
attn.append(make_attn(block_in, attn_type=attn_type))
|
263 |
+
down = nn.Module()
|
264 |
+
down.block = block
|
265 |
+
down.attn = attn
|
266 |
+
if i_level != self.num_resolutions-1:
|
267 |
+
down.downsample = Downsample(block_in, resamp_with_conv)
|
268 |
+
curr_res = curr_res // 2
|
269 |
+
self.down.append(down)
|
270 |
+
|
271 |
+
# middle
|
272 |
+
self.mid = nn.Module()
|
273 |
+
self.mid.block_1 = ResnetBlock(in_channels=block_in,
|
274 |
+
out_channels=block_in,
|
275 |
+
temb_channels=self.temb_ch,
|
276 |
+
dropout=dropout)
|
277 |
+
self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
|
278 |
+
self.mid.block_2 = ResnetBlock(in_channels=block_in,
|
279 |
+
out_channels=block_in,
|
280 |
+
temb_channels=self.temb_ch,
|
281 |
+
dropout=dropout)
|
282 |
+
|
283 |
+
# upsampling
|
284 |
+
self.up = nn.ModuleList()
|
285 |
+
for i_level in reversed(range(self.num_resolutions)):
|
286 |
+
block = nn.ModuleList()
|
287 |
+
attn = nn.ModuleList()
|
288 |
+
block_out = ch*ch_mult[i_level]
|
289 |
+
skip_in = ch*ch_mult[i_level]
|
290 |
+
for i_block in range(self.num_res_blocks+1):
|
291 |
+
if i_block == self.num_res_blocks:
|
292 |
+
skip_in = ch*in_ch_mult[i_level]
|
293 |
+
block.append(ResnetBlock(in_channels=block_in+skip_in,
|
294 |
+
out_channels=block_out,
|
295 |
+
temb_channels=self.temb_ch,
|
296 |
+
dropout=dropout))
|
297 |
+
block_in = block_out
|
298 |
+
if curr_res in attn_resolutions:
|
299 |
+
attn.append(make_attn(block_in, attn_type=attn_type))
|
300 |
+
up = nn.Module()
|
301 |
+
up.block = block
|
302 |
+
up.attn = attn
|
303 |
+
if i_level != 0:
|
304 |
+
up.upsample = Upsample(block_in, resamp_with_conv)
|
305 |
+
curr_res = curr_res * 2
|
306 |
+
self.up.insert(0, up) # prepend to get consistent order
|
307 |
+
|
308 |
+
# end
|
309 |
+
self.norm_out = Normalize(block_in)
|
310 |
+
self.conv_out = torch.nn.Conv2d(block_in,
|
311 |
+
out_ch,
|
312 |
+
kernel_size=3,
|
313 |
+
stride=1,
|
314 |
+
padding=1)
|
315 |
+
|
316 |
+
def forward(self, x, t=None, context=None):
|
317 |
+
#assert x.shape[2] == x.shape[3] == self.resolution
|
318 |
+
if context is not None:
|
319 |
+
# assume aligned context, cat along channel axis
|
320 |
+
x = torch.cat((x, context), dim=1)
|
321 |
+
if self.use_timestep:
|
322 |
+
# timestep embedding
|
323 |
+
assert t is not None
|
324 |
+
temb = get_timestep_embedding(t, self.ch)
|
325 |
+
temb = self.temb.dense[0](temb)
|
326 |
+
temb = nonlinearity(temb)
|
327 |
+
temb = self.temb.dense[1](temb)
|
328 |
+
else:
|
329 |
+
temb = None
|
330 |
+
|
331 |
+
# downsampling
|
332 |
+
hs = [self.conv_in(x)]
|
333 |
+
for i_level in range(self.num_resolutions):
|
334 |
+
for i_block in range(self.num_res_blocks):
|
335 |
+
h = self.down[i_level].block[i_block](hs[-1], temb)
|
336 |
+
if len(self.down[i_level].attn) > 0:
|
337 |
+
h = self.down[i_level].attn[i_block](h)
|
338 |
+
hs.append(h)
|
339 |
+
if i_level != self.num_resolutions-1:
|
340 |
+
hs.append(self.down[i_level].downsample(hs[-1]))
|
341 |
+
|
342 |
+
# middle
|
343 |
+
h = hs[-1]
|
344 |
+
h = self.mid.block_1(h, temb)
|
345 |
+
h = self.mid.attn_1(h)
|
346 |
+
h = self.mid.block_2(h, temb)
|
347 |
+
|
348 |
+
# upsampling
|
349 |
+
for i_level in reversed(range(self.num_resolutions)):
|
350 |
+
for i_block in range(self.num_res_blocks+1):
|
351 |
+
h = self.up[i_level].block[i_block](
|
352 |
+
torch.cat([h, hs.pop()], dim=1), temb)
|
353 |
+
if len(self.up[i_level].attn) > 0:
|
354 |
+
h = self.up[i_level].attn[i_block](h)
|
355 |
+
if i_level != 0:
|
356 |
+
h = self.up[i_level].upsample(h)
|
357 |
+
|
358 |
+
# end
|
359 |
+
h = self.norm_out(h)
|
360 |
+
h = nonlinearity(h)
|
361 |
+
h = self.conv_out(h)
|
362 |
+
return h
|
363 |
+
|
364 |
+
def get_last_layer(self):
|
365 |
+
return self.conv_out.weight
|
366 |
+
|
367 |
+
|
368 |
+
class Encoder(nn.Module):
|
369 |
+
def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
|
370 |
+
attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
|
371 |
+
resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
|
372 |
+
**ignore_kwargs):
|
373 |
+
super().__init__()
|
374 |
+
if use_linear_attn: attn_type = "linear"
|
375 |
+
self.ch = ch
|
376 |
+
self.temb_ch = 0
|
377 |
+
self.num_resolutions = len(ch_mult)
|
378 |
+
self.num_res_blocks = num_res_blocks
|
379 |
+
self.resolution = resolution
|
380 |
+
self.in_channels = in_channels
|
381 |
+
|
382 |
+
# downsampling
|
383 |
+
self.conv_in = torch.nn.Conv2d(in_channels,
|
384 |
+
self.ch,
|
385 |
+
kernel_size=3,
|
386 |
+
stride=1,
|
387 |
+
padding=1)
|
388 |
+
|
389 |
+
curr_res = resolution
|
390 |
+
in_ch_mult = (1,)+tuple(ch_mult)
|
391 |
+
self.in_ch_mult = in_ch_mult
|
392 |
+
self.down = nn.ModuleList()
|
393 |
+
for i_level in range(self.num_resolutions):
|
394 |
+
block = nn.ModuleList()
|
395 |
+
attn = nn.ModuleList()
|
396 |
+
block_in = ch*in_ch_mult[i_level]
|
397 |
+
block_out = ch*ch_mult[i_level]
|
398 |
+
for i_block in range(self.num_res_blocks):
|
399 |
+
block.append(ResnetBlock(in_channels=block_in,
|
400 |
+
out_channels=block_out,
|
401 |
+
temb_channels=self.temb_ch,
|
402 |
+
dropout=dropout))
|
403 |
+
block_in = block_out
|
404 |
+
if curr_res in attn_resolutions:
|
405 |
+
attn.append(make_attn(block_in, attn_type=attn_type))
|
406 |
+
down = nn.Module()
|
407 |
+
down.block = block
|
408 |
+
down.attn = attn
|
409 |
+
if i_level != self.num_resolutions-1:
|
410 |
+
down.downsample = Downsample(block_in, resamp_with_conv)
|
411 |
+
curr_res = curr_res // 2
|
412 |
+
self.down.append(down)
|
413 |
+
|
414 |
+
# middle
|
415 |
+
self.mid = nn.Module()
|
416 |
+
self.mid.block_1 = ResnetBlock(in_channels=block_in,
|
417 |
+
out_channels=block_in,
|
418 |
+
temb_channels=self.temb_ch,
|
419 |
+
dropout=dropout)
|
420 |
+
self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
|
421 |
+
self.mid.block_2 = ResnetBlock(in_channels=block_in,
|
422 |
+
out_channels=block_in,
|
423 |
+
temb_channels=self.temb_ch,
|
424 |
+
dropout=dropout)
|
425 |
+
|
426 |
+
# end
|
427 |
+
self.norm_out = Normalize(block_in)
|
428 |
+
self.conv_out = torch.nn.Conv2d(block_in,
|
429 |
+
2*z_channels if double_z else z_channels,
|
430 |
+
kernel_size=3,
|
431 |
+
stride=1,
|
432 |
+
padding=1)
|
433 |
+
|
434 |
+
def forward(self, x):
|
435 |
+
# timestep embedding
|
436 |
+
temb = None
|
437 |
+
|
438 |
+
# downsampling
|
439 |
+
hs = [self.conv_in(x)]
|
440 |
+
for i_level in range(self.num_resolutions):
|
441 |
+
for i_block in range(self.num_res_blocks):
|
442 |
+
h = self.down[i_level].block[i_block](hs[-1], temb)
|
443 |
+
if len(self.down[i_level].attn) > 0:
|
444 |
+
h = self.down[i_level].attn[i_block](h)
|
445 |
+
hs.append(h)
|
446 |
+
if i_level != self.num_resolutions-1:
|
447 |
+
hs.append(self.down[i_level].downsample(hs[-1]))
|
448 |
+
|
449 |
+
# middle
|
450 |
+
h = hs[-1]
|
451 |
+
h = self.mid.block_1(h, temb)
|
452 |
+
h = self.mid.attn_1(h)
|
453 |
+
h = self.mid.block_2(h, temb)
|
454 |
+
|
455 |
+
# end
|
456 |
+
h = self.norm_out(h)
|
457 |
+
h = nonlinearity(h)
|
458 |
+
h = self.conv_out(h)
|
459 |
+
return h
|
460 |
+
|
461 |
+
|
462 |
+
class Decoder(nn.Module):
|
463 |
+
def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
|
464 |
+
attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
|
465 |
+
resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
|
466 |
+
attn_type="vanilla", **ignorekwargs):
|
467 |
+
super().__init__()
|
468 |
+
if use_linear_attn: attn_type = "linear"
|
469 |
+
self.ch = ch
|
470 |
+
self.temb_ch = 0
|
471 |
+
self.num_resolutions = len(ch_mult)
|
472 |
+
self.num_res_blocks = num_res_blocks
|
473 |
+
self.resolution = resolution
|
474 |
+
self.in_channels = in_channels
|
475 |
+
self.give_pre_end = give_pre_end
|
476 |
+
self.tanh_out = tanh_out
|
477 |
+
|
478 |
+
# compute in_ch_mult, block_in and curr_res at lowest res
|
479 |
+
in_ch_mult = (1,)+tuple(ch_mult)
|
480 |
+
block_in = ch*ch_mult[self.num_resolutions-1]
|
481 |
+
curr_res = resolution // 2**(self.num_resolutions-1)
|
482 |
+
self.z_shape = (1,z_channels,curr_res,curr_res)
|
483 |
+
print("Working with z of shape {} = {} dimensions.".format(
|
484 |
+
self.z_shape, np.prod(self.z_shape)))
|
485 |
+
|
486 |
+
# z to block_in
|
487 |
+
self.conv_in = torch.nn.Conv2d(z_channels,
|
488 |
+
block_in,
|
489 |
+
kernel_size=3,
|
490 |
+
stride=1,
|
491 |
+
padding=1)
|
492 |
+
|
493 |
+
# middle
|
494 |
+
self.mid = nn.Module()
|
495 |
+
self.mid.block_1 = ResnetBlock(in_channels=block_in,
|
496 |
+
out_channels=block_in,
|
497 |
+
temb_channels=self.temb_ch,
|
498 |
+
dropout=dropout)
|
499 |
+
self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
|
500 |
+
self.mid.block_2 = ResnetBlock(in_channels=block_in,
|
501 |
+
out_channels=block_in,
|
502 |
+
temb_channels=self.temb_ch,
|
503 |
+
dropout=dropout)
|
504 |
+
|
505 |
+
# upsampling
|
506 |
+
self.up = nn.ModuleList()
|
507 |
+
for i_level in reversed(range(self.num_resolutions)):
|
508 |
+
block = nn.ModuleList()
|
509 |
+
attn = nn.ModuleList()
|
510 |
+
block_out = ch*ch_mult[i_level]
|
511 |
+
for i_block in range(self.num_res_blocks+1):
|
512 |
+
block.append(ResnetBlock(in_channels=block_in,
|
513 |
+
out_channels=block_out,
|
514 |
+
temb_channels=self.temb_ch,
|
515 |
+
dropout=dropout))
|
516 |
+
block_in = block_out
|
517 |
+
if curr_res in attn_resolutions:
|
518 |
+
attn.append(make_attn(block_in, attn_type=attn_type))
|
519 |
+
up = nn.Module()
|
520 |
+
up.block = block
|
521 |
+
up.attn = attn
|
522 |
+
if i_level != 0:
|
523 |
+
up.upsample = Upsample(block_in, resamp_with_conv)
|
524 |
+
curr_res = curr_res * 2
|
525 |
+
self.up.insert(0, up) # prepend to get consistent order
|
526 |
+
|
527 |
+
# end
|
528 |
+
self.norm_out = Normalize(block_in)
|
529 |
+
self.conv_out = torch.nn.Conv2d(block_in,
|
530 |
+
out_ch,
|
531 |
+
kernel_size=3,
|
532 |
+
stride=1,
|
533 |
+
padding=1)
|
534 |
+
|
535 |
+
def forward(self, z):
|
536 |
+
#assert z.shape[1:] == self.z_shape[1:]
|
537 |
+
self.last_z_shape = z.shape
|
538 |
+
|
539 |
+
# timestep embedding
|
540 |
+
temb = None
|
541 |
+
|
542 |
+
# z to block_in
|
543 |
+
h = self.conv_in(z)
|
544 |
+
|
545 |
+
# middle
|
546 |
+
h = self.mid.block_1(h, temb)
|
547 |
+
h = self.mid.attn_1(h)
|
548 |
+
h = self.mid.block_2(h, temb)
|
549 |
+
|
550 |
+
# upsampling
|
551 |
+
for i_level in reversed(range(self.num_resolutions)):
|
552 |
+
for i_block in range(self.num_res_blocks+1):
|
553 |
+
h = self.up[i_level].block[i_block](h, temb)
|
554 |
+
if len(self.up[i_level].attn) > 0:
|
555 |
+
h = self.up[i_level].attn[i_block](h)
|
556 |
+
if i_level != 0:
|
557 |
+
h = self.up[i_level].upsample(h)
|
558 |
+
|
559 |
+
# end
|
560 |
+
if self.give_pre_end:
|
561 |
+
return h
|
562 |
+
|
563 |
+
h = self.norm_out(h)
|
564 |
+
h = nonlinearity(h)
|
565 |
+
h = self.conv_out(h)
|
566 |
+
if self.tanh_out:
|
567 |
+
h = torch.tanh(h)
|
568 |
+
return h
|
569 |
+
|
570 |
+
|
571 |
+
class SimpleDecoder(nn.Module):
|
572 |
+
def __init__(self, in_channels, out_channels, *args, **kwargs):
|
573 |
+
super().__init__()
|
574 |
+
self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
|
575 |
+
ResnetBlock(in_channels=in_channels,
|
576 |
+
out_channels=2 * in_channels,
|
577 |
+
temb_channels=0, dropout=0.0),
|
578 |
+
ResnetBlock(in_channels=2 * in_channels,
|
579 |
+
out_channels=4 * in_channels,
|
580 |
+
temb_channels=0, dropout=0.0),
|
581 |
+
ResnetBlock(in_channels=4 * in_channels,
|
582 |
+
out_channels=2 * in_channels,
|
583 |
+
temb_channels=0, dropout=0.0),
|
584 |
+
nn.Conv2d(2*in_channels, in_channels, 1),
|
585 |
+
Upsample(in_channels, with_conv=True)])
|
586 |
+
# end
|
587 |
+
self.norm_out = Normalize(in_channels)
|
588 |
+
self.conv_out = torch.nn.Conv2d(in_channels,
|
589 |
+
out_channels,
|
590 |
+
kernel_size=3,
|
591 |
+
stride=1,
|
592 |
+
padding=1)
|
593 |
+
|
594 |
+
def forward(self, x):
|
595 |
+
for i, layer in enumerate(self.model):
|
596 |
+
if i in [1,2,3]:
|
597 |
+
x = layer(x, None)
|
598 |
+
else:
|
599 |
+
x = layer(x)
|
600 |
+
|
601 |
+
h = self.norm_out(x)
|
602 |
+
h = nonlinearity(h)
|
603 |
+
x = self.conv_out(h)
|
604 |
+
return x
|
605 |
+
|
606 |
+
|
607 |
+
class UpsampleDecoder(nn.Module):
|
608 |
+
def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
|
609 |
+
ch_mult=(2,2), dropout=0.0):
|
610 |
+
super().__init__()
|
611 |
+
# upsampling
|
612 |
+
self.temb_ch = 0
|
613 |
+
self.num_resolutions = len(ch_mult)
|
614 |
+
self.num_res_blocks = num_res_blocks
|
615 |
+
block_in = in_channels
|
616 |
+
curr_res = resolution // 2 ** (self.num_resolutions - 1)
|
617 |
+
self.res_blocks = nn.ModuleList()
|
618 |
+
self.upsample_blocks = nn.ModuleList()
|
619 |
+
for i_level in range(self.num_resolutions):
|
620 |
+
res_block = []
|
621 |
+
block_out = ch * ch_mult[i_level]
|
622 |
+
for i_block in range(self.num_res_blocks + 1):
|
623 |
+
res_block.append(ResnetBlock(in_channels=block_in,
|
624 |
+
out_channels=block_out,
|
625 |
+
temb_channels=self.temb_ch,
|
626 |
+
dropout=dropout))
|
627 |
+
block_in = block_out
|
628 |
+
self.res_blocks.append(nn.ModuleList(res_block))
|
629 |
+
if i_level != self.num_resolutions - 1:
|
630 |
+
self.upsample_blocks.append(Upsample(block_in, True))
|
631 |
+
curr_res = curr_res * 2
|
632 |
+
|
633 |
+
# end
|
634 |
+
self.norm_out = Normalize(block_in)
|
635 |
+
self.conv_out = torch.nn.Conv2d(block_in,
|
636 |
+
out_channels,
|
637 |
+
kernel_size=3,
|
638 |
+
stride=1,
|
639 |
+
padding=1)
|
640 |
+
|
641 |
+
def forward(self, x):
|
642 |
+
# upsampling
|
643 |
+
h = x
|
644 |
+
for k, i_level in enumerate(range(self.num_resolutions)):
|
645 |
+
for i_block in range(self.num_res_blocks + 1):
|
646 |
+
h = self.res_blocks[i_level][i_block](h, None)
|
647 |
+
if i_level != self.num_resolutions - 1:
|
648 |
+
h = self.upsample_blocks[k](h)
|
649 |
+
h = self.norm_out(h)
|
650 |
+
h = nonlinearity(h)
|
651 |
+
h = self.conv_out(h)
|
652 |
+
return h
|
653 |
+
|
654 |
+
|
655 |
+
class LatentRescaler(nn.Module):
|
656 |
+
def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
|
657 |
+
super().__init__()
|
658 |
+
# residual block, interpolate, residual block
|
659 |
+
self.factor = factor
|
660 |
+
self.conv_in = nn.Conv2d(in_channels,
|
661 |
+
mid_channels,
|
662 |
+
kernel_size=3,
|
663 |
+
stride=1,
|
664 |
+
padding=1)
|
665 |
+
self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
|
666 |
+
out_channels=mid_channels,
|
667 |
+
temb_channels=0,
|
668 |
+
dropout=0.0) for _ in range(depth)])
|
669 |
+
self.attn = AttnBlock(mid_channels)
|
670 |
+
self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
|
671 |
+
out_channels=mid_channels,
|
672 |
+
temb_channels=0,
|
673 |
+
dropout=0.0) for _ in range(depth)])
|
674 |
+
|
675 |
+
self.conv_out = nn.Conv2d(mid_channels,
|
676 |
+
out_channels,
|
677 |
+
kernel_size=1,
|
678 |
+
)
|
679 |
+
|
680 |
+
def forward(self, x):
|
681 |
+
x = self.conv_in(x)
|
682 |
+
for block in self.res_block1:
|
683 |
+
x = block(x, None)
|
684 |
+
x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor))))
|
685 |
+
x = self.attn(x)
|
686 |
+
for block in self.res_block2:
|
687 |
+
x = block(x, None)
|
688 |
+
x = self.conv_out(x)
|
689 |
+
return x
|
690 |
+
|
691 |
+
|
692 |
+
class MergedRescaleEncoder(nn.Module):
|
693 |
+
def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks,
|
694 |
+
attn_resolutions, dropout=0.0, resamp_with_conv=True,
|
695 |
+
ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1):
|
696 |
+
super().__init__()
|
697 |
+
intermediate_chn = ch * ch_mult[-1]
|
698 |
+
self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
|
699 |
+
z_channels=intermediate_chn, double_z=False, resolution=resolution,
|
700 |
+
attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv,
|
701 |
+
out_ch=None)
|
702 |
+
self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
|
703 |
+
mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
|
704 |
+
|
705 |
+
def forward(self, x):
|
706 |
+
x = self.encoder(x)
|
707 |
+
x = self.rescaler(x)
|
708 |
+
return x
|
709 |
+
|
710 |
+
|
711 |
+
class MergedRescaleDecoder(nn.Module):
|
712 |
+
def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8),
|
713 |
+
dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
|
714 |
+
super().__init__()
|
715 |
+
tmp_chn = z_channels*ch_mult[-1]
|
716 |
+
self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout,
|
717 |
+
resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
|
718 |
+
ch_mult=ch_mult, resolution=resolution, ch=ch)
|
719 |
+
self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
|
720 |
+
out_channels=tmp_chn, depth=rescale_module_depth)
|
721 |
+
|
722 |
+
def forward(self, x):
|
723 |
+
x = self.rescaler(x)
|
724 |
+
x = self.decoder(x)
|
725 |
+
return x
|
726 |
+
|
727 |
+
|
728 |
+
class Upsampler(nn.Module):
|
729 |
+
def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
|
730 |
+
super().__init__()
|
731 |
+
assert out_size >= in_size
|
732 |
+
num_blocks = int(np.log2(out_size//in_size))+1
|
733 |
+
factor_up = 1.+ (out_size % in_size)
|
734 |
+
print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
|
735 |
+
self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels,
|
736 |
+
out_channels=in_channels)
|
737 |
+
self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2,
|
738 |
+
attn_resolutions=[], in_channels=None, ch=in_channels,
|
739 |
+
ch_mult=[ch_mult for _ in range(num_blocks)])
|
740 |
+
|
741 |
+
def forward(self, x):
|
742 |
+
x = self.rescaler(x)
|
743 |
+
x = self.decoder(x)
|
744 |
+
return x
|
745 |
+
|
746 |
+
|
747 |
+
class Resize(nn.Module):
|
748 |
+
def __init__(self, in_channels=None, learned=False, mode="bilinear"):
|
749 |
+
super().__init__()
|
750 |
+
self.with_conv = learned
|
751 |
+
self.mode = mode
|
752 |
+
if self.with_conv:
|
753 |
+
print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
|
754 |
+
raise NotImplementedError()
|
755 |
+
assert in_channels is not None
|
756 |
+
# no asymmetric padding in torch conv, must do it ourselves
|
757 |
+
self.conv = torch.nn.Conv2d(in_channels,
|
758 |
+
in_channels,
|
759 |
+
kernel_size=4,
|
760 |
+
stride=2,
|
761 |
+
padding=1)
|
762 |
+
|
763 |
+
def forward(self, x, scale_factor=1.0):
|
764 |
+
if scale_factor==1.0:
|
765 |
+
return x
|
766 |
+
else:
|
767 |
+
x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
|
768 |
+
return x
|
769 |
+
|
770 |
+
class FirstStagePostProcessor(nn.Module):
|
771 |
+
|
772 |
+
def __init__(self, ch_mult:list, in_channels,
|
773 |
+
pretrained_model:nn.Module=None,
|
774 |
+
reshape=False,
|
775 |
+
n_channels=None,
|
776 |
+
dropout=0.,
|
777 |
+
pretrained_config=None):
|
778 |
+
super().__init__()
|
779 |
+
if pretrained_config is None:
|
780 |
+
assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
|
781 |
+
self.pretrained_model = pretrained_model
|
782 |
+
else:
|
783 |
+
assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
|
784 |
+
self.instantiate_pretrained(pretrained_config)
|
785 |
+
|
786 |
+
self.do_reshape = reshape
|
787 |
+
|
788 |
+
if n_channels is None:
|
789 |
+
n_channels = self.pretrained_model.encoder.ch
|
790 |
+
|
791 |
+
self.proj_norm = Normalize(in_channels,num_groups=in_channels//2)
|
792 |
+
self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3,
|
793 |
+
stride=1,padding=1)
|
794 |
+
|
795 |
+
blocks = []
|
796 |
+
downs = []
|
797 |
+
ch_in = n_channels
|
798 |
+
for m in ch_mult:
|
799 |
+
blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout))
|
800 |
+
ch_in = m * n_channels
|
801 |
+
downs.append(Downsample(ch_in, with_conv=False))
|
802 |
+
|
803 |
+
self.model = nn.ModuleList(blocks)
|
804 |
+
self.downsampler = nn.ModuleList(downs)
|
805 |
+
|
806 |
+
|
807 |
+
def instantiate_pretrained(self, config):
|
808 |
+
model = instantiate_from_config(config)
|
809 |
+
self.pretrained_model = model.eval()
|
810 |
+
# self.pretrained_model.train = False
|
811 |
+
for param in self.pretrained_model.parameters():
|
812 |
+
param.requires_grad = False
|
813 |
+
|
814 |
+
|
815 |
+
@torch.no_grad()
|
816 |
+
def encode_with_pretrained(self,x):
|
817 |
+
c = self.pretrained_model.encode(x)
|
818 |
+
if isinstance(c, DiagonalGaussianDistribution):
|
819 |
+
c = c.mode()
|
820 |
+
return c
|
821 |
+
|
822 |
+
def forward(self,x):
|
823 |
+
z_fs = self.encode_with_pretrained(x)
|
824 |
+
z = self.proj_norm(z_fs)
|
825 |
+
z = self.proj(z)
|
826 |
+
z = nonlinearity(z)
|
827 |
+
|
828 |
+
for submodel, downmodel in zip(self.model,self.downsampler):
|
829 |
+
z = submodel(z,temb=None)
|
830 |
+
z = downmodel(z)
|
831 |
+
|
832 |
+
if self.do_reshape:
|
833 |
+
z = rearrange(z,'b c h w -> b (h w) c')
|
834 |
+
return z
|
835 |
+
|
ldm/modules/diffusionmodules/openaimodel.py
ADDED
@@ -0,0 +1,996 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from abc import abstractmethod
|
2 |
+
from functools import partial
|
3 |
+
import math
|
4 |
+
from typing import Iterable
|
5 |
+
|
6 |
+
import numpy as np
|
7 |
+
import torch as th
|
8 |
+
import torch.nn as nn
|
9 |
+
import torch.nn.functional as F
|
10 |
+
|
11 |
+
from ldm.modules.diffusionmodules.util import (
|
12 |
+
checkpoint,
|
13 |
+
conv_nd,
|
14 |
+
linear,
|
15 |
+
avg_pool_nd,
|
16 |
+
zero_module,
|
17 |
+
normalization,
|
18 |
+
timestep_embedding,
|
19 |
+
)
|
20 |
+
from ldm.modules.attention import SpatialTransformer
|
21 |
+
from ldm.util import exists
|
22 |
+
|
23 |
+
|
24 |
+
# dummy replace
|
25 |
+
def convert_module_to_f16(x):
|
26 |
+
pass
|
27 |
+
|
28 |
+
def convert_module_to_f32(x):
|
29 |
+
pass
|
30 |
+
|
31 |
+
|
32 |
+
## go
|
33 |
+
class AttentionPool2d(nn.Module):
|
34 |
+
"""
|
35 |
+
Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
|
36 |
+
"""
|
37 |
+
|
38 |
+
def __init__(
|
39 |
+
self,
|
40 |
+
spacial_dim: int,
|
41 |
+
embed_dim: int,
|
42 |
+
num_heads_channels: int,
|
43 |
+
output_dim: int = None,
|
44 |
+
):
|
45 |
+
super().__init__()
|
46 |
+
self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
|
47 |
+
self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
|
48 |
+
self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
|
49 |
+
self.num_heads = embed_dim // num_heads_channels
|
50 |
+
self.attention = QKVAttention(self.num_heads)
|
51 |
+
|
52 |
+
def forward(self, x):
|
53 |
+
b, c, *_spatial = x.shape
|
54 |
+
x = x.reshape(b, c, -1) # NC(HW)
|
55 |
+
x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)
|
56 |
+
x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)
|
57 |
+
x = self.qkv_proj(x)
|
58 |
+
x = self.attention(x)
|
59 |
+
x = self.c_proj(x)
|
60 |
+
return x[:, :, 0]
|
61 |
+
|
62 |
+
|
63 |
+
class TimestepBlock(nn.Module):
|
64 |
+
"""
|
65 |
+
Any module where forward() takes timestep embeddings as a second argument.
|
66 |
+
"""
|
67 |
+
|
68 |
+
@abstractmethod
|
69 |
+
def forward(self, x, emb):
|
70 |
+
"""
|
71 |
+
Apply the module to `x` given `emb` timestep embeddings.
|
72 |
+
"""
|
73 |
+
|
74 |
+
|
75 |
+
class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
|
76 |
+
"""
|
77 |
+
A sequential module that passes timestep embeddings to the children that
|
78 |
+
support it as an extra input.
|
79 |
+
"""
|
80 |
+
|
81 |
+
def forward(self, x, emb, context=None):
|
82 |
+
for layer in self:
|
83 |
+
if isinstance(layer, TimestepBlock):
|
84 |
+
x = layer(x, emb)
|
85 |
+
elif isinstance(layer, SpatialTransformer):
|
86 |
+
x = layer(x, context)
|
87 |
+
else:
|
88 |
+
x = layer(x)
|
89 |
+
return x
|
90 |
+
|
91 |
+
|
92 |
+
class Upsample(nn.Module):
|
93 |
+
"""
|
94 |
+
An upsampling layer with an optional convolution.
|
95 |
+
:param channels: channels in the inputs and outputs.
|
96 |
+
:param use_conv: a bool determining if a convolution is applied.
|
97 |
+
:param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
|
98 |
+
upsampling occurs in the inner-two dimensions.
|
99 |
+
"""
|
100 |
+
|
101 |
+
def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
|
102 |
+
super().__init__()
|
103 |
+
self.channels = channels
|
104 |
+
self.out_channels = out_channels or channels
|
105 |
+
self.use_conv = use_conv
|
106 |
+
self.dims = dims
|
107 |
+
if use_conv:
|
108 |
+
self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
|
109 |
+
|
110 |
+
def forward(self, x):
|
111 |
+
assert x.shape[1] == self.channels
|
112 |
+
if self.dims == 3:
|
113 |
+
x = F.interpolate(
|
114 |
+
x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
|
115 |
+
)
|
116 |
+
else:
|
117 |
+
x = F.interpolate(x, scale_factor=2, mode="nearest")
|
118 |
+
if self.use_conv:
|
119 |
+
x = self.conv(x)
|
120 |
+
return x
|
121 |
+
|
122 |
+
class TransposedUpsample(nn.Module):
|
123 |
+
'Learned 2x upsampling without padding'
|
124 |
+
def __init__(self, channels, out_channels=None, ks=5):
|
125 |
+
super().__init__()
|
126 |
+
self.channels = channels
|
127 |
+
self.out_channels = out_channels or channels
|
128 |
+
|
129 |
+
self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
|
130 |
+
|
131 |
+
def forward(self,x):
|
132 |
+
return self.up(x)
|
133 |
+
|
134 |
+
|
135 |
+
class Downsample(nn.Module):
|
136 |
+
"""
|
137 |
+
A downsampling layer with an optional convolution.
|
138 |
+
:param channels: channels in the inputs and outputs.
|
139 |
+
:param use_conv: a bool determining if a convolution is applied.
|
140 |
+
:param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
|
141 |
+
downsampling occurs in the inner-two dimensions.
|
142 |
+
"""
|
143 |
+
|
144 |
+
def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1):
|
145 |
+
super().__init__()
|
146 |
+
self.channels = channels
|
147 |
+
self.out_channels = out_channels or channels
|
148 |
+
self.use_conv = use_conv
|
149 |
+
self.dims = dims
|
150 |
+
stride = 2 if dims != 3 else (1, 2, 2)
|
151 |
+
if use_conv:
|
152 |
+
self.op = conv_nd(
|
153 |
+
dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
|
154 |
+
)
|
155 |
+
else:
|
156 |
+
assert self.channels == self.out_channels
|
157 |
+
self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
|
158 |
+
|
159 |
+
def forward(self, x):
|
160 |
+
assert x.shape[1] == self.channels
|
161 |
+
return self.op(x)
|
162 |
+
|
163 |
+
|
164 |
+
class ResBlock(TimestepBlock):
|
165 |
+
"""
|
166 |
+
A residual block that can optionally change the number of channels.
|
167 |
+
:param channels: the number of input channels.
|
168 |
+
:param emb_channels: the number of timestep embedding channels.
|
169 |
+
:param dropout: the rate of dropout.
|
170 |
+
:param out_channels: if specified, the number of out channels.
|
171 |
+
:param use_conv: if True and out_channels is specified, use a spatial
|
172 |
+
convolution instead of a smaller 1x1 convolution to change the
|
173 |
+
channels in the skip connection.
|
174 |
+
:param dims: determines if the signal is 1D, 2D, or 3D.
|
175 |
+
:param use_checkpoint: if True, use gradient checkpointing on this module.
|
176 |
+
:param up: if True, use this block for upsampling.
|
177 |
+
:param down: if True, use this block for downsampling.
|
178 |
+
"""
|
179 |
+
|
180 |
+
def __init__(
|
181 |
+
self,
|
182 |
+
channels,
|
183 |
+
emb_channels,
|
184 |
+
dropout,
|
185 |
+
out_channels=None,
|
186 |
+
use_conv=False,
|
187 |
+
use_scale_shift_norm=False,
|
188 |
+
dims=2,
|
189 |
+
use_checkpoint=False,
|
190 |
+
up=False,
|
191 |
+
down=False,
|
192 |
+
):
|
193 |
+
super().__init__()
|
194 |
+
self.channels = channels
|
195 |
+
self.emb_channels = emb_channels
|
196 |
+
self.dropout = dropout
|
197 |
+
self.out_channels = out_channels or channels
|
198 |
+
self.use_conv = use_conv
|
199 |
+
self.use_checkpoint = use_checkpoint
|
200 |
+
self.use_scale_shift_norm = use_scale_shift_norm
|
201 |
+
|
202 |
+
self.in_layers = nn.Sequential(
|
203 |
+
normalization(channels),
|
204 |
+
nn.SiLU(),
|
205 |
+
conv_nd(dims, channels, self.out_channels, 3, padding=1),
|
206 |
+
)
|
207 |
+
|
208 |
+
self.updown = up or down
|
209 |
+
|
210 |
+
if up:
|
211 |
+
self.h_upd = Upsample(channels, False, dims)
|
212 |
+
self.x_upd = Upsample(channels, False, dims)
|
213 |
+
elif down:
|
214 |
+
self.h_upd = Downsample(channels, False, dims)
|
215 |
+
self.x_upd = Downsample(channels, False, dims)
|
216 |
+
else:
|
217 |
+
self.h_upd = self.x_upd = nn.Identity()
|
218 |
+
|
219 |
+
self.emb_layers = nn.Sequential(
|
220 |
+
nn.SiLU(),
|
221 |
+
linear(
|
222 |
+
emb_channels,
|
223 |
+
2 * self.out_channels if use_scale_shift_norm else self.out_channels,
|
224 |
+
),
|
225 |
+
)
|
226 |
+
self.out_layers = nn.Sequential(
|
227 |
+
normalization(self.out_channels),
|
228 |
+
nn.SiLU(),
|
229 |
+
nn.Dropout(p=dropout),
|
230 |
+
zero_module(
|
231 |
+
conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
|
232 |
+
),
|
233 |
+
)
|
234 |
+
|
235 |
+
if self.out_channels == channels:
|
236 |
+
self.skip_connection = nn.Identity()
|
237 |
+
elif use_conv:
|
238 |
+
self.skip_connection = conv_nd(
|
239 |
+
dims, channels, self.out_channels, 3, padding=1
|
240 |
+
)
|
241 |
+
else:
|
242 |
+
self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
|
243 |
+
|
244 |
+
def forward(self, x, emb):
|
245 |
+
"""
|
246 |
+
Apply the block to a Tensor, conditioned on a timestep embedding.
|
247 |
+
:param x: an [N x C x ...] Tensor of features.
|
248 |
+
:param emb: an [N x emb_channels] Tensor of timestep embeddings.
|
249 |
+
:return: an [N x C x ...] Tensor of outputs.
|
250 |
+
"""
|
251 |
+
return checkpoint(
|
252 |
+
self._forward, (x, emb), self.parameters(), self.use_checkpoint
|
253 |
+
)
|
254 |
+
|
255 |
+
|
256 |
+
def _forward(self, x, emb):
|
257 |
+
if self.updown:
|
258 |
+
in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
|
259 |
+
h = in_rest(x)
|
260 |
+
h = self.h_upd(h)
|
261 |
+
x = self.x_upd(x)
|
262 |
+
h = in_conv(h)
|
263 |
+
else:
|
264 |
+
h = self.in_layers(x)
|
265 |
+
emb_out = self.emb_layers(emb).type(h.dtype)
|
266 |
+
while len(emb_out.shape) < len(h.shape):
|
267 |
+
emb_out = emb_out[..., None]
|
268 |
+
if self.use_scale_shift_norm: # False
|
269 |
+
out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
|
270 |
+
scale, shift = th.chunk(emb_out, 2, dim=1)
|
271 |
+
h = out_norm(h) * (1 + scale) + shift
|
272 |
+
h = out_rest(h)
|
273 |
+
else:
|
274 |
+
h = h + emb_out
|
275 |
+
h = self.out_layers(h)
|
276 |
+
return self.skip_connection(x) + h
|
277 |
+
|
278 |
+
|
279 |
+
class AttentionBlock(nn.Module):
|
280 |
+
"""
|
281 |
+
An attention block that allows spatial positions to attend to each other.
|
282 |
+
Originally ported from here, but adapted to the N-d case.
|
283 |
+
https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
|
284 |
+
"""
|
285 |
+
|
286 |
+
def __init__(
|
287 |
+
self,
|
288 |
+
channels,
|
289 |
+
num_heads=1,
|
290 |
+
num_head_channels=-1,
|
291 |
+
use_checkpoint=False,
|
292 |
+
use_new_attention_order=False,
|
293 |
+
):
|
294 |
+
super().__init__()
|
295 |
+
self.channels = channels
|
296 |
+
if num_head_channels == -1:
|
297 |
+
self.num_heads = num_heads
|
298 |
+
else:
|
299 |
+
assert (
|
300 |
+
channels % num_head_channels == 0
|
301 |
+
), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
|
302 |
+
self.num_heads = channels // num_head_channels
|
303 |
+
self.use_checkpoint = use_checkpoint
|
304 |
+
self.norm = normalization(channels)
|
305 |
+
self.qkv = conv_nd(1, channels, channels * 3, 1)
|
306 |
+
if use_new_attention_order:
|
307 |
+
# split qkv before split heads
|
308 |
+
self.attention = QKVAttention(self.num_heads)
|
309 |
+
else:
|
310 |
+
# split heads before split qkv
|
311 |
+
self.attention = QKVAttentionLegacy(self.num_heads)
|
312 |
+
|
313 |
+
self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
|
314 |
+
|
315 |
+
def forward(self, x):
|
316 |
+
return checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
|
317 |
+
#return pt_checkpoint(self._forward, x) # pytorch
|
318 |
+
|
319 |
+
def _forward(self, x):
|
320 |
+
b, c, *spatial = x.shape
|
321 |
+
x = x.reshape(b, c, -1)
|
322 |
+
qkv = self.qkv(self.norm(x))
|
323 |
+
h = self.attention(qkv)
|
324 |
+
h = self.proj_out(h)
|
325 |
+
return (x + h).reshape(b, c, *spatial)
|
326 |
+
|
327 |
+
|
328 |
+
def count_flops_attn(model, _x, y):
|
329 |
+
"""
|
330 |
+
A counter for the `thop` package to count the operations in an
|
331 |
+
attention operation.
|
332 |
+
Meant to be used like:
|
333 |
+
macs, params = thop.profile(
|
334 |
+
model,
|
335 |
+
inputs=(inputs, timestamps),
|
336 |
+
custom_ops={QKVAttention: QKVAttention.count_flops},
|
337 |
+
)
|
338 |
+
"""
|
339 |
+
b, c, *spatial = y[0].shape
|
340 |
+
num_spatial = int(np.prod(spatial))
|
341 |
+
# We perform two matmuls with the same number of ops.
|
342 |
+
# The first computes the weight matrix, the second computes
|
343 |
+
# the combination of the value vectors.
|
344 |
+
matmul_ops = 2 * b * (num_spatial ** 2) * c
|
345 |
+
model.total_ops += th.DoubleTensor([matmul_ops])
|
346 |
+
|
347 |
+
|
348 |
+
class QKVAttentionLegacy(nn.Module):
|
349 |
+
"""
|
350 |
+
A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
|
351 |
+
"""
|
352 |
+
|
353 |
+
def __init__(self, n_heads):
|
354 |
+
super().__init__()
|
355 |
+
self.n_heads = n_heads
|
356 |
+
|
357 |
+
def forward(self, qkv):
|
358 |
+
"""
|
359 |
+
Apply QKV attention.
|
360 |
+
:param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
|
361 |
+
:return: an [N x (H * C) x T] tensor after attention.
|
362 |
+
"""
|
363 |
+
bs, width, length = qkv.shape
|
364 |
+
assert width % (3 * self.n_heads) == 0
|
365 |
+
ch = width // (3 * self.n_heads)
|
366 |
+
q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
|
367 |
+
scale = 1 / math.sqrt(math.sqrt(ch))
|
368 |
+
weight = th.einsum(
|
369 |
+
"bct,bcs->bts", q * scale, k * scale
|
370 |
+
) # More stable with f16 than dividing afterwards
|
371 |
+
weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
|
372 |
+
a = th.einsum("bts,bcs->bct", weight, v)
|
373 |
+
return a.reshape(bs, -1, length)
|
374 |
+
|
375 |
+
@staticmethod
|
376 |
+
def count_flops(model, _x, y):
|
377 |
+
return count_flops_attn(model, _x, y)
|
378 |
+
|
379 |
+
|
380 |
+
class QKVAttention(nn.Module):
|
381 |
+
"""
|
382 |
+
A module which performs QKV attention and splits in a different order.
|
383 |
+
"""
|
384 |
+
|
385 |
+
def __init__(self, n_heads):
|
386 |
+
super().__init__()
|
387 |
+
self.n_heads = n_heads
|
388 |
+
|
389 |
+
def forward(self, qkv):
|
390 |
+
"""
|
391 |
+
Apply QKV attention.
|
392 |
+
:param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
|
393 |
+
:return: an [N x (H * C) x T] tensor after attention.
|
394 |
+
"""
|
395 |
+
bs, width, length = qkv.shape
|
396 |
+
assert width % (3 * self.n_heads) == 0
|
397 |
+
ch = width // (3 * self.n_heads)
|
398 |
+
q, k, v = qkv.chunk(3, dim=1)
|
399 |
+
scale = 1 / math.sqrt(math.sqrt(ch))
|
400 |
+
weight = th.einsum(
|
401 |
+
"bct,bcs->bts",
|
402 |
+
(q * scale).view(bs * self.n_heads, ch, length),
|
403 |
+
(k * scale).view(bs * self.n_heads, ch, length),
|
404 |
+
) # More stable with f16 than dividing afterwards
|
405 |
+
weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
|
406 |
+
a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
|
407 |
+
return a.reshape(bs, -1, length)
|
408 |
+
|
409 |
+
@staticmethod
|
410 |
+
def count_flops(model, _x, y):
|
411 |
+
return count_flops_attn(model, _x, y)
|
412 |
+
|
413 |
+
|
414 |
+
class UNetModel(nn.Module):
|
415 |
+
"""
|
416 |
+
The full UNet model with attention and timestep embedding.
|
417 |
+
:param in_channels: channels in the input Tensor.
|
418 |
+
:param model_channels: base channel count for the model.
|
419 |
+
:param out_channels: channels in the output Tensor.
|
420 |
+
:param num_res_blocks: number of residual blocks per downsample.
|
421 |
+
:param attention_resolutions: a collection of downsample rates at which
|
422 |
+
attention will take place. May be a set, list, or tuple.
|
423 |
+
For example, if this contains 4, then at 4x downsampling, attention
|
424 |
+
will be used.
|
425 |
+
:param dropout: the dropout probability.
|
426 |
+
:param channel_mult: channel multiplier for each level of the UNet.
|
427 |
+
:param conv_resample: if True, use learned convolutions for upsampling and
|
428 |
+
downsampling.
|
429 |
+
:param dims: determines if the signal is 1D, 2D, or 3D.
|
430 |
+
:param num_classes: if specified (as an int), then this model will be
|
431 |
+
class-conditional with `num_classes` classes.
|
432 |
+
:param use_checkpoint: use gradient checkpointing to reduce memory usage.
|
433 |
+
:param num_heads: the number of attention heads in each attention layer.
|
434 |
+
:param num_heads_channels: if specified, ignore num_heads and instead use
|
435 |
+
a fixed channel width per attention head.
|
436 |
+
:param num_heads_upsample: works with num_heads to set a different number
|
437 |
+
of heads for upsampling. Deprecated.
|
438 |
+
:param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
|
439 |
+
:param resblock_updown: use residual blocks for up/downsampling.
|
440 |
+
:param use_new_attention_order: use a different attention pattern for potentially
|
441 |
+
increased efficiency.
|
442 |
+
"""
|
443 |
+
|
444 |
+
def __init__(
|
445 |
+
self,
|
446 |
+
image_size,
|
447 |
+
in_channels,
|
448 |
+
model_channels,
|
449 |
+
out_channels,
|
450 |
+
num_res_blocks,
|
451 |
+
attention_resolutions,
|
452 |
+
dropout=0,
|
453 |
+
channel_mult=(1, 2, 4, 8),
|
454 |
+
conv_resample=True,
|
455 |
+
dims=2,
|
456 |
+
num_classes=None,
|
457 |
+
use_checkpoint=False,
|
458 |
+
use_fp16=False,
|
459 |
+
num_heads=-1,
|
460 |
+
num_head_channels=-1,
|
461 |
+
num_heads_upsample=-1,
|
462 |
+
use_scale_shift_norm=False,
|
463 |
+
resblock_updown=False,
|
464 |
+
use_new_attention_order=False,
|
465 |
+
use_spatial_transformer=False, # custom transformer support
|
466 |
+
transformer_depth=1, # custom transformer support
|
467 |
+
context_dim=None, # custom transformer support
|
468 |
+
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
|
469 |
+
legacy=True,
|
470 |
+
disable_self_attentions=None,
|
471 |
+
num_attention_blocks=None
|
472 |
+
):
|
473 |
+
super().__init__()
|
474 |
+
if use_spatial_transformer:
|
475 |
+
assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
|
476 |
+
|
477 |
+
if context_dim is not None:
|
478 |
+
assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
|
479 |
+
from omegaconf.listconfig import ListConfig
|
480 |
+
if type(context_dim) == ListConfig:
|
481 |
+
context_dim = list(context_dim)
|
482 |
+
|
483 |
+
if num_heads_upsample == -1:
|
484 |
+
num_heads_upsample = num_heads
|
485 |
+
|
486 |
+
if num_heads == -1:
|
487 |
+
assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
|
488 |
+
|
489 |
+
if num_head_channels == -1:
|
490 |
+
assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
|
491 |
+
|
492 |
+
self.image_size = image_size
|
493 |
+
self.in_channels = in_channels
|
494 |
+
self.model_channels = model_channels
|
495 |
+
self.out_channels = out_channels
|
496 |
+
if isinstance(num_res_blocks, int):
|
497 |
+
self.num_res_blocks = len(channel_mult) * [num_res_blocks]
|
498 |
+
else:
|
499 |
+
if len(num_res_blocks) != len(channel_mult):
|
500 |
+
raise ValueError("provide num_res_blocks either as an int (globally constant) or "
|
501 |
+
"as a list/tuple (per-level) with the same length as channel_mult")
|
502 |
+
self.num_res_blocks = num_res_blocks
|
503 |
+
#self.num_res_blocks = num_res_blocks
|
504 |
+
if disable_self_attentions is not None:
|
505 |
+
# should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
|
506 |
+
assert len(disable_self_attentions) == len(channel_mult)
|
507 |
+
if num_attention_blocks is not None:
|
508 |
+
assert len(num_attention_blocks) == len(self.num_res_blocks)
|
509 |
+
assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
|
510 |
+
print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
|
511 |
+
f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
|
512 |
+
f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
|
513 |
+
f"attention will still not be set.") # todo: convert to warning
|
514 |
+
|
515 |
+
self.attention_resolutions = attention_resolutions
|
516 |
+
self.dropout = dropout
|
517 |
+
self.channel_mult = channel_mult
|
518 |
+
self.conv_resample = conv_resample
|
519 |
+
self.num_classes = num_classes
|
520 |
+
self.use_checkpoint = use_checkpoint
|
521 |
+
self.dtype = th.float16 if use_fp16 else th.float32
|
522 |
+
self.num_heads = num_heads
|
523 |
+
self.num_head_channels = num_head_channels
|
524 |
+
self.num_heads_upsample = num_heads_upsample
|
525 |
+
self.predict_codebook_ids = n_embed is not None
|
526 |
+
|
527 |
+
time_embed_dim = model_channels * 4
|
528 |
+
self.time_embed = nn.Sequential(
|
529 |
+
linear(model_channels, time_embed_dim),
|
530 |
+
nn.SiLU(),
|
531 |
+
linear(time_embed_dim, time_embed_dim),
|
532 |
+
)
|
533 |
+
|
534 |
+
if self.num_classes is not None:
|
535 |
+
self.label_emb = nn.Embedding(num_classes, time_embed_dim)
|
536 |
+
|
537 |
+
self.input_blocks = nn.ModuleList(
|
538 |
+
[
|
539 |
+
TimestepEmbedSequential(
|
540 |
+
conv_nd(dims, in_channels, model_channels, 3, padding=1)
|
541 |
+
)
|
542 |
+
]
|
543 |
+
) # 0
|
544 |
+
self._feature_size = model_channels
|
545 |
+
input_block_chans = [model_channels]
|
546 |
+
ch = model_channels
|
547 |
+
ds = 1
|
548 |
+
for level, mult in enumerate(channel_mult):
|
549 |
+
for nr in range(self.num_res_blocks[level]):
|
550 |
+
layers = [
|
551 |
+
ResBlock(
|
552 |
+
ch,
|
553 |
+
time_embed_dim,
|
554 |
+
dropout,
|
555 |
+
out_channels=mult * model_channels,
|
556 |
+
dims=dims,
|
557 |
+
use_checkpoint=use_checkpoint,
|
558 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
559 |
+
)
|
560 |
+
]
|
561 |
+
ch = mult * model_channels
|
562 |
+
if ds in attention_resolutions: # always True
|
563 |
+
if num_head_channels == -1:
|
564 |
+
dim_head = ch // num_heads
|
565 |
+
else:
|
566 |
+
num_heads = ch // num_head_channels
|
567 |
+
dim_head = num_head_channels
|
568 |
+
if legacy:
|
569 |
+
#num_heads = 1
|
570 |
+
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
|
571 |
+
if exists(disable_self_attentions):
|
572 |
+
disabled_sa = disable_self_attentions[level]
|
573 |
+
else:
|
574 |
+
disabled_sa = False
|
575 |
+
|
576 |
+
if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
|
577 |
+
layers.append(
|
578 |
+
AttentionBlock(
|
579 |
+
ch,
|
580 |
+
use_checkpoint=use_checkpoint,
|
581 |
+
num_heads=num_heads,
|
582 |
+
num_head_channels=dim_head,
|
583 |
+
use_new_attention_order=use_new_attention_order,
|
584 |
+
) if not use_spatial_transformer else SpatialTransformer(
|
585 |
+
ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
|
586 |
+
disable_self_attn=disabled_sa
|
587 |
+
)
|
588 |
+
)
|
589 |
+
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
590 |
+
self._feature_size += ch
|
591 |
+
input_block_chans.append(ch)
|
592 |
+
if level != len(channel_mult) - 1:
|
593 |
+
out_ch = ch
|
594 |
+
self.input_blocks.append(
|
595 |
+
TimestepEmbedSequential(
|
596 |
+
ResBlock(
|
597 |
+
ch,
|
598 |
+
time_embed_dim,
|
599 |
+
dropout,
|
600 |
+
out_channels=out_ch,
|
601 |
+
dims=dims,
|
602 |
+
use_checkpoint=use_checkpoint,
|
603 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
604 |
+
down=True,
|
605 |
+
)
|
606 |
+
if resblock_updown
|
607 |
+
else Downsample(
|
608 |
+
ch, conv_resample, dims=dims, out_channels=out_ch
|
609 |
+
)
|
610 |
+
)
|
611 |
+
)
|
612 |
+
ch = out_ch
|
613 |
+
input_block_chans.append(ch)
|
614 |
+
ds *= 2
|
615 |
+
self._feature_size += ch
|
616 |
+
|
617 |
+
if num_head_channels == -1:
|
618 |
+
dim_head = ch // num_heads
|
619 |
+
else:
|
620 |
+
num_heads = ch // num_head_channels
|
621 |
+
dim_head = num_head_channels
|
622 |
+
if legacy:
|
623 |
+
#num_heads = 1
|
624 |
+
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
|
625 |
+
self.middle_block = TimestepEmbedSequential(
|
626 |
+
ResBlock(
|
627 |
+
ch,
|
628 |
+
time_embed_dim,
|
629 |
+
dropout,
|
630 |
+
dims=dims,
|
631 |
+
use_checkpoint=use_checkpoint,
|
632 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
633 |
+
),
|
634 |
+
AttentionBlock(
|
635 |
+
ch,
|
636 |
+
use_checkpoint=use_checkpoint,
|
637 |
+
num_heads=num_heads,
|
638 |
+
num_head_channels=dim_head,
|
639 |
+
use_new_attention_order=use_new_attention_order,
|
640 |
+
) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn
|
641 |
+
ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim
|
642 |
+
),
|
643 |
+
ResBlock(
|
644 |
+
ch,
|
645 |
+
time_embed_dim,
|
646 |
+
dropout,
|
647 |
+
dims=dims,
|
648 |
+
use_checkpoint=use_checkpoint,
|
649 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
650 |
+
),
|
651 |
+
)
|
652 |
+
self._feature_size += ch
|
653 |
+
|
654 |
+
self.output_blocks = nn.ModuleList([])
|
655 |
+
for level, mult in list(enumerate(channel_mult))[::-1]:
|
656 |
+
for i in range(self.num_res_blocks[level] + 1):
|
657 |
+
ich = input_block_chans.pop()
|
658 |
+
layers = [
|
659 |
+
ResBlock(
|
660 |
+
ch + ich,
|
661 |
+
time_embed_dim,
|
662 |
+
dropout,
|
663 |
+
out_channels=model_channels * mult,
|
664 |
+
dims=dims,
|
665 |
+
use_checkpoint=use_checkpoint,
|
666 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
667 |
+
)
|
668 |
+
]
|
669 |
+
ch = model_channels * mult
|
670 |
+
if ds in attention_resolutions:
|
671 |
+
if num_head_channels == -1:
|
672 |
+
dim_head = ch // num_heads
|
673 |
+
else:
|
674 |
+
num_heads = ch // num_head_channels
|
675 |
+
dim_head = num_head_channels
|
676 |
+
if legacy:
|
677 |
+
#num_heads = 1
|
678 |
+
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
|
679 |
+
if exists(disable_self_attentions):
|
680 |
+
disabled_sa = disable_self_attentions[level]
|
681 |
+
else:
|
682 |
+
disabled_sa = False
|
683 |
+
|
684 |
+
if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
|
685 |
+
layers.append(
|
686 |
+
AttentionBlock(
|
687 |
+
ch,
|
688 |
+
use_checkpoint=use_checkpoint,
|
689 |
+
num_heads=num_heads_upsample,
|
690 |
+
num_head_channels=dim_head,
|
691 |
+
use_new_attention_order=use_new_attention_order,
|
692 |
+
) if not use_spatial_transformer else SpatialTransformer(
|
693 |
+
ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
|
694 |
+
disable_self_attn=disabled_sa
|
695 |
+
)
|
696 |
+
)
|
697 |
+
if level and i == self.num_res_blocks[level]:
|
698 |
+
out_ch = ch
|
699 |
+
layers.append(
|
700 |
+
ResBlock(
|
701 |
+
ch,
|
702 |
+
time_embed_dim,
|
703 |
+
dropout,
|
704 |
+
out_channels=out_ch,
|
705 |
+
dims=dims,
|
706 |
+
use_checkpoint=use_checkpoint,
|
707 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
708 |
+
up=True,
|
709 |
+
)
|
710 |
+
if resblock_updown
|
711 |
+
else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
|
712 |
+
)
|
713 |
+
ds //= 2
|
714 |
+
self.output_blocks.append(TimestepEmbedSequential(*layers))
|
715 |
+
self._feature_size += ch
|
716 |
+
|
717 |
+
self.out = nn.Sequential(
|
718 |
+
normalization(ch),
|
719 |
+
nn.SiLU(),
|
720 |
+
zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
|
721 |
+
)
|
722 |
+
if self.predict_codebook_ids:
|
723 |
+
self.id_predictor = nn.Sequential(
|
724 |
+
normalization(ch),
|
725 |
+
conv_nd(dims, model_channels, n_embed, 1),
|
726 |
+
#nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
|
727 |
+
)
|
728 |
+
|
729 |
+
def convert_to_fp16(self):
|
730 |
+
"""
|
731 |
+
Convert the torso of the model to float16.
|
732 |
+
"""
|
733 |
+
self.input_blocks.apply(convert_module_to_f16)
|
734 |
+
self.middle_block.apply(convert_module_to_f16)
|
735 |
+
self.output_blocks.apply(convert_module_to_f16)
|
736 |
+
|
737 |
+
def convert_to_fp32(self):
|
738 |
+
"""
|
739 |
+
Convert the torso of the model to float32.
|
740 |
+
"""
|
741 |
+
self.input_blocks.apply(convert_module_to_f32)
|
742 |
+
self.middle_block.apply(convert_module_to_f32)
|
743 |
+
self.output_blocks.apply(convert_module_to_f32)
|
744 |
+
|
745 |
+
def forward(self, x, timesteps=None, context=None, y=None,**kwargs):
|
746 |
+
"""
|
747 |
+
Apply the model to an input batch.
|
748 |
+
:param x: an [N x C x ...] Tensor of inputs.
|
749 |
+
:param timesteps: a 1-D batch of timesteps.
|
750 |
+
:param context: conditioning plugged in via crossattn
|
751 |
+
:param y: an [N] Tensor of labels, if class-conditional.
|
752 |
+
:return: an [N x C x ...] Tensor of outputs.
|
753 |
+
"""
|
754 |
+
assert (y is not None) == (
|
755 |
+
self.num_classes is not None
|
756 |
+
), "must specify y if and only if the model is class-conditional"
|
757 |
+
hs = []
|
758 |
+
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) # N
|
759 |
+
emb = self.time_embed(t_emb) #
|
760 |
+
|
761 |
+
if self.num_classes is not None:
|
762 |
+
assert y.shape == (x.shape[0],)
|
763 |
+
emb = emb + self.label_emb(y)
|
764 |
+
|
765 |
+
h = x.type(self.dtype)
|
766 |
+
for module in self.input_blocks:
|
767 |
+
h = module(h, emb, context) # conv
|
768 |
+
hs.append(h)
|
769 |
+
h = self.middle_block(h, emb, context)
|
770 |
+
for module in self.output_blocks:
|
771 |
+
h = th.cat([h, hs.pop()], dim=1)
|
772 |
+
h = module(h, emb, context)
|
773 |
+
h = h.type(x.dtype)
|
774 |
+
if self.predict_codebook_ids:
|
775 |
+
return self.id_predictor(h)
|
776 |
+
else:
|
777 |
+
return self.out(h)
|
778 |
+
|
779 |
+
|
780 |
+
class EncoderUNetModel(nn.Module):
|
781 |
+
"""
|
782 |
+
The half UNet model with attention and timestep embedding.
|
783 |
+
For usage, see UNet.
|
784 |
+
"""
|
785 |
+
|
786 |
+
def __init__(
|
787 |
+
self,
|
788 |
+
image_size,
|
789 |
+
in_channels,
|
790 |
+
model_channels,
|
791 |
+
out_channels,
|
792 |
+
num_res_blocks,
|
793 |
+
attention_resolutions,
|
794 |
+
dropout=0,
|
795 |
+
channel_mult=(1, 2, 4, 8),
|
796 |
+
conv_resample=True,
|
797 |
+
dims=2,
|
798 |
+
use_checkpoint=False,
|
799 |
+
use_fp16=False,
|
800 |
+
num_heads=1,
|
801 |
+
num_head_channels=-1,
|
802 |
+
num_heads_upsample=-1,
|
803 |
+
use_scale_shift_norm=False,
|
804 |
+
resblock_updown=False,
|
805 |
+
use_new_attention_order=False,
|
806 |
+
pool="adaptive",
|
807 |
+
*args,
|
808 |
+
**kwargs
|
809 |
+
):
|
810 |
+
super().__init__()
|
811 |
+
|
812 |
+
if num_heads_upsample == -1:
|
813 |
+
num_heads_upsample = num_heads
|
814 |
+
|
815 |
+
self.in_channels = in_channels
|
816 |
+
self.model_channels = model_channels
|
817 |
+
self.out_channels = out_channels
|
818 |
+
self.num_res_blocks = num_res_blocks
|
819 |
+
self.attention_resolutions = attention_resolutions
|
820 |
+
self.dropout = dropout
|
821 |
+
self.channel_mult = channel_mult
|
822 |
+
self.conv_resample = conv_resample
|
823 |
+
self.use_checkpoint = use_checkpoint
|
824 |
+
self.dtype = th.float16 if use_fp16 else th.float32
|
825 |
+
self.num_heads = num_heads
|
826 |
+
self.num_head_channels = num_head_channels
|
827 |
+
self.num_heads_upsample = num_heads_upsample
|
828 |
+
|
829 |
+
time_embed_dim = model_channels * 4
|
830 |
+
self.time_embed = nn.Sequential(
|
831 |
+
linear(model_channels, time_embed_dim),
|
832 |
+
nn.SiLU(),
|
833 |
+
linear(time_embed_dim, time_embed_dim),
|
834 |
+
)
|
835 |
+
|
836 |
+
self.input_blocks = nn.ModuleList(
|
837 |
+
[
|
838 |
+
TimestepEmbedSequential(
|
839 |
+
conv_nd(dims, in_channels, model_channels, 3, padding=1)
|
840 |
+
)
|
841 |
+
]
|
842 |
+
)
|
843 |
+
self._feature_size = model_channels
|
844 |
+
input_block_chans = [model_channels]
|
845 |
+
ch = model_channels
|
846 |
+
ds = 1
|
847 |
+
for level, mult in enumerate(channel_mult):
|
848 |
+
for _ in range(num_res_blocks):
|
849 |
+
layers = [
|
850 |
+
ResBlock(
|
851 |
+
ch,
|
852 |
+
time_embed_dim,
|
853 |
+
dropout,
|
854 |
+
out_channels=mult * model_channels,
|
855 |
+
dims=dims,
|
856 |
+
use_checkpoint=use_checkpoint,
|
857 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
858 |
+
)
|
859 |
+
]
|
860 |
+
ch = mult * model_channels
|
861 |
+
if ds in attention_resolutions:
|
862 |
+
layers.append(
|
863 |
+
AttentionBlock(
|
864 |
+
ch,
|
865 |
+
use_checkpoint=use_checkpoint,
|
866 |
+
num_heads=num_heads,
|
867 |
+
num_head_channels=num_head_channels,
|
868 |
+
use_new_attention_order=use_new_attention_order,
|
869 |
+
)
|
870 |
+
)
|
871 |
+
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
872 |
+
self._feature_size += ch
|
873 |
+
input_block_chans.append(ch)
|
874 |
+
if level != len(channel_mult) - 1:
|
875 |
+
out_ch = ch
|
876 |
+
self.input_blocks.append(
|
877 |
+
TimestepEmbedSequential(
|
878 |
+
ResBlock(
|
879 |
+
ch,
|
880 |
+
time_embed_dim,
|
881 |
+
dropout,
|
882 |
+
out_channels=out_ch,
|
883 |
+
dims=dims,
|
884 |
+
use_checkpoint=use_checkpoint,
|
885 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
886 |
+
down=True,
|
887 |
+
)
|
888 |
+
if resblock_updown
|
889 |
+
else Downsample(
|
890 |
+
ch, conv_resample, dims=dims, out_channels=out_ch
|
891 |
+
)
|
892 |
+
)
|
893 |
+
)
|
894 |
+
ch = out_ch
|
895 |
+
input_block_chans.append(ch)
|
896 |
+
ds *= 2
|
897 |
+
self._feature_size += ch
|
898 |
+
|
899 |
+
self.middle_block = TimestepEmbedSequential(
|
900 |
+
ResBlock(
|
901 |
+
ch,
|
902 |
+
time_embed_dim,
|
903 |
+
dropout,
|
904 |
+
dims=dims,
|
905 |
+
use_checkpoint=use_checkpoint,
|
906 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
907 |
+
),
|
908 |
+
AttentionBlock(
|
909 |
+
ch,
|
910 |
+
use_checkpoint=use_checkpoint,
|
911 |
+
num_heads=num_heads,
|
912 |
+
num_head_channels=num_head_channels,
|
913 |
+
use_new_attention_order=use_new_attention_order,
|
914 |
+
),
|
915 |
+
ResBlock(
|
916 |
+
ch,
|
917 |
+
time_embed_dim,
|
918 |
+
dropout,
|
919 |
+
dims=dims,
|
920 |
+
use_checkpoint=use_checkpoint,
|
921 |
+
use_scale_shift_norm=use_scale_shift_norm,
|
922 |
+
),
|
923 |
+
)
|
924 |
+
self._feature_size += ch
|
925 |
+
self.pool = pool
|
926 |
+
if pool == "adaptive":
|
927 |
+
self.out = nn.Sequential(
|
928 |
+
normalization(ch),
|
929 |
+
nn.SiLU(),
|
930 |
+
nn.AdaptiveAvgPool2d((1, 1)),
|
931 |
+
zero_module(conv_nd(dims, ch, out_channels, 1)),
|
932 |
+
nn.Flatten(),
|
933 |
+
)
|
934 |
+
elif pool == "attention":
|
935 |
+
assert num_head_channels != -1
|
936 |
+
self.out = nn.Sequential(
|
937 |
+
normalization(ch),
|
938 |
+
nn.SiLU(),
|
939 |
+
AttentionPool2d(
|
940 |
+
(image_size // ds), ch, num_head_channels, out_channels
|
941 |
+
),
|
942 |
+
)
|
943 |
+
elif pool == "spatial":
|
944 |
+
self.out = nn.Sequential(
|
945 |
+
nn.Linear(self._feature_size, 2048),
|
946 |
+
nn.ReLU(),
|
947 |
+
nn.Linear(2048, self.out_channels),
|
948 |
+
)
|
949 |
+
elif pool == "spatial_v2":
|
950 |
+
self.out = nn.Sequential(
|
951 |
+
nn.Linear(self._feature_size, 2048),
|
952 |
+
normalization(2048),
|
953 |
+
nn.SiLU(),
|
954 |
+
nn.Linear(2048, self.out_channels),
|
955 |
+
)
|
956 |
+
else:
|
957 |
+
raise NotImplementedError(f"Unexpected {pool} pooling")
|
958 |
+
|
959 |
+
def convert_to_fp16(self):
|
960 |
+
"""
|
961 |
+
Convert the torso of the model to float16.
|
962 |
+
"""
|
963 |
+
self.input_blocks.apply(convert_module_to_f16)
|
964 |
+
self.middle_block.apply(convert_module_to_f16)
|
965 |
+
|
966 |
+
def convert_to_fp32(self):
|
967 |
+
"""
|
968 |
+
Convert the torso of the model to float32.
|
969 |
+
"""
|
970 |
+
self.input_blocks.apply(convert_module_to_f32)
|
971 |
+
self.middle_block.apply(convert_module_to_f32)
|
972 |
+
|
973 |
+
def forward(self, x, timesteps):
|
974 |
+
"""
|
975 |
+
Apply the model to an input batch.
|
976 |
+
:param x: an [N x C x ...] Tensor of inputs.
|
977 |
+
:param timesteps: a 1-D batch of timesteps.
|
978 |
+
:return: an [N x K] Tensor of outputs.
|
979 |
+
"""
|
980 |
+
emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
|
981 |
+
|
982 |
+
results = []
|
983 |
+
h = x.type(self.dtype)
|
984 |
+
for module in self.input_blocks:
|
985 |
+
h = module(h, emb)
|
986 |
+
if self.pool.startswith("spatial"):
|
987 |
+
results.append(h.type(x.dtype).mean(dim=(2, 3)))
|
988 |
+
h = self.middle_block(h, emb)
|
989 |
+
if self.pool.startswith("spatial"):
|
990 |
+
results.append(h.type(x.dtype).mean(dim=(2, 3)))
|
991 |
+
h = th.cat(results, axis=-1)
|
992 |
+
return self.out(h)
|
993 |
+
else:
|
994 |
+
h = h.type(x.dtype)
|
995 |
+
return self.out(h)
|
996 |
+
|
ldm/modules/diffusionmodules/util.py
ADDED
@@ -0,0 +1,267 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# adopted from
|
2 |
+
# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
|
3 |
+
# and
|
4 |
+
# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
|
5 |
+
# and
|
6 |
+
# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
|
7 |
+
#
|
8 |
+
# thanks!
|
9 |
+
|
10 |
+
|
11 |
+
import os
|
12 |
+
import math
|
13 |
+
import torch
|
14 |
+
import torch.nn as nn
|
15 |
+
import numpy as np
|
16 |
+
from einops import repeat
|
17 |
+
|
18 |
+
from ldm.util import instantiate_from_config
|
19 |
+
|
20 |
+
|
21 |
+
def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
22 |
+
if schedule == "linear":
|
23 |
+
betas = (
|
24 |
+
torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
|
25 |
+
)
|
26 |
+
|
27 |
+
elif schedule == "cosine":
|
28 |
+
timesteps = (
|
29 |
+
torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
|
30 |
+
)
|
31 |
+
alphas = timesteps / (1 + cosine_s) * np.pi / 2
|
32 |
+
alphas = torch.cos(alphas).pow(2)
|
33 |
+
alphas = alphas / alphas[0]
|
34 |
+
betas = 1 - alphas[1:] / alphas[:-1]
|
35 |
+
betas = np.clip(betas, a_min=0, a_max=0.999)
|
36 |
+
|
37 |
+
elif schedule == "sqrt_linear":
|
38 |
+
betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
|
39 |
+
elif schedule == "sqrt":
|
40 |
+
betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
|
41 |
+
else:
|
42 |
+
raise ValueError(f"schedule '{schedule}' unknown.")
|
43 |
+
return betas.numpy()
|
44 |
+
|
45 |
+
|
46 |
+
def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
|
47 |
+
if ddim_discr_method == 'uniform':
|
48 |
+
c = num_ddpm_timesteps // num_ddim_timesteps
|
49 |
+
ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
|
50 |
+
elif ddim_discr_method == 'quad':
|
51 |
+
ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
|
52 |
+
else:
|
53 |
+
raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
|
54 |
+
|
55 |
+
# assert ddim_timesteps.shape[0] == num_ddim_timesteps
|
56 |
+
# add one to get the final alpha values right (the ones from first scale to data during sampling)
|
57 |
+
steps_out = ddim_timesteps + 1
|
58 |
+
if verbose:
|
59 |
+
print(f'Selected timesteps for ddim sampler: {steps_out}')
|
60 |
+
return steps_out
|
61 |
+
|
62 |
+
|
63 |
+
def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
|
64 |
+
# select alphas for computing the variance schedule
|
65 |
+
alphas = alphacums[ddim_timesteps]
|
66 |
+
alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
|
67 |
+
|
68 |
+
# according the the formula provided in https://arxiv.org/abs/2010.02502
|
69 |
+
sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
|
70 |
+
if verbose:
|
71 |
+
print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
|
72 |
+
print(f'For the chosen value of eta, which is {eta}, '
|
73 |
+
f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
|
74 |
+
return sigmas, alphas, alphas_prev
|
75 |
+
|
76 |
+
|
77 |
+
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
|
78 |
+
"""
|
79 |
+
Create a beta schedule that discretizes the given alpha_t_bar function,
|
80 |
+
which defines the cumulative product of (1-beta) over time from t = [0,1].
|
81 |
+
:param num_diffusion_timesteps: the number of betas to produce.
|
82 |
+
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
|
83 |
+
produces the cumulative product of (1-beta) up to that
|
84 |
+
part of the diffusion process.
|
85 |
+
:param max_beta: the maximum beta to use; use values lower than 1 to
|
86 |
+
prevent singularities.
|
87 |
+
"""
|
88 |
+
betas = []
|
89 |
+
for i in range(num_diffusion_timesteps):
|
90 |
+
t1 = i / num_diffusion_timesteps
|
91 |
+
t2 = (i + 1) / num_diffusion_timesteps
|
92 |
+
betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
|
93 |
+
return np.array(betas)
|
94 |
+
|
95 |
+
|
96 |
+
def extract_into_tensor(a, t, x_shape):
|
97 |
+
b, *_ = t.shape
|
98 |
+
out = a.gather(-1, t)
|
99 |
+
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
|
100 |
+
|
101 |
+
|
102 |
+
def checkpoint(func, inputs, params, flag):
|
103 |
+
"""
|
104 |
+
Evaluate a function without caching intermediate activations, allowing for
|
105 |
+
reduced memory at the expense of extra compute in the backward pass.
|
106 |
+
:param func: the function to evaluate.
|
107 |
+
:param inputs: the argument sequence to pass to `func`.
|
108 |
+
:param params: a sequence of parameters `func` depends on but does not
|
109 |
+
explicitly take as arguments.
|
110 |
+
:param flag: if False, disable gradient checkpointing.
|
111 |
+
"""
|
112 |
+
if flag:
|
113 |
+
args = tuple(inputs) + tuple(params)
|
114 |
+
return CheckpointFunction.apply(func, len(inputs), *args)
|
115 |
+
else:
|
116 |
+
return func(*inputs)
|
117 |
+
|
118 |
+
|
119 |
+
class CheckpointFunction(torch.autograd.Function):
|
120 |
+
@staticmethod
|
121 |
+
def forward(ctx, run_function, length, *args):
|
122 |
+
ctx.run_function = run_function
|
123 |
+
ctx.input_tensors = list(args[:length])
|
124 |
+
ctx.input_params = list(args[length:])
|
125 |
+
|
126 |
+
with torch.no_grad():
|
127 |
+
output_tensors = ctx.run_function(*ctx.input_tensors)
|
128 |
+
return output_tensors
|
129 |
+
|
130 |
+
@staticmethod
|
131 |
+
def backward(ctx, *output_grads):
|
132 |
+
ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
|
133 |
+
with torch.enable_grad():
|
134 |
+
# Fixes a bug where the first op in run_function modifies the
|
135 |
+
# Tensor storage in place, which is not allowed for detach()'d
|
136 |
+
# Tensors.
|
137 |
+
shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
|
138 |
+
output_tensors = ctx.run_function(*shallow_copies)
|
139 |
+
input_grads = torch.autograd.grad(
|
140 |
+
output_tensors,
|
141 |
+
ctx.input_tensors + ctx.input_params,
|
142 |
+
output_grads,
|
143 |
+
allow_unused=True,
|
144 |
+
)
|
145 |
+
del ctx.input_tensors
|
146 |
+
del ctx.input_params
|
147 |
+
del output_tensors
|
148 |
+
return (None, None) + input_grads
|
149 |
+
|
150 |
+
|
151 |
+
def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
|
152 |
+
"""
|
153 |
+
Create sinusoidal timestep embeddings.
|
154 |
+
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
155 |
+
These may be fractional.
|
156 |
+
:param dim: the dimension of the output.
|
157 |
+
:param max_period: controls the minimum frequency of the embeddings.
|
158 |
+
:return: an [N x dim] Tensor of positional embeddings.
|
159 |
+
"""
|
160 |
+
if not repeat_only:
|
161 |
+
half = dim // 2
|
162 |
+
freqs = torch.exp(
|
163 |
+
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
|
164 |
+
).to(device=timesteps.device)
|
165 |
+
args = timesteps[:, None].float() * freqs[None]
|
166 |
+
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
167 |
+
if dim % 2:
|
168 |
+
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
169 |
+
else:
|
170 |
+
embedding = repeat(timesteps, 'b -> b d', d=dim)
|
171 |
+
return embedding
|
172 |
+
|
173 |
+
|
174 |
+
def zero_module(module):
|
175 |
+
"""
|
176 |
+
Zero out the parameters of a module and return it.
|
177 |
+
"""
|
178 |
+
for p in module.parameters():
|
179 |
+
p.detach().zero_()
|
180 |
+
return module
|
181 |
+
|
182 |
+
|
183 |
+
def scale_module(module, scale):
|
184 |
+
"""
|
185 |
+
Scale the parameters of a module and return it.
|
186 |
+
"""
|
187 |
+
for p in module.parameters():
|
188 |
+
p.detach().mul_(scale)
|
189 |
+
return module
|
190 |
+
|
191 |
+
|
192 |
+
def mean_flat(tensor):
|
193 |
+
"""
|
194 |
+
Take the mean over all non-batch dimensions.
|
195 |
+
"""
|
196 |
+
return tensor.mean(dim=list(range(1, len(tensor.shape))))
|
197 |
+
|
198 |
+
|
199 |
+
def normalization(channels):
|
200 |
+
"""
|
201 |
+
Make a standard normalization layer.
|
202 |
+
:param channels: number of input channels.
|
203 |
+
:return: an nn.Module for normalization.
|
204 |
+
"""
|
205 |
+
return GroupNorm32(32, channels)
|
206 |
+
|
207 |
+
|
208 |
+
# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
|
209 |
+
class SiLU(nn.Module):
|
210 |
+
def forward(self, x):
|
211 |
+
return x * torch.sigmoid(x)
|
212 |
+
|
213 |
+
|
214 |
+
class GroupNorm32(nn.GroupNorm):
|
215 |
+
def forward(self, x):
|
216 |
+
return super().forward(x.float()).type(x.dtype)
|
217 |
+
|
218 |
+
def conv_nd(dims, *args, **kwargs):
|
219 |
+
"""
|
220 |
+
Create a 1D, 2D, or 3D convolution module.
|
221 |
+
"""
|
222 |
+
if dims == 1:
|
223 |
+
return nn.Conv1d(*args, **kwargs)
|
224 |
+
elif dims == 2:
|
225 |
+
return nn.Conv2d(*args, **kwargs)
|
226 |
+
elif dims == 3:
|
227 |
+
return nn.Conv3d(*args, **kwargs)
|
228 |
+
raise ValueError(f"unsupported dimensions: {dims}")
|
229 |
+
|
230 |
+
|
231 |
+
def linear(*args, **kwargs):
|
232 |
+
"""
|
233 |
+
Create a linear module.
|
234 |
+
"""
|
235 |
+
return nn.Linear(*args, **kwargs)
|
236 |
+
|
237 |
+
|
238 |
+
def avg_pool_nd(dims, *args, **kwargs):
|
239 |
+
"""
|
240 |
+
Create a 1D, 2D, or 3D average pooling module.
|
241 |
+
"""
|
242 |
+
if dims == 1:
|
243 |
+
return nn.AvgPool1d(*args, **kwargs)
|
244 |
+
elif dims == 2:
|
245 |
+
return nn.AvgPool2d(*args, **kwargs)
|
246 |
+
elif dims == 3:
|
247 |
+
return nn.AvgPool3d(*args, **kwargs)
|
248 |
+
raise ValueError(f"unsupported dimensions: {dims}")
|
249 |
+
|
250 |
+
|
251 |
+
class HybridConditioner(nn.Module):
|
252 |
+
|
253 |
+
def __init__(self, c_concat_config, c_crossattn_config):
|
254 |
+
super().__init__()
|
255 |
+
self.concat_conditioner = instantiate_from_config(c_concat_config)
|
256 |
+
self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
|
257 |
+
|
258 |
+
def forward(self, c_concat, c_crossattn):
|
259 |
+
c_concat = self.concat_conditioner(c_concat)
|
260 |
+
c_crossattn = self.crossattn_conditioner(c_crossattn)
|
261 |
+
return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
|
262 |
+
|
263 |
+
|
264 |
+
def noise_like(shape, device, repeat=False):
|
265 |
+
repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
|
266 |
+
noise = lambda: torch.randn(shape, device=device)
|
267 |
+
return repeat_noise() if repeat else noise()
|
ldm/modules/distributions/__init__.py
ADDED
File without changes
|
ldm/modules/distributions/distributions.py
ADDED
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import numpy as np
|
3 |
+
|
4 |
+
|
5 |
+
class AbstractDistribution:
|
6 |
+
def sample(self):
|
7 |
+
raise NotImplementedError()
|
8 |
+
|
9 |
+
def mode(self):
|
10 |
+
raise NotImplementedError()
|
11 |
+
|
12 |
+
|
13 |
+
class DiracDistribution(AbstractDistribution):
|
14 |
+
def __init__(self, value):
|
15 |
+
self.value = value
|
16 |
+
|
17 |
+
def sample(self):
|
18 |
+
return self.value
|
19 |
+
|
20 |
+
def mode(self):
|
21 |
+
return self.value
|
22 |
+
|
23 |
+
|
24 |
+
class DiagonalGaussianDistribution(object):
|
25 |
+
def __init__(self, parameters, deterministic=False):
|
26 |
+
self.parameters = parameters
|
27 |
+
self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
|
28 |
+
self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
|
29 |
+
self.deterministic = deterministic
|
30 |
+
self.std = torch.exp(0.5 * self.logvar)
|
31 |
+
self.var = torch.exp(self.logvar)
|
32 |
+
if self.deterministic:
|
33 |
+
self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
|
34 |
+
|
35 |
+
def sample(self):
|
36 |
+
x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
|
37 |
+
return x
|
38 |
+
|
39 |
+
def kl(self, other=None):
|
40 |
+
if self.deterministic:
|
41 |
+
return torch.Tensor([0.])
|
42 |
+
else:
|
43 |
+
if other is None:
|
44 |
+
return 0.5 * torch.sum(torch.pow(self.mean, 2)
|
45 |
+
+ self.var - 1.0 - self.logvar,
|
46 |
+
dim=[1, 2, 3])
|
47 |
+
else:
|
48 |
+
return 0.5 * torch.sum(
|
49 |
+
torch.pow(self.mean - other.mean, 2) / other.var
|
50 |
+
+ self.var / other.var - 1.0 - self.logvar + other.logvar,
|
51 |
+
dim=[1, 2, 3])
|
52 |
+
|
53 |
+
def nll(self, sample, dims=[1,2,3]):
|
54 |
+
if self.deterministic:
|
55 |
+
return torch.Tensor([0.])
|
56 |
+
logtwopi = np.log(2.0 * np.pi)
|
57 |
+
return 0.5 * torch.sum(
|
58 |
+
logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
|
59 |
+
dim=dims)
|
60 |
+
|
61 |
+
def mode(self):
|
62 |
+
return self.mean
|
63 |
+
|
64 |
+
|
65 |
+
def normal_kl(mean1, logvar1, mean2, logvar2):
|
66 |
+
"""
|
67 |
+
source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
|
68 |
+
Compute the KL divergence between two gaussians.
|
69 |
+
Shapes are automatically broadcasted, so batches can be compared to
|
70 |
+
scalars, among other use cases.
|
71 |
+
"""
|
72 |
+
tensor = None
|
73 |
+
for obj in (mean1, logvar1, mean2, logvar2):
|
74 |
+
if isinstance(obj, torch.Tensor):
|
75 |
+
tensor = obj
|
76 |
+
break
|
77 |
+
assert tensor is not None, "at least one argument must be a Tensor"
|
78 |
+
|
79 |
+
# Force variances to be Tensors. Broadcasting helps convert scalars to
|
80 |
+
# Tensors, but it does not work for torch.exp().
|
81 |
+
logvar1, logvar2 = [
|
82 |
+
x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
|
83 |
+
for x in (logvar1, logvar2)
|
84 |
+
]
|
85 |
+
|
86 |
+
return 0.5 * (
|
87 |
+
-1.0
|
88 |
+
+ logvar2
|
89 |
+
- logvar1
|
90 |
+
+ torch.exp(logvar1 - logvar2)
|
91 |
+
+ ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
|
92 |
+
)
|
ldm/modules/encoders/__init__.py
ADDED
File without changes
|
ldm/modules/encoders/modules.py
ADDED
@@ -0,0 +1,550 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
import numpy as np
|
4 |
+
from functools import partial
|
5 |
+
import kornia
|
6 |
+
|
7 |
+
from ldm.modules.x_transformer import Encoder, TransformerWrapper # TODO: can we directly rely on lucidrains code and simply add this as a reuirement? --> test
|
8 |
+
from ldm.util import default
|
9 |
+
import clip
|
10 |
+
|
11 |
+
|
12 |
+
class AbstractEncoder(nn.Module):
|
13 |
+
def __init__(self):
|
14 |
+
super().__init__()
|
15 |
+
|
16 |
+
def encode(self, *args, **kwargs):
|
17 |
+
raise NotImplementedError
|
18 |
+
|
19 |
+
class IdentityEncoder(AbstractEncoder):
|
20 |
+
|
21 |
+
def encode(self, x):
|
22 |
+
return x
|
23 |
+
|
24 |
+
class FaceClipEncoder(AbstractEncoder):
|
25 |
+
def __init__(self, augment=True, retreival_key=None):
|
26 |
+
super().__init__()
|
27 |
+
self.encoder = FrozenCLIPImageEmbedder()
|
28 |
+
self.augment = augment
|
29 |
+
self.retreival_key = retreival_key
|
30 |
+
|
31 |
+
def forward(self, img):
|
32 |
+
encodings = []
|
33 |
+
with torch.no_grad():
|
34 |
+
x_offset = 125
|
35 |
+
if self.retreival_key:
|
36 |
+
# Assumes retrieved image are packed into the second half of channels
|
37 |
+
face = img[:,3:,190:440,x_offset:(512-x_offset)]
|
38 |
+
other = img[:,:3,...].clone()
|
39 |
+
else:
|
40 |
+
face = img[:,:,190:440,x_offset:(512-x_offset)]
|
41 |
+
other = img.clone()
|
42 |
+
|
43 |
+
if self.augment:
|
44 |
+
face = K.RandomHorizontalFlip()(face)
|
45 |
+
|
46 |
+
other[:,:,190:440,x_offset:(512-x_offset)] *= 0
|
47 |
+
encodings = [
|
48 |
+
self.encoder.encode(face),
|
49 |
+
self.encoder.encode(other),
|
50 |
+
]
|
51 |
+
|
52 |
+
return torch.cat(encodings, dim=1)
|
53 |
+
|
54 |
+
def encode(self, img):
|
55 |
+
if isinstance(img, list):
|
56 |
+
# Uncondition
|
57 |
+
return torch.zeros((1, 2, 768), device=self.encoder.model.visual.conv1.weight.device)
|
58 |
+
|
59 |
+
return self(img)
|
60 |
+
|
61 |
+
class FaceIdClipEncoder(AbstractEncoder):
|
62 |
+
def __init__(self):
|
63 |
+
super().__init__()
|
64 |
+
self.encoder = FrozenCLIPImageEmbedder()
|
65 |
+
for p in self.encoder.parameters():
|
66 |
+
p.requires_grad = False
|
67 |
+
self.id = FrozenFaceEncoder("/home/jpinkney/code/stable-diffusion/model_ir_se50.pth", augment=True)
|
68 |
+
|
69 |
+
def forward(self, img):
|
70 |
+
encodings = []
|
71 |
+
with torch.no_grad():
|
72 |
+
face = kornia.geometry.resize(img, (256, 256),
|
73 |
+
interpolation='bilinear', align_corners=True)
|
74 |
+
|
75 |
+
other = img.clone()
|
76 |
+
other[:,:,184:452,122:396] *= 0
|
77 |
+
encodings = [
|
78 |
+
self.id.encode(face),
|
79 |
+
self.encoder.encode(other),
|
80 |
+
]
|
81 |
+
|
82 |
+
return torch.cat(encodings, dim=1)
|
83 |
+
|
84 |
+
def encode(self, img):
|
85 |
+
if isinstance(img, list):
|
86 |
+
# Uncondition
|
87 |
+
return torch.zeros((1, 2, 768), device=self.encoder.model.visual.conv1.weight.device)
|
88 |
+
|
89 |
+
return self(img)
|
90 |
+
|
91 |
+
class ClassEmbedder(nn.Module):
|
92 |
+
def __init__(self, embed_dim, n_classes=1000, key='class'):
|
93 |
+
super().__init__()
|
94 |
+
self.key = key
|
95 |
+
self.embedding = nn.Embedding(n_classes, embed_dim)
|
96 |
+
|
97 |
+
def forward(self, batch, key=None):
|
98 |
+
if key is None:
|
99 |
+
key = self.key
|
100 |
+
# this is for use in crossattn
|
101 |
+
c = batch[key][:, None]
|
102 |
+
c = self.embedding(c)
|
103 |
+
return c
|
104 |
+
|
105 |
+
|
106 |
+
class TransformerEmbedder(AbstractEncoder):
|
107 |
+
"""Some transformer encoder layers"""
|
108 |
+
def __init__(self, n_embed, n_layer, vocab_size, max_seq_len=77, device="cuda"):
|
109 |
+
super().__init__()
|
110 |
+
self.device = device
|
111 |
+
self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
|
112 |
+
attn_layers=Encoder(dim=n_embed, depth=n_layer))
|
113 |
+
|
114 |
+
def forward(self, tokens):
|
115 |
+
tokens = tokens.to(self.device) # meh
|
116 |
+
z = self.transformer(tokens, return_embeddings=True)
|
117 |
+
return z
|
118 |
+
|
119 |
+
def encode(self, x):
|
120 |
+
return self(x)
|
121 |
+
|
122 |
+
|
123 |
+
class BERTTokenizer(AbstractEncoder):
|
124 |
+
""" Uses a pretrained BERT tokenizer by huggingface. Vocab size: 30522 (?)"""
|
125 |
+
def __init__(self, device="cuda", vq_interface=True, max_length=77):
|
126 |
+
super().__init__()
|
127 |
+
from transformers import BertTokenizerFast # TODO: add to reuquirements
|
128 |
+
self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
|
129 |
+
self.device = device
|
130 |
+
self.vq_interface = vq_interface
|
131 |
+
self.max_length = max_length
|
132 |
+
|
133 |
+
def forward(self, text):
|
134 |
+
batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
|
135 |
+
return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
136 |
+
tokens = batch_encoding["input_ids"].to(self.device)
|
137 |
+
return tokens
|
138 |
+
|
139 |
+
@torch.no_grad()
|
140 |
+
def encode(self, text):
|
141 |
+
tokens = self(text)
|
142 |
+
if not self.vq_interface:
|
143 |
+
return tokens
|
144 |
+
return None, None, [None, None, tokens]
|
145 |
+
|
146 |
+
def decode(self, text):
|
147 |
+
return text
|
148 |
+
|
149 |
+
|
150 |
+
class BERTEmbedder(AbstractEncoder):
|
151 |
+
"""Uses the BERT tokenizr model and add some transformer encoder layers"""
|
152 |
+
def __init__(self, n_embed, n_layer, vocab_size=30522, max_seq_len=77,
|
153 |
+
device="cuda",use_tokenizer=True, embedding_dropout=0.0):
|
154 |
+
super().__init__()
|
155 |
+
self.use_tknz_fn = use_tokenizer
|
156 |
+
if self.use_tknz_fn:
|
157 |
+
self.tknz_fn = BERTTokenizer(vq_interface=False, max_length=max_seq_len)
|
158 |
+
self.device = device
|
159 |
+
self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
|
160 |
+
attn_layers=Encoder(dim=n_embed, depth=n_layer),
|
161 |
+
emb_dropout=embedding_dropout)
|
162 |
+
|
163 |
+
def forward(self, text):
|
164 |
+
if self.use_tknz_fn:
|
165 |
+
tokens = self.tknz_fn(text)#.to(self.device)
|
166 |
+
else:
|
167 |
+
tokens = text
|
168 |
+
z = self.transformer(tokens, return_embeddings=True)
|
169 |
+
return z
|
170 |
+
|
171 |
+
def encode(self, text):
|
172 |
+
# output of length 77
|
173 |
+
return self(text)
|
174 |
+
|
175 |
+
|
176 |
+
from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel
|
177 |
+
|
178 |
+
def disabled_train(self, mode=True):
|
179 |
+
"""Overwrite model.train with this function to make sure train/eval mode
|
180 |
+
does not change anymore."""
|
181 |
+
return self
|
182 |
+
|
183 |
+
|
184 |
+
class FrozenT5Embedder(AbstractEncoder):
|
185 |
+
"""Uses the T5 transformer encoder for text"""
|
186 |
+
def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77): # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
|
187 |
+
super().__init__()
|
188 |
+
self.tokenizer = T5Tokenizer.from_pretrained(version, cache_dir='/apdcephfs/private_rondyliu/projects/huggingface_models')
|
189 |
+
self.transformer = T5EncoderModel.from_pretrained(version, cache_dir='/apdcephfs/private_rondyliu/projects/huggingface_models')
|
190 |
+
self.device = device
|
191 |
+
self.max_length = max_length # TODO: typical value?
|
192 |
+
self.freeze()
|
193 |
+
|
194 |
+
def freeze(self):
|
195 |
+
self.transformer = self.transformer.eval()
|
196 |
+
#self.train = disabled_train
|
197 |
+
for param in self.parameters():
|
198 |
+
param.requires_grad = False
|
199 |
+
|
200 |
+
def forward(self, text):
|
201 |
+
batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
|
202 |
+
return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
203 |
+
tokens = batch_encoding["input_ids"].to(self.device)
|
204 |
+
outputs = self.transformer(input_ids=tokens)
|
205 |
+
|
206 |
+
z = outputs.last_hidden_state
|
207 |
+
return z
|
208 |
+
|
209 |
+
def encode(self, text):
|
210 |
+
return self(text)
|
211 |
+
|
212 |
+
from ldm.thirdp.psp.id_loss import IDFeatures
|
213 |
+
import kornia.augmentation as K
|
214 |
+
|
215 |
+
class FrozenFaceEncoder(AbstractEncoder):
|
216 |
+
def __init__(self, model_path, augment=False):
|
217 |
+
super().__init__()
|
218 |
+
self.loss_fn = IDFeatures(model_path)
|
219 |
+
# face encoder is frozen
|
220 |
+
for p in self.loss_fn.parameters():
|
221 |
+
p.requires_grad = False
|
222 |
+
# Mapper is trainable
|
223 |
+
self.mapper = torch.nn.Linear(512, 768)
|
224 |
+
p = 0.25
|
225 |
+
if augment:
|
226 |
+
self.augment = K.AugmentationSequential(
|
227 |
+
K.RandomHorizontalFlip(p=0.5),
|
228 |
+
K.RandomEqualize(p=p),
|
229 |
+
# K.RandomPlanckianJitter(p=p),
|
230 |
+
# K.RandomPlasmaBrightness(p=p),
|
231 |
+
# K.RandomPlasmaContrast(p=p),
|
232 |
+
# K.ColorJiggle(0.02, 0.2, 0.2, p=p),
|
233 |
+
)
|
234 |
+
else:
|
235 |
+
self.augment = False
|
236 |
+
|
237 |
+
def forward(self, img):
|
238 |
+
if isinstance(img, list):
|
239 |
+
# Uncondition
|
240 |
+
return torch.zeros((1, 1, 768), device=self.mapper.weight.device)
|
241 |
+
|
242 |
+
if self.augment is not None:
|
243 |
+
# Transforms require 0-1
|
244 |
+
img = self.augment((img + 1)/2)
|
245 |
+
img = 2*img - 1
|
246 |
+
|
247 |
+
feat = self.loss_fn(img, crop=True)
|
248 |
+
feat = self.mapper(feat.unsqueeze(1))
|
249 |
+
return feat
|
250 |
+
|
251 |
+
def encode(self, img):
|
252 |
+
return self(img)
|
253 |
+
|
254 |
+
class FrozenCLIPEmbedder(AbstractEncoder):
|
255 |
+
"""Uses the CLIP transformer encoder for text (from huggingface)"""
|
256 |
+
def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77): # clip-vit-base-patch32
|
257 |
+
super().__init__()
|
258 |
+
self.tokenizer = CLIPTokenizer.from_pretrained(version, cache_dir='/apdcephfs/private_rondyliu/projects/huggingface_models')
|
259 |
+
self.transformer = CLIPTextModel.from_pretrained(version, cache_dir='/apdcephfs/private_rondyliu/projects/huggingface_models')
|
260 |
+
self.device = device
|
261 |
+
self.max_length = max_length # TODO: typical value?
|
262 |
+
self.freeze()
|
263 |
+
|
264 |
+
def freeze(self):
|
265 |
+
self.transformer = self.transformer.eval()
|
266 |
+
#self.train = disabled_train
|
267 |
+
for param in self.parameters():
|
268 |
+
param.requires_grad = False
|
269 |
+
|
270 |
+
def forward(self, text):
|
271 |
+
batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
|
272 |
+
return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
273 |
+
tokens = batch_encoding["input_ids"].to(self.device)
|
274 |
+
outputs = self.transformer(input_ids=tokens)
|
275 |
+
|
276 |
+
z = outputs.last_hidden_state
|
277 |
+
return z
|
278 |
+
|
279 |
+
def encode(self, text):
|
280 |
+
return self(text)
|
281 |
+
|
282 |
+
import torch.nn.functional as F
|
283 |
+
from transformers import CLIPVisionModel
|
284 |
+
class ClipImageProjector(AbstractEncoder):
|
285 |
+
"""
|
286 |
+
Uses the CLIP image encoder.
|
287 |
+
"""
|
288 |
+
def __init__(self, version="openai/clip-vit-large-patch14", max_length=77): # clip-vit-base-patch32
|
289 |
+
super().__init__()
|
290 |
+
self.model = CLIPVisionModel.from_pretrained(version)
|
291 |
+
self.model.train()
|
292 |
+
self.max_length = max_length # TODO: typical value?
|
293 |
+
self.antialias = True
|
294 |
+
self.mapper = torch.nn.Linear(1024, 768)
|
295 |
+
self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
|
296 |
+
self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
|
297 |
+
null_cond = self.get_null_cond(version, max_length)
|
298 |
+
self.register_buffer('null_cond', null_cond)
|
299 |
+
|
300 |
+
@torch.no_grad()
|
301 |
+
def get_null_cond(self, version, max_length):
|
302 |
+
device = self.mean.device
|
303 |
+
embedder = FrozenCLIPEmbedder(version=version, device=device, max_length=max_length)
|
304 |
+
null_cond = embedder([""])
|
305 |
+
return null_cond
|
306 |
+
|
307 |
+
def preprocess(self, x):
|
308 |
+
# Expects inputs in the range -1, 1
|
309 |
+
x = kornia.geometry.resize(x, (224, 224),
|
310 |
+
interpolation='bicubic',align_corners=True,
|
311 |
+
antialias=self.antialias)
|
312 |
+
x = (x + 1.) / 2.
|
313 |
+
# renormalize according to clip
|
314 |
+
x = kornia.enhance.normalize(x, self.mean, self.std)
|
315 |
+
return x
|
316 |
+
|
317 |
+
def forward(self, x):
|
318 |
+
if isinstance(x, list):
|
319 |
+
return self.null_cond
|
320 |
+
# x is assumed to be in range [-1,1]
|
321 |
+
x = self.preprocess(x)
|
322 |
+
outputs = self.model(pixel_values=x)
|
323 |
+
last_hidden_state = outputs.last_hidden_state
|
324 |
+
last_hidden_state = self.mapper(last_hidden_state)
|
325 |
+
return F.pad(last_hidden_state, [0,0, 0,self.max_length-last_hidden_state.shape[1], 0,0])
|
326 |
+
|
327 |
+
def encode(self, im):
|
328 |
+
return self(im)
|
329 |
+
|
330 |
+
class ProjectedFrozenCLIPEmbedder(AbstractEncoder):
|
331 |
+
def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77): # clip-vit-base-patch32
|
332 |
+
super().__init__()
|
333 |
+
self.embedder = FrozenCLIPEmbedder(version=version, device=device, max_length=max_length)
|
334 |
+
self.projection = torch.nn.Linear(768, 768)
|
335 |
+
|
336 |
+
def forward(self, text):
|
337 |
+
z = self.embedder(text)
|
338 |
+
return self.projection(z)
|
339 |
+
|
340 |
+
def encode(self, text):
|
341 |
+
return self(text)
|
342 |
+
|
343 |
+
class FrozenCLIPImageEmbedder(AbstractEncoder):
|
344 |
+
"""
|
345 |
+
Uses the CLIP image encoder.
|
346 |
+
Not actually frozen... If you want that set cond_stage_trainable=False in cfg
|
347 |
+
"""
|
348 |
+
def __init__(
|
349 |
+
self,
|
350 |
+
model='ViT-L/14',
|
351 |
+
jit=False,
|
352 |
+
device='cpu',
|
353 |
+
antialias=False,
|
354 |
+
):
|
355 |
+
super().__init__()
|
356 |
+
self.model, _ = clip.load(name=model, device=device, jit=jit)
|
357 |
+
# We don't use the text part so delete it
|
358 |
+
del self.model.transformer
|
359 |
+
self.antialias = antialias
|
360 |
+
self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
|
361 |
+
self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
|
362 |
+
|
363 |
+
def preprocess(self, x):
|
364 |
+
# Expects inputs in the range -1, 1
|
365 |
+
x = kornia.geometry.resize(x, (224, 224),
|
366 |
+
interpolation='bicubic',align_corners=True,
|
367 |
+
antialias=self.antialias)
|
368 |
+
x = (x + 1.) / 2.
|
369 |
+
# renormalize according to clip
|
370 |
+
x = kornia.enhance.normalize(x, self.mean, self.std)
|
371 |
+
return x
|
372 |
+
|
373 |
+
def forward(self, x):
|
374 |
+
# x is assumed to be in range [-1,1]
|
375 |
+
if isinstance(x, list):
|
376 |
+
# [""] denotes condition dropout for ucg
|
377 |
+
device = self.model.visual.conv1.weight.device
|
378 |
+
return torch.zeros(1, 768, device=device)
|
379 |
+
return self.model.encode_image(self.preprocess(x)).float()
|
380 |
+
|
381 |
+
def encode(self, im):
|
382 |
+
return self(im).unsqueeze(1)
|
383 |
+
|
384 |
+
from torchvision import transforms
|
385 |
+
import random
|
386 |
+
|
387 |
+
class FrozenCLIPImageMutliEmbedder(AbstractEncoder):
|
388 |
+
"""
|
389 |
+
Uses the CLIP image encoder.
|
390 |
+
Not actually frozen... If you want that set cond_stage_trainable=False in cfg
|
391 |
+
"""
|
392 |
+
def __init__(
|
393 |
+
self,
|
394 |
+
model='ViT-L/14',
|
395 |
+
jit=False,
|
396 |
+
device='cpu',
|
397 |
+
antialias=True,
|
398 |
+
max_crops=5,
|
399 |
+
):
|
400 |
+
super().__init__()
|
401 |
+
self.model, _ = clip.load(name=model, device=device, jit=jit)
|
402 |
+
# We don't use the text part so delete it
|
403 |
+
del self.model.transformer
|
404 |
+
self.antialias = antialias
|
405 |
+
self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
|
406 |
+
self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
|
407 |
+
self.max_crops = max_crops
|
408 |
+
|
409 |
+
def preprocess(self, x):
|
410 |
+
|
411 |
+
# Expects inputs in the range -1, 1
|
412 |
+
randcrop = transforms.RandomResizedCrop(224, scale=(0.085, 1.0), ratio=(1,1))
|
413 |
+
max_crops = self.max_crops
|
414 |
+
patches = []
|
415 |
+
crops = [randcrop(x) for _ in range(max_crops)]
|
416 |
+
patches.extend(crops)
|
417 |
+
x = torch.cat(patches, dim=0)
|
418 |
+
x = (x + 1.) / 2.
|
419 |
+
# renormalize according to clip
|
420 |
+
x = kornia.enhance.normalize(x, self.mean, self.std)
|
421 |
+
return x
|
422 |
+
|
423 |
+
def forward(self, x):
|
424 |
+
# x is assumed to be in range [-1,1]
|
425 |
+
if isinstance(x, list):
|
426 |
+
# [""] denotes condition dropout for ucg
|
427 |
+
device = self.model.visual.conv1.weight.device
|
428 |
+
return torch.zeros(1, self.max_crops, 768, device=device)
|
429 |
+
batch_tokens = []
|
430 |
+
for im in x:
|
431 |
+
patches = self.preprocess(im.unsqueeze(0))
|
432 |
+
tokens = self.model.encode_image(patches).float()
|
433 |
+
for t in tokens:
|
434 |
+
if random.random() < 0.1:
|
435 |
+
t *= 0
|
436 |
+
batch_tokens.append(tokens.unsqueeze(0))
|
437 |
+
|
438 |
+
return torch.cat(batch_tokens, dim=0)
|
439 |
+
|
440 |
+
def encode(self, im):
|
441 |
+
return self(im)
|
442 |
+
|
443 |
+
class SpatialRescaler(nn.Module):
|
444 |
+
def __init__(self,
|
445 |
+
n_stages=1,
|
446 |
+
method='bilinear',
|
447 |
+
multiplier=0.5,
|
448 |
+
in_channels=3,
|
449 |
+
out_channels=None,
|
450 |
+
bias=False):
|
451 |
+
super().__init__()
|
452 |
+
self.n_stages = n_stages
|
453 |
+
assert self.n_stages >= 0
|
454 |
+
assert method in ['nearest','linear','bilinear','trilinear','bicubic','area']
|
455 |
+
self.multiplier = multiplier
|
456 |
+
self.interpolator = partial(torch.nn.functional.interpolate, mode=method)
|
457 |
+
self.remap_output = out_channels is not None
|
458 |
+
if self.remap_output:
|
459 |
+
print(f'Spatial Rescaler mapping from {in_channels} to {out_channels} channels after resizing.')
|
460 |
+
self.channel_mapper = nn.Conv2d(in_channels,out_channels,1,bias=bias)
|
461 |
+
|
462 |
+
def forward(self,x):
|
463 |
+
for stage in range(self.n_stages):
|
464 |
+
x = self.interpolator(x, scale_factor=self.multiplier)
|
465 |
+
|
466 |
+
|
467 |
+
if self.remap_output:
|
468 |
+
x = self.channel_mapper(x)
|
469 |
+
return x
|
470 |
+
|
471 |
+
def encode(self, x):
|
472 |
+
return self(x)
|
473 |
+
|
474 |
+
|
475 |
+
from ldm.util import instantiate_from_config
|
476 |
+
from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
|
477 |
+
|
478 |
+
|
479 |
+
class LowScaleEncoder(nn.Module):
|
480 |
+
def __init__(self, model_config, linear_start, linear_end, timesteps=1000, max_noise_level=250, output_size=64,
|
481 |
+
scale_factor=1.0):
|
482 |
+
super().__init__()
|
483 |
+
self.max_noise_level = max_noise_level
|
484 |
+
self.model = instantiate_from_config(model_config)
|
485 |
+
self.augmentation_schedule = self.register_schedule(timesteps=timesteps, linear_start=linear_start,
|
486 |
+
linear_end=linear_end)
|
487 |
+
self.out_size = output_size
|
488 |
+
self.scale_factor = scale_factor
|
489 |
+
|
490 |
+
def register_schedule(self, beta_schedule="linear", timesteps=1000,
|
491 |
+
linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
492 |
+
betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
|
493 |
+
cosine_s=cosine_s)
|
494 |
+
alphas = 1. - betas
|
495 |
+
alphas_cumprod = np.cumprod(alphas, axis=0)
|
496 |
+
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
|
497 |
+
|
498 |
+
timesteps, = betas.shape
|
499 |
+
self.num_timesteps = int(timesteps)
|
500 |
+
self.linear_start = linear_start
|
501 |
+
self.linear_end = linear_end
|
502 |
+
assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
|
503 |
+
|
504 |
+
to_torch = partial(torch.tensor, dtype=torch.float32)
|
505 |
+
|
506 |
+
self.register_buffer('betas', to_torch(betas))
|
507 |
+
self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
|
508 |
+
self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
|
509 |
+
|
510 |
+
# calculations for diffusion q(x_t | x_{t-1}) and others
|
511 |
+
self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
|
512 |
+
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
|
513 |
+
self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
|
514 |
+
self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
|
515 |
+
self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
|
516 |
+
|
517 |
+
def q_sample(self, x_start, t, noise=None):
|
518 |
+
noise = default(noise, lambda: torch.randn_like(x_start))
|
519 |
+
return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
|
520 |
+
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
|
521 |
+
|
522 |
+
def forward(self, x):
|
523 |
+
z = self.model.encode(x).sample()
|
524 |
+
z = z * self.scale_factor
|
525 |
+
noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
|
526 |
+
z = self.q_sample(z, noise_level)
|
527 |
+
if self.out_size is not None:
|
528 |
+
z = torch.nn.functional.interpolate(z, size=self.out_size, mode="nearest") # TODO: experiment with mode
|
529 |
+
# z = z.repeat_interleave(2, -2).repeat_interleave(2, -1)
|
530 |
+
return z, noise_level
|
531 |
+
|
532 |
+
def decode(self, z):
|
533 |
+
z = z / self.scale_factor
|
534 |
+
return self.model.decode(z)
|
535 |
+
|
536 |
+
|
537 |
+
if __name__ == "__main__":
|
538 |
+
from ldm.util import count_params
|
539 |
+
sentences = ["a hedgehog drinking a whiskey", "der mond ist aufgegangen", "Ein Satz mit vielen Sonderzeichen: äöü ß ?! : 'xx-y/@s'"]
|
540 |
+
model = FrozenT5Embedder(version="google/t5-v1_1-xl").cuda()
|
541 |
+
count_params(model, True)
|
542 |
+
z = model(sentences)
|
543 |
+
print(z.shape)
|
544 |
+
|
545 |
+
model = FrozenCLIPEmbedder().cuda()
|
546 |
+
count_params(model, True)
|
547 |
+
z = model(sentences)
|
548 |
+
print(z.shape)
|
549 |
+
|
550 |
+
print("done.")
|
ldm/modules/x_transformer.py
ADDED
@@ -0,0 +1,641 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""shout-out to https://github.com/lucidrains/x-transformers/tree/main/x_transformers"""
|
2 |
+
import torch
|
3 |
+
from torch import nn, einsum
|
4 |
+
import torch.nn.functional as F
|
5 |
+
from functools import partial
|
6 |
+
from inspect import isfunction
|
7 |
+
from collections import namedtuple
|
8 |
+
from einops import rearrange, repeat, reduce
|
9 |
+
|
10 |
+
# constants
|
11 |
+
|
12 |
+
DEFAULT_DIM_HEAD = 64
|
13 |
+
|
14 |
+
Intermediates = namedtuple('Intermediates', [
|
15 |
+
'pre_softmax_attn',
|
16 |
+
'post_softmax_attn'
|
17 |
+
])
|
18 |
+
|
19 |
+
LayerIntermediates = namedtuple('Intermediates', [
|
20 |
+
'hiddens',
|
21 |
+
'attn_intermediates'
|
22 |
+
])
|
23 |
+
|
24 |
+
|
25 |
+
class AbsolutePositionalEmbedding(nn.Module):
|
26 |
+
def __init__(self, dim, max_seq_len):
|
27 |
+
super().__init__()
|
28 |
+
self.emb = nn.Embedding(max_seq_len, dim)
|
29 |
+
self.init_()
|
30 |
+
|
31 |
+
def init_(self):
|
32 |
+
nn.init.normal_(self.emb.weight, std=0.02)
|
33 |
+
|
34 |
+
def forward(self, x):
|
35 |
+
n = torch.arange(x.shape[1], device=x.device)
|
36 |
+
return self.emb(n)[None, :, :]
|
37 |
+
|
38 |
+
|
39 |
+
class FixedPositionalEmbedding(nn.Module):
|
40 |
+
def __init__(self, dim):
|
41 |
+
super().__init__()
|
42 |
+
inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
|
43 |
+
self.register_buffer('inv_freq', inv_freq)
|
44 |
+
|
45 |
+
def forward(self, x, seq_dim=1, offset=0):
|
46 |
+
t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset
|
47 |
+
sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq)
|
48 |
+
emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
|
49 |
+
return emb[None, :, :]
|
50 |
+
|
51 |
+
|
52 |
+
# helpers
|
53 |
+
|
54 |
+
def exists(val):
|
55 |
+
return val is not None
|
56 |
+
|
57 |
+
|
58 |
+
def default(val, d):
|
59 |
+
if exists(val):
|
60 |
+
return val
|
61 |
+
return d() if isfunction(d) else d
|
62 |
+
|
63 |
+
|
64 |
+
def always(val):
|
65 |
+
def inner(*args, **kwargs):
|
66 |
+
return val
|
67 |
+
return inner
|
68 |
+
|
69 |
+
|
70 |
+
def not_equals(val):
|
71 |
+
def inner(x):
|
72 |
+
return x != val
|
73 |
+
return inner
|
74 |
+
|
75 |
+
|
76 |
+
def equals(val):
|
77 |
+
def inner(x):
|
78 |
+
return x == val
|
79 |
+
return inner
|
80 |
+
|
81 |
+
|
82 |
+
def max_neg_value(tensor):
|
83 |
+
return -torch.finfo(tensor.dtype).max
|
84 |
+
|
85 |
+
|
86 |
+
# keyword argument helpers
|
87 |
+
|
88 |
+
def pick_and_pop(keys, d):
|
89 |
+
values = list(map(lambda key: d.pop(key), keys))
|
90 |
+
return dict(zip(keys, values))
|
91 |
+
|
92 |
+
|
93 |
+
def group_dict_by_key(cond, d):
|
94 |
+
return_val = [dict(), dict()]
|
95 |
+
for key in d.keys():
|
96 |
+
match = bool(cond(key))
|
97 |
+
ind = int(not match)
|
98 |
+
return_val[ind][key] = d[key]
|
99 |
+
return (*return_val,)
|
100 |
+
|
101 |
+
|
102 |
+
def string_begins_with(prefix, str):
|
103 |
+
return str.startswith(prefix)
|
104 |
+
|
105 |
+
|
106 |
+
def group_by_key_prefix(prefix, d):
|
107 |
+
return group_dict_by_key(partial(string_begins_with, prefix), d)
|
108 |
+
|
109 |
+
|
110 |
+
def groupby_prefix_and_trim(prefix, d):
|
111 |
+
kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d)
|
112 |
+
kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
|
113 |
+
return kwargs_without_prefix, kwargs
|
114 |
+
|
115 |
+
|
116 |
+
# classes
|
117 |
+
class Scale(nn.Module):
|
118 |
+
def __init__(self, value, fn):
|
119 |
+
super().__init__()
|
120 |
+
self.value = value
|
121 |
+
self.fn = fn
|
122 |
+
|
123 |
+
def forward(self, x, **kwargs):
|
124 |
+
x, *rest = self.fn(x, **kwargs)
|
125 |
+
return (x * self.value, *rest)
|
126 |
+
|
127 |
+
|
128 |
+
class Rezero(nn.Module):
|
129 |
+
def __init__(self, fn):
|
130 |
+
super().__init__()
|
131 |
+
self.fn = fn
|
132 |
+
self.g = nn.Parameter(torch.zeros(1))
|
133 |
+
|
134 |
+
def forward(self, x, **kwargs):
|
135 |
+
x, *rest = self.fn(x, **kwargs)
|
136 |
+
return (x * self.g, *rest)
|
137 |
+
|
138 |
+
|
139 |
+
class ScaleNorm(nn.Module):
|
140 |
+
def __init__(self, dim, eps=1e-5):
|
141 |
+
super().__init__()
|
142 |
+
self.scale = dim ** -0.5
|
143 |
+
self.eps = eps
|
144 |
+
self.g = nn.Parameter(torch.ones(1))
|
145 |
+
|
146 |
+
def forward(self, x):
|
147 |
+
norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
|
148 |
+
return x / norm.clamp(min=self.eps) * self.g
|
149 |
+
|
150 |
+
|
151 |
+
class RMSNorm(nn.Module):
|
152 |
+
def __init__(self, dim, eps=1e-8):
|
153 |
+
super().__init__()
|
154 |
+
self.scale = dim ** -0.5
|
155 |
+
self.eps = eps
|
156 |
+
self.g = nn.Parameter(torch.ones(dim))
|
157 |
+
|
158 |
+
def forward(self, x):
|
159 |
+
norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
|
160 |
+
return x / norm.clamp(min=self.eps) * self.g
|
161 |
+
|
162 |
+
|
163 |
+
class Residual(nn.Module):
|
164 |
+
def forward(self, x, residual):
|
165 |
+
return x + residual
|
166 |
+
|
167 |
+
|
168 |
+
class GRUGating(nn.Module):
|
169 |
+
def __init__(self, dim):
|
170 |
+
super().__init__()
|
171 |
+
self.gru = nn.GRUCell(dim, dim)
|
172 |
+
|
173 |
+
def forward(self, x, residual):
|
174 |
+
gated_output = self.gru(
|
175 |
+
rearrange(x, 'b n d -> (b n) d'),
|
176 |
+
rearrange(residual, 'b n d -> (b n) d')
|
177 |
+
)
|
178 |
+
|
179 |
+
return gated_output.reshape_as(x)
|
180 |
+
|
181 |
+
|
182 |
+
# feedforward
|
183 |
+
|
184 |
+
class GEGLU(nn.Module):
|
185 |
+
def __init__(self, dim_in, dim_out):
|
186 |
+
super().__init__()
|
187 |
+
self.proj = nn.Linear(dim_in, dim_out * 2)
|
188 |
+
|
189 |
+
def forward(self, x):
|
190 |
+
x, gate = self.proj(x).chunk(2, dim=-1)
|
191 |
+
return x * F.gelu(gate)
|
192 |
+
|
193 |
+
|
194 |
+
class FeedForward(nn.Module):
|
195 |
+
def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
|
196 |
+
super().__init__()
|
197 |
+
inner_dim = int(dim * mult)
|
198 |
+
dim_out = default(dim_out, dim)
|
199 |
+
project_in = nn.Sequential(
|
200 |
+
nn.Linear(dim, inner_dim),
|
201 |
+
nn.GELU()
|
202 |
+
) if not glu else GEGLU(dim, inner_dim)
|
203 |
+
|
204 |
+
self.net = nn.Sequential(
|
205 |
+
project_in,
|
206 |
+
nn.Dropout(dropout),
|
207 |
+
nn.Linear(inner_dim, dim_out)
|
208 |
+
)
|
209 |
+
|
210 |
+
def forward(self, x):
|
211 |
+
return self.net(x)
|
212 |
+
|
213 |
+
|
214 |
+
# attention.
|
215 |
+
class Attention(nn.Module):
|
216 |
+
def __init__(
|
217 |
+
self,
|
218 |
+
dim,
|
219 |
+
dim_head=DEFAULT_DIM_HEAD,
|
220 |
+
heads=8,
|
221 |
+
causal=False,
|
222 |
+
mask=None,
|
223 |
+
talking_heads=False,
|
224 |
+
sparse_topk=None,
|
225 |
+
use_entmax15=False,
|
226 |
+
num_mem_kv=0,
|
227 |
+
dropout=0.,
|
228 |
+
on_attn=False
|
229 |
+
):
|
230 |
+
super().__init__()
|
231 |
+
if use_entmax15:
|
232 |
+
raise NotImplementedError("Check out entmax activation instead of softmax activation!")
|
233 |
+
self.scale = dim_head ** -0.5
|
234 |
+
self.heads = heads
|
235 |
+
self.causal = causal
|
236 |
+
self.mask = mask
|
237 |
+
|
238 |
+
inner_dim = dim_head * heads
|
239 |
+
|
240 |
+
self.to_q = nn.Linear(dim, inner_dim, bias=False)
|
241 |
+
self.to_k = nn.Linear(dim, inner_dim, bias=False)
|
242 |
+
self.to_v = nn.Linear(dim, inner_dim, bias=False)
|
243 |
+
self.dropout = nn.Dropout(dropout)
|
244 |
+
|
245 |
+
# talking heads
|
246 |
+
self.talking_heads = talking_heads
|
247 |
+
if talking_heads:
|
248 |
+
self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
|
249 |
+
self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
|
250 |
+
|
251 |
+
# explicit topk sparse attention
|
252 |
+
self.sparse_topk = sparse_topk
|
253 |
+
|
254 |
+
# entmax
|
255 |
+
#self.attn_fn = entmax15 if use_entmax15 else F.softmax
|
256 |
+
self.attn_fn = F.softmax
|
257 |
+
|
258 |
+
# add memory key / values
|
259 |
+
self.num_mem_kv = num_mem_kv
|
260 |
+
if num_mem_kv > 0:
|
261 |
+
self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
|
262 |
+
self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
|
263 |
+
|
264 |
+
# attention on attention
|
265 |
+
self.attn_on_attn = on_attn
|
266 |
+
self.to_out = nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(inner_dim, dim)
|
267 |
+
|
268 |
+
def forward(
|
269 |
+
self,
|
270 |
+
x,
|
271 |
+
context=None,
|
272 |
+
mask=None,
|
273 |
+
context_mask=None,
|
274 |
+
rel_pos=None,
|
275 |
+
sinusoidal_emb=None,
|
276 |
+
prev_attn=None,
|
277 |
+
mem=None
|
278 |
+
):
|
279 |
+
b, n, _, h, talking_heads, device = *x.shape, self.heads, self.talking_heads, x.device
|
280 |
+
kv_input = default(context, x)
|
281 |
+
|
282 |
+
q_input = x
|
283 |
+
k_input = kv_input
|
284 |
+
v_input = kv_input
|
285 |
+
|
286 |
+
if exists(mem):
|
287 |
+
k_input = torch.cat((mem, k_input), dim=-2)
|
288 |
+
v_input = torch.cat((mem, v_input), dim=-2)
|
289 |
+
|
290 |
+
if exists(sinusoidal_emb):
|
291 |
+
# in shortformer, the query would start at a position offset depending on the past cached memory
|
292 |
+
offset = k_input.shape[-2] - q_input.shape[-2]
|
293 |
+
q_input = q_input + sinusoidal_emb(q_input, offset=offset)
|
294 |
+
k_input = k_input + sinusoidal_emb(k_input)
|
295 |
+
|
296 |
+
q = self.to_q(q_input)
|
297 |
+
k = self.to_k(k_input)
|
298 |
+
v = self.to_v(v_input)
|
299 |
+
|
300 |
+
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
|
301 |
+
|
302 |
+
input_mask = None
|
303 |
+
if any(map(exists, (mask, context_mask))):
|
304 |
+
q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
|
305 |
+
k_mask = q_mask if not exists(context) else context_mask
|
306 |
+
k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool())
|
307 |
+
q_mask = rearrange(q_mask, 'b i -> b () i ()')
|
308 |
+
k_mask = rearrange(k_mask, 'b j -> b () () j')
|
309 |
+
input_mask = q_mask * k_mask
|
310 |
+
|
311 |
+
if self.num_mem_kv > 0:
|
312 |
+
mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v))
|
313 |
+
k = torch.cat((mem_k, k), dim=-2)
|
314 |
+
v = torch.cat((mem_v, v), dim=-2)
|
315 |
+
if exists(input_mask):
|
316 |
+
input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
|
317 |
+
|
318 |
+
dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
|
319 |
+
mask_value = max_neg_value(dots)
|
320 |
+
|
321 |
+
if exists(prev_attn):
|
322 |
+
dots = dots + prev_attn
|
323 |
+
|
324 |
+
pre_softmax_attn = dots
|
325 |
+
|
326 |
+
if talking_heads:
|
327 |
+
dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous()
|
328 |
+
|
329 |
+
if exists(rel_pos):
|
330 |
+
dots = rel_pos(dots)
|
331 |
+
|
332 |
+
if exists(input_mask):
|
333 |
+
dots.masked_fill_(~input_mask, mask_value)
|
334 |
+
del input_mask
|
335 |
+
|
336 |
+
if self.causal:
|
337 |
+
i, j = dots.shape[-2:]
|
338 |
+
r = torch.arange(i, device=device)
|
339 |
+
mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j')
|
340 |
+
mask = F.pad(mask, (j - i, 0), value=False)
|
341 |
+
dots.masked_fill_(mask, mask_value)
|
342 |
+
del mask
|
343 |
+
|
344 |
+
if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
|
345 |
+
top, _ = dots.topk(self.sparse_topk, dim=-1)
|
346 |
+
vk = top[..., -1].unsqueeze(-1).expand_as(dots)
|
347 |
+
mask = dots < vk
|
348 |
+
dots.masked_fill_(mask, mask_value)
|
349 |
+
del mask
|
350 |
+
|
351 |
+
attn = self.attn_fn(dots, dim=-1)
|
352 |
+
post_softmax_attn = attn
|
353 |
+
|
354 |
+
attn = self.dropout(attn)
|
355 |
+
|
356 |
+
if talking_heads:
|
357 |
+
attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous()
|
358 |
+
|
359 |
+
out = einsum('b h i j, b h j d -> b h i d', attn, v)
|
360 |
+
out = rearrange(out, 'b h n d -> b n (h d)')
|
361 |
+
|
362 |
+
intermediates = Intermediates(
|
363 |
+
pre_softmax_attn=pre_softmax_attn,
|
364 |
+
post_softmax_attn=post_softmax_attn
|
365 |
+
)
|
366 |
+
|
367 |
+
return self.to_out(out), intermediates
|
368 |
+
|
369 |
+
|
370 |
+
class AttentionLayers(nn.Module):
|
371 |
+
def __init__(
|
372 |
+
self,
|
373 |
+
dim,
|
374 |
+
depth,
|
375 |
+
heads=8,
|
376 |
+
causal=False,
|
377 |
+
cross_attend=False,
|
378 |
+
only_cross=False,
|
379 |
+
use_scalenorm=False,
|
380 |
+
use_rmsnorm=False,
|
381 |
+
use_rezero=False,
|
382 |
+
rel_pos_num_buckets=32,
|
383 |
+
rel_pos_max_distance=128,
|
384 |
+
position_infused_attn=False,
|
385 |
+
custom_layers=None,
|
386 |
+
sandwich_coef=None,
|
387 |
+
par_ratio=None,
|
388 |
+
residual_attn=False,
|
389 |
+
cross_residual_attn=False,
|
390 |
+
macaron=False,
|
391 |
+
pre_norm=True,
|
392 |
+
gate_residual=False,
|
393 |
+
**kwargs
|
394 |
+
):
|
395 |
+
super().__init__()
|
396 |
+
ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs)
|
397 |
+
attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs)
|
398 |
+
|
399 |
+
dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
|
400 |
+
|
401 |
+
self.dim = dim
|
402 |
+
self.depth = depth
|
403 |
+
self.layers = nn.ModuleList([])
|
404 |
+
|
405 |
+
self.has_pos_emb = position_infused_attn
|
406 |
+
self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None
|
407 |
+
self.rotary_pos_emb = always(None)
|
408 |
+
|
409 |
+
assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
|
410 |
+
self.rel_pos = None
|
411 |
+
|
412 |
+
self.pre_norm = pre_norm
|
413 |
+
|
414 |
+
self.residual_attn = residual_attn
|
415 |
+
self.cross_residual_attn = cross_residual_attn
|
416 |
+
|
417 |
+
norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
|
418 |
+
norm_class = RMSNorm if use_rmsnorm else norm_class
|
419 |
+
norm_fn = partial(norm_class, dim)
|
420 |
+
|
421 |
+
norm_fn = nn.Identity if use_rezero else norm_fn
|
422 |
+
branch_fn = Rezero if use_rezero else None
|
423 |
+
|
424 |
+
if cross_attend and not only_cross:
|
425 |
+
default_block = ('a', 'c', 'f')
|
426 |
+
elif cross_attend and only_cross:
|
427 |
+
default_block = ('c', 'f')
|
428 |
+
else:
|
429 |
+
default_block = ('a', 'f')
|
430 |
+
|
431 |
+
if macaron:
|
432 |
+
default_block = ('f',) + default_block
|
433 |
+
|
434 |
+
if exists(custom_layers):
|
435 |
+
layer_types = custom_layers
|
436 |
+
elif exists(par_ratio):
|
437 |
+
par_depth = depth * len(default_block)
|
438 |
+
assert 1 < par_ratio <= par_depth, 'par ratio out of range'
|
439 |
+
default_block = tuple(filter(not_equals('f'), default_block))
|
440 |
+
par_attn = par_depth // par_ratio
|
441 |
+
depth_cut = par_depth * 2 // 3 # 2 / 3 attention layer cutoff suggested by PAR paper
|
442 |
+
par_width = (depth_cut + depth_cut // par_attn) // par_attn
|
443 |
+
assert len(default_block) <= par_width, 'default block is too large for par_ratio'
|
444 |
+
par_block = default_block + ('f',) * (par_width - len(default_block))
|
445 |
+
par_head = par_block * par_attn
|
446 |
+
layer_types = par_head + ('f',) * (par_depth - len(par_head))
|
447 |
+
elif exists(sandwich_coef):
|
448 |
+
assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth'
|
449 |
+
layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef
|
450 |
+
else:
|
451 |
+
layer_types = default_block * depth
|
452 |
+
|
453 |
+
self.layer_types = layer_types
|
454 |
+
self.num_attn_layers = len(list(filter(equals('a'), layer_types)))
|
455 |
+
|
456 |
+
for layer_type in self.layer_types:
|
457 |
+
if layer_type == 'a':
|
458 |
+
layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
|
459 |
+
elif layer_type == 'c':
|
460 |
+
layer = Attention(dim, heads=heads, **attn_kwargs)
|
461 |
+
elif layer_type == 'f':
|
462 |
+
layer = FeedForward(dim, **ff_kwargs)
|
463 |
+
layer = layer if not macaron else Scale(0.5, layer)
|
464 |
+
else:
|
465 |
+
raise Exception(f'invalid layer type {layer_type}')
|
466 |
+
|
467 |
+
if isinstance(layer, Attention) and exists(branch_fn):
|
468 |
+
layer = branch_fn(layer)
|
469 |
+
|
470 |
+
if gate_residual:
|
471 |
+
residual_fn = GRUGating(dim)
|
472 |
+
else:
|
473 |
+
residual_fn = Residual()
|
474 |
+
|
475 |
+
self.layers.append(nn.ModuleList([
|
476 |
+
norm_fn(),
|
477 |
+
layer,
|
478 |
+
residual_fn
|
479 |
+
]))
|
480 |
+
|
481 |
+
def forward(
|
482 |
+
self,
|
483 |
+
x,
|
484 |
+
context=None,
|
485 |
+
mask=None,
|
486 |
+
context_mask=None,
|
487 |
+
mems=None,
|
488 |
+
return_hiddens=False
|
489 |
+
):
|
490 |
+
hiddens = []
|
491 |
+
intermediates = []
|
492 |
+
prev_attn = None
|
493 |
+
prev_cross_attn = None
|
494 |
+
|
495 |
+
mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
|
496 |
+
|
497 |
+
for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)):
|
498 |
+
is_last = ind == (len(self.layers) - 1)
|
499 |
+
|
500 |
+
if layer_type == 'a':
|
501 |
+
hiddens.append(x)
|
502 |
+
layer_mem = mems.pop(0)
|
503 |
+
|
504 |
+
residual = x
|
505 |
+
|
506 |
+
if self.pre_norm:
|
507 |
+
x = norm(x)
|
508 |
+
|
509 |
+
if layer_type == 'a':
|
510 |
+
out, inter = block(x, mask=mask, sinusoidal_emb=self.pia_pos_emb, rel_pos=self.rel_pos,
|
511 |
+
prev_attn=prev_attn, mem=layer_mem)
|
512 |
+
elif layer_type == 'c':
|
513 |
+
out, inter = block(x, context=context, mask=mask, context_mask=context_mask, prev_attn=prev_cross_attn)
|
514 |
+
elif layer_type == 'f':
|
515 |
+
out = block(x)
|
516 |
+
|
517 |
+
x = residual_fn(out, residual)
|
518 |
+
|
519 |
+
if layer_type in ('a', 'c'):
|
520 |
+
intermediates.append(inter)
|
521 |
+
|
522 |
+
if layer_type == 'a' and self.residual_attn:
|
523 |
+
prev_attn = inter.pre_softmax_attn
|
524 |
+
elif layer_type == 'c' and self.cross_residual_attn:
|
525 |
+
prev_cross_attn = inter.pre_softmax_attn
|
526 |
+
|
527 |
+
if not self.pre_norm and not is_last:
|
528 |
+
x = norm(x)
|
529 |
+
|
530 |
+
if return_hiddens:
|
531 |
+
intermediates = LayerIntermediates(
|
532 |
+
hiddens=hiddens,
|
533 |
+
attn_intermediates=intermediates
|
534 |
+
)
|
535 |
+
|
536 |
+
return x, intermediates
|
537 |
+
|
538 |
+
return x
|
539 |
+
|
540 |
+
|
541 |
+
class Encoder(AttentionLayers):
|
542 |
+
def __init__(self, **kwargs):
|
543 |
+
assert 'causal' not in kwargs, 'cannot set causality on encoder'
|
544 |
+
super().__init__(causal=False, **kwargs)
|
545 |
+
|
546 |
+
|
547 |
+
|
548 |
+
class TransformerWrapper(nn.Module):
|
549 |
+
def __init__(
|
550 |
+
self,
|
551 |
+
*,
|
552 |
+
num_tokens,
|
553 |
+
max_seq_len,
|
554 |
+
attn_layers,
|
555 |
+
emb_dim=None,
|
556 |
+
max_mem_len=0.,
|
557 |
+
emb_dropout=0.,
|
558 |
+
num_memory_tokens=None,
|
559 |
+
tie_embedding=False,
|
560 |
+
use_pos_emb=True
|
561 |
+
):
|
562 |
+
super().__init__()
|
563 |
+
assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
|
564 |
+
|
565 |
+
dim = attn_layers.dim
|
566 |
+
emb_dim = default(emb_dim, dim)
|
567 |
+
|
568 |
+
self.max_seq_len = max_seq_len
|
569 |
+
self.max_mem_len = max_mem_len
|
570 |
+
self.num_tokens = num_tokens
|
571 |
+
|
572 |
+
self.token_emb = nn.Embedding(num_tokens, emb_dim)
|
573 |
+
self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if (
|
574 |
+
use_pos_emb and not attn_layers.has_pos_emb) else always(0)
|
575 |
+
self.emb_dropout = nn.Dropout(emb_dropout)
|
576 |
+
|
577 |
+
self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
|
578 |
+
self.attn_layers = attn_layers
|
579 |
+
self.norm = nn.LayerNorm(dim)
|
580 |
+
|
581 |
+
self.init_()
|
582 |
+
|
583 |
+
self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t()
|
584 |
+
|
585 |
+
# memory tokens (like [cls]) from Memory Transformers paper
|
586 |
+
num_memory_tokens = default(num_memory_tokens, 0)
|
587 |
+
self.num_memory_tokens = num_memory_tokens
|
588 |
+
if num_memory_tokens > 0:
|
589 |
+
self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
|
590 |
+
|
591 |
+
# let funnel encoder know number of memory tokens, if specified
|
592 |
+
if hasattr(attn_layers, 'num_memory_tokens'):
|
593 |
+
attn_layers.num_memory_tokens = num_memory_tokens
|
594 |
+
|
595 |
+
def init_(self):
|
596 |
+
nn.init.normal_(self.token_emb.weight, std=0.02)
|
597 |
+
|
598 |
+
def forward(
|
599 |
+
self,
|
600 |
+
x,
|
601 |
+
return_embeddings=False,
|
602 |
+
mask=None,
|
603 |
+
return_mems=False,
|
604 |
+
return_attn=False,
|
605 |
+
mems=None,
|
606 |
+
**kwargs
|
607 |
+
):
|
608 |
+
b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens
|
609 |
+
x = self.token_emb(x)
|
610 |
+
x += self.pos_emb(x)
|
611 |
+
x = self.emb_dropout(x)
|
612 |
+
|
613 |
+
x = self.project_emb(x)
|
614 |
+
|
615 |
+
if num_mem > 0:
|
616 |
+
mem = repeat(self.memory_tokens, 'n d -> b n d', b=b)
|
617 |
+
x = torch.cat((mem, x), dim=1)
|
618 |
+
|
619 |
+
# auto-handle masking after appending memory tokens
|
620 |
+
if exists(mask):
|
621 |
+
mask = F.pad(mask, (num_mem, 0), value=True)
|
622 |
+
|
623 |
+
x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs)
|
624 |
+
x = self.norm(x)
|
625 |
+
|
626 |
+
mem, x = x[:, :num_mem], x[:, num_mem:]
|
627 |
+
|
628 |
+
out = self.to_logits(x) if not return_embeddings else x
|
629 |
+
|
630 |
+
if return_mems:
|
631 |
+
hiddens = intermediates.hiddens
|
632 |
+
new_mems = list(map(lambda pair: torch.cat(pair, dim=-2), zip(mems, hiddens))) if exists(mems) else hiddens
|
633 |
+
new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), new_mems))
|
634 |
+
return out, new_mems
|
635 |
+
|
636 |
+
if return_attn:
|
637 |
+
attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
|
638 |
+
return out, attn_maps
|
639 |
+
|
640 |
+
return out
|
641 |
+
|
ldm/thirdp/psp/helpers.py
ADDED
@@ -0,0 +1,121 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# https://github.com/eladrich/pixel2style2pixel
|
2 |
+
|
3 |
+
from collections import namedtuple
|
4 |
+
import torch
|
5 |
+
from torch.nn import Conv2d, BatchNorm2d, PReLU, ReLU, Sigmoid, MaxPool2d, AdaptiveAvgPool2d, Sequential, Module
|
6 |
+
|
7 |
+
"""
|
8 |
+
ArcFace implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch)
|
9 |
+
"""
|
10 |
+
|
11 |
+
|
12 |
+
class Flatten(Module):
|
13 |
+
def forward(self, input):
|
14 |
+
return input.view(input.size(0), -1)
|
15 |
+
|
16 |
+
|
17 |
+
def l2_norm(input, axis=1):
|
18 |
+
norm = torch.norm(input, 2, axis, True)
|
19 |
+
output = torch.div(input, norm)
|
20 |
+
return output
|
21 |
+
|
22 |
+
|
23 |
+
class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
|
24 |
+
""" A named tuple describing a ResNet block. """
|
25 |
+
|
26 |
+
|
27 |
+
def get_block(in_channel, depth, num_units, stride=2):
|
28 |
+
return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
|
29 |
+
|
30 |
+
|
31 |
+
def get_blocks(num_layers):
|
32 |
+
if num_layers == 50:
|
33 |
+
blocks = [
|
34 |
+
get_block(in_channel=64, depth=64, num_units=3),
|
35 |
+
get_block(in_channel=64, depth=128, num_units=4),
|
36 |
+
get_block(in_channel=128, depth=256, num_units=14),
|
37 |
+
get_block(in_channel=256, depth=512, num_units=3)
|
38 |
+
]
|
39 |
+
elif num_layers == 100:
|
40 |
+
blocks = [
|
41 |
+
get_block(in_channel=64, depth=64, num_units=3),
|
42 |
+
get_block(in_channel=64, depth=128, num_units=13),
|
43 |
+
get_block(in_channel=128, depth=256, num_units=30),
|
44 |
+
get_block(in_channel=256, depth=512, num_units=3)
|
45 |
+
]
|
46 |
+
elif num_layers == 152:
|
47 |
+
blocks = [
|
48 |
+
get_block(in_channel=64, depth=64, num_units=3),
|
49 |
+
get_block(in_channel=64, depth=128, num_units=8),
|
50 |
+
get_block(in_channel=128, depth=256, num_units=36),
|
51 |
+
get_block(in_channel=256, depth=512, num_units=3)
|
52 |
+
]
|
53 |
+
else:
|
54 |
+
raise ValueError("Invalid number of layers: {}. Must be one of [50, 100, 152]".format(num_layers))
|
55 |
+
return blocks
|
56 |
+
|
57 |
+
|
58 |
+
class SEModule(Module):
|
59 |
+
def __init__(self, channels, reduction):
|
60 |
+
super(SEModule, self).__init__()
|
61 |
+
self.avg_pool = AdaptiveAvgPool2d(1)
|
62 |
+
self.fc1 = Conv2d(channels, channels // reduction, kernel_size=1, padding=0, bias=False)
|
63 |
+
self.relu = ReLU(inplace=True)
|
64 |
+
self.fc2 = Conv2d(channels // reduction, channels, kernel_size=1, padding=0, bias=False)
|
65 |
+
self.sigmoid = Sigmoid()
|
66 |
+
|
67 |
+
def forward(self, x):
|
68 |
+
module_input = x
|
69 |
+
x = self.avg_pool(x)
|
70 |
+
x = self.fc1(x)
|
71 |
+
x = self.relu(x)
|
72 |
+
x = self.fc2(x)
|
73 |
+
x = self.sigmoid(x)
|
74 |
+
return module_input * x
|
75 |
+
|
76 |
+
|
77 |
+
class bottleneck_IR(Module):
|
78 |
+
def __init__(self, in_channel, depth, stride):
|
79 |
+
super(bottleneck_IR, self).__init__()
|
80 |
+
if in_channel == depth:
|
81 |
+
self.shortcut_layer = MaxPool2d(1, stride)
|
82 |
+
else:
|
83 |
+
self.shortcut_layer = Sequential(
|
84 |
+
Conv2d(in_channel, depth, (1, 1), stride, bias=False),
|
85 |
+
BatchNorm2d(depth)
|
86 |
+
)
|
87 |
+
self.res_layer = Sequential(
|
88 |
+
BatchNorm2d(in_channel),
|
89 |
+
Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), PReLU(depth),
|
90 |
+
Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth)
|
91 |
+
)
|
92 |
+
|
93 |
+
def forward(self, x):
|
94 |
+
shortcut = self.shortcut_layer(x)
|
95 |
+
res = self.res_layer(x)
|
96 |
+
return res + shortcut
|
97 |
+
|
98 |
+
|
99 |
+
class bottleneck_IR_SE(Module):
|
100 |
+
def __init__(self, in_channel, depth, stride):
|
101 |
+
super(bottleneck_IR_SE, self).__init__()
|
102 |
+
if in_channel == depth:
|
103 |
+
self.shortcut_layer = MaxPool2d(1, stride)
|
104 |
+
else:
|
105 |
+
self.shortcut_layer = Sequential(
|
106 |
+
Conv2d(in_channel, depth, (1, 1), stride, bias=False),
|
107 |
+
BatchNorm2d(depth)
|
108 |
+
)
|
109 |
+
self.res_layer = Sequential(
|
110 |
+
BatchNorm2d(in_channel),
|
111 |
+
Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
|
112 |
+
PReLU(depth),
|
113 |
+
Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
|
114 |
+
BatchNorm2d(depth),
|
115 |
+
SEModule(depth, 16)
|
116 |
+
)
|
117 |
+
|
118 |
+
def forward(self, x):
|
119 |
+
shortcut = self.shortcut_layer(x)
|
120 |
+
res = self.res_layer(x)
|
121 |
+
return res + shortcut
|
ldm/thirdp/psp/id_loss.py
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# https://github.com/eladrich/pixel2style2pixel
|
2 |
+
import torch
|
3 |
+
from torch import nn
|
4 |
+
from ldm.thirdp.psp.model_irse import Backbone
|
5 |
+
|
6 |
+
|
7 |
+
class IDFeatures(nn.Module):
|
8 |
+
def __init__(self, model_path):
|
9 |
+
super(IDFeatures, self).__init__()
|
10 |
+
print('Loading ResNet ArcFace')
|
11 |
+
self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
|
12 |
+
self.facenet.load_state_dict(torch.load(model_path, map_location="cpu"))
|
13 |
+
self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
|
14 |
+
self.facenet.eval()
|
15 |
+
|
16 |
+
def forward(self, x, crop=False):
|
17 |
+
# Not sure of the image range here
|
18 |
+
if crop:
|
19 |
+
x = torch.nn.functional.interpolate(x, (256, 256), mode="area")
|
20 |
+
x = x[:, :, 35:223, 32:220]
|
21 |
+
x = self.face_pool(x)
|
22 |
+
x_feats = self.facenet(x)
|
23 |
+
return x_feats
|
ldm/thirdp/psp/model_irse.py
ADDED
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# https://github.com/eladrich/pixel2style2pixel
|
2 |
+
|
3 |
+
from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Dropout, Sequential, Module
|
4 |
+
from ldm.thirdp.psp.helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE, l2_norm
|
5 |
+
|
6 |
+
"""
|
7 |
+
Modified Backbone implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch)
|
8 |
+
"""
|
9 |
+
|
10 |
+
|
11 |
+
class Backbone(Module):
|
12 |
+
def __init__(self, input_size, num_layers, mode='ir', drop_ratio=0.4, affine=True):
|
13 |
+
super(Backbone, self).__init__()
|
14 |
+
assert input_size in [112, 224], "input_size should be 112 or 224"
|
15 |
+
assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152"
|
16 |
+
assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se"
|
17 |
+
blocks = get_blocks(num_layers)
|
18 |
+
if mode == 'ir':
|
19 |
+
unit_module = bottleneck_IR
|
20 |
+
elif mode == 'ir_se':
|
21 |
+
unit_module = bottleneck_IR_SE
|
22 |
+
self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
|
23 |
+
BatchNorm2d(64),
|
24 |
+
PReLU(64))
|
25 |
+
if input_size == 112:
|
26 |
+
self.output_layer = Sequential(BatchNorm2d(512),
|
27 |
+
Dropout(drop_ratio),
|
28 |
+
Flatten(),
|
29 |
+
Linear(512 * 7 * 7, 512),
|
30 |
+
BatchNorm1d(512, affine=affine))
|
31 |
+
else:
|
32 |
+
self.output_layer = Sequential(BatchNorm2d(512),
|
33 |
+
Dropout(drop_ratio),
|
34 |
+
Flatten(),
|
35 |
+
Linear(512 * 14 * 14, 512),
|
36 |
+
BatchNorm1d(512, affine=affine))
|
37 |
+
|
38 |
+
modules = []
|
39 |
+
for block in blocks:
|
40 |
+
for bottleneck in block:
|
41 |
+
modules.append(unit_module(bottleneck.in_channel,
|
42 |
+
bottleneck.depth,
|
43 |
+
bottleneck.stride))
|
44 |
+
self.body = Sequential(*modules)
|
45 |
+
|
46 |
+
def forward(self, x):
|
47 |
+
x = self.input_layer(x)
|
48 |
+
x = self.body(x)
|
49 |
+
x = self.output_layer(x)
|
50 |
+
return l2_norm(x)
|
51 |
+
|
52 |
+
|
53 |
+
def IR_50(input_size):
|
54 |
+
"""Constructs a ir-50 model."""
|
55 |
+
model = Backbone(input_size, num_layers=50, mode='ir', drop_ratio=0.4, affine=False)
|
56 |
+
return model
|
57 |
+
|
58 |
+
|
59 |
+
def IR_101(input_size):
|
60 |
+
"""Constructs a ir-101 model."""
|
61 |
+
model = Backbone(input_size, num_layers=100, mode='ir', drop_ratio=0.4, affine=False)
|
62 |
+
return model
|
63 |
+
|
64 |
+
|
65 |
+
def IR_152(input_size):
|
66 |
+
"""Constructs a ir-152 model."""
|
67 |
+
model = Backbone(input_size, num_layers=152, mode='ir', drop_ratio=0.4, affine=False)
|
68 |
+
return model
|
69 |
+
|
70 |
+
|
71 |
+
def IR_SE_50(input_size):
|
72 |
+
"""Constructs a ir_se-50 model."""
|
73 |
+
model = Backbone(input_size, num_layers=50, mode='ir_se', drop_ratio=0.4, affine=False)
|
74 |
+
return model
|
75 |
+
|
76 |
+
|
77 |
+
def IR_SE_101(input_size):
|
78 |
+
"""Constructs a ir_se-101 model."""
|
79 |
+
model = Backbone(input_size, num_layers=100, mode='ir_se', drop_ratio=0.4, affine=False)
|
80 |
+
return model
|
81 |
+
|
82 |
+
|
83 |
+
def IR_SE_152(input_size):
|
84 |
+
"""Constructs a ir_se-152 model."""
|
85 |
+
model = Backbone(input_size, num_layers=152, mode='ir_se', drop_ratio=0.4, affine=False)
|
86 |
+
return model
|
ldm/util.py
ADDED
@@ -0,0 +1,302 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import importlib
|
2 |
+
|
3 |
+
import torchvision
|
4 |
+
import torch
|
5 |
+
from torch import optim
|
6 |
+
import numpy as np
|
7 |
+
|
8 |
+
from inspect import isfunction
|
9 |
+
from PIL import Image, ImageDraw, ImageFont
|
10 |
+
|
11 |
+
import os
|
12 |
+
import numpy as np
|
13 |
+
import matplotlib.pyplot as plt
|
14 |
+
from PIL import Image
|
15 |
+
import torch
|
16 |
+
import time
|
17 |
+
import cv2
|
18 |
+
import PIL
|
19 |
+
|
20 |
+
def pil_rectangle_crop(im):
|
21 |
+
width, height = im.size # Get dimensions
|
22 |
+
|
23 |
+
if width <= height:
|
24 |
+
left = 0
|
25 |
+
right = width
|
26 |
+
top = (height - width)/2
|
27 |
+
bottom = (height + width)/2
|
28 |
+
else:
|
29 |
+
|
30 |
+
top = 0
|
31 |
+
bottom = height
|
32 |
+
left = (width - height) / 2
|
33 |
+
bottom = (width + height) / 2
|
34 |
+
|
35 |
+
# Crop the center of the image
|
36 |
+
im = im.crop((left, top, right, bottom))
|
37 |
+
return im
|
38 |
+
|
39 |
+
def add_margin(pil_img, color=0, size=256):
|
40 |
+
width, height = pil_img.size
|
41 |
+
result = Image.new(pil_img.mode, (size, size), color)
|
42 |
+
result.paste(pil_img, ((size - width) // 2, (size - height) // 2))
|
43 |
+
return result
|
44 |
+
|
45 |
+
|
46 |
+
def create_carvekit_interface():
|
47 |
+
from carvekit.api.high import HiInterface
|
48 |
+
# Check doc strings for more information
|
49 |
+
interface = HiInterface(object_type="object", # Can be "object" or "hairs-like".
|
50 |
+
batch_size_seg=5,
|
51 |
+
batch_size_matting=1,
|
52 |
+
device='cuda' if torch.cuda.is_available() else 'cpu',
|
53 |
+
seg_mask_size=640, # Use 640 for Tracer B7 and 320 for U2Net
|
54 |
+
matting_mask_size=2048,
|
55 |
+
trimap_prob_threshold=231,
|
56 |
+
trimap_dilation=30,
|
57 |
+
trimap_erosion_iters=5,
|
58 |
+
fp16=False)
|
59 |
+
|
60 |
+
return interface
|
61 |
+
|
62 |
+
|
63 |
+
def load_and_preprocess(interface, input_im):
|
64 |
+
'''
|
65 |
+
:param input_im (PIL Image).
|
66 |
+
:return image (H, W, 3) array in [0, 1].
|
67 |
+
'''
|
68 |
+
# See https://github.com/Ir1d/image-background-remove-tool
|
69 |
+
image = input_im.convert('RGB')
|
70 |
+
|
71 |
+
image_without_background = interface([image])[0]
|
72 |
+
image_without_background = np.array(image_without_background)
|
73 |
+
est_seg = image_without_background > 127
|
74 |
+
image = np.array(image)
|
75 |
+
foreground = est_seg[:, : , -1].astype(np.bool_)
|
76 |
+
image[~foreground] = [255., 255., 255.]
|
77 |
+
x, y, w, h = cv2.boundingRect(foreground.astype(np.uint8))
|
78 |
+
image = image[y:y+h, x:x+w, :]
|
79 |
+
image = PIL.Image.fromarray(np.array(image))
|
80 |
+
|
81 |
+
# resize image such that long edge is 512
|
82 |
+
image.thumbnail([200, 200], Image.LANCZOS)
|
83 |
+
image = add_margin(image, (255, 255, 255), size=256)
|
84 |
+
image = np.array(image)
|
85 |
+
|
86 |
+
return image
|
87 |
+
|
88 |
+
|
89 |
+
def log_txt_as_img(wh, xc, size=10):
|
90 |
+
# wh a tuple of (width, height)
|
91 |
+
# xc a list of captions to plot
|
92 |
+
b = len(xc)
|
93 |
+
txts = list()
|
94 |
+
for bi in range(b):
|
95 |
+
txt = Image.new("RGB", wh, color="white")
|
96 |
+
draw = ImageDraw.Draw(txt)
|
97 |
+
font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
|
98 |
+
nc = int(40 * (wh[0] / 256))
|
99 |
+
lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
|
100 |
+
|
101 |
+
try:
|
102 |
+
draw.text((0, 0), lines, fill="black", font=font)
|
103 |
+
except UnicodeEncodeError:
|
104 |
+
print("Cant encode string for logging. Skipping.")
|
105 |
+
|
106 |
+
txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
|
107 |
+
txts.append(txt)
|
108 |
+
txts = np.stack(txts)
|
109 |
+
txts = torch.tensor(txts)
|
110 |
+
return txts
|
111 |
+
|
112 |
+
|
113 |
+
def ismap(x):
|
114 |
+
if not isinstance(x, torch.Tensor):
|
115 |
+
return False
|
116 |
+
return (len(x.shape) == 4) and (x.shape[1] > 3)
|
117 |
+
|
118 |
+
|
119 |
+
def isimage(x):
|
120 |
+
if not isinstance(x,torch.Tensor):
|
121 |
+
return False
|
122 |
+
return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
|
123 |
+
|
124 |
+
|
125 |
+
def exists(x):
|
126 |
+
return x is not None
|
127 |
+
|
128 |
+
|
129 |
+
def default(val, d):
|
130 |
+
if exists(val):
|
131 |
+
return val
|
132 |
+
return d() if isfunction(d) else d
|
133 |
+
|
134 |
+
|
135 |
+
def mean_flat(tensor):
|
136 |
+
"""
|
137 |
+
https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
|
138 |
+
Take the mean over all non-batch dimensions.
|
139 |
+
"""
|
140 |
+
return tensor.mean(dim=list(range(1, len(tensor.shape))))
|
141 |
+
|
142 |
+
|
143 |
+
def count_params(model, verbose=False):
|
144 |
+
total_params = sum(p.numel() for p in model.parameters())
|
145 |
+
if verbose:
|
146 |
+
print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
|
147 |
+
return total_params
|
148 |
+
|
149 |
+
|
150 |
+
def instantiate_from_config(config):
|
151 |
+
if not "target" in config:
|
152 |
+
if config == '__is_first_stage__':
|
153 |
+
return None
|
154 |
+
elif config == "__is_unconditional__":
|
155 |
+
return None
|
156 |
+
raise KeyError("Expected key `target` to instantiate.")
|
157 |
+
return get_obj_from_str(config["target"])(**config.get("params", dict()))
|
158 |
+
|
159 |
+
|
160 |
+
def get_obj_from_str(string, reload=False):
|
161 |
+
module, cls = string.rsplit(".", 1)
|
162 |
+
if reload:
|
163 |
+
module_imp = importlib.import_module(module)
|
164 |
+
importlib.reload(module_imp)
|
165 |
+
return getattr(importlib.import_module(module, package=None), cls)
|
166 |
+
|
167 |
+
|
168 |
+
class AdamWwithEMAandWings(optim.Optimizer):
|
169 |
+
# credit to https://gist.github.com/crowsonkb/65f7265353f403714fce3b2595e0b298
|
170 |
+
def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8, # TODO: check hyperparameters before using
|
171 |
+
weight_decay=1.e-2, amsgrad=False, ema_decay=0.9999, # ema decay to match previous code
|
172 |
+
ema_power=1., param_names=()):
|
173 |
+
"""AdamW that saves EMA versions of the parameters."""
|
174 |
+
if not 0.0 <= lr:
|
175 |
+
raise ValueError("Invalid learning rate: {}".format(lr))
|
176 |
+
if not 0.0 <= eps:
|
177 |
+
raise ValueError("Invalid epsilon value: {}".format(eps))
|
178 |
+
if not 0.0 <= betas[0] < 1.0:
|
179 |
+
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
|
180 |
+
if not 0.0 <= betas[1] < 1.0:
|
181 |
+
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
|
182 |
+
if not 0.0 <= weight_decay:
|
183 |
+
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
|
184 |
+
if not 0.0 <= ema_decay <= 1.0:
|
185 |
+
raise ValueError("Invalid ema_decay value: {}".format(ema_decay))
|
186 |
+
defaults = dict(lr=lr, betas=betas, eps=eps,
|
187 |
+
weight_decay=weight_decay, amsgrad=amsgrad, ema_decay=ema_decay,
|
188 |
+
ema_power=ema_power, param_names=param_names)
|
189 |
+
super().__init__(params, defaults)
|
190 |
+
|
191 |
+
def __setstate__(self, state):
|
192 |
+
super().__setstate__(state)
|
193 |
+
for group in self.param_groups:
|
194 |
+
group.setdefault('amsgrad', False)
|
195 |
+
|
196 |
+
@torch.no_grad()
|
197 |
+
def step(self, closure=None):
|
198 |
+
"""Performs a single optimization step.
|
199 |
+
Args:
|
200 |
+
closure (callable, optional): A closure that reevaluates the model
|
201 |
+
and returns the loss.
|
202 |
+
"""
|
203 |
+
loss = None
|
204 |
+
if closure is not None:
|
205 |
+
with torch.enable_grad():
|
206 |
+
loss = closure()
|
207 |
+
|
208 |
+
for group in self.param_groups:
|
209 |
+
params_with_grad = []
|
210 |
+
grads = []
|
211 |
+
exp_avgs = []
|
212 |
+
exp_avg_sqs = []
|
213 |
+
ema_params_with_grad = []
|
214 |
+
state_sums = []
|
215 |
+
max_exp_avg_sqs = []
|
216 |
+
state_steps = []
|
217 |
+
amsgrad = group['amsgrad']
|
218 |
+
beta1, beta2 = group['betas']
|
219 |
+
ema_decay = group['ema_decay']
|
220 |
+
ema_power = group['ema_power']
|
221 |
+
|
222 |
+
for p in group['params']:
|
223 |
+
if p.grad is None:
|
224 |
+
continue
|
225 |
+
params_with_grad.append(p)
|
226 |
+
if p.grad.is_sparse:
|
227 |
+
raise RuntimeError('AdamW does not support sparse gradients')
|
228 |
+
grads.append(p.grad)
|
229 |
+
|
230 |
+
state = self.state[p]
|
231 |
+
|
232 |
+
# State initialization
|
233 |
+
if len(state) == 0:
|
234 |
+
state['step'] = 0
|
235 |
+
# Exponential moving average of gradient values
|
236 |
+
state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
|
237 |
+
# Exponential moving average of squared gradient values
|
238 |
+
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
|
239 |
+
if amsgrad:
|
240 |
+
# Maintains max of all exp. moving avg. of sq. grad. values
|
241 |
+
state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
|
242 |
+
# Exponential moving average of parameter values
|
243 |
+
state['param_exp_avg'] = p.detach().float().clone()
|
244 |
+
|
245 |
+
exp_avgs.append(state['exp_avg'])
|
246 |
+
exp_avg_sqs.append(state['exp_avg_sq'])
|
247 |
+
ema_params_with_grad.append(state['param_exp_avg'])
|
248 |
+
|
249 |
+
if amsgrad:
|
250 |
+
max_exp_avg_sqs.append(state['max_exp_avg_sq'])
|
251 |
+
|
252 |
+
# update the steps for each param group update
|
253 |
+
state['step'] += 1
|
254 |
+
# record the step after step update
|
255 |
+
state_steps.append(state['step'])
|
256 |
+
|
257 |
+
optim._functional.adamw(params_with_grad,
|
258 |
+
grads,
|
259 |
+
exp_avgs,
|
260 |
+
exp_avg_sqs,
|
261 |
+
max_exp_avg_sqs,
|
262 |
+
state_steps,
|
263 |
+
amsgrad=amsgrad,
|
264 |
+
beta1=beta1,
|
265 |
+
beta2=beta2,
|
266 |
+
lr=group['lr'],
|
267 |
+
weight_decay=group['weight_decay'],
|
268 |
+
eps=group['eps'],
|
269 |
+
maximize=False)
|
270 |
+
|
271 |
+
cur_ema_decay = min(ema_decay, 1 - state['step'] ** -ema_power)
|
272 |
+
for param, ema_param in zip(params_with_grad, ema_params_with_grad):
|
273 |
+
ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)
|
274 |
+
|
275 |
+
return loss
|
276 |
+
|
277 |
+
def prepare_inputs(image_path, elevation_input, crop_size=-1, image_size=256):
|
278 |
+
image_input = Image.open(image_path)
|
279 |
+
|
280 |
+
if crop_size!=-1:
|
281 |
+
alpha_np = np.asarray(image_input)[:, :, 3]
|
282 |
+
coords = np.stack(np.nonzero(alpha_np), 1)[:, (1, 0)]
|
283 |
+
min_x, min_y = np.min(coords, 0)
|
284 |
+
max_x, max_y = np.max(coords, 0)
|
285 |
+
ref_img_ = image_input.crop((min_x, min_y, max_x, max_y))
|
286 |
+
h, w = ref_img_.height, ref_img_.width
|
287 |
+
scale = crop_size / max(h, w)
|
288 |
+
h_, w_ = int(scale * h), int(scale * w)
|
289 |
+
ref_img_ = ref_img_.resize((w_, h_), resample=Image.BICUBIC)
|
290 |
+
image_input = add_margin(ref_img_, size=image_size)
|
291 |
+
else:
|
292 |
+
image_input = add_margin(image_input, size=max(image_input.height, image_input.width))
|
293 |
+
image_input = image_input.resize((image_size, image_size), resample=Image.BICUBIC)
|
294 |
+
|
295 |
+
image_input = np.asarray(image_input)
|
296 |
+
image_input = image_input.astype(np.float32) / 255.0
|
297 |
+
ref_mask = image_input[:, :, 3:]
|
298 |
+
image_input[:, :, :3] = image_input[:, :, :3] * ref_mask + 1 - ref_mask # white background
|
299 |
+
image_input = image_input[:, :, :3] * 2.0 - 1.0
|
300 |
+
image_input = torch.from_numpy(image_input.astype(np.float32))
|
301 |
+
elevation_input = torch.from_numpy(np.asarray([np.deg2rad(elevation_input)], np.float32))
|
302 |
+
return {"input_image": image_input, "input_elevation": elevation_input}
|
meta_info/camera-16.pkl
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:d46537ff53982fd57a7b987673cd759e3a892f11b9f7ee44566f5612d6da6357
|
3 |
+
size 2142
|