jiangyzy commited on
Commit
28b27d8
1 Parent(s): d053b6e

Upload 35 files

Browse files
LICENSE ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright [yyyy] [name of copyright owner]
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
app.py ADDED
@@ -0,0 +1,293 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import os
3
+ import tempfile
4
+ from functools import partial
5
+ import cv2
6
+ import gradio as gr
7
+ import imageio
8
+ import numpy as np
9
+ import torch
10
+ import torchvision
11
+ from omegaconf import OmegaConf
12
+ from PIL import Image, ImageDraw
13
+ from pytorch_lightning import seed_everything
14
+ import sys
15
+ import copy
16
+ from utils.gradio_utils import load_preprocess_model, preprocess_image
17
+ from ldm.util import instantiate_from_config, img2tensor
18
+ from customnet.ddim import DDIMSampler
19
+ from einops import rearrange
20
+ import math
21
+
22
+
23
+
24
+ #### Description ####
25
+ title = r"""<h1 align="center">CustomNet: Object Customization with Variable-Viewpoints in Text-to-Image Diffusion Models</h1>"""
26
+
27
+ description = r"""
28
+ <b>Official Gradio demo</b> for <a href='https://github.com/TencentARC/CustomNet' target='_blank'><b>CustomNet: Object Customization with Variable-Viewpoints in Text-to-Image Diffusion Models</b></a>.<br>
29
+ 🔥 CustomNet is novel unified customization method that can generate harmonious customized images without
30
+ test-time optimization. CustomNet supports explicit viewpoint, location, text controls while ensuring
31
+ object identity preservation.<br>
32
+ 🤗 Try to customize the object gneration yourself!<br>
33
+ """
34
+ article = r"""
35
+ If CustomNet is helpful, please help to ⭐ the <a href='https://github.com/TencentARC/CustomNet' target='_blank'>Github Repo</a>. Thanks!
36
+ [![GitHub Stars](https://img.shields.io/github/stars/TencentARC%2FCustomNet)](https://github.com/TencentARC/CustomNet)
37
+
38
+ ---
39
+
40
+ 📝 **Citation**
41
+ <br>
42
+ If our work is useful for your research, please consider citing:
43
+ ```bibtex
44
+ @misc{yuan2023customnet,
45
+ title={CustomNet: Zero-shot Object Customization with Variable-Viewpoints in Text-to-Image Diffusion Models},
46
+ author={Ziyang Yuan and Mingdeng Cao and Xintao Wang and Zhongang Qi and Chun Yuan and Ying Shan},
47
+ year={2023},
48
+ eprint={2310.19784},
49
+ archivePrefix={arXiv},
50
+ primaryClass={cs.CV}
51
+ }
52
+ ```
53
+
54
+ 📧 **Contact**
55
+ <br>
56
+ If you have any questions, please feel free to reach me out at <b>[email protected]</b>.
57
+
58
+ """
59
+
60
+ # input_img = None
61
+ # concat_img = None
62
+ # T = None
63
+ # prompt = None
64
+ negtive_prompt = ""
65
+
66
+
67
+ def send_input_to_concat(input_image):
68
+ W, H = input_image.size
69
+ # image_array[:, 0, :] = image_array[:, 0, :]
70
+ draw = ImageDraw.Draw(input_image)
71
+ draw.rectangle([(0,0),(H-1, W-1)], outline="red", width=8)
72
+ return input_image
73
+
74
+ def preprocess_input(preprocess_model, input_image):
75
+ # global input_img
76
+ processed_image = preprocess_image(preprocess_model, input_image)
77
+ # input_img = (processed_image / 255.0).astype(np.float32)
78
+ return processed_image
79
+ # return processed_image, processed_image
80
+
81
+ def adjust_location(x0, y0, x1, y1, input_image):
82
+ x_0 = min(x0, x1)
83
+ x_1 = max(x0, x1)
84
+ y_0 = min(y0, y1)
85
+ y_1 = max(y0, y1)
86
+ print(x0, y0, x1, y1)
87
+ print(x_0, y_0, x_1, y_1)
88
+ new_size = (x_1-x_0, y_1-y_0)
89
+ input_image = input_image.resize(new_size)
90
+ img_array = np.array(input_image)
91
+ white_background = np.zeros((256, 256, 3))
92
+ white_background[y0:y1, x0:x1, :] = img_array
93
+ img_array = white_background.astype(np.uint8)
94
+ concat_img = Image.fromarray(img_array)
95
+ draw = ImageDraw.Draw(concat_img)
96
+ draw.rectangle([(x0,y0),(x1,y1)], outline="red", width=5)
97
+ return x_0, y_0, x_1, y_1, concat_img
98
+
99
+ def prepare_data(device, input_image, x0, y0, x1, y1, polar, azimuth, text):
100
+ if input_image.size[0] != 256 or input_image.size[1] != 256:
101
+ input_image = input_image.resize((256, 256))
102
+ input_image = np.array(input_image)
103
+
104
+ img_cond = img2tensor(input_image, bgr2rgb=False, float32=True).unsqueeze(0) / 255.
105
+ img_cond = img_cond*2-1
106
+
107
+ img_location = copy.deepcopy(img_cond)
108
+ input_im_padding = torch.ones_like(img_location)
109
+
110
+ x_0 = min(x0, x1)
111
+ x_1 = max(x0, x1)
112
+ y_0 = min(y0, y1)
113
+ y_1 = max(y0, y1)
114
+ print(x0, y0, x1, y1)
115
+ print(x_0, y_0, x_1, y_1)
116
+ img_location = torch.nn.functional.interpolate(img_location, (y_1-y_0, x_1-x_0), mode="bilinear")
117
+ input_im_padding[:,:, y_0:y_1, x_0:x_1] = img_location
118
+ img_location = input_im_padding
119
+
120
+ T = torch.tensor([[math.radians(polar), math.sin(math.radians(azimuth)), math.cos(math.radians(azimuth)), 0.0]]).unsqueeze(1)
121
+ batch = {
122
+ "image_cond": img_cond.to(device),
123
+ "image_location": img_location.to(device),
124
+ 'T': T.to(device),
125
+ 'text': [text],
126
+ }
127
+ return batch
128
+
129
+
130
+ @torch.no_grad()
131
+ def run_generation(sampler, model, device, input_image, x0, y0, x1, y1, polar, azimuth, text, seed):
132
+ seed_everything(seed)
133
+ batch = prepare_data(device, input_image, x0, y0, x1, y1, polar, azimuth, text)
134
+
135
+ c = model.get_learned_conditioning(batch["image_cond"])
136
+ c = torch.cat([c, batch["T"]], dim=-1)
137
+ c = model.cc_projection(c)
138
+
139
+ ## condition
140
+ cond = {}
141
+ cond['c_concat'] = [model.encode_first_stage((batch["image_location"])).mode().detach()]
142
+ cond['c_crossattn'] = [c]
143
+ text_embedding = model.text_encoder(batch["text"])
144
+ cond["c_crossattn"].append(text_embedding)
145
+
146
+ ## null-condition
147
+ uc = {}
148
+ neg_prompt = ""
149
+
150
+ uc['c_concat'] = [torch.zeros(1, 4, 32, 32).to(c.device)]
151
+ uc['c_crossattn'] = [torch.zeros_like(c).to(c.device)]
152
+ uc_text_embedding = model.text_encoder([neg_prompt])
153
+ uc['c_crossattn'].append(uc_text_embedding)
154
+
155
+ ## sample
156
+ shape = [4, 32, 32]
157
+ samples_latents, _ = sampler.sample(
158
+ S=50,
159
+ batch_size=1,
160
+ shape=shape,
161
+ verbose=False,
162
+ unconditional_guidance_scale=999, # useless
163
+ conditioning=cond,
164
+ unconditional_conditioning=uc,
165
+ cfg_type=0,
166
+ cfg_scale_dict={"img": 0., "text":0., "all": 3.0 }
167
+ )
168
+
169
+ x_samples = model.decode_first_stage(samples_latents)
170
+
171
+ x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0).cpu().numpy()
172
+ x_samples = rearrange(255.0 *x_samples[0], 'c h w -> h w c').astype(np.uint8)
173
+
174
+ output_image = Image.fromarray(x_samples)
175
+
176
+ return output_image
177
+
178
+
179
+ def load_example(input_image, x0, y0, x1, y1, polar, azimuth, prompt):
180
+ # print("AAAA")
181
+ # print(type(x0))
182
+ # print(type(polar))
183
+ return input_image, x0, y0, x1, y1, polar, azimuth, prompt
184
+
185
+ @torch.no_grad()
186
+ def main(args):
187
+ # load model
188
+ device = torch.device("cuda")
189
+ preprocess_model = load_preprocess_model()
190
+ config = OmegaConf.load("configs/config_customnet.yaml")
191
+ model = instantiate_from_config(config.model)
192
+
193
+ model_path='./customnet_v1.pt?download=true'
194
+ if not os.path.exists(model_path):
195
+ os.system(f'wget https://huggingface.co/TencentARC/CustomNet/resolve/main/customnet_v1.pt?download=true -P .')
196
+
197
+ ckpt = torch.load(model_path, map_location="cpu")
198
+ model.load_state_dict(ckpt)
199
+ del ckpt
200
+ model = model.to(device)
201
+ sampler = DDIMSampler(model, device=device)
202
+
203
+ # load demo
204
+ demo = gr.Blocks()
205
+ with demo:
206
+ gr.Markdown(title)
207
+ gr.Markdown(description)
208
+
209
+ with gr.Row():
210
+ ## Left column
211
+ with gr.Column():
212
+ ## step 1.
213
+ gr.Markdown("## Step 1: Upload an object image and process", show_label=False)
214
+ # with gr.Row(equal_height=True):
215
+ input_image = gr.Image(type="pil", interactive=True, elem_id="input_image", elem_classes='image', visible=True)
216
+ preprocess_botton = gr.Button(value="Need preprocess", visible=True)
217
+
218
+ ## step 2.
219
+ gr.Markdown("## Step 2: Set up different controls ", show_label=False, visible=True)
220
+ gr.Markdown("### 1: Object Location", show_label=False, visible=True)
221
+ with gr.Row():
222
+ with gr.Column():
223
+ with gr.Row():
224
+ x0 = gr.Slider(minimum=0, maximum=256, step=1, label="X_0", value=0, interactive=True, visible=True)
225
+ y0 = gr.Slider(minimum=0, maximum=256, step=1, label="Y_0", value=0, interactive=True, visible=True)
226
+ with gr.Row():
227
+ x1 = gr.Slider(minimum=0, maximum=256, step=1, label="X_1", value=256, interactive=True, visible=True)
228
+ y1 = gr.Slider(minimum=0, maximum=256, step=1, label="Y_1", value=256, interactive=True, visible=True)
229
+ location_botton = gr.Button(value="Update Location ", visible=True)
230
+
231
+ location_image = gr.Image(type="pil", interactive=True, elem_id="location", elem_classes='image', visible=True)
232
+ gr.Markdown("### 2: Object Viewpoint", show_label=False, visible=True)
233
+ with gr.Row():
234
+ polar = gr.Slider(minimum=-90, maximum=90, step=-0.5, label="Polar Angle", value=0.0, visible=True)
235
+ azimuth = gr.Slider(minimum=-60, maximum=90, step=-0.5, label="Azimuth angle", value=0.0, visible=True)
236
+ gr.Markdown("### 3: Text", show_label=False, visible=True)
237
+ prompt = gr.Textbox(value="on the seaside", label="Prompt", interactive=True, visible=True)
238
+
239
+ ## step 3.
240
+ gr.Markdown("## Step 3: Run Generation", show_label=False, visible=True)
241
+ seed = gr.Number(value=1234, precision=0, interactive=True, label="Seed", visible=True)
242
+ start = gr.Button(value="Run generation !", visible=True)
243
+
244
+
245
+
246
+ examples_full = [
247
+ ["examples/0.jpg", 50, 50, 256, 256, 0, -30, "a backpack in the office"],
248
+ ["examples/1.jpg", 20, 20, 256, 256, -25, -35, "a pair of shoes on dirt road"],
249
+ ["examples/2.jpg", 0, 0, 256, 256, -15, -20, "a car on the beach"],
250
+ ["examples/3.jpg", 0, 0, 256, 256, 0, 30, "in the jungle"],
251
+ ["examples/4.jpg", 0, 0, 256, 256, 0, -30, "in the snow"],
252
+ ["examples/5.jpg", 20, 20, 240, 240, 10, 20, "with mountain behind"],
253
+ ]
254
+
255
+ ## Right column
256
+ with gr.Column():
257
+ gr.Markdown("## Generation Results", show_label=False, visible=True)
258
+ output_image = gr.Image(type="pil", interactive=True, elem_id="output_image", elem_classes='image', visible=True)
259
+
260
+ gr.Examples(
261
+ examples=examples_full, # NOTE: elements must match inputs list!
262
+ fn=load_example,
263
+ inputs=[input_image, x0, y0, x1, y1, polar, azimuth, prompt],
264
+ outputs=[input_image, x0, y0, x1, y1, polar, azimuth, prompt],
265
+ cache_examples=False,
266
+ run_on_click=True,
267
+ )
268
+ gr.Markdown(article)
269
+
270
+ ## function
271
+ input_image.change(send_input_to_concat, inputs=input_image, outputs=location_image)
272
+ preprocess_botton.click(partial(preprocess_input, preprocess_model), inputs=input_image, outputs=input_image)
273
+ location_botton.click(adjust_location,
274
+ inputs=[x0, y0, x1, y1, input_image],
275
+ outputs=[x0, y0, x1, y1, location_image])
276
+
277
+ start.click(partial(run_generation, sampler, model, device),
278
+ inputs=[input_image, x0, y0, x1, y1, polar, azimuth, prompt, seed],
279
+ outputs=output_image)
280
+
281
+
282
+ # demo.launch(server_name='0.0.0.0', share=False, server_port=args.port)
283
+ demo.queue(concurrency_count=1, max_size=10)
284
+ demo.launch()
285
+
286
+
287
+
288
+ if __name__=="__main__":
289
+ parser = argparse.ArgumentParser()
290
+ parser.add_argument("--port", type=int, default=12345)
291
+ args = parser.parse_args()
292
+
293
+ main(args)
configs/config_customnet.yaml ADDED
@@ -0,0 +1,123 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ base_learning_rate: 2.0e-05
3
+ target: customnet.customnet.CustomNet
4
+ params:
5
+ linear_start: 0.00085
6
+ linear_end: 0.0120
7
+ num_timesteps_cond: 1
8
+ log_every_t: 200
9
+ timesteps: 1000
10
+ first_stage_key: "image_target"
11
+ cond_stage_key: "image_cond"
12
+ image_size: 32
13
+ channels: 4
14
+ cond_stage_trainable: false # Note: different from the one we trained before
15
+ conditioning_key: hybrid
16
+ monitor: val/loss_simple_ema
17
+ scale_factor: 0.18215
18
+ use_ema: false
19
+ use_cond_concat: true
20
+ use_bbox_mask: false
21
+ use_bg_inpainting: false
22
+ learning_rate_scale: 10
23
+ ucg_training:
24
+ txt: 0.15
25
+
26
+ sd_15_ckpt: #"v1-5-pruned-emaonly.ckpt"
27
+
28
+ unet_config:
29
+ target: customnet.openaimodel.UNetModel
30
+ params:
31
+ image_size: 32 # unused
32
+ in_channels: 8
33
+ out_channels: 4
34
+ model_channels: 320
35
+ attention_resolutions: [ 4, 2, 1 ]
36
+ num_res_blocks: 2
37
+ channel_mult: [ 1, 2, 4, 4 ]
38
+ num_heads: 8
39
+ use_spatial_transformer: True
40
+ transformer_depth: 1
41
+ context_dim: 768
42
+ use_checkpoint: True
43
+ legacy: False
44
+
45
+
46
+ first_stage_config:
47
+ target: ldm.models.autoencoder.AutoencoderKL
48
+ params:
49
+ embed_dim: 4
50
+ monitor: val/rec_loss
51
+ ddconfig:
52
+ double_z: true
53
+ z_channels: 4
54
+ resolution: 256
55
+ in_channels: 3
56
+ out_ch: 3
57
+ ch: 128
58
+ ch_mult:
59
+ - 1
60
+ - 2
61
+ - 4
62
+ - 4
63
+ num_res_blocks: 2
64
+ attn_resolutions: []
65
+ dropout: 0.0
66
+ lossconfig:
67
+ target: torch.nn.Identity
68
+
69
+ cond_stage_config:
70
+ target: ldm.modules.encoders.modules.FrozenCLIPImageEmbedder
71
+
72
+
73
+ text_encoder_config:
74
+ target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
75
+ params:
76
+ version: openai/clip-vit-large-patch14
77
+
78
+
79
+ ## this is a template dataset
80
+ train_data:
81
+ target: data.dataset.Dataset
82
+ params:
83
+ image_size: 256
84
+ root: examples/dataset/
85
+
86
+
87
+ train_dataloader:
88
+ batch_size: 12
89
+ num_workers: 8
90
+
91
+
92
+
93
+
94
+ lightning:
95
+ find_unused_parameters: false
96
+ metrics_over_trainsteps_checkpoint: True
97
+ modelcheckpoint:
98
+ params:
99
+ every_n_train_steps: 10000
100
+ save_top_k: -1
101
+ monitor: null
102
+ callbacks:
103
+ image_logger:
104
+ target: main.ImageLogger
105
+ params:
106
+ batch_frequency: 2500
107
+ max_images: 32
108
+ increase_log_steps: False
109
+ log_first_step: True
110
+ log_images_kwargs:
111
+ use_ema_scope: False
112
+ inpaint: False
113
+ plot_progressive_rows: False
114
+ plot_diffusion_rows: False
115
+ N: 32
116
+ unconditional_guidance_scale: 3.0
117
+ unconditional_guidance_label: [""]
118
+
119
+ trainer:
120
+ benchmark: True
121
+ limit_val_batches: 0
122
+ num_sanity_val_steps: 0
123
+ accumulate_grad_batches: 1
customnet/attention.py ADDED
@@ -0,0 +1,296 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
46
+
47
+ class FeedForward(nn.Module):
48
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
49
+ super().__init__()
50
+ inner_dim = int(dim * mult)
51
+ dim_out = default(dim_out, dim)
52
+ project_in = nn.Sequential(
53
+ nn.Linear(dim, inner_dim),
54
+ nn.GELU()
55
+ ) if not glu else GEGLU(dim, inner_dim)
56
+
57
+ self.net = nn.Sequential(
58
+ project_in,
59
+ nn.Dropout(dropout),
60
+ nn.Linear(inner_dim, dim_out)
61
+ )
62
+
63
+ def forward(self, x):
64
+ return self.net(x)
65
+
66
+
67
+ def zero_module(module):
68
+ """
69
+ Zero out the parameters of a module and return it.
70
+ """
71
+ for p in module.parameters():
72
+ p.detach().zero_()
73
+ return module
74
+
75
+
76
+ def Normalize(in_channels):
77
+ return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
78
+
79
+
80
+ class LinearAttention(nn.Module):
81
+ def __init__(self, dim, heads=4, dim_head=32):
82
+ super().__init__()
83
+ self.heads = heads
84
+ hidden_dim = dim_head * heads
85
+ self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
86
+ self.to_out = nn.Conv2d(hidden_dim, dim, 1)
87
+
88
+ def forward(self, x):
89
+ b, c, h, w = x.shape
90
+ qkv = self.to_qkv(x)
91
+ q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
92
+ k = k.softmax(dim=-1)
93
+ context = torch.einsum('bhdn,bhen->bhde', k, v)
94
+ out = torch.einsum('bhde,bhdn->bhen', context, q)
95
+ out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
96
+ return self.to_out(out)
97
+
98
+
99
+ class SpatialSelfAttention(nn.Module):
100
+ def __init__(self, in_channels):
101
+ super().__init__()
102
+ self.in_channels = in_channels
103
+
104
+ self.norm = Normalize(in_channels)
105
+ self.q = torch.nn.Conv2d(in_channels,
106
+ in_channels,
107
+ kernel_size=1,
108
+ stride=1,
109
+ padding=0)
110
+ self.k = torch.nn.Conv2d(in_channels,
111
+ in_channels,
112
+ kernel_size=1,
113
+ stride=1,
114
+ padding=0)
115
+ self.v = torch.nn.Conv2d(in_channels,
116
+ in_channels,
117
+ kernel_size=1,
118
+ stride=1,
119
+ padding=0)
120
+ self.proj_out = torch.nn.Conv2d(in_channels,
121
+ in_channels,
122
+ kernel_size=1,
123
+ stride=1,
124
+ padding=0)
125
+
126
+ def forward(self, x):
127
+ h_ = x
128
+ h_ = self.norm(h_)
129
+ q = self.q(h_)
130
+ k = self.k(h_)
131
+ v = self.v(h_)
132
+
133
+ # compute attention
134
+ b,c,h,w = q.shape
135
+ q = rearrange(q, 'b c h w -> b (h w) c')
136
+ k = rearrange(k, 'b c h w -> b c (h w)')
137
+ w_ = torch.einsum('bij,bjk->bik', q, k)
138
+
139
+ w_ = w_ * (int(c)**(-0.5))
140
+ w_ = torch.nn.functional.softmax(w_, dim=2)
141
+
142
+ # attend to values
143
+ v = rearrange(v, 'b c h w -> b c (h w)')
144
+ w_ = rearrange(w_, 'b i j -> b j i')
145
+ h_ = torch.einsum('bij,bjk->bik', v, w_)
146
+ h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
147
+ h_ = self.proj_out(h_)
148
+
149
+ return x+h_
150
+
151
+ ## Dual Cross-Attntion
152
+ class CrossAttention(nn.Module):
153
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., new_text_att=False,):
154
+ super().__init__()
155
+ inner_dim = dim_head * heads
156
+ context_dim = default(context_dim, query_dim)
157
+
158
+ self.scale = dim_head ** -0.5
159
+ self.heads = heads
160
+
161
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
162
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
163
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
164
+
165
+ self.to_out = nn.Sequential(
166
+ nn.Linear(inner_dim, query_dim),
167
+ nn.Dropout(dropout)
168
+ )
169
+
170
+ self.new_text_att = new_text_att
171
+ if new_text_att:
172
+ self.to_k_text = nn.Linear(context_dim, inner_dim, bias=False)
173
+ self.to_v_text = nn.Linear(context_dim, inner_dim, bias=False)
174
+ self.text_scale = 1.0
175
+
176
+
177
+
178
+ def forward(self, x, context=None, mask=None):
179
+ if self.new_text_att:
180
+ context_text = context[:,1:,:] ## text embedding
181
+ context = context[:,:1,:] ## image embedding
182
+
183
+ h = self.heads
184
+
185
+ q = self.to_q(x)
186
+ context = default(context, x)
187
+ k = self.to_k(context)
188
+ v = self.to_v(context)
189
+
190
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
191
+
192
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
193
+
194
+ if exists(mask):
195
+ mask = rearrange(mask, 'b ... -> b (...)')
196
+ max_neg_value = -torch.finfo(sim.dtype).max
197
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
198
+ sim.masked_fill_(~mask, max_neg_value)
199
+
200
+ # attention, what we cannot get enough of
201
+ attn = sim.softmax(dim=-1)
202
+
203
+ out = einsum('b i j, b j d -> b i d', attn, v)
204
+
205
+ ## text att
206
+ if self.new_text_att:
207
+ k_text = self.to_k_text(context_text)
208
+ v_text = self.to_v_text(context_text)
209
+ k_text, v_text = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (k_text, v_text))
210
+
211
+ sim_text = einsum('b i d, b j d -> b i j', q, k_text) * self.scale
212
+
213
+ attn_text = sim_text.softmax(dim=-1)
214
+ out_text = einsum('b i j, b j d -> b i d', attn_text, v_text)
215
+
216
+ out = out + out_text * self.text_scale
217
+
218
+ # print("no text attn")
219
+
220
+ out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
221
+ return self.to_out(out)
222
+
223
+
224
+ class BasicTransformerBlock(nn.Module):
225
+ def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
226
+ disable_self_attn=False):
227
+ super().__init__()
228
+ self.disable_self_attn = disable_self_attn
229
+ self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
230
+ context_dim=context_dim if self.disable_self_attn else None) # is a self-attention if not self.disable_self_attn
231
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
232
+ self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
233
+ heads=n_heads, dim_head=d_head, dropout=dropout,
234
+ new_text_att=True,
235
+ ) # is self-attn if context is none
236
+ self.norm1 = nn.LayerNorm(dim)
237
+ self.norm2 = nn.LayerNorm(dim)
238
+ self.norm3 = nn.LayerNorm(dim)
239
+ self.checkpoint = checkpoint
240
+
241
+ def forward(self, x, context=None):
242
+ return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
243
+
244
+ def _forward(self, x, context=None):
245
+ x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
246
+ x = self.attn2(self.norm2(x), context=context) + x
247
+ x = self.ff(self.norm3(x)) + x
248
+ return x
249
+
250
+
251
+ class SpatialTransformer(nn.Module):
252
+ """
253
+ Transformer block for image-like data.
254
+ First, project the input (aka embedding)
255
+ and reshape to b, t, d.
256
+ Then apply standard transformer action.
257
+ Finally, reshape to image
258
+ """
259
+ def __init__(self, in_channels, n_heads, d_head,
260
+ depth=1, dropout=0., context_dim=None,
261
+ disable_self_attn=False):
262
+ super().__init__()
263
+ self.in_channels = in_channels
264
+ inner_dim = n_heads * d_head
265
+ self.norm = Normalize(in_channels)
266
+
267
+ self.proj_in = nn.Conv2d(in_channels,
268
+ inner_dim,
269
+ kernel_size=1,
270
+ stride=1,
271
+ padding=0)
272
+
273
+ self.transformer_blocks = nn.ModuleList(
274
+ [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim,
275
+ disable_self_attn=disable_self_attn)
276
+ for d in range(depth)]
277
+ )
278
+
279
+ self.proj_out = zero_module(nn.Conv2d(inner_dim,
280
+ in_channels,
281
+ kernel_size=1,
282
+ stride=1,
283
+ padding=0))
284
+
285
+ def forward(self, x, context=None):
286
+ # note: if no context is given, cross-attention defaults to self-attention
287
+ b, c, h, w = x.shape
288
+ x_in = x
289
+ x = self.norm(x)
290
+ x = self.proj_in(x)
291
+ x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
292
+ for block in self.transformer_blocks:
293
+ x = block(x, context=context)
294
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
295
+ x = self.proj_out(x)
296
+ return x + x_in
customnet/customnet.py ADDED
@@ -0,0 +1,171 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import einops
3
+ import torch
4
+ import torch as th
5
+ import torch.nn as nn
6
+ import cv2
7
+ from pytorch_lightning.utilities.distributed import rank_zero_only
8
+ import numpy as np
9
+ from torch.optim.lr_scheduler import LambdaLR
10
+
11
+ from ldm.modules.diffusionmodules.util import (
12
+ conv_nd,
13
+ linear,
14
+ zero_module,
15
+ timestep_embedding,
16
+ )
17
+
18
+ from einops import rearrange, repeat
19
+ from torchvision.utils import make_grid
20
+ from ldm.modules.attention import SpatialTransformer
21
+ from ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
22
+ from ldm.models.diffusion.ddpm import LatentDiffusion
23
+ from ldm.util import log_txt_as_img, exists, instantiate_from_config
24
+
25
+ from ldm.models.diffusion.ddim import DDIMSampler
26
+ from ldm.util import load_state_dict
27
+
28
+
29
+
30
+
31
+
32
+
33
+ class CustomNet(LatentDiffusion):
34
+ def __init__(self,
35
+ text_encoder_config,
36
+ sd_15_ckpt=None,
37
+ use_cond_concat=False,
38
+ use_bbox_mask=False,
39
+ use_bg_inpainting=False,
40
+ learning_rate_scale=10,
41
+ *args, **kwargs):
42
+ super().__init__(*args, **kwargs)
43
+
44
+
45
+ self.text_encoder = instantiate_from_config(text_encoder_config)
46
+
47
+ if sd_15_ckpt is not None:
48
+ self.load_model_from_ckpt(ckpt=sd_15_ckpt)
49
+
50
+ self.use_cond_concat = use_cond_concat
51
+ self.use_bbox_mask = use_bbox_mask
52
+ self.use_bg_inpainting = use_bg_inpainting
53
+ self.learning_rate_scale = learning_rate_scale
54
+
55
+
56
+ def load_model_from_ckpt(self, ckpt, verbose=True):
57
+ print(" =========================== init Stable Diffusion pretrained checkpoint =========================== ")
58
+ print(f"Loading model from {ckpt}")
59
+ pl_sd = torch.load(ckpt, map_location="cpu")
60
+ if "global_step" in pl_sd:
61
+ print(f"Global Step: {pl_sd['global_step']}")
62
+ sd = pl_sd["state_dict"]
63
+ sd_keys = sd.keys()
64
+
65
+
66
+ missing = []
67
+ text_encoder_sd = self.text_encoder.state_dict()
68
+ for k in text_encoder_sd.keys():
69
+ sd_k = "cond_stage_model."+ k
70
+ if sd_k in sd_keys:
71
+ text_encoder_sd[k] = sd[sd_k]
72
+ else:
73
+ missing.append(k)
74
+
75
+ self.text_encoder.load_state_dict(text_encoder_sd)
76
+
77
+
78
+
79
+
80
+ def configure_optimizers(self):
81
+ lr = self.learning_rate
82
+ params = []
83
+ params += list(self.cc_projection.parameters())
84
+
85
+
86
+ params_dualattn = []
87
+ for k, v in self.model.named_parameters():
88
+ if "to_k_text" in k or "to_v_text" in k:
89
+ params_dualattn.append(v)
90
+ print("training weight: ", k)
91
+ else:
92
+ params.append(v)
93
+
94
+
95
+ opt = torch.optim.AdamW([
96
+ {'params':params_dualattn, 'lr': lr*self.learning_rate_scale},
97
+ {'params': params, 'lr': lr}
98
+ ])
99
+
100
+
101
+ if self.use_scheduler:
102
+ assert 'target' in self.scheduler_config
103
+ scheduler = instantiate_from_config(self.scheduler_config)
104
+
105
+ print("Setting up LambdaLR scheduler...")
106
+ scheduler = [
107
+ {
108
+ 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
109
+ 'interval': 'step',
110
+ 'frequency': 1
111
+ }]
112
+ return [opt], scheduler
113
+ return opt
114
+
115
+ def training_step(self, batch, batch_idx):
116
+ loss, loss_dict = self.shared_step(batch)
117
+
118
+ self.log_dict(loss_dict, prog_bar=True,
119
+ logger=True, on_step=True, on_epoch=True)
120
+
121
+ self.log("global_step", self.global_step,
122
+ prog_bar=True, logger=True, on_step=True, on_epoch=False)
123
+
124
+ if self.use_scheduler:
125
+ lr = self.optimizers().param_groups[0]['lr']
126
+ self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
127
+
128
+ return loss
129
+
130
+
131
+
132
+
133
+ def shared_step(self, batch, **kwargs):
134
+
135
+ if 'txt' in self.ucg_training:
136
+ k = 'txt'
137
+ p = self.ucg_training[k]
138
+ for i in range(len(batch[k])):
139
+ if self.ucg_prng.choice(2, p=[1 - p, p]):
140
+ if isinstance(batch[k], list):
141
+ batch[k][i] = ""
142
+
143
+ with torch.no_grad():
144
+ text = batch['txt']
145
+ text_embedding = self.text_encoder(text)
146
+
147
+
148
+ x, c = self.get_input(batch, self.first_stage_key)
149
+
150
+ c["c_crossattn"].append(text_embedding)
151
+ loss = self(x, c,)
152
+ return loss
153
+
154
+
155
+ def apply_model(self, x_noisy, t, cond, return_ids=False, **kwargs):
156
+
157
+ if isinstance(cond, dict):
158
+ # hybrid case, cond is exptected to be a dict
159
+ pass
160
+ else:
161
+ if not isinstance(cond, list):
162
+ cond = [cond]
163
+ key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
164
+ cond = {key: cond}
165
+
166
+ x_recon = self.model(x_noisy, t, **cond)
167
+
168
+ if isinstance(x_recon, tuple) and not return_ids:
169
+ return x_recon[0]
170
+ else:
171
+ return x_recon
customnet/ddim.py ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+ from functools import partial
7
+ from einops import rearrange
8
+
9
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
10
+ from ldm.models.diffusion.sampling_util import renorm_thresholding, norm_thresholding, spatial_norm_thresholding
11
+
12
+
13
+ class DDIMSampler(object):
14
+ def __init__(self, model, schedule="linear", **kwargs):
15
+ super().__init__()
16
+ self.model = model
17
+ self.ddpm_num_timesteps = model.num_timesteps
18
+ self.schedule = schedule
19
+
20
+ def to(self, device):
21
+ """Same as to in torch module
22
+ Don't really underestand why this isn't a module in the first place"""
23
+ for k, v in self.__dict__.items():
24
+ if isinstance(v, torch.Tensor):
25
+ new_v = getattr(self, k).to(device)
26
+ setattr(self, k, new_v)
27
+
28
+
29
+ def register_buffer(self, name, attr):
30
+ if type(attr) == torch.Tensor:
31
+ if attr.device != torch.device("cuda"):
32
+ attr = attr.to(torch.device("cuda"))
33
+ setattr(self, name, attr)
34
+
35
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
36
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
37
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
38
+ alphas_cumprod = self.model.alphas_cumprod
39
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
40
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
41
+
42
+ self.register_buffer('betas', to_torch(self.model.betas))
43
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
44
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
45
+
46
+ # calculations for diffusion q(x_t | x_{t-1}) and others
47
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
48
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
49
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
50
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
51
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
52
+
53
+ # ddim sampling parameters
54
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
55
+ ddim_timesteps=self.ddim_timesteps,
56
+ eta=ddim_eta,verbose=verbose)
57
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
58
+ self.register_buffer('ddim_alphas', ddim_alphas)
59
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
60
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
61
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
62
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
63
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
64
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
65
+
66
+ @torch.no_grad()
67
+ def sample(self,
68
+ S,
69
+ batch_size,
70
+ shape,
71
+ conditioning=None,
72
+ callback=None,
73
+ normals_sequence=None,
74
+ img_callback=None,
75
+ quantize_x0=False,
76
+ eta=0.,
77
+ mask=None,
78
+ x0=None,
79
+ temperature=1.,
80
+ noise_dropout=0.,
81
+ score_corrector=None,
82
+ corrector_kwargs=None,
83
+ verbose=True,
84
+ x_T=None,
85
+ log_every_t=100,
86
+ unconditional_guidance_scale=1.,
87
+ unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
88
+ dynamic_threshold=None,
89
+ **kwargs
90
+ ):
91
+ if conditioning is not None:
92
+ if isinstance(conditioning, dict):
93
+ ctmp = conditioning[list(conditioning.keys())[0]]
94
+ while isinstance(ctmp, list): ctmp = ctmp[0]
95
+ cbs = ctmp.shape[0]
96
+ if cbs != batch_size:
97
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
98
+
99
+ else:
100
+ if conditioning.shape[0] != batch_size:
101
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
102
+
103
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
104
+ # sampling
105
+ C, H, W = shape
106
+ size = (batch_size, C, H, W)
107
+ print(f'Data shape for DDIM sampling is {size}, eta {eta}')
108
+
109
+ samples, intermediates = self.ddim_sampling(conditioning, size,
110
+ callback=callback,
111
+ img_callback=img_callback,
112
+ quantize_denoised=quantize_x0,
113
+ mask=mask, x0=x0,
114
+ ddim_use_original_steps=False,
115
+ noise_dropout=noise_dropout,
116
+ temperature=temperature,
117
+ score_corrector=score_corrector,
118
+ corrector_kwargs=corrector_kwargs,
119
+ x_T=x_T,
120
+ log_every_t=log_every_t,
121
+ unconditional_guidance_scale=unconditional_guidance_scale,
122
+ unconditional_conditioning=unconditional_conditioning,
123
+ dynamic_threshold=dynamic_threshold,
124
+ **kwargs)
125
+ return samples, intermediates
126
+
127
+ @torch.no_grad()
128
+ def ddim_sampling(self, cond, shape,
129
+ x_T=None, ddim_use_original_steps=False,
130
+ callback=None, timesteps=None, quantize_denoised=False,
131
+ mask=None, x0=None, img_callback=None, log_every_t=100,
132
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
133
+ unconditional_guidance_scale=1., unconditional_conditioning=None, dynamic_threshold=None,
134
+ t_start=-1,
135
+ **kwargs):
136
+ device = self.model.betas.device
137
+ b = shape[0]
138
+ if x_T is None:
139
+ img = torch.randn(shape, device=device)
140
+ else:
141
+ img = x_T
142
+
143
+ if timesteps is None:
144
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
145
+ elif timesteps is not None and not ddim_use_original_steps:
146
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
147
+ timesteps = self.ddim_timesteps[:subset_end]
148
+
149
+ timesteps = timesteps[:t_start]
150
+
151
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
152
+ time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
153
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
154
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
155
+
156
+ iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
157
+
158
+ for i, step in enumerate(iterator):
159
+ index = total_steps - i - 1
160
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
161
+
162
+ if mask is not None:
163
+ assert x0 is not None
164
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
165
+ img = img_orig * mask + (1. - mask) * img
166
+
167
+ outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
168
+ quantize_denoised=quantize_denoised, temperature=temperature,
169
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
170
+ corrector_kwargs=corrector_kwargs,
171
+ unconditional_guidance_scale=unconditional_guidance_scale,
172
+ unconditional_conditioning=unconditional_conditioning,
173
+ dynamic_threshold=dynamic_threshold,
174
+ **kwargs)
175
+ img, pred_x0 = outs
176
+ if callback:
177
+ img = callback(i, img, pred_x0)
178
+ if img_callback: img_callback(pred_x0, i)
179
+
180
+ if index % log_every_t == 0 or index == total_steps - 1:
181
+ intermediates['x_inter'].append(img)
182
+ intermediates['pred_x0'].append(pred_x0)
183
+
184
+ return img, intermediates
185
+
186
+ @torch.no_grad()
187
+ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
188
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
189
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
190
+ dynamic_threshold=None, **kwargs):
191
+ b, *_, device = *x.shape, x.device
192
+
193
+
194
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
195
+ e_t = self.model.apply_model(x, t, c)
196
+ else:
197
+ cfg_type = kwargs.get("cfg_type", 0)
198
+ cfg = kwargs.get("cfg_scale_dict", {"img": 0., "text":0., "all": 3.0 })
199
+ truncate_step = kwargs.get("cfg_type", 250) ## 250 -1
200
+
201
+ if cfg_type == 0 :
202
+ model_t = self.model.apply_model(x, t, c)
203
+
204
+ model_uc = self.model.apply_model(x, t, unconditional_conditioning)
205
+
206
+ e_t = model_uc + cfg["all"]*(model_t - model_uc)
207
+
208
+ elif cfg_type==1:
209
+
210
+ model_t = self.model.apply_model(x, t, c)
211
+
212
+ c_text = {"c_concat": unconditional_conditioning["c_concat"],
213
+ "c_crossattn": [unconditional_conditioning["c_crossattn"][0],
214
+ c["c_crossattn"][1],]}
215
+
216
+ c_img = {"c_concat": c["c_concat"],
217
+ "c_crossattn": [c["c_crossattn"][0],
218
+ unconditional_conditioning["c_crossattn"][1],]}
219
+
220
+ model_t_text = self.model.apply_model(x, t, c_text)
221
+ model_t_img = self.model.apply_model(x, t, c_img)
222
+ model_uc = self.model.apply_model(x, t, unconditional_conditioning)
223
+
224
+ e_t = model_uc + cfg["all"]*(model_t - model_uc)
225
+ if t[0] > truncate_step:
226
+ print(1, t[0])
227
+ e_t += cfg["img"] * (model_t_img - model_uc)
228
+ e_t += cfg["text"] * (model_t_text - model_uc)
229
+
230
+
231
+ if score_corrector is not None:
232
+ assert self.model.parameterization == "eps"
233
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
234
+
235
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
236
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
237
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
238
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
239
+ # select parameters corresponding to the currently considered timestep
240
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
241
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
242
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
243
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
244
+
245
+ # current prediction for x_0
246
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
247
+ if quantize_denoised:
248
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
249
+
250
+ if dynamic_threshold is not None:
251
+ pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
252
+
253
+ # direction pointing to x_t
254
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
255
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
256
+ if noise_dropout > 0.:
257
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
258
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
259
+ return x_prev, pred_x0
260
+
261
+ @torch.no_grad()
262
+ def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
263
+ unconditional_guidance_scale=1.0, unconditional_conditioning=None):
264
+ num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
265
+
266
+ assert t_enc <= num_reference_steps
267
+ num_steps = t_enc
268
+
269
+ if use_original_steps:
270
+ alphas_next = self.alphas_cumprod[:num_steps]
271
+ alphas = self.alphas_cumprod_prev[:num_steps]
272
+ else:
273
+ alphas_next = self.ddim_alphas[:num_steps]
274
+ alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
275
+
276
+ x_next = x0
277
+ intermediates = []
278
+ inter_steps = []
279
+ for i in tqdm(range(num_steps), desc='Encoding Image'):
280
+ t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)
281
+ if unconditional_guidance_scale == 1.:
282
+ noise_pred = self.model.apply_model(x_next, t, c)
283
+ else:
284
+ assert unconditional_conditioning is not None
285
+ e_t_uncond, noise_pred = torch.chunk(
286
+ self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),
287
+ torch.cat((unconditional_conditioning, c))), 2)
288
+ noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)
289
+
290
+ xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
291
+ weighted_noise_pred = alphas_next[i].sqrt() * (
292
+ (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
293
+ x_next = xt_weighted + weighted_noise_pred
294
+ if return_intermediates and i % (
295
+ num_steps // return_intermediates) == 0 and i < num_steps - 1:
296
+ intermediates.append(x_next)
297
+ inter_steps.append(i)
298
+ elif return_intermediates and i >= num_steps - 2:
299
+ intermediates.append(x_next)
300
+ inter_steps.append(i)
301
+
302
+ out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
303
+ if return_intermediates:
304
+ out.update({'intermediates': intermediates})
305
+ return x_next, out
306
+
307
+ @torch.no_grad()
308
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
309
+ # fast, but does not allow for exact reconstruction
310
+ # t serves as an index to gather the correct alphas
311
+ if use_original_steps:
312
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
313
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
314
+ else:
315
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
316
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
317
+
318
+ if noise is None:
319
+ noise = torch.randn_like(x0)
320
+ return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
321
+ extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
322
+
323
+ @torch.no_grad()
324
+ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
325
+ use_original_steps=False):
326
+
327
+ timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
328
+ timesteps = timesteps[:t_start]
329
+
330
+ time_range = np.flip(timesteps)
331
+ total_steps = timesteps.shape[0]
332
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
333
+
334
+ iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
335
+ x_dec = x_latent
336
+ for i, step in enumerate(iterator):
337
+ index = total_steps - i - 1
338
+ ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
339
+ x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
340
+ unconditional_guidance_scale=unconditional_guidance_scale,
341
+ unconditional_conditioning=unconditional_conditioning)
342
+ return x_dec
343
+
customnet/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 customnet.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:
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
+ )
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:
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)
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)
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
+
examples/0.jpg ADDED
examples/1.jpg ADDED
examples/2.jpg ADDED
examples/3.jpg ADDED
examples/4.jpg ADDED
examples/5.jpg ADDED
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/classifier.py ADDED
@@ -0,0 +1,267 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import pytorch_lightning as pl
4
+ from omegaconf import OmegaConf
5
+ from torch.nn import functional as F
6
+ from torch.optim import AdamW
7
+ from torch.optim.lr_scheduler import LambdaLR
8
+ from copy import deepcopy
9
+ from einops import rearrange
10
+ from glob import glob
11
+ from natsort import natsorted
12
+
13
+ from ldm.modules.diffusionmodules.openaimodel import EncoderUNetModel, UNetModel
14
+ from ldm.util import log_txt_as_img, default, ismap, instantiate_from_config
15
+
16
+ __models__ = {
17
+ 'class_label': EncoderUNetModel,
18
+ 'segmentation': UNetModel
19
+ }
20
+
21
+
22
+ def disabled_train(self, mode=True):
23
+ """Overwrite model.train with this function to make sure train/eval mode
24
+ does not change anymore."""
25
+ return self
26
+
27
+
28
+ class NoisyLatentImageClassifier(pl.LightningModule):
29
+
30
+ def __init__(self,
31
+ diffusion_path,
32
+ num_classes,
33
+ ckpt_path=None,
34
+ pool='attention',
35
+ label_key=None,
36
+ diffusion_ckpt_path=None,
37
+ scheduler_config=None,
38
+ weight_decay=1.e-2,
39
+ log_steps=10,
40
+ monitor='val/loss',
41
+ *args,
42
+ **kwargs):
43
+ super().__init__(*args, **kwargs)
44
+ self.num_classes = num_classes
45
+ # get latest config of diffusion model
46
+ diffusion_config = natsorted(glob(os.path.join(diffusion_path, 'configs', '*-project.yaml')))[-1]
47
+ self.diffusion_config = OmegaConf.load(diffusion_config).model
48
+ self.diffusion_config.params.ckpt_path = diffusion_ckpt_path
49
+ self.load_diffusion()
50
+
51
+ self.monitor = monitor
52
+ self.numd = self.diffusion_model.first_stage_model.encoder.num_resolutions - 1
53
+ self.log_time_interval = self.diffusion_model.num_timesteps // log_steps
54
+ self.log_steps = log_steps
55
+
56
+ self.label_key = label_key if not hasattr(self.diffusion_model, 'cond_stage_key') \
57
+ else self.diffusion_model.cond_stage_key
58
+
59
+ assert self.label_key is not None, 'label_key neither in diffusion model nor in model.params'
60
+
61
+ if self.label_key not in __models__:
62
+ raise NotImplementedError()
63
+
64
+ self.load_classifier(ckpt_path, pool)
65
+
66
+ self.scheduler_config = scheduler_config
67
+ self.use_scheduler = self.scheduler_config is not None
68
+ self.weight_decay = weight_decay
69
+
70
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
71
+ sd = torch.load(path, map_location="cpu")
72
+ if "state_dict" in list(sd.keys()):
73
+ sd = sd["state_dict"]
74
+ keys = list(sd.keys())
75
+ for k in keys:
76
+ for ik in ignore_keys:
77
+ if k.startswith(ik):
78
+ print("Deleting key {} from state_dict.".format(k))
79
+ del sd[k]
80
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
81
+ sd, strict=False)
82
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
83
+ if len(missing) > 0:
84
+ print(f"Missing Keys: {missing}")
85
+ if len(unexpected) > 0:
86
+ print(f"Unexpected Keys: {unexpected}")
87
+
88
+ def load_diffusion(self):
89
+ model = instantiate_from_config(self.diffusion_config)
90
+ self.diffusion_model = model.eval()
91
+ self.diffusion_model.train = disabled_train
92
+ for param in self.diffusion_model.parameters():
93
+ param.requires_grad = False
94
+
95
+ def load_classifier(self, ckpt_path, pool):
96
+ model_config = deepcopy(self.diffusion_config.params.unet_config.params)
97
+ model_config.in_channels = self.diffusion_config.params.unet_config.params.out_channels
98
+ model_config.out_channels = self.num_classes
99
+ if self.label_key == 'class_label':
100
+ model_config.pool = pool
101
+
102
+ self.model = __models__[self.label_key](**model_config)
103
+ if ckpt_path is not None:
104
+ print('#####################################################################')
105
+ print(f'load from ckpt "{ckpt_path}"')
106
+ print('#####################################################################')
107
+ self.init_from_ckpt(ckpt_path)
108
+
109
+ @torch.no_grad()
110
+ def get_x_noisy(self, x, t, noise=None):
111
+ noise = default(noise, lambda: torch.randn_like(x))
112
+ continuous_sqrt_alpha_cumprod = None
113
+ if self.diffusion_model.use_continuous_noise:
114
+ continuous_sqrt_alpha_cumprod = self.diffusion_model.sample_continuous_noise_level(x.shape[0], t + 1)
115
+ # todo: make sure t+1 is correct here
116
+
117
+ return self.diffusion_model.q_sample(x_start=x, t=t, noise=noise,
118
+ continuous_sqrt_alpha_cumprod=continuous_sqrt_alpha_cumprod)
119
+
120
+ def forward(self, x_noisy, t, *args, **kwargs):
121
+ return self.model(x_noisy, t)
122
+
123
+ @torch.no_grad()
124
+ def get_input(self, batch, k):
125
+ x = batch[k]
126
+ if len(x.shape) == 3:
127
+ x = x[..., None]
128
+ x = rearrange(x, 'b h w c -> b c h w')
129
+ x = x.to(memory_format=torch.contiguous_format).float()
130
+ return x
131
+
132
+ @torch.no_grad()
133
+ def get_conditioning(self, batch, k=None):
134
+ if k is None:
135
+ k = self.label_key
136
+ assert k is not None, 'Needs to provide label key'
137
+
138
+ targets = batch[k].to(self.device)
139
+
140
+ if self.label_key == 'segmentation':
141
+ targets = rearrange(targets, 'b h w c -> b c h w')
142
+ for down in range(self.numd):
143
+ h, w = targets.shape[-2:]
144
+ targets = F.interpolate(targets, size=(h // 2, w // 2), mode='nearest')
145
+
146
+ # targets = rearrange(targets,'b c h w -> b h w c')
147
+
148
+ return targets
149
+
150
+ def compute_top_k(self, logits, labels, k, reduction="mean"):
151
+ _, top_ks = torch.topk(logits, k, dim=1)
152
+ if reduction == "mean":
153
+ return (top_ks == labels[:, None]).float().sum(dim=-1).mean().item()
154
+ elif reduction == "none":
155
+ return (top_ks == labels[:, None]).float().sum(dim=-1)
156
+
157
+ def on_train_epoch_start(self):
158
+ # save some memory
159
+ self.diffusion_model.model.to('cpu')
160
+
161
+ @torch.no_grad()
162
+ def write_logs(self, loss, logits, targets):
163
+ log_prefix = 'train' if self.training else 'val'
164
+ log = {}
165
+ log[f"{log_prefix}/loss"] = loss.mean()
166
+ log[f"{log_prefix}/acc@1"] = self.compute_top_k(
167
+ logits, targets, k=1, reduction="mean"
168
+ )
169
+ log[f"{log_prefix}/acc@5"] = self.compute_top_k(
170
+ logits, targets, k=5, reduction="mean"
171
+ )
172
+
173
+ self.log_dict(log, prog_bar=False, logger=True, on_step=self.training, on_epoch=True)
174
+ self.log('loss', log[f"{log_prefix}/loss"], prog_bar=True, logger=False)
175
+ self.log('global_step', self.global_step, logger=False, on_epoch=False, prog_bar=True)
176
+ lr = self.optimizers().param_groups[0]['lr']
177
+ self.log('lr_abs', lr, on_step=True, logger=True, on_epoch=False, prog_bar=True)
178
+
179
+ def shared_step(self, batch, t=None):
180
+ x, *_ = self.diffusion_model.get_input(batch, k=self.diffusion_model.first_stage_key)
181
+ targets = self.get_conditioning(batch)
182
+ if targets.dim() == 4:
183
+ targets = targets.argmax(dim=1)
184
+ if t is None:
185
+ t = torch.randint(0, self.diffusion_model.num_timesteps, (x.shape[0],), device=self.device).long()
186
+ else:
187
+ t = torch.full(size=(x.shape[0],), fill_value=t, device=self.device).long()
188
+ x_noisy = self.get_x_noisy(x, t)
189
+ logits = self(x_noisy, t)
190
+
191
+ loss = F.cross_entropy(logits, targets, reduction='none')
192
+
193
+ self.write_logs(loss.detach(), logits.detach(), targets.detach())
194
+
195
+ loss = loss.mean()
196
+ return loss, logits, x_noisy, targets
197
+
198
+ def training_step(self, batch, batch_idx):
199
+ loss, *_ = self.shared_step(batch)
200
+ return loss
201
+
202
+ def reset_noise_accs(self):
203
+ self.noisy_acc = {t: {'acc@1': [], 'acc@5': []} for t in
204
+ range(0, self.diffusion_model.num_timesteps, self.diffusion_model.log_every_t)}
205
+
206
+ def on_validation_start(self):
207
+ self.reset_noise_accs()
208
+
209
+ @torch.no_grad()
210
+ def validation_step(self, batch, batch_idx):
211
+ loss, *_ = self.shared_step(batch)
212
+
213
+ for t in self.noisy_acc:
214
+ _, logits, _, targets = self.shared_step(batch, t)
215
+ self.noisy_acc[t]['acc@1'].append(self.compute_top_k(logits, targets, k=1, reduction='mean'))
216
+ self.noisy_acc[t]['acc@5'].append(self.compute_top_k(logits, targets, k=5, reduction='mean'))
217
+
218
+ return loss
219
+
220
+ def configure_optimizers(self):
221
+ optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
222
+
223
+ if self.use_scheduler:
224
+ scheduler = instantiate_from_config(self.scheduler_config)
225
+
226
+ print("Setting up LambdaLR scheduler...")
227
+ scheduler = [
228
+ {
229
+ 'scheduler': LambdaLR(optimizer, lr_lambda=scheduler.schedule),
230
+ 'interval': 'step',
231
+ 'frequency': 1
232
+ }]
233
+ return [optimizer], scheduler
234
+
235
+ return optimizer
236
+
237
+ @torch.no_grad()
238
+ def log_images(self, batch, N=8, *args, **kwargs):
239
+ log = dict()
240
+ x = self.get_input(batch, self.diffusion_model.first_stage_key)
241
+ log['inputs'] = x
242
+
243
+ y = self.get_conditioning(batch)
244
+
245
+ if self.label_key == 'class_label':
246
+ y = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"])
247
+ log['labels'] = y
248
+
249
+ if ismap(y):
250
+ log['labels'] = self.diffusion_model.to_rgb(y)
251
+
252
+ for step in range(self.log_steps):
253
+ current_time = step * self.log_time_interval
254
+
255
+ _, logits, x_noisy, _ = self.shared_step(batch, t=current_time)
256
+
257
+ log[f'inputs@t{current_time}'] = x_noisy
258
+
259
+ pred = F.one_hot(logits.argmax(dim=1), num_classes=self.num_classes)
260
+ pred = rearrange(pred, 'b h w c -> b c h w')
261
+
262
+ log[f'pred@t{current_time}'] = self.diffusion_model.to_rgb(pred)
263
+
264
+ for key in log:
265
+ log[key] = log[key][:N]
266
+
267
+ return log
ldm/models/diffusion/ddim.py ADDED
@@ -0,0 +1,325 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+ from functools import partial
7
+ from einops import rearrange
8
+
9
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
10
+ from ldm.models.diffusion.sampling_util import renorm_thresholding, norm_thresholding, spatial_norm_thresholding
11
+
12
+
13
+ class DDIMSampler(object):
14
+ def __init__(self, model, schedule="linear", **kwargs):
15
+ super().__init__()
16
+ self.model = model
17
+ self.ddpm_num_timesteps = model.num_timesteps
18
+ self.schedule = schedule
19
+
20
+ def to(self, device):
21
+ """Same as to in torch module
22
+ Don't really underestand why this isn't a module in the first place"""
23
+ for k, v in self.__dict__.items():
24
+ if isinstance(v, torch.Tensor):
25
+ new_v = getattr(self, k).to(device)
26
+ setattr(self, k, new_v)
27
+
28
+
29
+ def register_buffer(self, name, attr):
30
+ if type(attr) == torch.Tensor:
31
+ if attr.device != torch.device("cuda"):
32
+ attr = attr.to(torch.device("cuda"))
33
+ setattr(self, name, attr)
34
+
35
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
36
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
37
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
38
+ alphas_cumprod = self.model.alphas_cumprod
39
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
40
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
41
+
42
+ self.register_buffer('betas', to_torch(self.model.betas))
43
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
44
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
45
+
46
+ # calculations for diffusion q(x_t | x_{t-1}) and others
47
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
48
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
49
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
50
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
51
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
52
+
53
+ # ddim sampling parameters
54
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
55
+ ddim_timesteps=self.ddim_timesteps,
56
+ eta=ddim_eta,verbose=verbose)
57
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
58
+ self.register_buffer('ddim_alphas', ddim_alphas)
59
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
60
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
61
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
62
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
63
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
64
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
65
+
66
+ @torch.no_grad()
67
+ def sample(self,
68
+ S,
69
+ batch_size,
70
+ shape,
71
+ conditioning=None,
72
+ callback=None,
73
+ normals_sequence=None,
74
+ img_callback=None,
75
+ quantize_x0=False,
76
+ eta=0.,
77
+ mask=None,
78
+ x0=None,
79
+ temperature=1.,
80
+ noise_dropout=0.,
81
+ score_corrector=None,
82
+ corrector_kwargs=None,
83
+ verbose=True,
84
+ x_T=None,
85
+ log_every_t=100,
86
+ unconditional_guidance_scale=1.,
87
+ unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
88
+ dynamic_threshold=None,
89
+ **kwargs
90
+ ):
91
+ if conditioning is not None:
92
+ if isinstance(conditioning, dict):
93
+ ctmp = conditioning[list(conditioning.keys())[0]]
94
+ while isinstance(ctmp, list): ctmp = ctmp[0]
95
+ cbs = ctmp.shape[0]
96
+ if cbs != batch_size:
97
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
98
+
99
+ else:
100
+ if conditioning.shape[0] != batch_size:
101
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
102
+
103
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
104
+ # sampling
105
+ C, H, W = shape
106
+ size = (batch_size, C, H, W)
107
+ print(f'Data shape for DDIM sampling is {size}, eta {eta}')
108
+
109
+ samples, intermediates = self.ddim_sampling(conditioning, size,
110
+ callback=callback,
111
+ img_callback=img_callback,
112
+ quantize_denoised=quantize_x0,
113
+ mask=mask, x0=x0,
114
+ ddim_use_original_steps=False,
115
+ noise_dropout=noise_dropout,
116
+ temperature=temperature,
117
+ score_corrector=score_corrector,
118
+ corrector_kwargs=corrector_kwargs,
119
+ x_T=x_T,
120
+ log_every_t=log_every_t,
121
+ unconditional_guidance_scale=unconditional_guidance_scale,
122
+ unconditional_conditioning=unconditional_conditioning,
123
+ dynamic_threshold=dynamic_threshold,
124
+ )
125
+ return samples, intermediates
126
+
127
+ @torch.no_grad()
128
+ def ddim_sampling(self, cond, shape,
129
+ x_T=None, ddim_use_original_steps=False,
130
+ callback=None, timesteps=None, quantize_denoised=False,
131
+ mask=None, x0=None, img_callback=None, log_every_t=100,
132
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
133
+ unconditional_guidance_scale=1., unconditional_conditioning=None, dynamic_threshold=None,
134
+ t_start=-1):
135
+ device = self.model.betas.device
136
+ b = shape[0]
137
+ if x_T is None:
138
+ img = torch.randn(shape, device=device)
139
+ else:
140
+ img = x_T
141
+
142
+ if timesteps is None:
143
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
144
+ elif timesteps is not None and not ddim_use_original_steps:
145
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
146
+ timesteps = self.ddim_timesteps[:subset_end]
147
+
148
+ timesteps = timesteps[:t_start]
149
+
150
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
151
+ time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
152
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
153
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
154
+
155
+ iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
156
+
157
+ for i, step in enumerate(iterator):
158
+ index = total_steps - i - 1
159
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
160
+
161
+ if mask is not None:
162
+ assert x0 is not None
163
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
164
+ img = img_orig * mask + (1. - mask) * img
165
+
166
+ outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
167
+ quantize_denoised=quantize_denoised, temperature=temperature,
168
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
169
+ corrector_kwargs=corrector_kwargs,
170
+ unconditional_guidance_scale=unconditional_guidance_scale,
171
+ unconditional_conditioning=unconditional_conditioning,
172
+ dynamic_threshold=dynamic_threshold)
173
+ img, pred_x0 = outs
174
+ if callback:
175
+ img = callback(i, img, pred_x0)
176
+ if img_callback: img_callback(pred_x0, i)
177
+
178
+ if index % log_every_t == 0 or index == total_steps - 1:
179
+ intermediates['x_inter'].append(img)
180
+ intermediates['pred_x0'].append(pred_x0)
181
+
182
+ return img, intermediates
183
+
184
+ @torch.no_grad()
185
+ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
186
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
187
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
188
+ dynamic_threshold=None, **kwargs):
189
+ b, *_, device = *x.shape, x.device
190
+
191
+
192
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
193
+ e_t = self.model.apply_model(x, t, c)
194
+ else:
195
+ x_in = torch.cat([x] * 2)
196
+ t_in = torch.cat([t] * 2)
197
+ if isinstance(c, dict):
198
+ assert isinstance(unconditional_conditioning, dict)
199
+ c_in = dict()
200
+ for k in c:
201
+ if isinstance(c[k], list):
202
+ c_in[k] = [torch.cat([
203
+ unconditional_conditioning[k][i],
204
+ c[k][i]]) for i in range(len(c[k]))]
205
+ else:
206
+ c_in[k] = torch.cat([
207
+ unconditional_conditioning[k],
208
+ c[k]])
209
+ else:
210
+ c_in = torch.cat([unconditional_conditioning, c])
211
+ e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
212
+ e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
213
+
214
+ if score_corrector is not None:
215
+ assert self.model.parameterization == "eps"
216
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
217
+
218
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
219
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
220
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
221
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
222
+ # select parameters corresponding to the currently considered timestep
223
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
224
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
225
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
226
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
227
+
228
+ # current prediction for x_0
229
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
230
+ if quantize_denoised:
231
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
232
+
233
+ if dynamic_threshold is not None:
234
+ pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
235
+
236
+ # direction pointing to x_t
237
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
238
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
239
+ if noise_dropout > 0.:
240
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
241
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
242
+ return x_prev, pred_x0
243
+
244
+ @torch.no_grad()
245
+ def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
246
+ unconditional_guidance_scale=1.0, unconditional_conditioning=None):
247
+ num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
248
+
249
+ assert t_enc <= num_reference_steps
250
+ num_steps = t_enc
251
+
252
+ if use_original_steps:
253
+ alphas_next = self.alphas_cumprod[:num_steps]
254
+ alphas = self.alphas_cumprod_prev[:num_steps]
255
+ else:
256
+ alphas_next = self.ddim_alphas[:num_steps]
257
+ alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
258
+
259
+ x_next = x0
260
+ intermediates = []
261
+ inter_steps = []
262
+ for i in tqdm(range(num_steps), desc='Encoding Image'):
263
+ t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)
264
+ if unconditional_guidance_scale == 1.:
265
+ noise_pred = self.model.apply_model(x_next, t, c)
266
+ else:
267
+ assert unconditional_conditioning is not None
268
+ e_t_uncond, noise_pred = torch.chunk(
269
+ self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),
270
+ torch.cat((unconditional_conditioning, c))), 2)
271
+ noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)
272
+
273
+ xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
274
+ weighted_noise_pred = alphas_next[i].sqrt() * (
275
+ (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
276
+ x_next = xt_weighted + weighted_noise_pred
277
+ if return_intermediates and i % (
278
+ num_steps // return_intermediates) == 0 and i < num_steps - 1:
279
+ intermediates.append(x_next)
280
+ inter_steps.append(i)
281
+ elif return_intermediates and i >= num_steps - 2:
282
+ intermediates.append(x_next)
283
+ inter_steps.append(i)
284
+
285
+ out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
286
+ if return_intermediates:
287
+ out.update({'intermediates': intermediates})
288
+ return x_next, out
289
+
290
+ @torch.no_grad()
291
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
292
+ # fast, but does not allow for exact reconstruction
293
+ # t serves as an index to gather the correct alphas
294
+ if use_original_steps:
295
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
296
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
297
+ else:
298
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
299
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
300
+
301
+ if noise is None:
302
+ noise = torch.randn_like(x0)
303
+ return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
304
+ extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
305
+
306
+ @torch.no_grad()
307
+ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
308
+ use_original_steps=False):
309
+
310
+ timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
311
+ timesteps = timesteps[:t_start]
312
+
313
+ time_range = np.flip(timesteps)
314
+ total_steps = timesteps.shape[0]
315
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
316
+
317
+ iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
318
+ x_dec = x_latent
319
+ for i, step in enumerate(iterator):
320
+ index = total_steps - i - 1
321
+ ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
322
+ x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
323
+ unconditional_guidance_scale=unconditional_guidance_scale,
324
+ unconditional_conditioning=unconditional_conditioning)
325
+ return x_dec
ldm/models/diffusion/ddpm.py ADDED
@@ -0,0 +1,1502 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ wild mixture of
3
+ https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
4
+ https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py
5
+ https://github.com/CompVis/taming-transformers
6
+ -- merci
7
+ """
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import numpy as np
12
+ import pytorch_lightning as pl
13
+ from torch.optim.lr_scheduler import LambdaLR
14
+ from einops import rearrange, repeat
15
+ from contextlib import contextmanager, nullcontext
16
+ from functools import partial
17
+ import itertools
18
+ from tqdm import tqdm
19
+ from torchvision.utils import make_grid
20
+ from pytorch_lightning.utilities.distributed import rank_zero_only
21
+ from omegaconf import ListConfig
22
+
23
+ from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
24
+ from ldm.modules.ema import LitEma
25
+ from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
26
+ from ldm.models.autoencoder import VQModelInterface, IdentityFirstStage, AutoencoderKL
27
+ from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
28
+ from ldm.models.diffusion.ddim import DDIMSampler
29
+ from ldm.modules.attention import CrossAttention
30
+
31
+
32
+ __conditioning_keys__ = {'concat': 'c_concat',
33
+ 'crossattn': 'c_crossattn',
34
+ 'adm': 'y'}
35
+
36
+
37
+ def disabled_train(self, mode=True):
38
+ """Overwrite model.train with this function to make sure train/eval mode
39
+ does not change anymore."""
40
+ return self
41
+
42
+
43
+ def uniform_on_device(r1, r2, shape, device):
44
+ return (r1 - r2) * torch.rand(*shape, device=device) + r2
45
+
46
+
47
+ class DDPM(pl.LightningModule):
48
+ # classic DDPM with Gaussian diffusion, in image space
49
+ def __init__(self,
50
+ unet_config,
51
+ timesteps=1000,
52
+ beta_schedule="linear",
53
+ loss_type="l2",
54
+ ckpt_path=None,
55
+ ignore_keys=[],
56
+ load_only_unet=False,
57
+ monitor="val/loss",
58
+ use_ema=True,
59
+ first_stage_key="image",
60
+ image_size=256,
61
+ channels=3,
62
+ log_every_t=100,
63
+ clip_denoised=True,
64
+ linear_start=1e-4,
65
+ linear_end=2e-2,
66
+ cosine_s=8e-3,
67
+ given_betas=None,
68
+ original_elbo_weight=0.,
69
+ v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
70
+ l_simple_weight=1.,
71
+ conditioning_key=None,
72
+ parameterization="eps", # all assuming fixed variance schedules
73
+ scheduler_config=None,
74
+ use_positional_encodings=False,
75
+ learn_logvar=False,
76
+ logvar_init=0.,
77
+ make_it_fit=False,
78
+ ucg_training=None,
79
+ ):
80
+ super().__init__()
81
+ assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"'
82
+ self.parameterization = parameterization
83
+ print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
84
+ self.cond_stage_model = None
85
+ self.clip_denoised = clip_denoised
86
+ self.log_every_t = log_every_t
87
+ self.first_stage_key = first_stage_key
88
+ self.image_size = image_size # try conv?
89
+ self.channels = channels
90
+ self.use_positional_encodings = use_positional_encodings
91
+ self.model = DiffusionWrapper(unet_config, conditioning_key)
92
+ count_params(self.model, verbose=True)
93
+ self.use_ema = use_ema
94
+ if self.use_ema:
95
+ self.model_ema = LitEma(self.model)
96
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
97
+
98
+ self.use_scheduler = scheduler_config is not None
99
+ if self.use_scheduler:
100
+ self.scheduler_config = scheduler_config
101
+
102
+ self.v_posterior = v_posterior
103
+ self.original_elbo_weight = original_elbo_weight
104
+ self.l_simple_weight = l_simple_weight
105
+
106
+ if monitor is not None:
107
+ self.monitor = monitor
108
+ self.make_it_fit = make_it_fit
109
+ if ckpt_path is not None:
110
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
111
+
112
+ self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
113
+ linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
114
+
115
+ self.loss_type = loss_type
116
+
117
+ self.learn_logvar = learn_logvar
118
+ self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
119
+ if self.learn_logvar:
120
+ self.logvar = nn.Parameter(self.logvar, requires_grad=True)
121
+
122
+ self.ucg_training = ucg_training or dict()
123
+ if self.ucg_training:
124
+ self.ucg_prng = np.random.RandomState()
125
+
126
+ def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
127
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
128
+ if exists(given_betas):
129
+ betas = given_betas
130
+ else:
131
+ betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
132
+ cosine_s=cosine_s)
133
+ alphas = 1. - betas
134
+ alphas_cumprod = np.cumprod(alphas, axis=0)
135
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
136
+
137
+ timesteps, = betas.shape
138
+ self.num_timesteps = int(timesteps)
139
+ self.linear_start = linear_start
140
+ self.linear_end = linear_end
141
+ assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
142
+
143
+ to_torch = partial(torch.tensor, dtype=torch.float32)
144
+
145
+ self.register_buffer('betas', to_torch(betas))
146
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
147
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
148
+
149
+ # calculations for diffusion q(x_t | x_{t-1}) and others
150
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
151
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
152
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
153
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
154
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
155
+
156
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
157
+ posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
158
+ 1. - alphas_cumprod) + self.v_posterior * betas
159
+ # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
160
+ self.register_buffer('posterior_variance', to_torch(posterior_variance))
161
+ # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
162
+ self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
163
+ self.register_buffer('posterior_mean_coef1', to_torch(
164
+ betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
165
+ self.register_buffer('posterior_mean_coef2', to_torch(
166
+ (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
167
+
168
+ if self.parameterization == "eps":
169
+ lvlb_weights = self.betas ** 2 / (
170
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))
171
+ elif self.parameterization == "x0":
172
+ lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod))
173
+ else:
174
+ raise NotImplementedError("mu not supported")
175
+ # TODO how to choose this term
176
+ lvlb_weights[0] = lvlb_weights[1]
177
+ self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
178
+ assert not torch.isnan(self.lvlb_weights).all()
179
+
180
+ @contextmanager
181
+ def ema_scope(self, context=None):
182
+ if self.use_ema:
183
+ self.model_ema.store(self.model.parameters())
184
+ self.model_ema.copy_to(self.model)
185
+ if context is not None:
186
+ print(f"{context}: Switched to EMA weights")
187
+ try:
188
+ yield None
189
+ finally:
190
+ if self.use_ema:
191
+ self.model_ema.restore(self.model.parameters())
192
+ if context is not None:
193
+ print(f"{context}: Restored training weights")
194
+
195
+ @torch.no_grad()
196
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
197
+ sd = torch.load(path, map_location="cpu")
198
+ if "state_dict" in list(sd.keys()):
199
+ sd = sd["state_dict"]
200
+ keys = list(sd.keys())
201
+
202
+ if self.make_it_fit:
203
+ n_params = len([name for name, _ in
204
+ itertools.chain(self.named_parameters(),
205
+ self.named_buffers())])
206
+ for name, param in tqdm(
207
+ itertools.chain(self.named_parameters(),
208
+ self.named_buffers()),
209
+ desc="Fitting old weights to new weights",
210
+ total=n_params
211
+ ):
212
+ if not name in sd:
213
+ continue
214
+ old_shape = sd[name].shape
215
+ new_shape = param.shape
216
+ assert len(old_shape)==len(new_shape)
217
+ if len(new_shape) > 2:
218
+ # we only modify first two axes
219
+ assert new_shape[2:] == old_shape[2:]
220
+ # assumes first axis corresponds to output dim
221
+ if not new_shape == old_shape:
222
+ new_param = param.clone()
223
+ old_param = sd[name]
224
+ if len(new_shape) == 1:
225
+ for i in range(new_param.shape[0]):
226
+ new_param[i] = old_param[i % old_shape[0]]
227
+ elif len(new_shape) >= 2:
228
+ for i in range(new_param.shape[0]):
229
+ for j in range(new_param.shape[1]):
230
+ new_param[i, j] = old_param[i % old_shape[0], j % old_shape[1]]
231
+
232
+ n_used_old = torch.ones(old_shape[1])
233
+ for j in range(new_param.shape[1]):
234
+ n_used_old[j % old_shape[1]] += 1
235
+ n_used_new = torch.zeros(new_shape[1])
236
+ for j in range(new_param.shape[1]):
237
+ n_used_new[j] = n_used_old[j % old_shape[1]]
238
+
239
+ n_used_new = n_used_new[None, :]
240
+ while len(n_used_new.shape) < len(new_shape):
241
+ n_used_new = n_used_new.unsqueeze(-1)
242
+ new_param /= n_used_new
243
+
244
+ sd[name] = new_param
245
+
246
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
247
+ sd, strict=False)
248
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
249
+ if len(missing) > 0:
250
+ print(f"Missing Keys: {missing}")
251
+ if len(unexpected) > 0:
252
+ print(f"Unexpected Keys: {unexpected}")
253
+
254
+ def q_mean_variance(self, x_start, t):
255
+ """
256
+ Get the distribution q(x_t | x_0).
257
+ :param x_start: the [N x C x ...] tensor of noiseless inputs.
258
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
259
+ :return: A tuple (mean, variance, log_variance), all of x_start's shape.
260
+ """
261
+ mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start)
262
+ variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
263
+ log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
264
+ return mean, variance, log_variance
265
+
266
+ def predict_start_from_noise(self, x_t, t, noise):
267
+ return (
268
+ extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
269
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
270
+ )
271
+
272
+ def q_posterior(self, x_start, x_t, t):
273
+ posterior_mean = (
274
+ extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
275
+ extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
276
+ )
277
+ posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
278
+ posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
279
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
280
+
281
+ def p_mean_variance(self, x, t, clip_denoised: bool):
282
+ model_out = self.model(x, t)
283
+ if self.parameterization == "eps":
284
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
285
+ elif self.parameterization == "x0":
286
+ x_recon = model_out
287
+ if clip_denoised:
288
+ x_recon.clamp_(-1., 1.)
289
+
290
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
291
+ return model_mean, posterior_variance, posterior_log_variance
292
+
293
+ @torch.no_grad()
294
+ def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
295
+ b, *_, device = *x.shape, x.device
296
+ model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
297
+ noise = noise_like(x.shape, device, repeat_noise)
298
+ # no noise when t == 0
299
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
300
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
301
+
302
+ @torch.no_grad()
303
+ def p_sample_loop(self, shape, return_intermediates=False):
304
+ device = self.betas.device
305
+ b = shape[0]
306
+ img = torch.randn(shape, device=device)
307
+ intermediates = [img]
308
+ for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
309
+ img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
310
+ clip_denoised=self.clip_denoised)
311
+ if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
312
+ intermediates.append(img)
313
+ if return_intermediates:
314
+ return img, intermediates
315
+ return img
316
+
317
+ @torch.no_grad()
318
+ def sample(self, batch_size=16, return_intermediates=False):
319
+ image_size = self.image_size
320
+ channels = self.channels
321
+ return self.p_sample_loop((batch_size, channels, image_size, image_size),
322
+ return_intermediates=return_intermediates)
323
+
324
+ def q_sample(self, x_start, t, noise=None):
325
+ noise = default(noise, lambda: torch.randn_like(x_start))
326
+ return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
327
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
328
+
329
+ def get_loss(self, pred, target, mean=True):
330
+ if self.loss_type == 'l1':
331
+ loss = (target - pred).abs()
332
+ if mean:
333
+ loss = loss.mean()
334
+ elif self.loss_type == 'l2':
335
+ if mean:
336
+ loss = torch.nn.functional.mse_loss(target, pred)
337
+ else:
338
+ loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
339
+ else:
340
+ raise NotImplementedError("unknown loss type '{loss_type}'")
341
+
342
+ return loss
343
+
344
+ def p_losses(self, x_start, t, noise=None):
345
+ noise = default(noise, lambda: torch.randn_like(x_start))
346
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
347
+ model_out = self.model(x_noisy, t)
348
+
349
+ loss_dict = {}
350
+ if self.parameterization == "eps":
351
+ target = noise
352
+ elif self.parameterization == "x0":
353
+ target = x_start
354
+ else:
355
+ raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported")
356
+
357
+ loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
358
+
359
+ log_prefix = 'train' if self.training else 'val'
360
+
361
+ loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
362
+ loss_simple = loss.mean() * self.l_simple_weight
363
+
364
+ loss_vlb = (self.lvlb_weights[t] * loss).mean()
365
+ loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
366
+
367
+ loss = loss_simple + self.original_elbo_weight * loss_vlb
368
+
369
+ loss_dict.update({f'{log_prefix}/loss': loss})
370
+
371
+ return loss, loss_dict
372
+
373
+ def forward(self, x, *args, **kwargs):
374
+ # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
375
+ # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
376
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
377
+ return self.p_losses(x, t, *args, **kwargs)
378
+
379
+ def get_input(self, batch, k):
380
+ x = batch[k]
381
+ if len(x.shape) == 3:
382
+ x = x[..., None]
383
+ x = rearrange(x, 'b h w c -> b c h w')
384
+ x = x.to(memory_format=torch.contiguous_format).float()
385
+ return x
386
+
387
+ def shared_step(self, batch):
388
+ x = self.get_input(batch, self.first_stage_key)
389
+ loss, loss_dict = self(x)
390
+ return loss, loss_dict
391
+
392
+ def training_step(self, batch, batch_idx):
393
+ for k in self.ucg_training:
394
+ p = self.ucg_training[k]["p"]
395
+ val = self.ucg_training[k]["val"]
396
+ if val is None:
397
+ val = ""
398
+ for i in range(len(batch[k])):
399
+ if self.ucg_prng.choice(2, p=[1-p, p]):
400
+ batch[k][i] = val
401
+
402
+ loss, loss_dict = self.shared_step(batch)
403
+
404
+ self.log_dict(loss_dict, prog_bar=True,
405
+ logger=True, on_step=True, on_epoch=True)
406
+
407
+ self.log("global_step", self.global_step,
408
+ prog_bar=True, logger=True, on_step=True, on_epoch=False)
409
+
410
+ if self.use_scheduler:
411
+ lr = self.optimizers().param_groups[0]['lr']
412
+ self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
413
+
414
+ return loss
415
+
416
+ @torch.no_grad()
417
+ def validation_step(self, batch, batch_idx):
418
+ _, loss_dict_no_ema = self.shared_step(batch)
419
+ with self.ema_scope():
420
+ _, loss_dict_ema = self.shared_step(batch)
421
+ loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
422
+ self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
423
+ self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
424
+
425
+ def on_train_batch_end(self, *args, **kwargs):
426
+ if self.use_ema:
427
+ self.model_ema(self.model)
428
+
429
+ def _get_rows_from_list(self, samples):
430
+ n_imgs_per_row = len(samples)
431
+ denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
432
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
433
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
434
+ return denoise_grid
435
+
436
+ @torch.no_grad()
437
+ def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
438
+ log = dict()
439
+ x = self.get_input(batch, self.first_stage_key)
440
+ N = min(x.shape[0], N)
441
+ n_row = min(x.shape[0], n_row)
442
+ x = x.to(self.device)[:N]
443
+ log["inputs"] = x
444
+
445
+ # get diffusion row
446
+ diffusion_row = list()
447
+ x_start = x[:n_row]
448
+
449
+ for t in range(self.num_timesteps):
450
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
451
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
452
+ t = t.to(self.device).long()
453
+ noise = torch.randn_like(x_start)
454
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
455
+ diffusion_row.append(x_noisy)
456
+
457
+ log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
458
+
459
+ if sample:
460
+ # get denoise row
461
+ with self.ema_scope("Plotting"):
462
+ samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
463
+
464
+ log["samples"] = samples
465
+ log["denoise_row"] = self._get_rows_from_list(denoise_row)
466
+
467
+ if return_keys:
468
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
469
+ return log
470
+ else:
471
+ return {key: log[key] for key in return_keys}
472
+ return log
473
+
474
+ def configure_optimizers(self):
475
+ lr = self.learning_rate
476
+ params = list(self.model.parameters())
477
+ if self.learn_logvar:
478
+ params = params + [self.logvar]
479
+ opt = torch.optim.AdamW(params, lr=lr)
480
+ return opt
481
+
482
+
483
+ class LatentDiffusion(DDPM):
484
+ """main class"""
485
+ def __init__(self,
486
+ first_stage_config,
487
+ cond_stage_config,
488
+ num_timesteps_cond=None,
489
+ cond_stage_key="image",
490
+ cond_stage_trainable=False,
491
+ concat_mode=True,
492
+ cond_stage_forward=None,
493
+ conditioning_key=None,
494
+ scale_factor=1.0,
495
+ scale_by_std=False,
496
+ unet_trainable=True,
497
+ *args, **kwargs):
498
+ self.num_timesteps_cond = default(num_timesteps_cond, 1)
499
+ self.scale_by_std = scale_by_std
500
+ assert self.num_timesteps_cond <= kwargs['timesteps']
501
+ # for backwards compatibility after implementation of DiffusionWrapper
502
+ if conditioning_key is None:
503
+ conditioning_key = 'concat' if concat_mode else 'crossattn'
504
+ if cond_stage_config == '__is_unconditional__':
505
+ conditioning_key = None
506
+ ckpt_path = kwargs.pop("ckpt_path", None)
507
+ ignore_keys = kwargs.pop("ignore_keys", [])
508
+ super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
509
+ self.concat_mode = concat_mode
510
+ self.cond_stage_trainable = cond_stage_trainable
511
+ self.unet_trainable = unet_trainable
512
+ self.cond_stage_key = cond_stage_key
513
+ try:
514
+ self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
515
+ except:
516
+ self.num_downs = 0
517
+ if not scale_by_std:
518
+ self.scale_factor = scale_factor
519
+ else:
520
+ self.register_buffer('scale_factor', torch.tensor(scale_factor))
521
+ self.instantiate_first_stage(first_stage_config)
522
+ self.instantiate_cond_stage(cond_stage_config)
523
+ self.cond_stage_forward = cond_stage_forward
524
+
525
+ # construct linear projection layer for concatenating image CLIP embedding and RT
526
+ self.cc_projection = nn.Linear(772, 768)
527
+ nn.init.eye_(list(self.cc_projection.parameters())[0][:768, :768])
528
+ nn.init.zeros_(list(self.cc_projection.parameters())[1])
529
+ self.cc_projection.requires_grad_(True)
530
+
531
+ self.clip_denoised = False
532
+ self.bbox_tokenizer = None
533
+
534
+ self.restarted_from_ckpt = False
535
+ if ckpt_path is not None:
536
+ self.init_from_ckpt(ckpt_path, ignore_keys)
537
+ self.restarted_from_ckpt = True
538
+
539
+ def make_cond_schedule(self, ):
540
+ self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
541
+ ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
542
+ self.cond_ids[:self.num_timesteps_cond] = ids
543
+
544
+ @rank_zero_only
545
+ @torch.no_grad()
546
+ def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
547
+ # only for very first batch
548
+ if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
549
+ assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
550
+ # set rescale weight to 1./std of encodings
551
+ print("### USING STD-RESCALING ###")
552
+ x = super().get_input(batch, self.first_stage_key)
553
+ x = x.to(self.device)
554
+ encoder_posterior = self.encode_first_stage(x)
555
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
556
+ del self.scale_factor
557
+ self.register_buffer('scale_factor', 1. / z.flatten().std())
558
+ print(f"setting self.scale_factor to {self.scale_factor}")
559
+ print("### USING STD-RESCALING ###")
560
+
561
+ def register_schedule(self,
562
+ given_betas=None, beta_schedule="linear", timesteps=1000,
563
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
564
+ super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
565
+
566
+ self.shorten_cond_schedule = self.num_timesteps_cond > 1
567
+ if self.shorten_cond_schedule:
568
+ self.make_cond_schedule()
569
+
570
+ def instantiate_first_stage(self, config):
571
+ model = instantiate_from_config(config)
572
+ self.first_stage_model = model.eval()
573
+ self.first_stage_model.train = disabled_train
574
+ for param in self.first_stage_model.parameters():
575
+ param.requires_grad = False
576
+
577
+ def instantiate_cond_stage(self, config):
578
+ if not self.cond_stage_trainable:
579
+ if config == "__is_first_stage__":
580
+ print("Using first stage also as cond stage.")
581
+ self.cond_stage_model = self.first_stage_model
582
+ elif config == "__is_unconditional__":
583
+ print(f"Training {self.__class__.__name__} as an unconditional model.")
584
+ self.cond_stage_model = None
585
+ # self.be_unconditional = True
586
+ else:
587
+ model = instantiate_from_config(config)
588
+ self.cond_stage_model = model.eval()
589
+ self.cond_stage_model.train = disabled_train
590
+ for param in self.cond_stage_model.parameters():
591
+ param.requires_grad = False
592
+ else:
593
+ assert config != '__is_first_stage__'
594
+ assert config != '__is_unconditional__'
595
+ model = instantiate_from_config(config)
596
+ self.cond_stage_model = model
597
+
598
+ def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False):
599
+ denoise_row = []
600
+ for zd in tqdm(samples, desc=desc):
601
+ denoise_row.append(self.decode_first_stage(zd.to(self.device),
602
+ force_not_quantize=force_no_decoder_quantization))
603
+ n_imgs_per_row = len(denoise_row)
604
+ denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W
605
+ denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
606
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
607
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
608
+ return denoise_grid
609
+
610
+ def get_first_stage_encoding(self, encoder_posterior):
611
+ if isinstance(encoder_posterior, DiagonalGaussianDistribution):
612
+ z = encoder_posterior.sample()
613
+ elif isinstance(encoder_posterior, torch.Tensor):
614
+ z = encoder_posterior
615
+ else:
616
+ raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
617
+ return self.scale_factor * z
618
+
619
+ def get_learned_conditioning(self, c):
620
+ if self.cond_stage_forward is None:
621
+ if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
622
+ c = self.cond_stage_model.encode(c)
623
+ if isinstance(c, DiagonalGaussianDistribution):
624
+ c = c.mode()
625
+ else:
626
+ c = self.cond_stage_model(c)
627
+ else:
628
+ assert hasattr(self.cond_stage_model, self.cond_stage_forward)
629
+ c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
630
+ return c
631
+
632
+
633
+ def meshgrid(self, h, w):
634
+ y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
635
+ x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
636
+
637
+ arr = torch.cat([y, x], dim=-1)
638
+ return arr
639
+
640
+ def delta_border(self, h, w):
641
+ """
642
+ :param h: height
643
+ :param w: width
644
+ :return: normalized distance to image border,
645
+ wtith min distance = 0 at border and max dist = 0.5 at image center
646
+ """
647
+ lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
648
+ arr = self.meshgrid(h, w) / lower_right_corner
649
+ dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
650
+ dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
651
+ edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
652
+ return edge_dist
653
+
654
+ def get_weighting(self, h, w, Ly, Lx, device):
655
+ weighting = self.delta_border(h, w)
656
+ weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"],
657
+ self.split_input_params["clip_max_weight"], )
658
+ weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
659
+
660
+ if self.split_input_params["tie_braker"]:
661
+ L_weighting = self.delta_border(Ly, Lx)
662
+ L_weighting = torch.clip(L_weighting,
663
+ self.split_input_params["clip_min_tie_weight"],
664
+ self.split_input_params["clip_max_tie_weight"])
665
+
666
+ L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
667
+ weighting = weighting * L_weighting
668
+ return weighting
669
+
670
+ def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code
671
+ """
672
+ :param x: img of size (bs, c, h, w)
673
+ :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
674
+ """
675
+ bs, nc, h, w = x.shape
676
+
677
+ # number of crops in image
678
+ Ly = (h - kernel_size[0]) // stride[0] + 1
679
+ Lx = (w - kernel_size[1]) // stride[1] + 1
680
+
681
+ if uf == 1 and df == 1:
682
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
683
+ unfold = torch.nn.Unfold(**fold_params)
684
+
685
+ fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
686
+
687
+ weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype)
688
+ normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap
689
+ weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
690
+
691
+ elif uf > 1 and df == 1:
692
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
693
+ unfold = torch.nn.Unfold(**fold_params)
694
+
695
+ fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
696
+ dilation=1, padding=0,
697
+ stride=(stride[0] * uf, stride[1] * uf))
698
+ fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
699
+
700
+ weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype)
701
+ normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap
702
+ weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
703
+
704
+ elif df > 1 and uf == 1:
705
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
706
+ unfold = torch.nn.Unfold(**fold_params)
707
+
708
+ fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
709
+ dilation=1, padding=0,
710
+ stride=(stride[0] // df, stride[1] // df))
711
+ fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
712
+
713
+ weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype)
714
+ normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap
715
+ weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
716
+
717
+ else:
718
+ raise NotImplementedError
719
+
720
+ return fold, unfold, normalization, weighting
721
+
722
+
723
+
724
+ @torch.no_grad()
725
+ def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
726
+ cond_key=None, return_original_cond=False, bs=None, uncond=0.05, ):
727
+ x = super().get_input(batch, k)
728
+ T = batch['T'].to(memory_format=torch.contiguous_format).float().to(self.device)
729
+
730
+ if bs is not None:
731
+ x = x[:bs]
732
+ T = T[:bs].to(self.device)
733
+
734
+ x = x.to(self.device)
735
+ encoder_posterior = self.encode_first_stage(x)
736
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
737
+ cond_key = cond_key or self.cond_stage_key
738
+ xc = super().get_input(batch, cond_key).to(self.device)
739
+ if bs is not None:
740
+ xc = xc[:bs]
741
+ cond = {}
742
+
743
+ # To support classifier-free guidance, randomly drop out only text conditioning 5%, only image conditioning 5%, and both 5%.
744
+ random = torch.rand(x.size(0), device=x.device)
745
+ prompt_mask = rearrange(random < 2 * uncond, "n -> n 1 1")
746
+ input_mask = 1 - rearrange((random >= uncond).float() * (random < 3 * uncond).float(), "n -> n 1 1 1")
747
+ null_prompt = self.get_learned_conditioning([""])
748
+
749
+ # z.shape: [8, 4, 32, 32]; c.shape: [8, 1, 768]
750
+ # print('=========== xc shape ===========', xc.shape)
751
+ with torch.enable_grad():
752
+ clip_emb = self.get_learned_conditioning(xc).detach()
753
+ null_prompt = self.get_learned_conditioning([""]).detach()
754
+ cond["c_crossattn"] = [self.cc_projection(torch.cat([torch.where(prompt_mask, null_prompt, clip_emb), T[:, None, :]], dim=-1))]
755
+
756
+ text_only = batch["text_only"] # (B,1)
757
+ cond["c_crossattn"][0] = torch.einsum("ijk, ij->ijk", cond["c_crossattn"][0], (1. - text_only))
758
+
759
+ if self.use_cond_concat:
760
+ xc_concat = super().get_input(batch, "image_cond_concat") #.to(self.device)
761
+ cond["c_concat"] = [input_mask * self.encode_first_stage((xc_concat.to(self.device))).mode().detach()]
762
+ else:
763
+ cond["c_concat"] = [input_mask * self.encode_first_stage((xc.to(self.device))).mode().detach()]
764
+
765
+ if self.use_bbox_mask:
766
+ ## bbox_mask
767
+ bbox_mask = batch["bbox_mask"].to(self.device)
768
+ cond["c_concat"].append(bbox_mask)
769
+ ###
770
+
771
+ if self.use_bg_inpainting:
772
+ bg_img = super().get_input(batch, "bg_concat")
773
+ cond["c_concat"].append(self.encode_first_stage((bg_img.to(self.device))).mode().detach())
774
+
775
+
776
+ out = [z, cond]
777
+ if return_first_stage_outputs:
778
+ xrec = self.decode_first_stage(z)
779
+ out.extend([x, xrec])
780
+ if return_original_cond:
781
+ out.append(xc)
782
+ return out
783
+
784
+
785
+
786
+
787
+ # @torch.no_grad()
788
+ def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
789
+ if predict_cids:
790
+ if z.dim() == 4:
791
+ z = torch.argmax(z.exp(), dim=1).long()
792
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
793
+ z = rearrange(z, 'b h w c -> b c h w').contiguous()
794
+
795
+ z = 1. / self.scale_factor * z
796
+
797
+ if hasattr(self, "split_input_params"):
798
+ if self.split_input_params["patch_distributed_vq"]:
799
+ ks = self.split_input_params["ks"] # eg. (128, 128)
800
+ stride = self.split_input_params["stride"] # eg. (64, 64)
801
+ uf = self.split_input_params["vqf"]
802
+ bs, nc, h, w = z.shape
803
+ if ks[0] > h or ks[1] > w:
804
+ ks = (min(ks[0], h), min(ks[1], w))
805
+ print("reducing Kernel")
806
+
807
+ if stride[0] > h or stride[1] > w:
808
+ stride = (min(stride[0], h), min(stride[1], w))
809
+ print("reducing stride")
810
+
811
+ fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf)
812
+
813
+ z = unfold(z) # (bn, nc * prod(**ks), L)
814
+ # 1. Reshape to img shape
815
+ z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
816
+
817
+ # 2. apply model loop over last dim
818
+ if isinstance(self.first_stage_model, VQModelInterface):
819
+ output_list = [self.first_stage_model.decode(z[:, :, :, :, i],
820
+ force_not_quantize=predict_cids or force_not_quantize)
821
+ for i in range(z.shape[-1])]
822
+ else:
823
+
824
+ output_list = [self.first_stage_model.decode(z[:, :, :, :, i])
825
+ for i in range(z.shape[-1])]
826
+
827
+ o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L)
828
+ o = o * weighting
829
+ # Reverse 1. reshape to img shape
830
+ o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L)
831
+ # stitch crops together
832
+ decoded = fold(o)
833
+ decoded = decoded / normalization # norm is shape (1, 1, h, w)
834
+ return decoded
835
+ else:
836
+ if isinstance(self.first_stage_model, VQModelInterface):
837
+ return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
838
+ else:
839
+ return self.first_stage_model.decode(z)
840
+
841
+ else:
842
+ if isinstance(self.first_stage_model, VQModelInterface):
843
+ return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
844
+ else:
845
+ return self.first_stage_model.decode(z)
846
+
847
+ @torch.no_grad()
848
+ def encode_first_stage(self, x):
849
+ if hasattr(self, "split_input_params"):
850
+ if self.split_input_params["patch_distributed_vq"]:
851
+ ks = self.split_input_params["ks"] # eg. (128, 128)
852
+ stride = self.split_input_params["stride"] # eg. (64, 64)
853
+ df = self.split_input_params["vqf"]
854
+ self.split_input_params['original_image_size'] = x.shape[-2:]
855
+ bs, nc, h, w = x.shape
856
+ if ks[0] > h or ks[1] > w:
857
+ ks = (min(ks[0], h), min(ks[1], w))
858
+ print("reducing Kernel")
859
+
860
+ if stride[0] > h or stride[1] > w:
861
+ stride = (min(stride[0], h), min(stride[1], w))
862
+ print("reducing stride")
863
+
864
+ fold, unfold, normalization, weighting = self.get_fold_unfold(x, ks, stride, df=df)
865
+ z = unfold(x) # (bn, nc * prod(**ks), L)
866
+ # Reshape to img shape
867
+ z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
868
+
869
+ output_list = [self.first_stage_model.encode(z[:, :, :, :, i])
870
+ for i in range(z.shape[-1])]
871
+
872
+ o = torch.stack(output_list, axis=-1)
873
+ o = o * weighting
874
+
875
+ # Reverse reshape to img shape
876
+ o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L)
877
+ # stitch crops together
878
+ decoded = fold(o)
879
+ decoded = decoded / normalization
880
+ return decoded
881
+
882
+ else:
883
+ return self.first_stage_model.encode(x)
884
+ else:
885
+ return self.first_stage_model.encode(x)
886
+
887
+ def shared_step(self, batch, **kwargs):
888
+ x, c = self.get_input(batch, self.first_stage_key)
889
+ loss = self(x, c)
890
+ return loss
891
+
892
+ def forward(self, x, c, *args, **kwargs):
893
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
894
+ if self.model.conditioning_key is not None:
895
+ assert c is not None
896
+ # if self.cond_stage_trainable:
897
+ # c = self.get_learned_conditioning(c)
898
+ if self.shorten_cond_schedule: # TODO: drop this option
899
+ tc = self.cond_ids[t].to(self.device)
900
+ c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
901
+ return self.p_losses(x, c, t, *args, **kwargs)
902
+
903
+ def _rescale_annotations(self, bboxes, crop_coordinates): # TODO: move to dataset
904
+ def rescale_bbox(bbox):
905
+ x0 = clamp((bbox[0] - crop_coordinates[0]) / crop_coordinates[2])
906
+ y0 = clamp((bbox[1] - crop_coordinates[1]) / crop_coordinates[3])
907
+ w = min(bbox[2] / crop_coordinates[2], 1 - x0)
908
+ h = min(bbox[3] / crop_coordinates[3], 1 - y0)
909
+ return x0, y0, w, h
910
+
911
+ return [rescale_bbox(b) for b in bboxes]
912
+
913
+ def apply_model(self, x_noisy, t, cond, return_ids=False):
914
+
915
+ if isinstance(cond, dict):
916
+ # hybrid case, cond is exptected to be a dict
917
+ pass
918
+ else:
919
+ if not isinstance(cond, list):
920
+ cond = [cond]
921
+ key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
922
+ cond = {key: cond}
923
+
924
+ if hasattr(self, "split_input_params"):
925
+ assert len(cond) == 1 # todo can only deal with one conditioning atm
926
+ assert not return_ids
927
+ ks = self.split_input_params["ks"] # eg. (128, 128)
928
+ stride = self.split_input_params["stride"] # eg. (64, 64)
929
+
930
+ h, w = x_noisy.shape[-2:]
931
+
932
+ fold, unfold, normalization, weighting = self.get_fold_unfold(x_noisy, ks, stride)
933
+
934
+ z = unfold(x_noisy) # (bn, nc * prod(**ks), L)
935
+ # Reshape to img shape
936
+ z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
937
+ z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])]
938
+
939
+ if self.cond_stage_key in ["image", "LR_image", "segmentation",
940
+ 'bbox_img'] and self.model.conditioning_key: # todo check for completeness
941
+ c_key = next(iter(cond.keys())) # get key
942
+ c = next(iter(cond.values())) # get value
943
+ assert (len(c) == 1) # todo extend to list with more than one elem
944
+ c = c[0] # get element
945
+
946
+ c = unfold(c)
947
+ c = c.view((c.shape[0], -1, ks[0], ks[1], c.shape[-1])) # (bn, nc, ks[0], ks[1], L )
948
+
949
+ cond_list = [{c_key: [c[:, :, :, :, i]]} for i in range(c.shape[-1])]
950
+
951
+ elif self.cond_stage_key == 'coordinates_bbox':
952
+ assert 'original_image_size' in self.split_input_params, 'BoudingBoxRescaling is missing original_image_size'
953
+
954
+ # assuming padding of unfold is always 0 and its dilation is always 1
955
+ n_patches_per_row = int((w - ks[0]) / stride[0] + 1)
956
+ full_img_h, full_img_w = self.split_input_params['original_image_size']
957
+ # as we are operating on latents, we need the factor from the original image size to the
958
+ # spatial latent size to properly rescale the crops for regenerating the bbox annotations
959
+ num_downs = self.first_stage_model.encoder.num_resolutions - 1
960
+ rescale_latent = 2 ** (num_downs)
961
+
962
+ # get top left postions of patches as conforming for the bbbox tokenizer, therefore we
963
+ # need to rescale the tl patch coordinates to be in between (0,1)
964
+ tl_patch_coordinates = [(rescale_latent * stride[0] * (patch_nr % n_patches_per_row) / full_img_w,
965
+ rescale_latent * stride[1] * (patch_nr // n_patches_per_row) / full_img_h)
966
+ for patch_nr in range(z.shape[-1])]
967
+
968
+ # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w)
969
+ patch_limits = [(x_tl, y_tl,
970
+ rescale_latent * ks[0] / full_img_w,
971
+ rescale_latent * ks[1] / full_img_h) for x_tl, y_tl in tl_patch_coordinates]
972
+ # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates]
973
+
974
+ # tokenize crop coordinates for the bounding boxes of the respective patches
975
+ patch_limits_tknzd = [torch.LongTensor(self.bbox_tokenizer._crop_encoder(bbox))[None].to(self.device)
976
+ for bbox in patch_limits] # list of length l with tensors of shape (1, 2)
977
+ # cut tknzd crop position from conditioning
978
+ assert isinstance(cond, dict), 'cond must be dict to be fed into model'
979
+ cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device)
980
+
981
+ adapted_cond = torch.stack([torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd])
982
+ adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n')
983
+ adapted_cond = self.get_learned_conditioning(adapted_cond)
984
+ adapted_cond = rearrange(adapted_cond, '(l b) n d -> l b n d', l=z.shape[-1])
985
+
986
+ cond_list = [{'c_crossattn': [e]} for e in adapted_cond]
987
+
988
+ else:
989
+ cond_list = [cond for i in range(z.shape[-1])] # Todo make this more efficient
990
+
991
+ # apply model by loop over crops
992
+ output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])]
993
+ assert not isinstance(output_list[0],
994
+ tuple) # todo cant deal with multiple model outputs check this never happens
995
+
996
+ o = torch.stack(output_list, axis=-1)
997
+ o = o * weighting
998
+ # Reverse reshape to img shape
999
+ o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L)
1000
+ # stitch crops together
1001
+ x_recon = fold(o) / normalization
1002
+
1003
+ else:
1004
+ x_recon = self.model(x_noisy, t, **cond)
1005
+
1006
+ if isinstance(x_recon, tuple) and not return_ids:
1007
+ return x_recon[0]
1008
+ else:
1009
+ return x_recon
1010
+
1011
+ def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
1012
+ return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
1013
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
1014
+
1015
+ def _prior_bpd(self, x_start):
1016
+ """
1017
+ Get the prior KL term for the variational lower-bound, measured in
1018
+ bits-per-dim.
1019
+ This term can't be optimized, as it only depends on the encoder.
1020
+ :param x_start: the [N x C x ...] tensor of inputs.
1021
+ :return: a batch of [N] KL values (in bits), one per batch element.
1022
+ """
1023
+ batch_size = x_start.shape[0]
1024
+ t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
1025
+ qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
1026
+ kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
1027
+ return mean_flat(kl_prior) / np.log(2.0)
1028
+
1029
+ def p_losses(self, x_start, cond, t, noise=None, **kwargs):
1030
+ noise = default(noise, lambda: torch.randn_like(x_start))
1031
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
1032
+ model_output = self.apply_model(x_noisy, t, cond, **kwargs)
1033
+
1034
+ loss_dict = {}
1035
+ prefix = 'train' if self.training else 'val'
1036
+
1037
+ if self.parameterization == "x0":
1038
+ target = x_start
1039
+ elif self.parameterization == "eps":
1040
+ target = noise
1041
+ else:
1042
+ raise NotImplementedError()
1043
+
1044
+ loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
1045
+ loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
1046
+
1047
+ logvar_t = self.logvar[t].to(self.device)
1048
+ loss = loss_simple / torch.exp(logvar_t) + logvar_t
1049
+ # loss = loss_simple / torch.exp(self.logvar) + self.logvar
1050
+ if self.learn_logvar:
1051
+ loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
1052
+ loss_dict.update({'logvar': self.logvar.data.mean()})
1053
+
1054
+ loss = self.l_simple_weight * loss.mean()
1055
+
1056
+ loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3))
1057
+ loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
1058
+ loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
1059
+ loss += (self.original_elbo_weight * loss_vlb)
1060
+ loss_dict.update({f'{prefix}/loss': loss})
1061
+
1062
+ return loss, loss_dict
1063
+
1064
+ def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False,
1065
+ return_x0=False, score_corrector=None, corrector_kwargs=None):
1066
+ t_in = t
1067
+ model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
1068
+
1069
+ if score_corrector is not None:
1070
+ assert self.parameterization == "eps"
1071
+ model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs)
1072
+
1073
+ if return_codebook_ids:
1074
+ model_out, logits = model_out
1075
+
1076
+ if self.parameterization == "eps":
1077
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
1078
+ elif self.parameterization == "x0":
1079
+ x_recon = model_out
1080
+ else:
1081
+ raise NotImplementedError()
1082
+
1083
+ if clip_denoised:
1084
+ x_recon.clamp_(-1., 1.)
1085
+ if quantize_denoised:
1086
+ x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
1087
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
1088
+ if return_codebook_ids:
1089
+ return model_mean, posterior_variance, posterior_log_variance, logits
1090
+ elif return_x0:
1091
+ return model_mean, posterior_variance, posterior_log_variance, x_recon
1092
+ else:
1093
+ return model_mean, posterior_variance, posterior_log_variance
1094
+
1095
+ @torch.no_grad()
1096
+ def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
1097
+ return_codebook_ids=False, quantize_denoised=False, return_x0=False,
1098
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None):
1099
+ b, *_, device = *x.shape, x.device
1100
+ outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
1101
+ return_codebook_ids=return_codebook_ids,
1102
+ quantize_denoised=quantize_denoised,
1103
+ return_x0=return_x0,
1104
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
1105
+ if return_codebook_ids:
1106
+ raise DeprecationWarning("Support dropped.")
1107
+ model_mean, _, model_log_variance, logits = outputs
1108
+ elif return_x0:
1109
+ model_mean, _, model_log_variance, x0 = outputs
1110
+ else:
1111
+ model_mean, _, model_log_variance = outputs
1112
+
1113
+ noise = noise_like(x.shape, device, repeat_noise) * temperature
1114
+ if noise_dropout > 0.:
1115
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
1116
+ # no noise when t == 0
1117
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
1118
+
1119
+ if return_codebook_ids:
1120
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
1121
+ if return_x0:
1122
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
1123
+ else:
1124
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
1125
+
1126
+ @torch.no_grad()
1127
+ def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False,
1128
+ img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
1129
+ score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
1130
+ log_every_t=None):
1131
+ if not log_every_t:
1132
+ log_every_t = self.log_every_t
1133
+ timesteps = self.num_timesteps
1134
+ if batch_size is not None:
1135
+ b = batch_size if batch_size is not None else shape[0]
1136
+ shape = [batch_size] + list(shape)
1137
+ else:
1138
+ b = batch_size = shape[0]
1139
+ if x_T is None:
1140
+ img = torch.randn(shape, device=self.device)
1141
+ else:
1142
+ img = x_T
1143
+ intermediates = []
1144
+ if cond is not None:
1145
+ if isinstance(cond, dict):
1146
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
1147
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
1148
+ else:
1149
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
1150
+
1151
+ if start_T is not None:
1152
+ timesteps = min(timesteps, start_T)
1153
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation',
1154
+ total=timesteps) if verbose else reversed(
1155
+ range(0, timesteps))
1156
+ if type(temperature) == float:
1157
+ temperature = [temperature] * timesteps
1158
+
1159
+ for i in iterator:
1160
+ ts = torch.full((b,), i, device=self.device, dtype=torch.long)
1161
+ if self.shorten_cond_schedule:
1162
+ assert self.model.conditioning_key != 'hybrid'
1163
+ tc = self.cond_ids[ts].to(cond.device)
1164
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1165
+
1166
+ img, x0_partial = self.p_sample(img, cond, ts,
1167
+ clip_denoised=self.clip_denoised,
1168
+ quantize_denoised=quantize_denoised, return_x0=True,
1169
+ temperature=temperature[i], noise_dropout=noise_dropout,
1170
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
1171
+ if mask is not None:
1172
+ assert x0 is not None
1173
+ img_orig = self.q_sample(x0, ts)
1174
+ img = img_orig * mask + (1. - mask) * img
1175
+
1176
+ if i % log_every_t == 0 or i == timesteps - 1:
1177
+ intermediates.append(x0_partial)
1178
+ if callback: callback(i)
1179
+ if img_callback: img_callback(img, i)
1180
+ return img, intermediates
1181
+
1182
+ @torch.no_grad()
1183
+ def p_sample_loop(self, cond, shape, return_intermediates=False,
1184
+ x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False,
1185
+ mask=None, x0=None, img_callback=None, start_T=None,
1186
+ log_every_t=None):
1187
+
1188
+ if not log_every_t:
1189
+ log_every_t = self.log_every_t
1190
+ device = self.betas.device
1191
+ b = shape[0]
1192
+ if x_T is None:
1193
+ img = torch.randn(shape, device=device)
1194
+ else:
1195
+ img = x_T
1196
+
1197
+ intermediates = [img]
1198
+ if timesteps is None:
1199
+ timesteps = self.num_timesteps
1200
+
1201
+ if start_T is not None:
1202
+ timesteps = min(timesteps, start_T)
1203
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(
1204
+ range(0, timesteps))
1205
+
1206
+ if mask is not None:
1207
+ assert x0 is not None
1208
+ assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match
1209
+
1210
+ for i in iterator:
1211
+ ts = torch.full((b,), i, device=device, dtype=torch.long)
1212
+ if self.shorten_cond_schedule:
1213
+ assert self.model.conditioning_key != 'hybrid'
1214
+ tc = self.cond_ids[ts].to(cond.device)
1215
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1216
+
1217
+ img = self.p_sample(img, cond, ts,
1218
+ clip_denoised=self.clip_denoised,
1219
+ quantize_denoised=quantize_denoised)
1220
+ if mask is not None:
1221
+ img_orig = self.q_sample(x0, ts)
1222
+ img = img_orig * mask + (1. - mask) * img
1223
+
1224
+ if i % log_every_t == 0 or i == timesteps - 1:
1225
+ intermediates.append(img)
1226
+ if callback: callback(i)
1227
+ if img_callback: img_callback(img, i)
1228
+
1229
+ if return_intermediates:
1230
+ return img, intermediates
1231
+ return img
1232
+
1233
+ @torch.no_grad()
1234
+ def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
1235
+ verbose=True, timesteps=None, quantize_denoised=False,
1236
+ mask=None, x0=None, shape=None,**kwargs):
1237
+ if shape is None:
1238
+ shape = (batch_size, self.channels, self.image_size, self.image_size)
1239
+ if cond is not None:
1240
+ if isinstance(cond, dict):
1241
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
1242
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
1243
+ else:
1244
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
1245
+ return self.p_sample_loop(cond,
1246
+ shape,
1247
+ return_intermediates=return_intermediates, x_T=x_T,
1248
+ verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
1249
+ mask=mask, x0=x0)
1250
+
1251
+ @torch.no_grad()
1252
+ def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
1253
+ if ddim:
1254
+ ddim_sampler = DDIMSampler(self)
1255
+ shape = (self.channels, self.image_size, self.image_size)
1256
+ samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size,
1257
+ shape, cond, verbose=False, **kwargs)
1258
+
1259
+ else:
1260
+ samples, intermediates = self.sample(cond=cond, batch_size=batch_size,
1261
+ return_intermediates=True, **kwargs)
1262
+
1263
+ return samples, intermediates
1264
+
1265
+ @torch.no_grad()
1266
+ def get_unconditional_conditioning(self, batch_size, null_label=None, image_size=512):
1267
+ if null_label is not None:
1268
+ xc = null_label
1269
+ if isinstance(xc, ListConfig):
1270
+ xc = list(xc)
1271
+ if isinstance(xc, dict) or isinstance(xc, list):
1272
+ c = self.get_learned_conditioning(xc)
1273
+ else:
1274
+ if hasattr(xc, "to"):
1275
+ xc = xc.to(self.device)
1276
+ c = self.get_learned_conditioning(xc)
1277
+ else:
1278
+ # todo: get null label from cond_stage_model
1279
+ raise NotImplementedError()
1280
+ c = repeat(c, '1 ... -> b ...', b=batch_size).to(self.device)
1281
+ cond = {}
1282
+ cond["c_crossattn"] = [c]
1283
+ cond["c_concat"] = [torch.zeros([batch_size, 4, image_size // 8, image_size // 8]).to(self.device)]
1284
+ return cond
1285
+
1286
+ @torch.no_grad()
1287
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
1288
+ quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
1289
+ plot_diffusion_rows=True, unconditional_guidance_scale=1., unconditional_guidance_label=None,
1290
+ use_ema_scope=True,
1291
+ **kwargs):
1292
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1293
+ use_ddim = ddim_steps is not None
1294
+
1295
+ log = dict()
1296
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key,
1297
+ return_first_stage_outputs=True,
1298
+ force_c_encode=True,
1299
+ return_original_cond=True,
1300
+ bs=N)
1301
+
1302
+
1303
+ N = min(x.shape[0], N)
1304
+ n_row = min(x.shape[0], n_row)
1305
+ log["inputs"] = x
1306
+ log["reconstruction"] = xrec
1307
+ if self.model.conditioning_key is not None:
1308
+ if hasattr(self.cond_stage_model, "decode"):
1309
+ xc = self.cond_stage_model.decode(c)
1310
+ log["conditioning"] = xc
1311
+ elif self.cond_stage_key in ["caption", "txt"]:
1312
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2]//25)
1313
+ log["conditioning"] = xc
1314
+ elif self.cond_stage_key == 'class_label':
1315
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2]//25)
1316
+ log['conditioning'] = xc
1317
+ elif isimage(xc):
1318
+ log["conditioning"] = xc
1319
+ if ismap(xc):
1320
+ log["original_conditioning"] = self.to_rgb(xc)
1321
+
1322
+ if plot_diffusion_rows:
1323
+ # get diffusion row
1324
+ diffusion_row = list()
1325
+ z_start = z[:n_row]
1326
+ for t in range(self.num_timesteps):
1327
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1328
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1329
+ t = t.to(self.device).long()
1330
+ noise = torch.randn_like(z_start)
1331
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1332
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1333
+
1334
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1335
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1336
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1337
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1338
+ log["diffusion_row"] = diffusion_grid
1339
+
1340
+ if sample:
1341
+ # get denoise row
1342
+ with ema_scope("Sampling"):
1343
+ samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim,
1344
+ ddim_steps=ddim_steps,eta=ddim_eta)
1345
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1346
+ x_samples = self.decode_first_stage(samples)
1347
+ log["samples"] = x_samples
1348
+ if plot_denoise_rows:
1349
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1350
+ log["denoise_row"] = denoise_grid
1351
+
1352
+ if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance(
1353
+ self.first_stage_model, IdentityFirstStage):
1354
+ # also display when quantizing x0 while sampling
1355
+ with ema_scope("Plotting Quantized Denoised"):
1356
+ samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim,
1357
+ ddim_steps=ddim_steps,eta=ddim_eta,
1358
+ quantize_denoised=True)
1359
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True,
1360
+ # quantize_denoised=True)
1361
+ x_samples = self.decode_first_stage(samples.to(self.device))
1362
+ log["samples_x0_quantized"] = x_samples
1363
+
1364
+ if unconditional_guidance_scale > 1.0:
1365
+ uc = self.get_unconditional_conditioning(N, unconditional_guidance_label, image_size=x.shape[-1])
1366
+ # uc = torch.zeros_like(c)
1367
+ with ema_scope("Sampling with classifier-free guidance"):
1368
+ samples_cfg, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1369
+ ddim_steps=ddim_steps, eta=ddim_eta,
1370
+ unconditional_guidance_scale=unconditional_guidance_scale,
1371
+ unconditional_conditioning=uc,
1372
+ )
1373
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1374
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1375
+
1376
+ if inpaint:
1377
+ # make a simple center square
1378
+ b, h, w = z.shape[0], z.shape[2], z.shape[3]
1379
+ mask = torch.ones(N, h, w).to(self.device)
1380
+ # zeros will be filled in
1381
+ mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
1382
+ mask = mask[:, None, ...]
1383
+ with ema_scope("Plotting Inpaint"):
1384
+
1385
+ samples, _ = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, eta=ddim_eta,
1386
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1387
+ x_samples = self.decode_first_stage(samples.to(self.device))
1388
+ log["samples_inpainting"] = x_samples
1389
+ log["mask"] = mask
1390
+
1391
+ # outpaint
1392
+ mask = 1. - mask
1393
+ with ema_scope("Plotting Outpaint"):
1394
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,eta=ddim_eta,
1395
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1396
+ x_samples = self.decode_first_stage(samples.to(self.device))
1397
+ log["samples_outpainting"] = x_samples
1398
+
1399
+ if plot_progressive_rows:
1400
+ with ema_scope("Plotting Progressives"):
1401
+ img, progressives = self.progressive_denoising(c,
1402
+ shape=(self.channels, self.image_size, self.image_size),
1403
+ batch_size=N)
1404
+ prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
1405
+ log["progressive_row"] = prog_row
1406
+
1407
+ if return_keys:
1408
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
1409
+ return log
1410
+ else:
1411
+ return {key: log[key] for key in return_keys}
1412
+ return log
1413
+
1414
+ def configure_optimizers(self):
1415
+ lr = self.learning_rate
1416
+ params = []
1417
+ if self.unet_trainable == "attn":
1418
+ print("Training only unet attention layers")
1419
+ for n, m in self.model.named_modules():
1420
+ if isinstance(m, CrossAttention) and n.endswith('attn2'):
1421
+ params.extend(m.parameters())
1422
+ if self.unet_trainable == "conv_in":
1423
+ print("Training only unet input conv layers")
1424
+ params = list(self.model.diffusion_model.input_blocks[0][0].parameters())
1425
+ elif self.unet_trainable is True or self.unet_trainable == "all":
1426
+ print("Training the full unet")
1427
+ params = list(self.model.parameters())
1428
+ else:
1429
+ raise ValueError(f"Unrecognised setting for unet_trainable: {self.unet_trainable}")
1430
+
1431
+ if self.cond_stage_trainable:
1432
+ print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
1433
+ params = params + list(self.cond_stage_model.parameters())
1434
+ if self.learn_logvar:
1435
+ print('Diffusion model optimizing logvar')
1436
+ params.append(self.logvar)
1437
+
1438
+ if self.cc_projection is not None:
1439
+ params = params + list(self.cc_projection.parameters())
1440
+ print('========== optimizing for cc projection weight ==========')
1441
+
1442
+ opt = torch.optim.AdamW([{"params": self.model.parameters(), "lr": lr},
1443
+ {"params": self.cc_projection.parameters(), "lr": 10. * lr}], lr=lr)
1444
+ if self.use_scheduler:
1445
+ assert 'target' in self.scheduler_config
1446
+ scheduler = instantiate_from_config(self.scheduler_config)
1447
+
1448
+ print("Setting up LambdaLR scheduler...")
1449
+ scheduler = [
1450
+ {
1451
+ 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
1452
+ 'interval': 'step',
1453
+ 'frequency': 1
1454
+ }]
1455
+ return [opt], scheduler
1456
+ return opt
1457
+
1458
+ @torch.no_grad()
1459
+ def to_rgb(self, x):
1460
+ x = x.float()
1461
+ if not hasattr(self, "colorize"):
1462
+ self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
1463
+ x = nn.functional.conv2d(x, weight=self.colorize)
1464
+ x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
1465
+ return x
1466
+
1467
+
1468
+ class DiffusionWrapper(pl.LightningModule):
1469
+ def __init__(self, diff_model_config, conditioning_key):
1470
+ super().__init__()
1471
+ self.diffusion_model = instantiate_from_config(diff_model_config)
1472
+ self.conditioning_key = conditioning_key
1473
+ assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'hybrid-adm']
1474
+
1475
+ def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None, **kwargs):
1476
+ if self.conditioning_key is None:
1477
+ out = self.diffusion_model(x, t)
1478
+ elif self.conditioning_key == 'concat':
1479
+ xc = torch.cat([x] + c_concat, dim=1)
1480
+ out = self.diffusion_model(xc, t)
1481
+ elif self.conditioning_key == 'crossattn':
1482
+ # c_crossattn dimension: torch.Size([8, 1, 768]) 1
1483
+ # cc dimension: torch.Size([8, 1, 768]
1484
+ cc = torch.cat(c_crossattn, 1)
1485
+ out = self.diffusion_model(x, t, context=cc)
1486
+ elif self.conditioning_key == 'hybrid': ### This One
1487
+ xc = torch.cat([x] + c_concat, dim=1)
1488
+ cc = torch.cat(c_crossattn, 1)
1489
+ out = self.diffusion_model(xc, t, context=cc, **kwargs)
1490
+ elif self.conditioning_key == 'hybrid-adm':
1491
+ assert c_adm is not None
1492
+ xc = torch.cat([x] + c_concat, dim=1)
1493
+ cc = torch.cat(c_crossattn, 1)
1494
+ out = self.diffusion_model(xc, t, context=cc, y=c_adm)
1495
+ elif self.conditioning_key == 'adm':
1496
+ cc = c_crossattn[0]
1497
+ out = self.diffusion_model(x, t, y=cc)
1498
+ else:
1499
+ raise NotImplementedError()
1500
+
1501
+ return out
1502
+
ldm/models/diffusion/plms.py ADDED
@@ -0,0 +1,259 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+ from functools import partial
7
+
8
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
9
+ from ldm.models.diffusion.sampling_util import norm_thresholding
10
+
11
+
12
+ class PLMSSampler(object):
13
+ def __init__(self, model, schedule="linear", **kwargs):
14
+ super().__init__()
15
+ self.model = model
16
+ self.ddpm_num_timesteps = model.num_timesteps
17
+ self.schedule = schedule
18
+
19
+ def register_buffer(self, name, attr):
20
+ if type(attr) == torch.Tensor:
21
+ if attr.device != torch.device("cuda"):
22
+ attr = attr.to(torch.device("cuda"))
23
+ setattr(self, name, attr)
24
+
25
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
26
+ if ddim_eta != 0:
27
+ raise ValueError('ddim_eta must be 0 for PLMS')
28
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
29
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
30
+ alphas_cumprod = self.model.alphas_cumprod
31
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
32
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
33
+
34
+ self.register_buffer('betas', to_torch(self.model.betas))
35
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
36
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
37
+
38
+ # calculations for diffusion q(x_t | x_{t-1}) and others
39
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
40
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
41
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
42
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
43
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
44
+
45
+ # ddim sampling parameters
46
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
47
+ ddim_timesteps=self.ddim_timesteps,
48
+ eta=ddim_eta,verbose=verbose)
49
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
50
+ self.register_buffer('ddim_alphas', ddim_alphas)
51
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
52
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
53
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
54
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
55
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
56
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
57
+
58
+ @torch.no_grad()
59
+ def sample(self,
60
+ S,
61
+ batch_size,
62
+ shape,
63
+ conditioning=None,
64
+ callback=None,
65
+ normals_sequence=None,
66
+ img_callback=None,
67
+ quantize_x0=False,
68
+ eta=0.,
69
+ mask=None,
70
+ x0=None,
71
+ temperature=1.,
72
+ noise_dropout=0.,
73
+ score_corrector=None,
74
+ corrector_kwargs=None,
75
+ verbose=True,
76
+ x_T=None,
77
+ log_every_t=100,
78
+ unconditional_guidance_scale=1.,
79
+ unconditional_conditioning=None,
80
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
81
+ dynamic_threshold=None,
82
+ **kwargs
83
+ ):
84
+ if conditioning is not None:
85
+ if isinstance(conditioning, dict):
86
+ ctmp = conditioning[list(conditioning.keys())[0]]
87
+ while isinstance(ctmp, list): ctmp = ctmp[0]
88
+ cbs = ctmp.shape[0]
89
+ if cbs != batch_size:
90
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
91
+ else:
92
+ if conditioning.shape[0] != batch_size:
93
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
94
+
95
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
96
+ # sampling
97
+ C, H, W = shape
98
+ size = (batch_size, C, H, W)
99
+ print(f'Data shape for PLMS sampling is {size}')
100
+
101
+ samples, intermediates = self.plms_sampling(conditioning, size,
102
+ callback=callback,
103
+ img_callback=img_callback,
104
+ quantize_denoised=quantize_x0,
105
+ mask=mask, x0=x0,
106
+ ddim_use_original_steps=False,
107
+ noise_dropout=noise_dropout,
108
+ temperature=temperature,
109
+ score_corrector=score_corrector,
110
+ corrector_kwargs=corrector_kwargs,
111
+ x_T=x_T,
112
+ log_every_t=log_every_t,
113
+ unconditional_guidance_scale=unconditional_guidance_scale,
114
+ unconditional_conditioning=unconditional_conditioning,
115
+ dynamic_threshold=dynamic_threshold,
116
+ )
117
+ return samples, intermediates
118
+
119
+ @torch.no_grad()
120
+ def plms_sampling(self, cond, shape,
121
+ x_T=None, ddim_use_original_steps=False,
122
+ callback=None, timesteps=None, quantize_denoised=False,
123
+ mask=None, x0=None, img_callback=None, log_every_t=100,
124
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
125
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
126
+ dynamic_threshold=None):
127
+ device = self.model.betas.device
128
+ b = shape[0]
129
+ if x_T is None:
130
+ img = torch.randn(shape, device=device)
131
+ else:
132
+ img = x_T
133
+
134
+ if timesteps is None:
135
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
136
+ elif timesteps is not None and not ddim_use_original_steps:
137
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
138
+ timesteps = self.ddim_timesteps[:subset_end]
139
+
140
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
141
+ time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps)
142
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
143
+ print(f"Running PLMS Sampling with {total_steps} timesteps")
144
+
145
+ iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
146
+ old_eps = []
147
+
148
+ for i, step in enumerate(iterator):
149
+ index = total_steps - i - 1
150
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
151
+ ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long)
152
+
153
+ if mask is not None:
154
+ assert x0 is not None
155
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
156
+ img = img_orig * mask + (1. - mask) * img
157
+
158
+ outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
159
+ quantize_denoised=quantize_denoised, temperature=temperature,
160
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
161
+ corrector_kwargs=corrector_kwargs,
162
+ unconditional_guidance_scale=unconditional_guidance_scale,
163
+ unconditional_conditioning=unconditional_conditioning,
164
+ old_eps=old_eps, t_next=ts_next,
165
+ dynamic_threshold=dynamic_threshold)
166
+ img, pred_x0, e_t = outs
167
+ old_eps.append(e_t)
168
+ if len(old_eps) >= 4:
169
+ old_eps.pop(0)
170
+ if callback: callback(i)
171
+ if img_callback: img_callback(pred_x0, i)
172
+
173
+ if index % log_every_t == 0 or index == total_steps - 1:
174
+ intermediates['x_inter'].append(img)
175
+ intermediates['pred_x0'].append(pred_x0)
176
+
177
+ return img, intermediates
178
+
179
+ @torch.no_grad()
180
+ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
181
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
182
+ unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None,
183
+ dynamic_threshold=None):
184
+ b, *_, device = *x.shape, x.device
185
+
186
+ def get_model_output(x, t):
187
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
188
+ e_t = self.model.apply_model(x, t, c)
189
+ else:
190
+ x_in = torch.cat([x] * 2)
191
+ t_in = torch.cat([t] * 2)
192
+ if isinstance(c, dict):
193
+ assert isinstance(unconditional_conditioning, dict)
194
+ c_in = dict()
195
+ for k in c:
196
+ if isinstance(c[k], list):
197
+ c_in[k] = [torch.cat([
198
+ unconditional_conditioning[k][i],
199
+ c[k][i]]) for i in range(len(c[k]))]
200
+ else:
201
+ c_in[k] = torch.cat([
202
+ unconditional_conditioning[k],
203
+ c[k]])
204
+ else:
205
+ c_in = torch.cat([unconditional_conditioning, c])
206
+ e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
207
+ e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
208
+
209
+ if score_corrector is not None:
210
+ assert self.model.parameterization == "eps"
211
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
212
+
213
+ return e_t
214
+
215
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
216
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
217
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
218
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
219
+
220
+ def get_x_prev_and_pred_x0(e_t, index):
221
+ # select parameters corresponding to the currently considered timestep
222
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
223
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
224
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
225
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
226
+
227
+ # current prediction for x_0
228
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
229
+ if quantize_denoised:
230
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
231
+ if dynamic_threshold is not None:
232
+ pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
233
+ # direction pointing to x_t
234
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
235
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
236
+ if noise_dropout > 0.:
237
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
238
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
239
+ return x_prev, pred_x0
240
+
241
+ e_t = get_model_output(x, t)
242
+ if len(old_eps) == 0:
243
+ # Pseudo Improved Euler (2nd order)
244
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
245
+ e_t_next = get_model_output(x_prev, t_next)
246
+ e_t_prime = (e_t + e_t_next) / 2
247
+ elif len(old_eps) == 1:
248
+ # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
249
+ e_t_prime = (3 * e_t - old_eps[-1]) / 2
250
+ elif len(old_eps) == 2:
251
+ # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
252
+ e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
253
+ elif len(old_eps) >= 3:
254
+ # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
255
+ e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
256
+
257
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
258
+
259
+ return x_prev, pred_x0, e_t
ldm/models/diffusion/sampling_util.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+
4
+
5
+ def append_dims(x, target_dims):
6
+ """Appends dimensions to the end of a tensor until it has target_dims dimensions.
7
+ From https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/utils.py"""
8
+ dims_to_append = target_dims - x.ndim
9
+ if dims_to_append < 0:
10
+ raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
11
+ return x[(...,) + (None,) * dims_to_append]
12
+
13
+
14
+ def renorm_thresholding(x0, value):
15
+ # renorm
16
+ pred_max = x0.max()
17
+ pred_min = x0.min()
18
+ pred_x0 = (x0 - pred_min) / (pred_max - pred_min) # 0 ... 1
19
+ pred_x0 = 2 * pred_x0 - 1. # -1 ... 1
20
+
21
+ s = torch.quantile(
22
+ rearrange(pred_x0, 'b ... -> b (...)').abs(),
23
+ value,
24
+ dim=-1
25
+ )
26
+ s.clamp_(min=1.0)
27
+ s = s.view(-1, *((1,) * (pred_x0.ndim - 1)))
28
+
29
+ # clip by threshold
30
+ # pred_x0 = pred_x0.clamp(-s, s) / s # needs newer pytorch # TODO bring back to pure-gpu with min/max
31
+
32
+ # temporary hack: numpy on cpu
33
+ pred_x0 = np.clip(pred_x0.cpu().numpy(), -s.cpu().numpy(), s.cpu().numpy()) / s.cpu().numpy()
34
+ pred_x0 = torch.tensor(pred_x0).to(self.model.device)
35
+
36
+ # re.renorm
37
+ pred_x0 = (pred_x0 + 1.) / 2. # 0 ... 1
38
+ pred_x0 = (pred_max - pred_min) * pred_x0 + pred_min # orig range
39
+ return pred_x0
40
+
41
+
42
+ def norm_thresholding(x0, value):
43
+ s = append_dims(x0.pow(2).flatten(1).mean(1).sqrt().clamp(min=value), x0.ndim)
44
+ return x0 * (value / s)
45
+
46
+
47
+ def spatial_norm_thresholding(x0, value):
48
+ # b c h w
49
+ s = x0.pow(2).mean(1, keepdim=True).sqrt().clamp(min=value)
50
+ return x0 * (value / s)
ldm/modules/attention.py ADDED
@@ -0,0 +1,266 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
46
+
47
+ class FeedForward(nn.Module):
48
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
49
+ super().__init__()
50
+ inner_dim = int(dim * mult)
51
+ dim_out = default(dim_out, dim)
52
+ project_in = nn.Sequential(
53
+ nn.Linear(dim, inner_dim),
54
+ nn.GELU()
55
+ ) if not glu else GEGLU(dim, inner_dim)
56
+
57
+ self.net = nn.Sequential(
58
+ project_in,
59
+ nn.Dropout(dropout),
60
+ nn.Linear(inner_dim, dim_out)
61
+ )
62
+
63
+ def forward(self, x):
64
+ return self.net(x)
65
+
66
+
67
+ def zero_module(module):
68
+ """
69
+ Zero out the parameters of a module and return it.
70
+ """
71
+ for p in module.parameters():
72
+ p.detach().zero_()
73
+ return module
74
+
75
+
76
+ def Normalize(in_channels):
77
+ return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
78
+
79
+
80
+ class LinearAttention(nn.Module):
81
+ def __init__(self, dim, heads=4, dim_head=32):
82
+ super().__init__()
83
+ self.heads = heads
84
+ hidden_dim = dim_head * heads
85
+ self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
86
+ self.to_out = nn.Conv2d(hidden_dim, dim, 1)
87
+
88
+ def forward(self, x):
89
+ b, c, h, w = x.shape
90
+ qkv = self.to_qkv(x)
91
+ q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
92
+ k = k.softmax(dim=-1)
93
+ context = torch.einsum('bhdn,bhen->bhde', k, v)
94
+ out = torch.einsum('bhde,bhdn->bhen', context, q)
95
+ out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
96
+ return self.to_out(out)
97
+
98
+
99
+ class SpatialSelfAttention(nn.Module):
100
+ def __init__(self, in_channels):
101
+ super().__init__()
102
+ self.in_channels = in_channels
103
+
104
+ self.norm = Normalize(in_channels)
105
+ self.q = torch.nn.Conv2d(in_channels,
106
+ in_channels,
107
+ kernel_size=1,
108
+ stride=1,
109
+ padding=0)
110
+ self.k = torch.nn.Conv2d(in_channels,
111
+ in_channels,
112
+ kernel_size=1,
113
+ stride=1,
114
+ padding=0)
115
+ self.v = torch.nn.Conv2d(in_channels,
116
+ in_channels,
117
+ kernel_size=1,
118
+ stride=1,
119
+ padding=0)
120
+ self.proj_out = torch.nn.Conv2d(in_channels,
121
+ in_channels,
122
+ kernel_size=1,
123
+ stride=1,
124
+ padding=0)
125
+
126
+ def forward(self, x):
127
+ h_ = x
128
+ h_ = self.norm(h_)
129
+ q = self.q(h_)
130
+ k = self.k(h_)
131
+ v = self.v(h_)
132
+
133
+ # compute attention
134
+ b,c,h,w = q.shape
135
+ q = rearrange(q, 'b c h w -> b (h w) c')
136
+ k = rearrange(k, 'b c h w -> b c (h w)')
137
+ w_ = torch.einsum('bij,bjk->bik', q, k)
138
+
139
+ w_ = w_ * (int(c)**(-0.5))
140
+ w_ = torch.nn.functional.softmax(w_, dim=2)
141
+
142
+ # attend to values
143
+ v = rearrange(v, 'b c h w -> b c (h w)')
144
+ w_ = rearrange(w_, 'b i j -> b j i')
145
+ h_ = torch.einsum('bij,bjk->bik', v, w_)
146
+ h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
147
+ h_ = self.proj_out(h_)
148
+
149
+ return x+h_
150
+
151
+
152
+ class CrossAttention(nn.Module):
153
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
154
+ super().__init__()
155
+ inner_dim = dim_head * heads
156
+ context_dim = default(context_dim, query_dim)
157
+
158
+ self.scale = dim_head ** -0.5
159
+ self.heads = heads
160
+
161
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
162
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
163
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
164
+
165
+ self.to_out = nn.Sequential(
166
+ nn.Linear(inner_dim, query_dim),
167
+ nn.Dropout(dropout)
168
+ )
169
+
170
+ def forward(self, x, context=None, mask=None):
171
+ h = self.heads
172
+
173
+ q = self.to_q(x)
174
+ context = default(context, x)
175
+ k = self.to_k(context)
176
+ v = self.to_v(context)
177
+
178
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
179
+
180
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
181
+
182
+ if exists(mask):
183
+ mask = rearrange(mask, 'b ... -> b (...)')
184
+ max_neg_value = -torch.finfo(sim.dtype).max
185
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
186
+ sim.masked_fill_(~mask, max_neg_value)
187
+
188
+ # attention, what we cannot get enough of
189
+ attn = sim.softmax(dim=-1)
190
+
191
+ out = einsum('b i j, b j d -> b i d', attn, v)
192
+ out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
193
+ return self.to_out(out)
194
+
195
+
196
+ class BasicTransformerBlock(nn.Module):
197
+ def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
198
+ disable_self_attn=False):
199
+ super().__init__()
200
+ self.disable_self_attn = disable_self_attn
201
+ self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
202
+ context_dim=context_dim if self.disable_self_attn else None) # is a self-attention if not self.disable_self_attn
203
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
204
+ self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
205
+ heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none
206
+ self.norm1 = nn.LayerNorm(dim)
207
+ self.norm2 = nn.LayerNorm(dim)
208
+ self.norm3 = nn.LayerNorm(dim)
209
+ self.checkpoint = checkpoint
210
+
211
+ def forward(self, x, context=None):
212
+ return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
213
+
214
+ def _forward(self, x, context=None):
215
+ x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
216
+ x = self.attn2(self.norm2(x), context=context) + x
217
+ x = self.ff(self.norm3(x)) + x
218
+ return x
219
+
220
+
221
+ class SpatialTransformer(nn.Module):
222
+ """
223
+ Transformer block for image-like data.
224
+ First, project the input (aka embedding)
225
+ and reshape to b, t, d.
226
+ Then apply standard transformer action.
227
+ Finally, reshape to image
228
+ """
229
+ def __init__(self, in_channels, n_heads, d_head,
230
+ depth=1, dropout=0., context_dim=None,
231
+ disable_self_attn=False):
232
+ super().__init__()
233
+ self.in_channels = in_channels
234
+ inner_dim = n_heads * d_head
235
+ self.norm = Normalize(in_channels)
236
+
237
+ self.proj_in = nn.Conv2d(in_channels,
238
+ inner_dim,
239
+ kernel_size=1,
240
+ stride=1,
241
+ padding=0)
242
+
243
+ self.transformer_blocks = nn.ModuleList(
244
+ [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim,
245
+ disable_self_attn=disable_self_attn)
246
+ for d in range(depth)]
247
+ )
248
+
249
+ self.proj_out = zero_module(nn.Conv2d(inner_dim,
250
+ in_channels,
251
+ kernel_size=1,
252
+ stride=1,
253
+ padding=0))
254
+
255
+ def forward(self, x, context=None):
256
+ # note: if no context is given, cross-attention defaults to self-attention
257
+ b, c, h, w = x.shape
258
+ x_in = x
259
+ x = self.norm(x)
260
+ x = self.proj_in(x)
261
+ x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
262
+ for block in self.transformer_blocks:
263
+ x = block(x, context=context)
264
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
265
+ x = self.proj_out(x)
266
+ 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:
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
+ )
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:
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)
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)
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/ema.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch import nn
3
+
4
+
5
+ class LitEma(nn.Module):
6
+ def __init__(self, model, decay=0.9999, use_num_upates=True):
7
+ super().__init__()
8
+ if decay < 0.0 or decay > 1.0:
9
+ raise ValueError('Decay must be between 0 and 1')
10
+
11
+ self.m_name2s_name = {}
12
+ self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32))
13
+ self.register_buffer('num_updates', torch.tensor(0,dtype=torch.int) if use_num_upates
14
+ else torch.tensor(-1,dtype=torch.int))
15
+
16
+ for name, p in model.named_parameters():
17
+ if p.requires_grad:
18
+ #remove as '.'-character is not allowed in buffers
19
+ s_name = name.replace('.','')
20
+ self.m_name2s_name.update({name:s_name})
21
+ self.register_buffer(s_name,p.clone().detach().data)
22
+
23
+ self.collected_params = []
24
+
25
+ def forward(self,model):
26
+ decay = self.decay
27
+
28
+ if self.num_updates >= 0:
29
+ self.num_updates += 1
30
+ decay = min(self.decay,(1 + self.num_updates) / (10 + self.num_updates))
31
+
32
+ one_minus_decay = 1.0 - decay
33
+
34
+ with torch.no_grad():
35
+ m_param = dict(model.named_parameters())
36
+ shadow_params = dict(self.named_buffers())
37
+
38
+ for key in m_param:
39
+ if m_param[key].requires_grad:
40
+ sname = self.m_name2s_name[key]
41
+ shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
42
+ shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
43
+ else:
44
+ assert not key in self.m_name2s_name
45
+
46
+ def copy_to(self, model):
47
+ m_param = dict(model.named_parameters())
48
+ shadow_params = dict(self.named_buffers())
49
+ for key in m_param:
50
+ if m_param[key].requires_grad:
51
+ m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
52
+ else:
53
+ assert not key in self.m_name2s_name
54
+
55
+ def store(self, parameters):
56
+ """
57
+ Save the current parameters for restoring later.
58
+ Args:
59
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
60
+ temporarily stored.
61
+ """
62
+ self.collected_params = [param.clone() for param in parameters]
63
+
64
+ def restore(self, parameters):
65
+ """
66
+ Restore the parameters stored with the `store` method.
67
+ Useful to validate the model with EMA parameters without affecting the
68
+ original optimization process. Store the parameters before the
69
+ `copy_to` method. After validation (or model saving), use this to
70
+ restore the former parameters.
71
+ Args:
72
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
73
+ updated with the stored parameters.
74
+ """
75
+ for c_param, param in zip(self.collected_params, parameters):
76
+ param.data.copy_(c_param.data)
ldm/modules/encoders/__init__.py ADDED
File without changes
ldm/modules/encoders/modules.py ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import torch.nn as nn
4
+ import numpy as np
5
+ from functools import partial
6
+ import kornia
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 transformers import CLIPTokenizer, CLIPTextModel
9
+ import torch.nn.functional as F
10
+ from torchvision import transforms
11
+ import random
12
+ from ldm.util import default, instantiate_from_config
13
+ from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
14
+ import clip
15
+
16
+ class AbstractEncoder(nn.Module):
17
+ def __init__(self):
18
+ super().__init__()
19
+
20
+ def encode(self, *args, **kwargs):
21
+ raise NotImplementedError
22
+
23
+
24
+
25
+ def disabled_train(self, mode=True):
26
+ """Overwrite model.train with this function to make sure train/eval mode
27
+ does not change anymore."""
28
+ return self
29
+
30
+
31
+ class FrozenCLIPEmbedder(AbstractEncoder):
32
+ """Uses the CLIP transformer encoder for text (from huggingface)"""
33
+ def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77): # clip-vit-base-patch32
34
+ super().__init__()
35
+ self.tokenizer = CLIPTokenizer.from_pretrained(version)
36
+ self.transformer = CLIPTextModel.from_pretrained(version)
37
+ self.device = device
38
+ self.max_length = max_length # TODO: typical value?
39
+ self.freeze()
40
+
41
+ def freeze(self):
42
+ self.transformer = self.transformer.eval()
43
+ #self.train = disabled_train
44
+ for param in self.parameters():
45
+ param.requires_grad = False
46
+
47
+ def forward(self, text, return_pool=False):
48
+ batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
49
+ return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
50
+ tokens = batch_encoding["input_ids"].to(self.device)
51
+ outputs = self.transformer(input_ids=tokens)
52
+
53
+ z = outputs.last_hidden_state
54
+ if return_pool:
55
+ return z, outputs.pooler_output
56
+ else:
57
+ return z
58
+
59
+ def encode(self, text):
60
+ return self(text)
61
+
62
+
63
+ class FrozenCLIPImageEmbedder(AbstractEncoder):
64
+ """
65
+ Uses the CLIP image encoder.
66
+ Not actually frozen... If you want that set cond_stage_trainable=False in cfg
67
+ """
68
+ def __init__(
69
+ self,
70
+ model='ViT-L/14',
71
+ jit=False,
72
+ device='cpu',
73
+ antialias=False,
74
+ ):
75
+ super().__init__()
76
+ self.model, _ = clip.load(name=model, device=device, jit=jit,)
77
+ # We don't use the text part so delete it
78
+ del self.model.transformer
79
+ self.antialias = antialias
80
+ self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
81
+ self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
82
+
83
+ def preprocess(self, x):
84
+ # Expects inputs in the range -1, 1
85
+ # x = kornia.geometry.resize(x, (224, 224),
86
+ # interpolation='bicubic',align_corners=True,
87
+ # antialias=self.antialias)
88
+
89
+ x = kornia.geometry.resize(x, (224, 224),
90
+ interpolation='bicubic',align_corners=True)
91
+
92
+ x = (x + 1.) / 2.
93
+ # renormalize according to clip
94
+ x = kornia.enhance.normalize(x, self.mean, self.std)
95
+ return x
96
+
97
+ def forward(self, x):
98
+ # x is assumed to be in range [-1,1]
99
+ if isinstance(x, list):
100
+ # [""] denotes condition dropout for ucg
101
+ device = self.model.visual.conv1.weight.device
102
+ return torch.zeros(1, 768, device=device)
103
+ return self.model.encode_image(self.preprocess(x)).float()
104
+
105
+ def encode(self, im):
106
+ return self(im).unsqueeze(1)
107
+
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/util.py ADDED
@@ -0,0 +1,289 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ from carvekit.api.high import HiInterface
19
+ import PIL
20
+
21
+ def pil_rectangle_crop(im):
22
+ width, height = im.size # Get dimensions
23
+
24
+ if width <= height:
25
+ left = 0
26
+ right = width
27
+ top = (height - width)/2
28
+ bottom = (height + width)/2
29
+ else:
30
+
31
+ top = 0
32
+ bottom = height
33
+ left = (width - height) / 2
34
+ bottom = (width + height) / 2
35
+
36
+ # Crop the center of the image
37
+ im = im.crop((left, top, right, bottom))
38
+ return im
39
+
40
+ def add_margin(pil_img, color, size=256):
41
+ width, height = pil_img.size
42
+ result = Image.new(pil_img.mode, (size, size), color)
43
+ result.paste(pil_img, ((size - width) // 2, (size - height) // 2))
44
+ return result
45
+
46
+
47
+ def create_carvekit_interface():
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.Resampling.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 get_state_dict(d):
278
+ return d.get('state_dict', d)
279
+
280
+ def load_state_dict(ckpt_path, location='cpu'):
281
+ _, extension = os.path.splitext(ckpt_path)
282
+ if extension.lower() == ".safetensors":
283
+ import safetensors.torch
284
+ state_dict = safetensors.torch.load_file(ckpt_path, device=location)
285
+ else:
286
+ state_dict = get_state_dict(torch.load(ckpt_path, map_location=torch.device(location)))
287
+ state_dict = get_state_dict(state_dict)
288
+ print(f'Loaded state_dict from [{ckpt_path}]')
289
+ return state_dict
requirements.txt ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
2
+ torch==1.12.1
3
+ torchvision==0.13.1
4
+ albumentations==0.4.3
5
+ opencv-python
6
+ pudb==2019.2
7
+ imageio==2.9.0
8
+ imageio-ffmpeg==0.4.2
9
+ pytorch-lightning==1.4.2
10
+ omegaconf==2.1.1
11
+ test-tube>=0.7.5
12
+ streamlit>=0.73.1
13
+ einops==0.3.0
14
+ torch-fidelity==0.3.0
15
+ transformers==4.22.2
16
+ kornia==0.6
17
+ webdataset==0.2.5
18
+ torchmetrics==0.6.0
19
+ fire==0.4.0
20
+ diffusers==0.12.1
21
+ datasets[vision]==2.4.0
22
+ carvekit-colab==4.1.0
23
+ rich>=13.3.2
24
+ lovely-numpy>=0.2.8
25
+ lovely-tensors>=0.1.14
26
+ plotly==5.13.1
27
+ Pillow==9.5.0
28
+ gradio==3.37.0
29
+ ultralytics==8.0.120
30
+ -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
31
+ -e git+https://github.com/openai/CLIP.git@main#egg=clip
utils/gradio_utils.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ldm.util import create_carvekit_interface, load_and_preprocess
2
+
3
+ def load_preprocess_model():
4
+ carvekit = create_carvekit_interface()
5
+ return carvekit
6
+
7
+ def preprocess_image(models, input_im):
8
+ '''
9
+ :param input_im (PIL Image).
10
+ :return input_im (H, W, 3) array.
11
+ '''
12
+ input_im = load_and_preprocess(models, input_im)
13
+ return input_im