MC-E commited on
Commit
c05d22e
β€’
1 Parent(s): 6d6d488

first push

Browse files
This view is limited to 50 files because it contains too many changes. Β  See raw diff
Files changed (50) hide show
  1. LICENSE +201 -0
  2. README.md +8 -11
  3. app.py +230 -0
  4. configs/mm/faster_rcnn_r50_fpn_coco.py +182 -0
  5. configs/mm/hrnet_w48_coco_256x192.py +169 -0
  6. configs/stable-diffusion/sd-v1-inference.yaml +65 -0
  7. dist_util.py +91 -0
  8. ldm/__pycache__/inference_base.cpython-38.pyc +0 -0
  9. ldm/__pycache__/util.cpython-38.pyc +0 -0
  10. ldm/data/__init__.py +0 -0
  11. ldm/data/dataset_coco.py +36 -0
  12. ldm/data/dataset_depth.py +35 -0
  13. ldm/data/dataset_laion.py +130 -0
  14. ldm/data/dataset_wikiart.py +67 -0
  15. ldm/data/utils.py +60 -0
  16. ldm/inference_base.py +292 -0
  17. ldm/lr_scheduler.py +98 -0
  18. ldm/models/__pycache__/autoencoder.cpython-38.pyc +0 -0
  19. ldm/models/autoencoder.py +211 -0
  20. ldm/models/diffusion/__init__.py +0 -0
  21. ldm/models/diffusion/__pycache__/__init__.cpython-38.pyc +0 -0
  22. ldm/models/diffusion/__pycache__/ddim.cpython-38.pyc +0 -0
  23. ldm/models/diffusion/__pycache__/ddpm.cpython-38.pyc +0 -0
  24. ldm/models/diffusion/__pycache__/plms.cpython-38.pyc +0 -0
  25. ldm/models/diffusion/ddim.py +293 -0
  26. ldm/models/diffusion/ddpm.py +1329 -0
  27. ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  28. ldm/models/diffusion/dpm_solver/dpm_solver.py +1217 -0
  29. ldm/models/diffusion/dpm_solver/sampler.py +87 -0
  30. ldm/models/diffusion/plms.py +243 -0
  31. ldm/modules/__pycache__/attention.cpython-38.pyc +0 -0
  32. ldm/modules/__pycache__/ema.cpython-38.pyc +0 -0
  33. ldm/modules/attention.py +344 -0
  34. ldm/modules/diffusionmodules/__init__.py +0 -0
  35. ldm/modules/diffusionmodules/__pycache__/__init__.cpython-38.pyc +0 -0
  36. ldm/modules/diffusionmodules/__pycache__/model.cpython-38.pyc +0 -0
  37. ldm/modules/diffusionmodules/__pycache__/openaimodel.cpython-38.pyc +0 -0
  38. ldm/modules/diffusionmodules/__pycache__/util.cpython-38.pyc +0 -0
  39. ldm/modules/diffusionmodules/model.py +852 -0
  40. ldm/modules/diffusionmodules/openaimodel.py +798 -0
  41. ldm/modules/diffusionmodules/util.py +270 -0
  42. ldm/modules/distributions/__init__.py +0 -0
  43. ldm/modules/distributions/__pycache__/__init__.cpython-38.pyc +0 -0
  44. ldm/modules/distributions/__pycache__/distributions.cpython-38.pyc +0 -0
  45. ldm/modules/distributions/distributions.py +92 -0
  46. ldm/modules/ema.py +80 -0
  47. ldm/modules/encoders/__init__.py +0 -0
  48. ldm/modules/encoders/__pycache__/__init__.cpython-38.pyc +0 -0
  49. ldm/modules/encoders/__pycache__/adapter.cpython-38.pyc +0 -0
  50. ldm/modules/encoders/__pycache__/modules.cpython-38.pyc +0 -0
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.
README.md CHANGED
@@ -1,13 +1,10 @@
1
- ---
2
- title: T2I Coadapter
3
- emoji: 🏒
4
- colorFrom: green
5
- colorTo: gray
6
  sdk: gradio
7
- sdk_version: 3.21.0
8
- app_file: app.py
 
 
9
  pinned: false
10
- license: openrail
11
- ---
12
-
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
1
+ license: openrail
2
+ title: T2I-Adapter
 
 
 
3
  sdk: gradio
4
+ sdk_version: 3.19.1
5
+ emoji: 😻
6
+ colorFrom: pink
7
+ colorTo: blue
8
  pinned: false
9
+ python_version: 3.8.16
10
+ app_file: app.py
 
 
app.py ADDED
@@ -0,0 +1,230 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # demo inspired by https://huggingface.co/spaces/lambdalabs/image-mixer-demo
2
+ import argparse
3
+ import copy
4
+ import gradio as gr
5
+ import torch
6
+ from functools import partial
7
+ from itertools import chain
8
+ from torch import autocast
9
+ from pytorch_lightning import seed_everything
10
+
11
+ from basicsr.utils import tensor2img
12
+ from ldm.inference_base import DEFAULT_NEGATIVE_PROMPT, diffusion_inference, get_adapters, get_sd_models
13
+ from ldm.modules.extra_condition import api
14
+ from ldm.modules.extra_condition.api import ExtraCondition, get_cond_model
15
+ from ldm.modules.encoders.adapter import CoAdapterFuser
16
+ import os
17
+ from huggingface_hub import hf_hub_url
18
+ import subprocess
19
+ import shlex
20
+
21
+ torch.set_grad_enabled(False)
22
+
23
+ urls = {
24
+ 'TencentARC/T2I-Adapter':[
25
+ 'third-party-models/body_pose_model.pth', 'third-party-models/table5_pidinet.pth',
26
+ 'models/coadapter-canny-sd15v1.pth',
27
+ 'models/coadapter-color-sd15v1.pth',
28
+ 'models/coadapter-sketch-sd15v1.pth',
29
+ 'models/coadapter-style-sd15v1.pth',
30
+ 'models/coadapter-depth-sd15v1.pth',
31
+ 'models/coadapter-fuser-sd15v1.pth',
32
+
33
+ ],
34
+ 'runwayml/stable-diffusion-v1-5': ['v1-5-pruned-emaonly.ckpt'],
35
+ 'andite/anything-v4.0': ['anything-v4.5-pruned.ckpt', 'anything-v4.0.vae.pt'],
36
+ }
37
+
38
+ if os.path.exists('models') == False:
39
+ os.mkdir('models')
40
+ for repo in urls:
41
+ files = urls[repo]
42
+ for file in files:
43
+ url = hf_hub_url(repo, file)
44
+ name_ckp = url.split('/')[-1]
45
+ save_path = os.path.join('models',name_ckp)
46
+ if os.path.exists(save_path) == False:
47
+ subprocess.run(shlex.split(f'wget {url} -O {save_path}'))
48
+
49
+ supported_cond = ['style', 'color', 'sketch', 'depth', 'canny']
50
+
51
+ # config
52
+ parser = argparse.ArgumentParser()
53
+ parser.add_argument(
54
+ '--sd_ckpt',
55
+ type=str,
56
+ default='models/v1-5-pruned-emaonly.ckpt',
57
+ help='path to checkpoint of stable diffusion model, both .ckpt and .safetensor are supported',
58
+ )
59
+ parser.add_argument(
60
+ '--vae_ckpt',
61
+ type=str,
62
+ default=None,
63
+ help='vae checkpoint, anime SD models usually have seperate vae ckpt that need to be loaded',
64
+ )
65
+ global_opt = parser.parse_args()
66
+ global_opt.config = 'configs/stable-diffusion/sd-v1-inference.yaml'
67
+ for cond_name in supported_cond:
68
+ setattr(global_opt, f'{cond_name}_adapter_ckpt', f'models/coadapter-{cond_name}-sd15v1.pth')
69
+ global_opt.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
70
+ global_opt.max_resolution = 512 * 512
71
+ global_opt.sampler = 'ddim'
72
+ global_opt.cond_weight = 1.0
73
+ global_opt.C = 4
74
+ global_opt.f = 8
75
+ #TODO: expose style_cond_tau to users
76
+ global_opt.style_cond_tau = 1.0
77
+
78
+ # stable-diffusion model
79
+ sd_model, sampler = get_sd_models(global_opt)
80
+ # adapters and models to processing condition inputs
81
+ adapters = {}
82
+ cond_models = {}
83
+
84
+ torch.cuda.empty_cache()
85
+
86
+ # fuser is indispensable
87
+ coadapter_fuser = CoAdapterFuser(unet_channels=[320, 640, 1280, 1280], width=768, num_head=8, n_layes=3)
88
+ coadapter_fuser.load_state_dict(torch.load(f'models/coadapter-fuser-sd15v1.pth'))
89
+ coadapter_fuser = coadapter_fuser.to(global_opt.device)
90
+
91
+
92
+ def run(*args):
93
+ with torch.inference_mode(), \
94
+ sd_model.ema_scope(), \
95
+ autocast('cuda'):
96
+
97
+ inps = []
98
+ for i in range(0, len(args) - 8, len(supported_cond)):
99
+ inps.append(args[i:i + len(supported_cond)])
100
+
101
+ opt = copy.deepcopy(global_opt)
102
+ opt.prompt, opt.neg_prompt, opt.scale, opt.n_samples, opt.seed, opt.steps, opt.resize_short_edge, opt.cond_tau \
103
+ = args[-8:]
104
+
105
+ conds = []
106
+ activated_conds = []
107
+ for idx, (b, im1, im2, cond_weight) in enumerate(zip(*inps)):
108
+ cond_name = supported_cond[idx]
109
+ if b == 'Nothing':
110
+ if cond_name in adapters:
111
+ adapters[cond_name]['model'] = adapters[cond_name]['model'].cpu()
112
+ else:
113
+ activated_conds.append(cond_name)
114
+ if cond_name in adapters:
115
+ adapters[cond_name]['model'] = adapters[cond_name]['model'].to(opt.device)
116
+ else:
117
+ adapters[cond_name] = get_adapters(opt, getattr(ExtraCondition, cond_name))
118
+ adapters[cond_name]['cond_weight'] = cond_weight
119
+
120
+ process_cond_module = getattr(api, f'get_cond_{cond_name}')
121
+
122
+ if b == 'Image':
123
+ if cond_name not in cond_models:
124
+ cond_models[cond_name] = get_cond_model(opt, getattr(ExtraCondition, cond_name))
125
+ conds.append(process_cond_module(opt, im1, 'image', cond_models[cond_name]))
126
+ else:
127
+ conds.append(process_cond_module(opt, im2, cond_name, None))
128
+
129
+ features = dict()
130
+ for idx, cond_name in enumerate(activated_conds):
131
+ cur_feats = adapters[cond_name]['model'](conds[idx])
132
+ if isinstance(cur_feats, list):
133
+ for i in range(len(cur_feats)):
134
+ cur_feats[i] *= adapters[cond_name]['cond_weight']
135
+ else:
136
+ cur_feats *= adapters[cond_name]['cond_weight']
137
+ features[cond_name] = cur_feats
138
+
139
+ adapter_features, append_to_context = coadapter_fuser(features)
140
+
141
+ output_conds = []
142
+ for cond in conds:
143
+ output_conds.append(tensor2img(cond, rgb2bgr=False))
144
+
145
+ ims = []
146
+ seed_everything(opt.seed)
147
+ for _ in range(opt.n_samples):
148
+ result = diffusion_inference(opt, sd_model, sampler, adapter_features, append_to_context)
149
+ ims.append(tensor2img(result, rgb2bgr=False))
150
+
151
+ # Clear GPU memory cache so less likely to OOM
152
+ torch.cuda.empty_cache()
153
+ return ims, output_conds
154
+
155
+
156
+ def change_visible(im1, im2, val):
157
+ outputs = {}
158
+ if val == "Image":
159
+ outputs[im1] = gr.update(visible=True)
160
+ outputs[im2] = gr.update(visible=False)
161
+ elif val == "Nothing":
162
+ outputs[im1] = gr.update(visible=False)
163
+ outputs[im2] = gr.update(visible=False)
164
+ else:
165
+ outputs[im1] = gr.update(visible=False)
166
+ outputs[im2] = gr.update(visible=True)
167
+ return outputs
168
+
169
+
170
+ DESCRIPTION = '''# CoAdapter
171
+ [Paper](https://arxiv.org/abs/2302.08453) [GitHub](https://github.com/TencentARC/T2I-Adapter)
172
+
173
+ This gradio demo is for a simple experience of CoAdapter:
174
+ '''
175
+ with gr.Blocks(title="CoAdapter", css=".gr-box {border-color: #8136e2}") as demo:
176
+ gr.Markdown(DESCRIPTION)
177
+
178
+ btns = []
179
+ ims1 = []
180
+ ims2 = []
181
+ cond_weights = []
182
+
183
+ with gr.Row():
184
+ for cond_name in supported_cond:
185
+ with gr.Box():
186
+ with gr.Column():
187
+ btn1 = gr.Radio(
188
+ choices=["Image", cond_name, "Nothing"],
189
+ label=f"Input type for {cond_name}",
190
+ interactive=True,
191
+ value="Nothing",
192
+ )
193
+ im1 = gr.Image(source='upload', label="Image", interactive=True, visible=False, type="numpy")
194
+ im2 = gr.Image(source='upload', label=cond_name, interactive=True, visible=False, type="numpy")
195
+ cond_weight = gr.Slider(
196
+ label="Condition weight", minimum=0, maximum=5, step=0.05, value=1, interactive=True)
197
+
198
+ fn = partial(change_visible, im1, im2)
199
+ btn1.change(fn=fn, inputs=[btn1], outputs=[im1, im2], queue=False)
200
+
201
+ btns.append(btn1)
202
+ ims1.append(im1)
203
+ ims2.append(im2)
204
+ cond_weights.append(cond_weight)
205
+
206
+ with gr.Column():
207
+ prompt = gr.Textbox(label="Prompt")
208
+ neg_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
209
+ scale = gr.Slider(label="Guidance Scale (Classifier free guidance)", value=7.5, minimum=1, maximum=20, step=0.1)
210
+ n_samples = gr.Slider(label="Num samples", value=1, minimum=1, maximum=8, step=1)
211
+ seed = gr.Slider(label="Seed", value=42, minimum=0, maximum=10000, step=1)
212
+ steps = gr.Slider(label="Steps", value=50, minimum=10, maximum=100, step=1)
213
+ resize_short_edge = gr.Slider(label="Image resolution", value=512, minimum=320, maximum=1024, step=1)
214
+ cond_tau = gr.Slider(
215
+ label="timestamp parameter that determines until which step the adapter is applied",
216
+ value=1.0,
217
+ minimum=0.1,
218
+ maximum=1.0,
219
+ step=0.05)
220
+
221
+ with gr.Row():
222
+ submit = gr.Button("Generate")
223
+ output = gr.Gallery().style(grid=2, height='auto')
224
+ cond = gr.Gallery().style(grid=2, height='auto')
225
+
226
+ inps = list(chain(btns, ims1, ims2, cond_weights))
227
+ inps.extend([prompt, neg_prompt, scale, n_samples, seed, steps, resize_short_edge, cond_tau])
228
+ submit.click(fn=run, inputs=inps, outputs=[output, cond])
229
+ # demo.launch()
230
+ demo.launch(server_port=43343, server_name='0.0.0.0')
configs/mm/faster_rcnn_r50_fpn_coco.py ADDED
@@ -0,0 +1,182 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ checkpoint_config = dict(interval=1)
2
+ # yapf:disable
3
+ log_config = dict(
4
+ interval=50,
5
+ hooks=[
6
+ dict(type='TextLoggerHook'),
7
+ # dict(type='TensorboardLoggerHook')
8
+ ])
9
+ # yapf:enable
10
+ dist_params = dict(backend='nccl')
11
+ log_level = 'INFO'
12
+ load_from = None
13
+ resume_from = None
14
+ workflow = [('train', 1)]
15
+ # optimizer
16
+ optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
17
+ optimizer_config = dict(grad_clip=None)
18
+ # learning policy
19
+ lr_config = dict(
20
+ policy='step',
21
+ warmup='linear',
22
+ warmup_iters=500,
23
+ warmup_ratio=0.001,
24
+ step=[8, 11])
25
+ total_epochs = 12
26
+
27
+ model = dict(
28
+ type='FasterRCNN',
29
+ pretrained='torchvision://resnet50',
30
+ backbone=dict(
31
+ type='ResNet',
32
+ depth=50,
33
+ num_stages=4,
34
+ out_indices=(0, 1, 2, 3),
35
+ frozen_stages=1,
36
+ norm_cfg=dict(type='BN', requires_grad=True),
37
+ norm_eval=True,
38
+ style='pytorch'),
39
+ neck=dict(
40
+ type='FPN',
41
+ in_channels=[256, 512, 1024, 2048],
42
+ out_channels=256,
43
+ num_outs=5),
44
+ rpn_head=dict(
45
+ type='RPNHead',
46
+ in_channels=256,
47
+ feat_channels=256,
48
+ anchor_generator=dict(
49
+ type='AnchorGenerator',
50
+ scales=[8],
51
+ ratios=[0.5, 1.0, 2.0],
52
+ strides=[4, 8, 16, 32, 64]),
53
+ bbox_coder=dict(
54
+ type='DeltaXYWHBBoxCoder',
55
+ target_means=[.0, .0, .0, .0],
56
+ target_stds=[1.0, 1.0, 1.0, 1.0]),
57
+ loss_cls=dict(
58
+ type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
59
+ loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
60
+ roi_head=dict(
61
+ type='StandardRoIHead',
62
+ bbox_roi_extractor=dict(
63
+ type='SingleRoIExtractor',
64
+ roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
65
+ out_channels=256,
66
+ featmap_strides=[4, 8, 16, 32]),
67
+ bbox_head=dict(
68
+ type='Shared2FCBBoxHead',
69
+ in_channels=256,
70
+ fc_out_channels=1024,
71
+ roi_feat_size=7,
72
+ num_classes=80,
73
+ bbox_coder=dict(
74
+ type='DeltaXYWHBBoxCoder',
75
+ target_means=[0., 0., 0., 0.],
76
+ target_stds=[0.1, 0.1, 0.2, 0.2]),
77
+ reg_class_agnostic=False,
78
+ loss_cls=dict(
79
+ type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
80
+ loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
81
+ # model training and testing settings
82
+ train_cfg=dict(
83
+ rpn=dict(
84
+ assigner=dict(
85
+ type='MaxIoUAssigner',
86
+ pos_iou_thr=0.7,
87
+ neg_iou_thr=0.3,
88
+ min_pos_iou=0.3,
89
+ match_low_quality=True,
90
+ ignore_iof_thr=-1),
91
+ sampler=dict(
92
+ type='RandomSampler',
93
+ num=256,
94
+ pos_fraction=0.5,
95
+ neg_pos_ub=-1,
96
+ add_gt_as_proposals=False),
97
+ allowed_border=-1,
98
+ pos_weight=-1,
99
+ debug=False),
100
+ rpn_proposal=dict(
101
+ nms_pre=2000,
102
+ max_per_img=1000,
103
+ nms=dict(type='nms', iou_threshold=0.7),
104
+ min_bbox_size=0),
105
+ rcnn=dict(
106
+ assigner=dict(
107
+ type='MaxIoUAssigner',
108
+ pos_iou_thr=0.5,
109
+ neg_iou_thr=0.5,
110
+ min_pos_iou=0.5,
111
+ match_low_quality=False,
112
+ ignore_iof_thr=-1),
113
+ sampler=dict(
114
+ type='RandomSampler',
115
+ num=512,
116
+ pos_fraction=0.25,
117
+ neg_pos_ub=-1,
118
+ add_gt_as_proposals=True),
119
+ pos_weight=-1,
120
+ debug=False)),
121
+ test_cfg=dict(
122
+ rpn=dict(
123
+ nms_pre=1000,
124
+ max_per_img=1000,
125
+ nms=dict(type='nms', iou_threshold=0.7),
126
+ min_bbox_size=0),
127
+ rcnn=dict(
128
+ score_thr=0.05,
129
+ nms=dict(type='nms', iou_threshold=0.5),
130
+ max_per_img=100)
131
+ # soft-nms is also supported for rcnn testing
132
+ # e.g., nms=dict(type='soft_nms', iou_threshold=0.5, min_score=0.05)
133
+ ))
134
+
135
+ dataset_type = 'CocoDataset'
136
+ data_root = 'data/coco'
137
+ img_norm_cfg = dict(
138
+ mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
139
+ train_pipeline = [
140
+ dict(type='LoadImageFromFile'),
141
+ dict(type='LoadAnnotations', with_bbox=True),
142
+ dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
143
+ dict(type='RandomFlip', flip_ratio=0.5),
144
+ dict(type='Normalize', **img_norm_cfg),
145
+ dict(type='Pad', size_divisor=32),
146
+ dict(type='DefaultFormatBundle'),
147
+ dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),
148
+ ]
149
+ test_pipeline = [
150
+ dict(type='LoadImageFromFile'),
151
+ dict(
152
+ type='MultiScaleFlipAug',
153
+ img_scale=(1333, 800),
154
+ flip=False,
155
+ transforms=[
156
+ dict(type='Resize', keep_ratio=True),
157
+ dict(type='RandomFlip'),
158
+ dict(type='Normalize', **img_norm_cfg),
159
+ dict(type='Pad', size_divisor=32),
160
+ dict(type='DefaultFormatBundle'),
161
+ dict(type='Collect', keys=['img']),
162
+ ])
163
+ ]
164
+ data = dict(
165
+ samples_per_gpu=2,
166
+ workers_per_gpu=2,
167
+ train=dict(
168
+ type=dataset_type,
169
+ ann_file=f'{data_root}/annotations/instances_train2017.json',
170
+ img_prefix=f'{data_root}/train2017/',
171
+ pipeline=train_pipeline),
172
+ val=dict(
173
+ type=dataset_type,
174
+ ann_file=f'{data_root}/annotations/instances_val2017.json',
175
+ img_prefix=f'{data_root}/val2017/',
176
+ pipeline=test_pipeline),
177
+ test=dict(
178
+ type=dataset_type,
179
+ ann_file=f'{data_root}/annotations/instances_val2017.json',
180
+ img_prefix=f'{data_root}/val2017/',
181
+ pipeline=test_pipeline))
182
+ evaluation = dict(interval=1, metric='bbox')
configs/mm/hrnet_w48_coco_256x192.py ADDED
@@ -0,0 +1,169 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # _base_ = [
2
+ # '../../../../_base_/default_runtime.py',
3
+ # '../../../../_base_/datasets/coco.py'
4
+ # ]
5
+ evaluation = dict(interval=10, metric='mAP', save_best='AP')
6
+
7
+ optimizer = dict(
8
+ type='Adam',
9
+ lr=5e-4,
10
+ )
11
+ optimizer_config = dict(grad_clip=None)
12
+ # learning policy
13
+ lr_config = dict(
14
+ policy='step',
15
+ warmup='linear',
16
+ warmup_iters=500,
17
+ warmup_ratio=0.001,
18
+ step=[170, 200])
19
+ total_epochs = 210
20
+ channel_cfg = dict(
21
+ num_output_channels=17,
22
+ dataset_joints=17,
23
+ dataset_channel=[
24
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
25
+ ],
26
+ inference_channel=[
27
+ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
28
+ ])
29
+
30
+ # model settings
31
+ model = dict(
32
+ type='TopDown',
33
+ pretrained='https://download.openmmlab.com/mmpose/'
34
+ 'pretrain_models/hrnet_w48-8ef0771d.pth',
35
+ backbone=dict(
36
+ type='HRNet',
37
+ in_channels=3,
38
+ extra=dict(
39
+ stage1=dict(
40
+ num_modules=1,
41
+ num_branches=1,
42
+ block='BOTTLENECK',
43
+ num_blocks=(4, ),
44
+ num_channels=(64, )),
45
+ stage2=dict(
46
+ num_modules=1,
47
+ num_branches=2,
48
+ block='BASIC',
49
+ num_blocks=(4, 4),
50
+ num_channels=(48, 96)),
51
+ stage3=dict(
52
+ num_modules=4,
53
+ num_branches=3,
54
+ block='BASIC',
55
+ num_blocks=(4, 4, 4),
56
+ num_channels=(48, 96, 192)),
57
+ stage4=dict(
58
+ num_modules=3,
59
+ num_branches=4,
60
+ block='BASIC',
61
+ num_blocks=(4, 4, 4, 4),
62
+ num_channels=(48, 96, 192, 384))),
63
+ ),
64
+ keypoint_head=dict(
65
+ type='TopdownHeatmapSimpleHead',
66
+ in_channels=48,
67
+ out_channels=channel_cfg['num_output_channels'],
68
+ num_deconv_layers=0,
69
+ extra=dict(final_conv_kernel=1, ),
70
+ loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),
71
+ train_cfg=dict(),
72
+ test_cfg=dict(
73
+ flip_test=True,
74
+ post_process='default',
75
+ shift_heatmap=True,
76
+ modulate_kernel=11))
77
+
78
+ data_cfg = dict(
79
+ image_size=[192, 256],
80
+ heatmap_size=[48, 64],
81
+ num_output_channels=channel_cfg['num_output_channels'],
82
+ num_joints=channel_cfg['dataset_joints'],
83
+ dataset_channel=channel_cfg['dataset_channel'],
84
+ inference_channel=channel_cfg['inference_channel'],
85
+ soft_nms=False,
86
+ nms_thr=1.0,
87
+ oks_thr=0.9,
88
+ vis_thr=0.2,
89
+ use_gt_bbox=False,
90
+ det_bbox_thr=0.0,
91
+ bbox_file='data/coco/person_detection_results/'
92
+ 'COCO_val2017_detections_AP_H_56_person.json',
93
+ )
94
+
95
+ train_pipeline = [
96
+ dict(type='LoadImageFromFile'),
97
+ dict(type='TopDownGetBboxCenterScale', padding=1.25),
98
+ dict(type='TopDownRandomShiftBboxCenter', shift_factor=0.16, prob=0.3),
99
+ dict(type='TopDownRandomFlip', flip_prob=0.5),
100
+ dict(
101
+ type='TopDownHalfBodyTransform',
102
+ num_joints_half_body=8,
103
+ prob_half_body=0.3),
104
+ dict(
105
+ type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5),
106
+ dict(type='TopDownAffine'),
107
+ dict(type='ToTensor'),
108
+ dict(
109
+ type='NormalizeTensor',
110
+ mean=[0.485, 0.456, 0.406],
111
+ std=[0.229, 0.224, 0.225]),
112
+ dict(type='TopDownGenerateTarget', sigma=2),
113
+ dict(
114
+ type='Collect',
115
+ keys=['img', 'target', 'target_weight'],
116
+ meta_keys=[
117
+ 'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale',
118
+ 'rotation', 'bbox_score', 'flip_pairs'
119
+ ]),
120
+ ]
121
+
122
+ val_pipeline = [
123
+ dict(type='LoadImageFromFile'),
124
+ dict(type='TopDownGetBboxCenterScale', padding=1.25),
125
+ dict(type='TopDownAffine'),
126
+ dict(type='ToTensor'),
127
+ dict(
128
+ type='NormalizeTensor',
129
+ mean=[0.485, 0.456, 0.406],
130
+ std=[0.229, 0.224, 0.225]),
131
+ dict(
132
+ type='Collect',
133
+ keys=['img'],
134
+ meta_keys=[
135
+ 'image_file', 'center', 'scale', 'rotation', 'bbox_score',
136
+ 'flip_pairs'
137
+ ]),
138
+ ]
139
+
140
+ test_pipeline = val_pipeline
141
+
142
+ data_root = 'data/coco'
143
+ data = dict(
144
+ samples_per_gpu=32,
145
+ workers_per_gpu=2,
146
+ val_dataloader=dict(samples_per_gpu=32),
147
+ test_dataloader=dict(samples_per_gpu=32),
148
+ train=dict(
149
+ type='TopDownCocoDataset',
150
+ ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',
151
+ img_prefix=f'{data_root}/train2017/',
152
+ data_cfg=data_cfg,
153
+ pipeline=train_pipeline,
154
+ dataset_info={{_base_.dataset_info}}),
155
+ val=dict(
156
+ type='TopDownCocoDataset',
157
+ ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',
158
+ img_prefix=f'{data_root}/val2017/',
159
+ data_cfg=data_cfg,
160
+ pipeline=val_pipeline,
161
+ dataset_info={{_base_.dataset_info}}),
162
+ test=dict(
163
+ type='TopDownCocoDataset',
164
+ ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',
165
+ img_prefix=f'{data_root}/val2017/',
166
+ data_cfg=data_cfg,
167
+ pipeline=test_pipeline,
168
+ dataset_info={{_base_.dataset_info}}),
169
+ )
configs/stable-diffusion/sd-v1-inference.yaml ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ base_learning_rate: 1.0e-04
3
+ target: ldm.models.diffusion.ddpm.LatentDiffusion
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: "jpg"
11
+ cond_stage_key: "txt"
12
+ image_size: 64
13
+ channels: 4
14
+ cond_stage_trainable: false # Note: different from the one we trained before
15
+ conditioning_key: crossattn
16
+ monitor: val/loss_simple_ema
17
+ scale_factor: 0.18215
18
+ use_ema: False
19
+
20
+ unet_config:
21
+ target: ldm.modules.diffusionmodules.openaimodel.UNetModel
22
+ params:
23
+ use_fp16: True
24
+ image_size: 32 # unused
25
+ in_channels: 4
26
+ out_channels: 4
27
+ model_channels: 320
28
+ attention_resolutions: [ 4, 2, 1 ]
29
+ num_res_blocks: 2
30
+ channel_mult: [ 1, 2, 4, 4 ]
31
+ num_heads: 8
32
+ use_spatial_transformer: True
33
+ transformer_depth: 1
34
+ context_dim: 768
35
+ use_checkpoint: True
36
+ legacy: False
37
+
38
+ first_stage_config:
39
+ target: ldm.models.autoencoder.AutoencoderKL
40
+ params:
41
+ embed_dim: 4
42
+ monitor: val/rec_loss
43
+ ddconfig:
44
+ double_z: true
45
+ z_channels: 4
46
+ resolution: 512
47
+ in_channels: 3
48
+ out_ch: 3
49
+ ch: 128
50
+ ch_mult:
51
+ - 1
52
+ - 2
53
+ - 4
54
+ - 4
55
+ num_res_blocks: 2
56
+ attn_resolutions: []
57
+ dropout: 0.0
58
+ lossconfig:
59
+ target: torch.nn.Identity
60
+
61
+ cond_stage_config:
62
+ target: ldm.modules.encoders.modules.WebUIFrozenCLIPEmebedder
63
+ params:
64
+ version: openai/clip-vit-large-patch14
65
+ layer: last
dist_util.py ADDED
@@ -0,0 +1,91 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py # noqa: E501
2
+ import functools
3
+ import os
4
+ import subprocess
5
+ import torch
6
+ import torch.distributed as dist
7
+ import torch.multiprocessing as mp
8
+ from torch.nn.parallel import DataParallel, DistributedDataParallel
9
+
10
+
11
+ def init_dist(launcher, backend='nccl', **kwargs):
12
+ if mp.get_start_method(allow_none=True) is None:
13
+ mp.set_start_method('spawn')
14
+ if launcher == 'pytorch':
15
+ _init_dist_pytorch(backend, **kwargs)
16
+ elif launcher == 'slurm':
17
+ _init_dist_slurm(backend, **kwargs)
18
+ else:
19
+ raise ValueError(f'Invalid launcher type: {launcher}')
20
+
21
+
22
+ def _init_dist_pytorch(backend, **kwargs):
23
+ rank = int(os.environ['RANK'])
24
+ num_gpus = torch.cuda.device_count()
25
+ torch.cuda.set_device(rank % num_gpus)
26
+ dist.init_process_group(backend=backend, **kwargs)
27
+
28
+
29
+ def _init_dist_slurm(backend, port=None):
30
+ """Initialize slurm distributed training environment.
31
+
32
+ If argument ``port`` is not specified, then the master port will be system
33
+ environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
34
+ environment variable, then a default port ``29500`` will be used.
35
+
36
+ Args:
37
+ backend (str): Backend of torch.distributed.
38
+ port (int, optional): Master port. Defaults to None.
39
+ """
40
+ proc_id = int(os.environ['SLURM_PROCID'])
41
+ ntasks = int(os.environ['SLURM_NTASKS'])
42
+ node_list = os.environ['SLURM_NODELIST']
43
+ num_gpus = torch.cuda.device_count()
44
+ torch.cuda.set_device(proc_id % num_gpus)
45
+ addr = subprocess.getoutput(f'scontrol show hostname {node_list} | head -n1')
46
+ # specify master port
47
+ if port is not None:
48
+ os.environ['MASTER_PORT'] = str(port)
49
+ elif 'MASTER_PORT' in os.environ:
50
+ pass # use MASTER_PORT in the environment variable
51
+ else:
52
+ # 29500 is torch.distributed default port
53
+ os.environ['MASTER_PORT'] = '29500'
54
+ os.environ['MASTER_ADDR'] = addr
55
+ os.environ['WORLD_SIZE'] = str(ntasks)
56
+ os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
57
+ os.environ['RANK'] = str(proc_id)
58
+ dist.init_process_group(backend=backend)
59
+
60
+
61
+ def get_dist_info():
62
+ if dist.is_available():
63
+ initialized = dist.is_initialized()
64
+ else:
65
+ initialized = False
66
+ if initialized:
67
+ rank = dist.get_rank()
68
+ world_size = dist.get_world_size()
69
+ else:
70
+ rank = 0
71
+ world_size = 1
72
+ return rank, world_size
73
+
74
+
75
+ def master_only(func):
76
+
77
+ @functools.wraps(func)
78
+ def wrapper(*args, **kwargs):
79
+ rank, _ = get_dist_info()
80
+ if rank == 0:
81
+ return func(*args, **kwargs)
82
+
83
+ return wrapper
84
+
85
+ def get_bare_model(net):
86
+ """Get bare model, especially under wrapping with
87
+ DistributedDataParallel or DataParallel.
88
+ """
89
+ if isinstance(net, (DataParallel, DistributedDataParallel)):
90
+ net = net.module
91
+ return net
ldm/__pycache__/inference_base.cpython-38.pyc ADDED
Binary file (6.28 kB). View file
 
ldm/__pycache__/util.cpython-38.pyc ADDED
Binary file (6.23 kB). View file
 
ldm/data/__init__.py ADDED
File without changes
ldm/data/dataset_coco.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import cv2
3
+ import os
4
+ from basicsr.utils import img2tensor
5
+
6
+
7
+ class dataset_coco_mask_color():
8
+ def __init__(self, path_json, root_path_im, root_path_mask, image_size):
9
+ super(dataset_coco_mask_color, self).__init__()
10
+ with open(path_json, 'r', encoding='utf-8') as fp:
11
+ data = json.load(fp)
12
+ data = data['annotations']
13
+ self.files = []
14
+ self.root_path_im = root_path_im
15
+ self.root_path_mask = root_path_mask
16
+ for file in data:
17
+ name = "%012d.png" % file['image_id']
18
+ self.files.append({'name': name, 'sentence': file['caption']})
19
+
20
+ def __getitem__(self, idx):
21
+ file = self.files[idx]
22
+ name = file['name']
23
+ # print(os.path.join(self.root_path_im, name))
24
+ im = cv2.imread(os.path.join(self.root_path_im, name.replace('.png', '.jpg')))
25
+ im = cv2.resize(im, (512, 512))
26
+ im = img2tensor(im, bgr2rgb=True, float32=True) / 255.
27
+
28
+ mask = cv2.imread(os.path.join(self.root_path_mask, name)) # [:,:,0]
29
+ mask = cv2.resize(mask, (512, 512))
30
+ mask = img2tensor(mask, bgr2rgb=True, float32=True) / 255. # [0].unsqueeze(0)#/255.
31
+
32
+ sentence = file['sentence']
33
+ return {'im': im, 'mask': mask, 'sentence': sentence}
34
+
35
+ def __len__(self):
36
+ return len(self.files)
ldm/data/dataset_depth.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import cv2
3
+ import os
4
+ from basicsr.utils import img2tensor
5
+
6
+
7
+ class DepthDataset():
8
+ def __init__(self, meta_file):
9
+ super(DepthDataset, self).__init__()
10
+
11
+ self.files = []
12
+ with open(meta_file, 'r') as f:
13
+ lines = f.readlines()
14
+ for line in lines:
15
+ img_path = line.strip()
16
+ depth_img_path = img_path.rsplit('.', 1)[0] + '.depth.png'
17
+ txt_path = img_path.rsplit('.', 1)[0] + '.txt'
18
+ self.files.append({'img_path': img_path, 'depth_img_path': depth_img_path, 'txt_path': txt_path})
19
+
20
+ def __getitem__(self, idx):
21
+ file = self.files[idx]
22
+
23
+ im = cv2.imread(file['img_path'])
24
+ im = img2tensor(im, bgr2rgb=True, float32=True) / 255.
25
+
26
+ depth = cv2.imread(file['depth_img_path']) # [:,:,0]
27
+ depth = img2tensor(depth, bgr2rgb=True, float32=True) / 255. # [0].unsqueeze(0)#/255.
28
+
29
+ with open(file['txt_path'], 'r') as fs:
30
+ sentence = fs.readline().strip()
31
+
32
+ return {'im': im, 'depth': depth, 'sentence': sentence}
33
+
34
+ def __len__(self):
35
+ return len(self.files)
ldm/data/dataset_laion.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import numpy as np
4
+ import os
5
+ import pytorch_lightning as pl
6
+ import torch
7
+ import webdataset as wds
8
+ from torchvision.transforms import transforms
9
+
10
+ from ldm.util import instantiate_from_config
11
+
12
+
13
+ def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True):
14
+ """Take a list of samples (as dictionary) and create a batch, preserving the keys.
15
+ If `tensors` is True, `ndarray` objects are combined into
16
+ tensor batches.
17
+ :param dict samples: list of samples
18
+ :param bool tensors: whether to turn lists of ndarrays into a single ndarray
19
+ :returns: single sample consisting of a batch
20
+ :rtype: dict
21
+ """
22
+ keys = set.intersection(*[set(sample.keys()) for sample in samples])
23
+ batched = {key: [] for key in keys}
24
+
25
+ for s in samples:
26
+ [batched[key].append(s[key]) for key in batched]
27
+
28
+ result = {}
29
+ for key in batched:
30
+ if isinstance(batched[key][0], (int, float)):
31
+ if combine_scalars:
32
+ result[key] = np.array(list(batched[key]))
33
+ elif isinstance(batched[key][0], torch.Tensor):
34
+ if combine_tensors:
35
+ result[key] = torch.stack(list(batched[key]))
36
+ elif isinstance(batched[key][0], np.ndarray):
37
+ if combine_tensors:
38
+ result[key] = np.array(list(batched[key]))
39
+ else:
40
+ result[key] = list(batched[key])
41
+ return result
42
+
43
+
44
+ class WebDataModuleFromConfig(pl.LightningDataModule):
45
+
46
+ def __init__(self,
47
+ tar_base,
48
+ batch_size,
49
+ train=None,
50
+ validation=None,
51
+ test=None,
52
+ num_workers=4,
53
+ multinode=True,
54
+ min_size=None,
55
+ max_pwatermark=1.0,
56
+ **kwargs):
57
+ super().__init__()
58
+ print(f'Setting tar base to {tar_base}')
59
+ self.tar_base = tar_base
60
+ self.batch_size = batch_size
61
+ self.num_workers = num_workers
62
+ self.train = train
63
+ self.validation = validation
64
+ self.test = test
65
+ self.multinode = multinode
66
+ self.min_size = min_size # filter out very small images
67
+ self.max_pwatermark = max_pwatermark # filter out watermarked images
68
+
69
+ def make_loader(self, dataset_config):
70
+ image_transforms = [instantiate_from_config(tt) for tt in dataset_config.image_transforms]
71
+ image_transforms = transforms.Compose(image_transforms)
72
+
73
+ process = instantiate_from_config(dataset_config['process'])
74
+
75
+ shuffle = dataset_config.get('shuffle', 0)
76
+ shardshuffle = shuffle > 0
77
+
78
+ nodesplitter = wds.shardlists.split_by_node if self.multinode else wds.shardlists.single_node_only
79
+
80
+ tars = os.path.join(self.tar_base, dataset_config.shards)
81
+
82
+ dset = wds.WebDataset(
83
+ tars, nodesplitter=nodesplitter, shardshuffle=shardshuffle,
84
+ handler=wds.warn_and_continue).repeat().shuffle(shuffle)
85
+ print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.')
86
+
87
+ dset = (
88
+ dset.select(self.filter_keys).decode('pil',
89
+ handler=wds.warn_and_continue).select(self.filter_size).map_dict(
90
+ jpg=image_transforms, handler=wds.warn_and_continue).map(process))
91
+ dset = (dset.batched(self.batch_size, partial=False, collation_fn=dict_collation_fn))
92
+
93
+ loader = wds.WebLoader(dset, batch_size=None, shuffle=False, num_workers=self.num_workers)
94
+
95
+ return loader
96
+
97
+ def filter_size(self, x):
98
+ if self.min_size is None:
99
+ return True
100
+ try:
101
+ return x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size and x[
102
+ 'json']['pwatermark'] <= self.max_pwatermark
103
+ except Exception:
104
+ return False
105
+
106
+ def filter_keys(self, x):
107
+ try:
108
+ return ("jpg" in x) and ("txt" in x)
109
+ except Exception:
110
+ return False
111
+
112
+ def train_dataloader(self):
113
+ return self.make_loader(self.train)
114
+
115
+ def val_dataloader(self):
116
+ return None
117
+
118
+ def test_dataloader(self):
119
+ return None
120
+
121
+
122
+ if __name__ == '__main__':
123
+ from omegaconf import OmegaConf
124
+ config = OmegaConf.load("configs/stable-diffusion/train_canny_sd_v1.yaml")
125
+ datamod = WebDataModuleFromConfig(**config["data"]["params"])
126
+ dataloader = datamod.train_dataloader()
127
+
128
+ for batch in dataloader:
129
+ print(batch.keys())
130
+ print(batch['jpg'].shape)
ldm/data/dataset_wikiart.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import os.path
3
+
4
+ from PIL import Image
5
+ from torch.utils.data import DataLoader
6
+
7
+ from transformers import CLIPProcessor
8
+ from torchvision.transforms import transforms
9
+
10
+ import pytorch_lightning as pl
11
+
12
+
13
+ class WikiArtDataset():
14
+ def __init__(self, meta_file):
15
+ super(WikiArtDataset, self).__init__()
16
+
17
+ self.files = []
18
+ with open(meta_file, 'r') as f:
19
+ js = json.load(f)
20
+ for img_path in js:
21
+ img_name = os.path.splitext(os.path.basename(img_path))[0]
22
+ caption = img_name.split('_')[-1]
23
+ caption = caption.split('-')
24
+ j = len(caption) - 1
25
+ while j >= 0:
26
+ if not caption[j].isdigit():
27
+ break
28
+ j -= 1
29
+ if j < 0:
30
+ continue
31
+ sentence = ' '.join(caption[:j + 1])
32
+ self.files.append({'img_path': os.path.join('datasets/wikiart', img_path), 'sentence': sentence})
33
+
34
+ version = 'openai/clip-vit-large-patch14'
35
+ self.processor = CLIPProcessor.from_pretrained(version)
36
+
37
+ self.jpg_transform = transforms.Compose([
38
+ transforms.Resize(512),
39
+ transforms.RandomCrop(512),
40
+ transforms.ToTensor(),
41
+ ])
42
+
43
+ def __getitem__(self, idx):
44
+ file = self.files[idx]
45
+
46
+ im = Image.open(file['img_path'])
47
+
48
+ im_tensor = self.jpg_transform(im)
49
+
50
+ clip_im = self.processor(images=im, return_tensors="pt")['pixel_values'][0]
51
+
52
+ return {'jpg': im_tensor, 'style': clip_im, 'txt': file['sentence']}
53
+
54
+ def __len__(self):
55
+ return len(self.files)
56
+
57
+
58
+ class WikiArtDataModule(pl.LightningDataModule):
59
+ def __init__(self, meta_file, batch_size, num_workers):
60
+ super(WikiArtDataModule, self).__init__()
61
+ self.train_dataset = WikiArtDataset(meta_file)
62
+ self.batch_size = batch_size
63
+ self.num_workers = num_workers
64
+
65
+ def train_dataloader(self):
66
+ return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers,
67
+ pin_memory=True)
ldm/data/utils.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import cv2
4
+ import numpy as np
5
+ from torchvision.transforms import transforms
6
+ from torchvision.transforms.functional import to_tensor
7
+ from transformers import CLIPProcessor
8
+
9
+ from basicsr.utils import img2tensor
10
+
11
+
12
+ class AddCannyFreezeThreshold(object):
13
+
14
+ def __init__(self, low_threshold=100, high_threshold=200):
15
+ self.low_threshold = low_threshold
16
+ self.high_threshold = high_threshold
17
+
18
+ def __call__(self, sample):
19
+ # sample['jpg'] is PIL image
20
+ x = sample['jpg']
21
+ img = cv2.cvtColor(np.array(x), cv2.COLOR_RGB2BGR)
22
+ canny = cv2.Canny(img, self.low_threshold, self.high_threshold)[..., None]
23
+ sample['canny'] = img2tensor(canny, bgr2rgb=True, float32=True) / 255.
24
+ sample['jpg'] = to_tensor(x)
25
+ return sample
26
+
27
+
28
+ class AddCannyRandomThreshold(object):
29
+
30
+ def __init__(self, low_threshold=100, high_threshold=200, shift_range=50):
31
+ self.low_threshold = low_threshold
32
+ self.high_threshold = high_threshold
33
+ self.threshold_prng = np.random.RandomState()
34
+ self.shift_range = shift_range
35
+
36
+ def __call__(self, sample):
37
+ # sample['jpg'] is PIL image
38
+ x = sample['jpg']
39
+ img = cv2.cvtColor(np.array(x), cv2.COLOR_RGB2BGR)
40
+ low_threshold = self.low_threshold + self.threshold_prng.randint(-self.shift_range, self.shift_range)
41
+ high_threshold = self.high_threshold + self.threshold_prng.randint(-self.shift_range, self.shift_range)
42
+ canny = cv2.Canny(img, low_threshold, high_threshold)[..., None]
43
+ sample['canny'] = img2tensor(canny, bgr2rgb=True, float32=True) / 255.
44
+ sample['jpg'] = to_tensor(x)
45
+ return sample
46
+
47
+
48
+ class AddStyle(object):
49
+
50
+ def __init__(self, version):
51
+ self.processor = CLIPProcessor.from_pretrained(version)
52
+ self.pil_to_tensor = transforms.ToTensor()
53
+
54
+ def __call__(self, sample):
55
+ # sample['jpg'] is PIL image
56
+ x = sample['jpg']
57
+ style = self.processor(images=x, return_tensors="pt")['pixel_values'][0]
58
+ sample['style'] = style
59
+ sample['jpg'] = to_tensor(x)
60
+ return sample
ldm/inference_base.py ADDED
@@ -0,0 +1,292 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import torch
3
+ from omegaconf import OmegaConf
4
+
5
+ from ldm.models.diffusion.ddim import DDIMSampler
6
+ from ldm.models.diffusion.plms import PLMSSampler
7
+ from ldm.modules.encoders.adapter import Adapter, StyleAdapter, Adapter_light
8
+ from ldm.modules.extra_condition.api import ExtraCondition
9
+ from ldm.util import fix_cond_shapes, load_model_from_config, read_state_dict
10
+
11
+ DEFAULT_NEGATIVE_PROMPT = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, ' \
12
+ 'fewer digits, cropped, worst quality, low quality'
13
+
14
+
15
+ def get_base_argument_parser() -> argparse.ArgumentParser:
16
+ """get the base argument parser for inference scripts"""
17
+ parser = argparse.ArgumentParser()
18
+ parser.add_argument(
19
+ '--outdir',
20
+ type=str,
21
+ help='dir to write results to',
22
+ default=None,
23
+ )
24
+
25
+ parser.add_argument(
26
+ '--prompt',
27
+ type=str,
28
+ nargs='?',
29
+ default=None,
30
+ help='positive prompt',
31
+ )
32
+
33
+ parser.add_argument(
34
+ '--neg_prompt',
35
+ type=str,
36
+ default=DEFAULT_NEGATIVE_PROMPT,
37
+ help='negative prompt',
38
+ )
39
+
40
+ parser.add_argument(
41
+ '--cond_path',
42
+ type=str,
43
+ default=None,
44
+ help='condition image path',
45
+ )
46
+
47
+ parser.add_argument(
48
+ '--cond_inp_type',
49
+ type=str,
50
+ default='image',
51
+ help='the type of the input condition image, take depth T2I as example, the input can be raw image, '
52
+ 'which depth will be calculated, or the input can be a directly a depth map image',
53
+ )
54
+
55
+ parser.add_argument(
56
+ '--sampler',
57
+ type=str,
58
+ default='ddim',
59
+ choices=['ddim', 'plms'],
60
+ help='sampling algorithm, currently, only ddim and plms are supported, more are on the way',
61
+ )
62
+
63
+ parser.add_argument(
64
+ '--steps',
65
+ type=int,
66
+ default=50,
67
+ help='number of sampling steps',
68
+ )
69
+
70
+ parser.add_argument(
71
+ '--sd_ckpt',
72
+ type=str,
73
+ default='models/sd-v1-4.ckpt',
74
+ help='path to checkpoint of stable diffusion model, both .ckpt and .safetensor are supported',
75
+ )
76
+
77
+ parser.add_argument(
78
+ '--vae_ckpt',
79
+ type=str,
80
+ default=None,
81
+ help='vae checkpoint, anime SD models usually have seperate vae ckpt that need to be loaded',
82
+ )
83
+
84
+ parser.add_argument(
85
+ '--adapter_ckpt',
86
+ type=str,
87
+ default=None,
88
+ help='path to checkpoint of adapter',
89
+ )
90
+
91
+ parser.add_argument(
92
+ '--config',
93
+ type=str,
94
+ default='configs/stable-diffusion/sd-v1-inference.yaml',
95
+ help='path to config which constructs SD model',
96
+ )
97
+
98
+ parser.add_argument(
99
+ '--max_resolution',
100
+ type=float,
101
+ default=512 * 512,
102
+ help='max image height * width, only for computer with limited vram',
103
+ )
104
+
105
+ parser.add_argument(
106
+ '--resize_short_edge',
107
+ type=int,
108
+ default=None,
109
+ help='resize short edge of the input image, if this arg is set, max_resolution will not be used',
110
+ )
111
+
112
+ parser.add_argument(
113
+ '--C',
114
+ type=int,
115
+ default=4,
116
+ help='latent channels',
117
+ )
118
+
119
+ parser.add_argument(
120
+ '--f',
121
+ type=int,
122
+ default=8,
123
+ help='downsampling factor',
124
+ )
125
+
126
+ parser.add_argument(
127
+ '--scale',
128
+ type=float,
129
+ default=7.5,
130
+ help='unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))',
131
+ )
132
+
133
+ parser.add_argument(
134
+ '--cond_tau',
135
+ type=float,
136
+ default=1.0,
137
+ help='timestamp parameter that determines until which step the adapter is applied, '
138
+ 'similar as Prompt-to-Prompt tau',
139
+ )
140
+
141
+ parser.add_argument(
142
+ '--style_cond_tau',
143
+ type=float,
144
+ default=1.0,
145
+ help='timestamp parameter that determines until which step the adapter is applied, '
146
+ 'similar as Prompt-to-Prompt tau',
147
+ )
148
+
149
+ parser.add_argument(
150
+ '--cond_weight',
151
+ type=float,
152
+ default=1.0,
153
+ help='the adapter features are multiplied by the cond_weight. The larger the cond_weight, the more aligned '
154
+ 'the generated image and condition will be, but the generated quality may be reduced',
155
+ )
156
+
157
+ parser.add_argument(
158
+ '--seed',
159
+ type=int,
160
+ default=42,
161
+ )
162
+
163
+ parser.add_argument(
164
+ '--n_samples',
165
+ type=int,
166
+ default=4,
167
+ help='# of samples to generate',
168
+ )
169
+
170
+ return parser
171
+
172
+
173
+ def get_sd_models(opt):
174
+ """
175
+ build stable diffusion model, sampler
176
+ """
177
+ # SD
178
+ config = OmegaConf.load(f"{opt.config}")
179
+ model = load_model_from_config(config, opt.sd_ckpt, opt.vae_ckpt)
180
+ sd_model = model.to(opt.device)
181
+
182
+ # sampler
183
+ if opt.sampler == 'plms':
184
+ sampler = PLMSSampler(model)
185
+ elif opt.sampler == 'ddim':
186
+ sampler = DDIMSampler(model)
187
+ else:
188
+ raise NotImplementedError
189
+
190
+ return sd_model, sampler
191
+
192
+
193
+ def get_t2i_adapter_models(opt):
194
+ config = OmegaConf.load(f"{opt.config}")
195
+ model = load_model_from_config(config, opt.sd_ckpt, opt.vae_ckpt)
196
+ adapter_ckpt_path = getattr(opt, f'{opt.which_cond}_adapter_ckpt', None)
197
+ if adapter_ckpt_path is None:
198
+ adapter_ckpt_path = getattr(opt, 'adapter_ckpt')
199
+ adapter_ckpt = read_state_dict(adapter_ckpt_path)
200
+ new_state_dict = {}
201
+ for k, v in adapter_ckpt.items():
202
+ if not k.startswith('adapter.'):
203
+ new_state_dict[f'adapter.{k}'] = v
204
+ else:
205
+ new_state_dict[k] = v
206
+ m, u = model.load_state_dict(new_state_dict, strict=False)
207
+ if len(u) > 0:
208
+ print(f"unexpected keys in loading adapter ckpt {adapter_ckpt_path}:")
209
+ print(u)
210
+
211
+ model = model.to(opt.device)
212
+
213
+ # sampler
214
+ if opt.sampler == 'plms':
215
+ sampler = PLMSSampler(model)
216
+ elif opt.sampler == 'ddim':
217
+ sampler = DDIMSampler(model)
218
+ else:
219
+ raise NotImplementedError
220
+
221
+ return model, sampler
222
+
223
+
224
+ def get_cond_ch(cond_type: ExtraCondition):
225
+ if cond_type == ExtraCondition.sketch or cond_type == ExtraCondition.canny:
226
+ return 1
227
+ return 3
228
+
229
+
230
+ def get_adapters(opt, cond_type: ExtraCondition):
231
+ adapter = {}
232
+ cond_weight = getattr(opt, f'{cond_type.name}_weight', None)
233
+ if cond_weight is None:
234
+ cond_weight = getattr(opt, 'cond_weight')
235
+ adapter['cond_weight'] = cond_weight
236
+
237
+ if cond_type == ExtraCondition.style:
238
+ adapter['model'] = StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8).to(opt.device)
239
+ elif cond_type == ExtraCondition.color:
240
+ adapter['model'] = Adapter_light(
241
+ cin=64 * get_cond_ch(cond_type),
242
+ channels=[320, 640, 1280, 1280],
243
+ nums_rb=4).to(opt.device)
244
+ else:
245
+ adapter['model'] = Adapter(
246
+ cin=64 * get_cond_ch(cond_type),
247
+ channels=[320, 640, 1280, 1280][:4],
248
+ nums_rb=2,
249
+ ksize=1,
250
+ sk=True,
251
+ use_conv=False).to(opt.device)
252
+ ckpt_path = getattr(opt, f'{cond_type.name}_adapter_ckpt', None)
253
+ if ckpt_path is None:
254
+ ckpt_path = getattr(opt, 'adapter_ckpt')
255
+ adapter['model'].load_state_dict(torch.load(ckpt_path))
256
+
257
+ return adapter
258
+
259
+
260
+ def diffusion_inference(opt, model, sampler, adapter_features, append_to_context=None):
261
+ # get text embedding
262
+ c = model.get_learned_conditioning([opt.prompt])
263
+ if opt.scale != 1.0:
264
+ uc = model.get_learned_conditioning([opt.neg_prompt])
265
+ else:
266
+ uc = None
267
+ c, uc = fix_cond_shapes(model, c, uc)
268
+
269
+ if not hasattr(opt, 'H'):
270
+ opt.H = 512
271
+ opt.W = 512
272
+ shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
273
+
274
+ samples_latents, _ = sampler.sample(
275
+ S=opt.steps,
276
+ conditioning=c,
277
+ batch_size=1,
278
+ shape=shape,
279
+ verbose=False,
280
+ unconditional_guidance_scale=opt.scale,
281
+ unconditional_conditioning=uc,
282
+ x_T=None,
283
+ features_adapter=adapter_features,
284
+ append_to_context=append_to_context,
285
+ cond_tau=opt.cond_tau,
286
+ style_cond_tau=opt.style_cond_tau,
287
+ )
288
+
289
+ x_samples = model.decode_first_stage(samples_latents)
290
+ x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
291
+
292
+ return x_samples
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/__pycache__/autoencoder.cpython-38.pyc ADDED
Binary file (7.48 kB). View file
 
ldm/models/autoencoder.py ADDED
@@ -0,0 +1,211 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import pytorch_lightning as pl
3
+ import torch.nn.functional as F
4
+ import torch.nn as nn
5
+ from contextlib import contextmanager
6
+
7
+ from ldm.modules.diffusionmodules.model import Encoder, Decoder
8
+ from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
9
+
10
+ from ldm.util import instantiate_from_config
11
+ from ldm.modules.ema import LitEma
12
+
13
+
14
+ class AutoencoderKL(pl.LightningModule):
15
+ def __init__(self,
16
+ ddconfig,
17
+ lossconfig,
18
+ embed_dim,
19
+ ckpt_path=None,
20
+ ignore_keys=[],
21
+ image_key="image",
22
+ colorize_nlabels=None,
23
+ monitor=None,
24
+ ema_decay=None,
25
+ learn_logvar=False
26
+ ):
27
+ super().__init__()
28
+ self.learn_logvar = learn_logvar
29
+ self.image_key = image_key
30
+ self.encoder = Encoder(**ddconfig)
31
+ self.decoder = Decoder(**ddconfig)
32
+ self.loss = instantiate_from_config(lossconfig)
33
+ assert ddconfig["double_z"]
34
+ self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
35
+ self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
36
+ self.embed_dim = embed_dim
37
+ if colorize_nlabels is not None:
38
+ assert type(colorize_nlabels)==int
39
+ self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
40
+ if monitor is not None:
41
+ self.monitor = monitor
42
+
43
+ self.use_ema = ema_decay is not None
44
+ if self.use_ema:
45
+ self.ema_decay = ema_decay
46
+ assert 0. < ema_decay < 1.
47
+ self.model_ema = LitEma(self, decay=ema_decay)
48
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
49
+
50
+ if ckpt_path is not None:
51
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
52
+
53
+ def init_from_ckpt(self, path, ignore_keys=list()):
54
+ sd = torch.load(path, map_location="cpu")["state_dict"]
55
+ keys = list(sd.keys())
56
+ for k in keys:
57
+ for ik in ignore_keys:
58
+ if k.startswith(ik):
59
+ print("Deleting key {} from state_dict.".format(k))
60
+ del sd[k]
61
+ self.load_state_dict(sd, strict=False)
62
+ print(f"Restored from {path}")
63
+
64
+ @contextmanager
65
+ def ema_scope(self, context=None):
66
+ if self.use_ema:
67
+ self.model_ema.store(self.parameters())
68
+ self.model_ema.copy_to(self)
69
+ if context is not None:
70
+ print(f"{context}: Switched to EMA weights")
71
+ try:
72
+ yield None
73
+ finally:
74
+ if self.use_ema:
75
+ self.model_ema.restore(self.parameters())
76
+ if context is not None:
77
+ print(f"{context}: Restored training weights")
78
+
79
+ def on_train_batch_end(self, *args, **kwargs):
80
+ if self.use_ema:
81
+ self.model_ema(self)
82
+
83
+ def encode(self, x):
84
+ h = self.encoder(x)
85
+ moments = self.quant_conv(h)
86
+ posterior = DiagonalGaussianDistribution(moments)
87
+ return posterior
88
+
89
+ def decode(self, z):
90
+ z = self.post_quant_conv(z)
91
+ dec = self.decoder(z)
92
+ return dec
93
+
94
+ def forward(self, input, sample_posterior=True):
95
+ posterior = self.encode(input)
96
+ if sample_posterior:
97
+ z = posterior.sample()
98
+ else:
99
+ z = posterior.mode()
100
+ dec = self.decode(z)
101
+ return dec, posterior
102
+
103
+ def get_input(self, batch, k):
104
+ x = batch[k]
105
+ if len(x.shape) == 3:
106
+ x = x[..., None]
107
+ x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
108
+ return x
109
+
110
+ def training_step(self, batch, batch_idx, optimizer_idx):
111
+ inputs = self.get_input(batch, self.image_key)
112
+ reconstructions, posterior = self(inputs)
113
+
114
+ if optimizer_idx == 0:
115
+ # train encoder+decoder+logvar
116
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
117
+ last_layer=self.get_last_layer(), split="train")
118
+ self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
119
+ self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
120
+ return aeloss
121
+
122
+ if optimizer_idx == 1:
123
+ # train the discriminator
124
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
125
+ last_layer=self.get_last_layer(), split="train")
126
+
127
+ self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
128
+ self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
129
+ return discloss
130
+
131
+ def validation_step(self, batch, batch_idx):
132
+ log_dict = self._validation_step(batch, batch_idx)
133
+ with self.ema_scope():
134
+ log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
135
+ return log_dict
136
+
137
+ def _validation_step(self, batch, batch_idx, postfix=""):
138
+ inputs = self.get_input(batch, self.image_key)
139
+ reconstructions, posterior = self(inputs)
140
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
141
+ last_layer=self.get_last_layer(), split="val"+postfix)
142
+
143
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
144
+ last_layer=self.get_last_layer(), split="val"+postfix)
145
+
146
+ self.log(f"val{postfix}/rec_loss", log_dict_ae[f"val{postfix}/rec_loss"])
147
+ self.log_dict(log_dict_ae)
148
+ self.log_dict(log_dict_disc)
149
+ return self.log_dict
150
+
151
+ def configure_optimizers(self):
152
+ lr = self.learning_rate
153
+ ae_params_list = list(self.encoder.parameters()) + list(self.decoder.parameters()) + list(
154
+ self.quant_conv.parameters()) + list(self.post_quant_conv.parameters())
155
+ if self.learn_logvar:
156
+ print(f"{self.__class__.__name__}: Learning logvar")
157
+ ae_params_list.append(self.loss.logvar)
158
+ opt_ae = torch.optim.Adam(ae_params_list,
159
+ lr=lr, betas=(0.5, 0.9))
160
+ opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
161
+ lr=lr, betas=(0.5, 0.9))
162
+ return [opt_ae, opt_disc], []
163
+
164
+ def get_last_layer(self):
165
+ return self.decoder.conv_out.weight
166
+
167
+ @torch.no_grad()
168
+ def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
169
+ log = dict()
170
+ x = self.get_input(batch, self.image_key)
171
+ x = x.to(self.device)
172
+ if not only_inputs:
173
+ xrec, posterior = self(x)
174
+ if x.shape[1] > 3:
175
+ # colorize with random projection
176
+ assert xrec.shape[1] > 3
177
+ x = self.to_rgb(x)
178
+ xrec = self.to_rgb(xrec)
179
+ log["samples"] = self.decode(torch.randn_like(posterior.sample()))
180
+ log["reconstructions"] = xrec
181
+ log["inputs"] = x
182
+ return log
183
+
184
+ def to_rgb(self, x):
185
+ assert self.image_key == "segmentation"
186
+ if not hasattr(self, "colorize"):
187
+ self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
188
+ x = F.conv2d(x, weight=self.colorize)
189
+ x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
190
+ return x
191
+
192
+
193
+ class IdentityFirstStage(nn.Module):
194
+ def __init__(self, *args, vq_interface=False, **kwargs):
195
+ self.vq_interface = vq_interface
196
+ super().__init__()
197
+
198
+ def encode(self, x, *args, **kwargs):
199
+ return x
200
+
201
+ def decode(self, x, *args, **kwargs):
202
+ return x
203
+
204
+ def quantize(self, x, *args, **kwargs):
205
+ if self.vq_interface:
206
+ return x, None, [None, None, None]
207
+ return x
208
+
209
+ def forward(self, x, *args, **kwargs):
210
+ return x
211
+
ldm/models/diffusion/__init__.py ADDED
File without changes
ldm/models/diffusion/__pycache__/__init__.cpython-38.pyc ADDED
Binary file (206 Bytes). View file
 
ldm/models/diffusion/__pycache__/ddim.cpython-38.pyc ADDED
Binary file (8.35 kB). View file
 
ldm/models/diffusion/__pycache__/ddpm.cpython-38.pyc ADDED
Binary file (39.5 kB). View file
 
ldm/models/diffusion/__pycache__/plms.cpython-38.pyc ADDED
Binary file (7.53 kB). View file
 
ldm/models/diffusion/ddim.py ADDED
@@ -0,0 +1,293 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+
7
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, \
8
+ extract_into_tensor
9
+
10
+
11
+ class DDIMSampler(object):
12
+ def __init__(self, model, schedule="linear", **kwargs):
13
+ super().__init__()
14
+ self.model = model
15
+ self.ddpm_num_timesteps = model.num_timesteps
16
+ self.schedule = schedule
17
+
18
+ def register_buffer(self, name, attr):
19
+ if type(attr) == torch.Tensor:
20
+ if attr.device != torch.device("cuda"):
21
+ attr = attr.to(torch.device("cuda"))
22
+ setattr(self, name, attr)
23
+
24
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
25
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
26
+ num_ddpm_timesteps=self.ddpm_num_timesteps, verbose=verbose)
27
+ alphas_cumprod = self.model.alphas_cumprod
28
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
29
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
30
+
31
+ self.register_buffer('betas', to_torch(self.model.betas))
32
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
33
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
34
+
35
+ # calculations for diffusion q(x_t | x_{t-1}) and others
36
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
37
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
38
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
39
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
40
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
41
+
42
+ # ddim sampling parameters
43
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
44
+ ddim_timesteps=self.ddim_timesteps,
45
+ eta=ddim_eta, verbose=verbose)
46
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
47
+ self.register_buffer('ddim_alphas', ddim_alphas)
48
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
49
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
50
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
51
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
52
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
53
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
54
+
55
+ @torch.no_grad()
56
+ def sample(self,
57
+ S,
58
+ batch_size,
59
+ shape,
60
+ conditioning=None,
61
+ callback=None,
62
+ normals_sequence=None,
63
+ img_callback=None,
64
+ quantize_x0=False,
65
+ eta=0.,
66
+ mask=None,
67
+ x0=None,
68
+ temperature=1.,
69
+ noise_dropout=0.,
70
+ score_corrector=None,
71
+ corrector_kwargs=None,
72
+ verbose=True,
73
+ x_T=None,
74
+ log_every_t=100,
75
+ unconditional_guidance_scale=1.,
76
+ unconditional_conditioning=None,
77
+ features_adapter=None,
78
+ append_to_context=None,
79
+ cond_tau=0.4,
80
+ style_cond_tau=1.0,
81
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
82
+ **kwargs
83
+ ):
84
+ if conditioning is not None:
85
+ if isinstance(conditioning, dict):
86
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
87
+ if cbs != batch_size:
88
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
89
+ else:
90
+ if conditioning.shape[0] != batch_size:
91
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
92
+
93
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
94
+ # sampling
95
+ C, H, W = shape
96
+ size = (batch_size, C, H, W)
97
+ print(f'Data shape for DDIM sampling is {size}, eta {eta}')
98
+
99
+ samples, intermediates = self.ddim_sampling(conditioning, size,
100
+ callback=callback,
101
+ img_callback=img_callback,
102
+ quantize_denoised=quantize_x0,
103
+ mask=mask, x0=x0,
104
+ ddim_use_original_steps=False,
105
+ noise_dropout=noise_dropout,
106
+ temperature=temperature,
107
+ score_corrector=score_corrector,
108
+ corrector_kwargs=corrector_kwargs,
109
+ x_T=x_T,
110
+ log_every_t=log_every_t,
111
+ unconditional_guidance_scale=unconditional_guidance_scale,
112
+ unconditional_conditioning=unconditional_conditioning,
113
+ features_adapter=features_adapter,
114
+ append_to_context=append_to_context,
115
+ cond_tau=cond_tau,
116
+ style_cond_tau=style_cond_tau,
117
+ )
118
+ return samples, intermediates
119
+
120
+ @torch.no_grad()
121
+ def ddim_sampling(self, cond, shape,
122
+ x_T=None, ddim_use_original_steps=False,
123
+ callback=None, timesteps=None, quantize_denoised=False,
124
+ mask=None, x0=None, img_callback=None, log_every_t=100,
125
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
126
+ unconditional_guidance_scale=1., unconditional_conditioning=None, features_adapter=None,
127
+ append_to_context=None, cond_tau=0.4, style_cond_tau=1.0):
128
+ device = self.model.betas.device
129
+ b = shape[0]
130
+ if x_T is None:
131
+ img = torch.randn(shape, device=device)
132
+ else:
133
+ img = x_T
134
+
135
+ if timesteps is None:
136
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
137
+ elif timesteps is not None and not ddim_use_original_steps:
138
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
139
+ timesteps = self.ddim_timesteps[:subset_end]
140
+
141
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
142
+ time_range = reversed(range(0, timesteps)) if ddim_use_original_steps else np.flip(timesteps)
143
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
144
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
145
+
146
+ iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
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
+
152
+ if mask is not None:
153
+ assert x0 is not None
154
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
155
+ img = img_orig * mask + (1. - mask) * img
156
+
157
+ outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
158
+ quantize_denoised=quantize_denoised, temperature=temperature,
159
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
160
+ corrector_kwargs=corrector_kwargs,
161
+ unconditional_guidance_scale=unconditional_guidance_scale,
162
+ unconditional_conditioning=unconditional_conditioning,
163
+ features_adapter=None if index < int(
164
+ (1 - cond_tau) * total_steps) else features_adapter,
165
+ append_to_context=None if index < int(
166
+ (1 - style_cond_tau) * total_steps) else append_to_context,
167
+ )
168
+ img, pred_x0 = outs
169
+ if callback: callback(i)
170
+ if img_callback: img_callback(pred_x0, i)
171
+
172
+ if index % log_every_t == 0 or index == total_steps - 1:
173
+ intermediates['x_inter'].append(img)
174
+ intermediates['pred_x0'].append(pred_x0)
175
+
176
+ return img, intermediates
177
+
178
+ @torch.no_grad()
179
+ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
180
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
181
+ unconditional_guidance_scale=1., unconditional_conditioning=None, features_adapter=None,
182
+ append_to_context=None):
183
+ b, *_, device = *x.shape, x.device
184
+
185
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
186
+ if append_to_context is not None:
187
+ model_output = self.model.apply_model(x, t, torch.cat([c, append_to_context], dim=1),
188
+ features_adapter=features_adapter)
189
+ else:
190
+ model_output = self.model.apply_model(x, t, c, features_adapter=features_adapter)
191
+ else:
192
+ x_in = torch.cat([x] * 2)
193
+ t_in = torch.cat([t] * 2)
194
+ if isinstance(c, dict):
195
+ assert isinstance(unconditional_conditioning, dict)
196
+ c_in = dict()
197
+ for k in c:
198
+ if isinstance(c[k], list):
199
+ c_in[k] = [torch.cat([
200
+ unconditional_conditioning[k][i],
201
+ c[k][i]]) for i in range(len(c[k]))]
202
+ else:
203
+ c_in[k] = torch.cat([
204
+ unconditional_conditioning[k],
205
+ c[k]])
206
+ elif isinstance(c, list):
207
+ c_in = list()
208
+ assert isinstance(unconditional_conditioning, list)
209
+ for i in range(len(c)):
210
+ c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))
211
+ else:
212
+ if append_to_context is not None:
213
+ pad_len = append_to_context.size(1)
214
+ new_unconditional_conditioning = torch.cat(
215
+ [unconditional_conditioning, unconditional_conditioning[:, -pad_len:, :]], dim=1)
216
+ new_c = torch.cat([c, append_to_context], dim=1)
217
+ c_in = torch.cat([new_unconditional_conditioning, new_c])
218
+ else:
219
+ c_in = torch.cat([unconditional_conditioning, c])
220
+ model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in, features_adapter=features_adapter).chunk(2)
221
+ model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)
222
+
223
+ if self.model.parameterization == "v":
224
+ e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
225
+ else:
226
+ e_t = model_output
227
+
228
+ if score_corrector is not None:
229
+ assert self.model.parameterization == "eps", 'not implemented'
230
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
231
+
232
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
233
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
234
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
235
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
236
+ # select parameters corresponding to the currently considered timestep
237
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
238
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
239
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
240
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index], device=device)
241
+
242
+ # current prediction for x_0
243
+ if self.model.parameterization != "v":
244
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
245
+ else:
246
+ pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
247
+
248
+ if quantize_denoised:
249
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
250
+ # direction pointing to x_t
251
+ dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
252
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
253
+ if noise_dropout > 0.:
254
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
255
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
256
+ return x_prev, pred_x0
257
+
258
+ @torch.no_grad()
259
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
260
+ # fast, but does not allow for exact reconstruction
261
+ # t serves as an index to gather the correct alphas
262
+ if use_original_steps:
263
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
264
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
265
+ else:
266
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
267
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
268
+
269
+ if noise is None:
270
+ noise = torch.randn_like(x0)
271
+ return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
272
+ extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
273
+
274
+ @torch.no_grad()
275
+ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
276
+ use_original_steps=False):
277
+
278
+ timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
279
+ timesteps = timesteps[:t_start]
280
+
281
+ time_range = np.flip(timesteps)
282
+ total_steps = timesteps.shape[0]
283
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
284
+
285
+ iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
286
+ x_dec = x_latent
287
+ for i, step in enumerate(iterator):
288
+ index = total_steps - i - 1
289
+ ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
290
+ x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
291
+ unconditional_guidance_scale=unconditional_guidance_scale,
292
+ unconditional_conditioning=unconditional_conditioning)
293
+ return x_dec
ldm/models/diffusion/ddpm.py ADDED
@@ -0,0 +1,1329 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 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
+
30
+
31
+ __conditioning_keys__ = {'concat': 'c_concat',
32
+ 'crossattn': 'c_crossattn',
33
+ 'adm': 'y'}
34
+
35
+
36
+ def disabled_train(self, mode=True):
37
+ """Overwrite model.train with this function to make sure train/eval mode
38
+ does not change anymore."""
39
+ return self
40
+
41
+
42
+ def uniform_on_device(r1, r2, shape, device):
43
+ return (r1 - r2) * torch.rand(*shape, device=device) + r2
44
+
45
+
46
+ class DDPM(pl.LightningModule):
47
+ # classic DDPM with Gaussian diffusion, in image space
48
+ def __init__(self,
49
+ unet_config,
50
+ timesteps=1000,
51
+ beta_schedule="linear",
52
+ loss_type="l2",
53
+ ckpt_path=None,
54
+ ignore_keys=[],
55
+ load_only_unet=False,
56
+ monitor="val/loss",
57
+ use_ema=True,
58
+ first_stage_key="image",
59
+ image_size=256,
60
+ channels=3,
61
+ log_every_t=100,
62
+ clip_denoised=True,
63
+ linear_start=1e-4,
64
+ linear_end=2e-2,
65
+ cosine_s=8e-3,
66
+ given_betas=None,
67
+ original_elbo_weight=0.,
68
+ v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
69
+ l_simple_weight=1.,
70
+ conditioning_key=None,
71
+ parameterization="eps", # all assuming fixed variance schedules
72
+ scheduler_config=None,
73
+ use_positional_encodings=False,
74
+ learn_logvar=False,
75
+ logvar_init=0.,
76
+ make_it_fit=False,
77
+ ucg_training=None,
78
+ reset_ema=False,
79
+ reset_num_ema_updates=False,
80
+ ):
81
+ super().__init__()
82
+ assert parameterization in ["eps", "x0", "v"], 'currently only supporting "eps" and "x0" and "v"'
83
+ self.parameterization = parameterization
84
+ print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
85
+ self.cond_stage_model = None
86
+ self.clip_denoised = clip_denoised
87
+ self.log_every_t = log_every_t
88
+ self.first_stage_key = first_stage_key
89
+ self.image_size = image_size # try conv?
90
+ self.channels = channels
91
+ self.use_positional_encodings = use_positional_encodings
92
+ self.model = DiffusionWrapper(unet_config, conditioning_key)
93
+ count_params(self.model, verbose=True)
94
+ self.use_ema = use_ema
95
+ if self.use_ema:
96
+ self.model_ema = LitEma(self.model)
97
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
98
+
99
+ self.use_scheduler = scheduler_config is not None
100
+ if self.use_scheduler:
101
+ self.scheduler_config = scheduler_config
102
+
103
+ self.v_posterior = v_posterior
104
+ self.original_elbo_weight = original_elbo_weight
105
+ self.l_simple_weight = l_simple_weight
106
+
107
+ if monitor is not None:
108
+ self.monitor = monitor
109
+ self.make_it_fit = make_it_fit
110
+ if reset_ema: assert exists(ckpt_path)
111
+ if ckpt_path is not None:
112
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
113
+ if reset_ema:
114
+ assert self.use_ema
115
+ print(f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint.")
116
+ self.model_ema = LitEma(self.model)
117
+ if reset_num_ema_updates:
118
+ print(" +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ ")
119
+ assert self.use_ema
120
+ self.model_ema.reset_num_updates()
121
+
122
+ self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
123
+ linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
124
+
125
+ self.loss_type = loss_type
126
+
127
+ self.learn_logvar = learn_logvar
128
+ self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
129
+ if self.learn_logvar:
130
+ self.logvar = nn.Parameter(self.logvar, requires_grad=True)
131
+
132
+ self.ucg_training = ucg_training or dict()
133
+ if self.ucg_training:
134
+ self.ucg_prng = np.random.RandomState()
135
+
136
+ def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
137
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
138
+ if exists(given_betas):
139
+ betas = given_betas
140
+ else:
141
+ betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
142
+ cosine_s=cosine_s)
143
+ alphas = 1. - betas
144
+ alphas_cumprod = np.cumprod(alphas, axis=0)
145
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
146
+
147
+ timesteps, = betas.shape
148
+ self.num_timesteps = int(timesteps)
149
+ self.linear_start = linear_start
150
+ self.linear_end = linear_end
151
+ assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
152
+
153
+ to_torch = partial(torch.tensor, dtype=torch.float32)
154
+
155
+ self.register_buffer('betas', to_torch(betas))
156
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
157
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
158
+
159
+ # calculations for diffusion q(x_t | x_{t-1}) and others
160
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
161
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
162
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
163
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
164
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
165
+
166
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
167
+ posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
168
+ 1. - alphas_cumprod) + self.v_posterior * betas
169
+ # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
170
+ self.register_buffer('posterior_variance', to_torch(posterior_variance))
171
+ # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
172
+ self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
173
+ self.register_buffer('posterior_mean_coef1', to_torch(
174
+ betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
175
+ self.register_buffer('posterior_mean_coef2', to_torch(
176
+ (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
177
+
178
+ if self.parameterization == "eps":
179
+ lvlb_weights = self.betas ** 2 / (
180
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))
181
+ elif self.parameterization == "x0":
182
+ lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod))
183
+ elif self.parameterization == "v":
184
+ lvlb_weights = torch.ones_like(self.betas ** 2 / (
185
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)))
186
+ else:
187
+ raise NotImplementedError("mu not supported")
188
+ lvlb_weights[0] = lvlb_weights[1]
189
+ self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
190
+ assert not torch.isnan(self.lvlb_weights).all()
191
+
192
+ @contextmanager
193
+ def ema_scope(self, context=None):
194
+ if self.use_ema:
195
+ self.model_ema.store(self.model.parameters())
196
+ self.model_ema.copy_to(self.model)
197
+ if context is not None:
198
+ print(f"{context}: Switched to EMA weights")
199
+ try:
200
+ yield None
201
+ finally:
202
+ if self.use_ema:
203
+ self.model_ema.restore(self.model.parameters())
204
+ if context is not None:
205
+ print(f"{context}: Restored training weights")
206
+
207
+ @torch.no_grad()
208
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
209
+ sd = torch.load(path, map_location="cpu")
210
+ if "state_dict" in list(sd.keys()):
211
+ sd = sd["state_dict"]
212
+ keys = list(sd.keys())
213
+ for k in keys:
214
+ for ik in ignore_keys:
215
+ if k.startswith(ik):
216
+ print("Deleting key {} from state_dict.".format(k))
217
+ del sd[k]
218
+ if self.make_it_fit:
219
+ n_params = len([name for name, _ in
220
+ itertools.chain(self.named_parameters(),
221
+ self.named_buffers())])
222
+ for name, param in tqdm(
223
+ itertools.chain(self.named_parameters(),
224
+ self.named_buffers()),
225
+ desc="Fitting old weights to new weights",
226
+ total=n_params
227
+ ):
228
+ if not name in sd:
229
+ continue
230
+ old_shape = sd[name].shape
231
+ new_shape = param.shape
232
+ assert len(old_shape) == len(new_shape)
233
+ if len(new_shape) > 2:
234
+ # we only modify first two axes
235
+ assert new_shape[2:] == old_shape[2:]
236
+ # assumes first axis corresponds to output dim
237
+ if not new_shape == old_shape:
238
+ new_param = param.clone()
239
+ old_param = sd[name]
240
+ if len(new_shape) == 1:
241
+ for i in range(new_param.shape[0]):
242
+ new_param[i] = old_param[i % old_shape[0]]
243
+ elif len(new_shape) >= 2:
244
+ for i in range(new_param.shape[0]):
245
+ for j in range(new_param.shape[1]):
246
+ new_param[i, j] = old_param[i % old_shape[0], j % old_shape[1]]
247
+
248
+ n_used_old = torch.ones(old_shape[1])
249
+ for j in range(new_param.shape[1]):
250
+ n_used_old[j % old_shape[1]] += 1
251
+ n_used_new = torch.zeros(new_shape[1])
252
+ for j in range(new_param.shape[1]):
253
+ n_used_new[j] = n_used_old[j % old_shape[1]]
254
+
255
+ n_used_new = n_used_new[None, :]
256
+ while len(n_used_new.shape) < len(new_shape):
257
+ n_used_new = n_used_new.unsqueeze(-1)
258
+ new_param /= n_used_new
259
+
260
+ sd[name] = new_param
261
+
262
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
263
+ sd, strict=False)
264
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
265
+ if len(missing) > 0:
266
+ print(f"Missing Keys:\n {missing}")
267
+ if len(unexpected) > 0:
268
+ print(f"\nUnexpected Keys:\n {unexpected}")
269
+
270
+ def q_mean_variance(self, x_start, t):
271
+ """
272
+ Get the distribution q(x_t | x_0).
273
+ :param x_start: the [N x C x ...] tensor of noiseless inputs.
274
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
275
+ :return: A tuple (mean, variance, log_variance), all of x_start's shape.
276
+ """
277
+ mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start)
278
+ variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
279
+ log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
280
+ return mean, variance, log_variance
281
+
282
+ def predict_start_from_noise(self, x_t, t, noise):
283
+ return (
284
+ extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
285
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
286
+ )
287
+
288
+ def predict_start_from_z_and_v(self, x_t, t, v):
289
+ # self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
290
+ # self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
291
+ return (
292
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -
293
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
294
+ )
295
+
296
+ def predict_eps_from_z_and_v(self, x_t, t, v):
297
+ return (
298
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * v +
299
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * x_t
300
+ )
301
+
302
+ def q_posterior(self, x_start, x_t, t):
303
+ posterior_mean = (
304
+ extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
305
+ extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
306
+ )
307
+ posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
308
+ posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
309
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
310
+
311
+ def p_mean_variance(self, x, t, clip_denoised: bool):
312
+ model_out = self.model(x, t)
313
+ if self.parameterization == "eps":
314
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
315
+ elif self.parameterization == "x0":
316
+ x_recon = model_out
317
+ if clip_denoised:
318
+ x_recon.clamp_(-1., 1.)
319
+
320
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
321
+ return model_mean, posterior_variance, posterior_log_variance
322
+
323
+ @torch.no_grad()
324
+ def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
325
+ b, *_, device = *x.shape, x.device
326
+ model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
327
+ noise = noise_like(x.shape, device, repeat_noise)
328
+ # no noise when t == 0
329
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
330
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
331
+
332
+ @torch.no_grad()
333
+ def p_sample_loop(self, shape, return_intermediates=False):
334
+ device = self.betas.device
335
+ b = shape[0]
336
+ img = torch.randn(shape, device=device)
337
+ intermediates = [img]
338
+ for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
339
+ img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
340
+ clip_denoised=self.clip_denoised)
341
+ if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
342
+ intermediates.append(img)
343
+ if return_intermediates:
344
+ return img, intermediates
345
+ return img
346
+
347
+ @torch.no_grad()
348
+ def sample(self, batch_size=16, return_intermediates=False):
349
+ image_size = self.image_size
350
+ channels = self.channels
351
+ return self.p_sample_loop((batch_size, channels, image_size, image_size),
352
+ return_intermediates=return_intermediates)
353
+
354
+ def q_sample(self, x_start, t, noise=None):
355
+ noise = default(noise, lambda: torch.randn_like(x_start))
356
+ return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
357
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
358
+
359
+ def get_v(self, x, noise, t):
360
+ return (
361
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -
362
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
363
+ )
364
+
365
+ def get_loss(self, pred, target, mean=True):
366
+ if self.loss_type == 'l1':
367
+ loss = (target - pred).abs()
368
+ if mean:
369
+ loss = loss.mean()
370
+ elif self.loss_type == 'l2':
371
+ if mean:
372
+ loss = torch.nn.functional.mse_loss(target, pred)
373
+ else:
374
+ loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
375
+ else:
376
+ raise NotImplementedError("unknown loss type '{loss_type}'")
377
+
378
+ return loss
379
+
380
+ def p_losses(self, x_start, t, noise=None):
381
+ noise = default(noise, lambda: torch.randn_like(x_start))
382
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
383
+ model_out = self.model(x_noisy, t)
384
+
385
+ loss_dict = {}
386
+ if self.parameterization == "eps":
387
+ target = noise
388
+ elif self.parameterization == "x0":
389
+ target = x_start
390
+ elif self.parameterization == "v":
391
+ target = self.get_v(x_start, noise, t)
392
+ else:
393
+ raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported")
394
+
395
+ loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
396
+
397
+ log_prefix = 'train' if self.training else 'val'
398
+
399
+ loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
400
+ loss_simple = loss.mean() * self.l_simple_weight
401
+
402
+ loss_vlb = (self.lvlb_weights[t] * loss).mean()
403
+ loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
404
+
405
+ loss = loss_simple + self.original_elbo_weight * loss_vlb
406
+
407
+ loss_dict.update({f'{log_prefix}/loss': loss})
408
+
409
+ return loss, loss_dict
410
+
411
+ def forward(self, x, *args, **kwargs):
412
+ # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
413
+ # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
414
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
415
+ return self.p_losses(x, t, *args, **kwargs)
416
+
417
+ def get_input(self, batch, k):
418
+ x = batch[k]
419
+ # if len(x.shape) == 3:
420
+ # x = x[..., None]
421
+ # x = rearrange(x, 'b h w c -> b c h w')
422
+ # x = x.to(memory_format=torch.contiguous_format).float()
423
+ return x
424
+
425
+ def shared_step(self, batch):
426
+ x = self.get_input(batch, self.first_stage_key)
427
+ loss, loss_dict = self(x)
428
+ return loss, loss_dict
429
+
430
+ def training_step(self, batch, batch_idx):
431
+ loss, loss_dict = self.shared_step(batch)
432
+
433
+ self.log_dict(loss_dict, prog_bar=True,
434
+ logger=True, on_step=True, on_epoch=True)
435
+
436
+ self.log("global_step", self.global_step,
437
+ prog_bar=True, logger=True, on_step=True, on_epoch=False)
438
+
439
+ if self.use_scheduler:
440
+ lr = self.optimizers().param_groups[0]['lr']
441
+ self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
442
+
443
+ return loss
444
+
445
+ @torch.no_grad()
446
+ def validation_step(self, batch, batch_idx):
447
+ _, loss_dict_no_ema = self.shared_step(batch)
448
+ with self.ema_scope():
449
+ _, loss_dict_ema = self.shared_step(batch)
450
+ loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
451
+ self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
452
+ self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
453
+
454
+ def on_train_batch_end(self, *args, **kwargs):
455
+ if self.use_ema:
456
+ self.model_ema(self.model)
457
+
458
+ def _get_rows_from_list(self, samples):
459
+ n_imgs_per_row = len(samples)
460
+ denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
461
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
462
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
463
+ return denoise_grid
464
+
465
+ @torch.no_grad()
466
+ def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
467
+ log = dict()
468
+ x = self.get_input(batch, self.first_stage_key)
469
+ N = min(x.shape[0], N)
470
+ n_row = min(x.shape[0], n_row)
471
+ x = x.to(self.device)[:N]
472
+ log["inputs"] = x
473
+
474
+ # get diffusion row
475
+ diffusion_row = list()
476
+ x_start = x[:n_row]
477
+
478
+ for t in range(self.num_timesteps):
479
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
480
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
481
+ t = t.to(self.device).long()
482
+ noise = torch.randn_like(x_start)
483
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
484
+ diffusion_row.append(x_noisy)
485
+
486
+ log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
487
+
488
+ if sample:
489
+ # get denoise row
490
+ with self.ema_scope("Plotting"):
491
+ samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
492
+
493
+ log["samples"] = samples
494
+ log["denoise_row"] = self._get_rows_from_list(denoise_row)
495
+
496
+ if return_keys:
497
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
498
+ return log
499
+ else:
500
+ return {key: log[key] for key in return_keys}
501
+ return log
502
+
503
+ def configure_optimizers(self):
504
+ lr = self.learning_rate
505
+ params = list(self.model.parameters())
506
+ if self.learn_logvar:
507
+ params = params + [self.logvar]
508
+ opt = torch.optim.AdamW(params, lr=lr)
509
+ return opt
510
+
511
+
512
+ class LatentDiffusion(DDPM):
513
+ """main class"""
514
+
515
+ def __init__(self,
516
+ first_stage_config,
517
+ cond_stage_config,
518
+ num_timesteps_cond=None,
519
+ cond_stage_key="image",
520
+ cond_stage_trainable=False,
521
+ concat_mode=True,
522
+ cond_stage_forward=None,
523
+ conditioning_key=None,
524
+ scale_factor=1.0,
525
+ scale_by_std=False,
526
+ *args, **kwargs):
527
+ self.num_timesteps_cond = default(num_timesteps_cond, 1)
528
+ self.scale_by_std = scale_by_std
529
+ assert self.num_timesteps_cond <= kwargs['timesteps']
530
+ # for backwards compatibility after implementation of DiffusionWrapper
531
+ if conditioning_key is None:
532
+ conditioning_key = 'concat' if concat_mode else 'crossattn'
533
+ if cond_stage_config == '__is_unconditional__':
534
+ conditioning_key = None
535
+ ckpt_path = kwargs.pop("ckpt_path", None)
536
+ reset_ema = kwargs.pop("reset_ema", False)
537
+ reset_num_ema_updates = kwargs.pop("reset_num_ema_updates", False)
538
+ ignore_keys = kwargs.pop("ignore_keys", [])
539
+ super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
540
+ self.concat_mode = concat_mode
541
+ self.cond_stage_trainable = cond_stage_trainable
542
+ self.cond_stage_key = cond_stage_key
543
+ try:
544
+ self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
545
+ except:
546
+ self.num_downs = 0
547
+ if not scale_by_std:
548
+ self.scale_factor = scale_factor
549
+ else:
550
+ self.register_buffer('scale_factor', torch.tensor(scale_factor))
551
+ self.instantiate_first_stage(first_stage_config)
552
+ self.instantiate_cond_stage(cond_stage_config)
553
+ self.cond_stage_forward = cond_stage_forward
554
+ self.clip_denoised = False
555
+ self.bbox_tokenizer = None
556
+
557
+ self.restarted_from_ckpt = False
558
+ if ckpt_path is not None:
559
+ self.init_from_ckpt(ckpt_path, ignore_keys)
560
+ self.restarted_from_ckpt = True
561
+ if reset_ema:
562
+ assert self.use_ema
563
+ print(
564
+ f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint.")
565
+ self.model_ema = LitEma(self.model)
566
+ if reset_num_ema_updates:
567
+ print(" +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ ")
568
+ assert self.use_ema
569
+ self.model_ema.reset_num_updates()
570
+
571
+ def make_cond_schedule(self, ):
572
+ self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
573
+ ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
574
+ self.cond_ids[:self.num_timesteps_cond] = ids
575
+
576
+ def register_schedule(self,
577
+ given_betas=None, beta_schedule="linear", timesteps=1000,
578
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
579
+ super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
580
+
581
+ self.shorten_cond_schedule = self.num_timesteps_cond > 1
582
+ if self.shorten_cond_schedule:
583
+ self.make_cond_schedule()
584
+
585
+ def instantiate_first_stage(self, config):
586
+ model = instantiate_from_config(config)
587
+ self.first_stage_model = model.eval()
588
+ self.first_stage_model.train = disabled_train
589
+ for param in self.first_stage_model.parameters():
590
+ param.requires_grad = False
591
+
592
+ def instantiate_cond_stage(self, config):
593
+ if not self.cond_stage_trainable:
594
+ if config == "__is_first_stage__":
595
+ print("Using first stage also as cond stage.")
596
+ self.cond_stage_model = self.first_stage_model
597
+ elif config == "__is_unconditional__":
598
+ print(f"Training {self.__class__.__name__} as an unconditional model.")
599
+ self.cond_stage_model = None
600
+ # self.be_unconditional = True
601
+ else:
602
+ model = instantiate_from_config(config)
603
+ self.cond_stage_model = model.eval()
604
+ self.cond_stage_model.train = disabled_train
605
+ for param in self.cond_stage_model.parameters():
606
+ param.requires_grad = False
607
+ else:
608
+ assert config != '__is_first_stage__'
609
+ assert config != '__is_unconditional__'
610
+ model = instantiate_from_config(config)
611
+ self.cond_stage_model = model
612
+
613
+ def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False):
614
+ denoise_row = []
615
+ for zd in tqdm(samples, desc=desc):
616
+ denoise_row.append(self.decode_first_stage(zd.to(self.device),
617
+ force_not_quantize=force_no_decoder_quantization))
618
+ n_imgs_per_row = len(denoise_row)
619
+ denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W
620
+ denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
621
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
622
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
623
+ return denoise_grid
624
+
625
+ def get_first_stage_encoding(self, encoder_posterior):
626
+ if isinstance(encoder_posterior, DiagonalGaussianDistribution):
627
+ z = encoder_posterior.sample()
628
+ elif isinstance(encoder_posterior, torch.Tensor):
629
+ z = encoder_posterior
630
+ else:
631
+ raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
632
+ return self.scale_factor * z
633
+
634
+ def get_learned_conditioning(self, c):
635
+ if self.cond_stage_forward is None:
636
+ if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
637
+ c = self.cond_stage_model.encode(c)
638
+ if isinstance(c, DiagonalGaussianDistribution):
639
+ c = c.mode()
640
+ else:
641
+ c = self.cond_stage_model(c)
642
+ else:
643
+ assert hasattr(self.cond_stage_model, self.cond_stage_forward)
644
+ c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
645
+ return c
646
+
647
+ def meshgrid(self, h, w):
648
+ y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
649
+ x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
650
+
651
+ arr = torch.cat([y, x], dim=-1)
652
+ return arr
653
+
654
+ def delta_border(self, h, w):
655
+ """
656
+ :param h: height
657
+ :param w: width
658
+ :return: normalized distance to image border,
659
+ wtith min distance = 0 at border and max dist = 0.5 at image center
660
+ """
661
+ lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
662
+ arr = self.meshgrid(h, w) / lower_right_corner
663
+ dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
664
+ dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
665
+ edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
666
+ return edge_dist
667
+
668
+ def get_weighting(self, h, w, Ly, Lx, device):
669
+ weighting = self.delta_border(h, w)
670
+ weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"],
671
+ self.split_input_params["clip_max_weight"], )
672
+ weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
673
+
674
+ if self.split_input_params["tie_braker"]:
675
+ L_weighting = self.delta_border(Ly, Lx)
676
+ L_weighting = torch.clip(L_weighting,
677
+ self.split_input_params["clip_min_tie_weight"],
678
+ self.split_input_params["clip_max_tie_weight"])
679
+
680
+ L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
681
+ weighting = weighting * L_weighting
682
+ return weighting
683
+
684
+ def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code
685
+ """
686
+ :param x: img of size (bs, c, h, w)
687
+ :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
688
+ """
689
+ bs, nc, h, w = x.shape
690
+
691
+ # number of crops in image
692
+ Ly = (h - kernel_size[0]) // stride[0] + 1
693
+ Lx = (w - kernel_size[1]) // stride[1] + 1
694
+
695
+ if uf == 1 and df == 1:
696
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
697
+ unfold = torch.nn.Unfold(**fold_params)
698
+
699
+ fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
700
+
701
+ weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype)
702
+ normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap
703
+ weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
704
+
705
+ elif uf > 1 and df == 1:
706
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
707
+ unfold = torch.nn.Unfold(**fold_params)
708
+
709
+ fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
710
+ dilation=1, padding=0,
711
+ stride=(stride[0] * uf, stride[1] * uf))
712
+ fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
713
+
714
+ weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype)
715
+ normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap
716
+ weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
717
+
718
+ elif df > 1 and uf == 1:
719
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
720
+ unfold = torch.nn.Unfold(**fold_params)
721
+
722
+ fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
723
+ dilation=1, padding=0,
724
+ stride=(stride[0] // df, stride[1] // df))
725
+ fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
726
+
727
+ weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype)
728
+ normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap
729
+ weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
730
+
731
+ else:
732
+ raise NotImplementedError
733
+
734
+ return fold, unfold, normalization, weighting
735
+
736
+ @torch.no_grad()
737
+ def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
738
+ cond_key=None, return_original_cond=False, bs=None):
739
+ x = super().get_input(batch, k)
740
+ if bs is not None:
741
+ x = x[:bs]
742
+ x = x.to(self.device)
743
+ encoder_posterior = self.encode_first_stage(x)
744
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
745
+
746
+ if self.model.conditioning_key is not None:
747
+ if cond_key is None:
748
+ cond_key = self.cond_stage_key
749
+ if cond_key != self.first_stage_key:
750
+ if cond_key in ['caption', 'coordinates_bbox', "txt"]:
751
+ xc = batch[cond_key]
752
+ elif cond_key in ['class_label', 'cls']:
753
+ xc = batch
754
+ else:
755
+ xc = super().get_input(batch, cond_key).to(self.device)
756
+ else:
757
+ xc = x
758
+ if not self.cond_stage_trainable or force_c_encode:
759
+ if isinstance(xc, dict) or isinstance(xc, list):
760
+ # import pudb; pudb.set_trace()
761
+ c = self.get_learned_conditioning(xc)
762
+ else:
763
+ c = self.get_learned_conditioning(xc.to(self.device))
764
+ else:
765
+ c = xc
766
+ if bs is not None:
767
+ c = c[:bs]
768
+
769
+ if self.use_positional_encodings:
770
+ pos_x, pos_y = self.compute_latent_shifts(batch)
771
+ ckey = __conditioning_keys__[self.model.conditioning_key]
772
+ c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y}
773
+
774
+ else:
775
+ c = None
776
+ xc = None
777
+ if self.use_positional_encodings:
778
+ pos_x, pos_y = self.compute_latent_shifts(batch)
779
+ c = {'pos_x': pos_x, 'pos_y': pos_y}
780
+ out = [z, c]
781
+ if return_first_stage_outputs:
782
+ xrec = self.decode_first_stage(z)
783
+ out.extend([x, xrec])
784
+ if return_original_cond:
785
+ out.append(xc)
786
+ return out
787
+
788
+ @torch.no_grad()
789
+ def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
790
+ if predict_cids:
791
+ if z.dim() == 4:
792
+ z = torch.argmax(z.exp(), dim=1).long()
793
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
794
+ z = rearrange(z, 'b h w c -> b c h w').contiguous()
795
+
796
+ z = 1. / self.scale_factor * z
797
+ return self.first_stage_model.decode(z)
798
+
799
+ @torch.no_grad()
800
+ def encode_first_stage(self, x):
801
+ return self.first_stage_model.encode(x)
802
+
803
+ def shared_step(self, batch, **kwargs):
804
+ x, c = self.get_input(batch, self.first_stage_key)
805
+ loss = self(x, c, **kwargs)
806
+ return loss
807
+
808
+ def get_time_with_schedule(self, scheduler, bs):
809
+ if scheduler == 'linear':
810
+ t = torch.randint(0, self.num_timesteps, (bs,), device=self.device).long()
811
+ elif scheduler == 'cosine':
812
+ t = torch.rand((bs, ), device=self.device)
813
+ t = torch.cos(torch.pi / 2. * t) * self.num_timesteps
814
+ t = t.long()
815
+ elif scheduler == 'cubic':
816
+ t = torch.rand((bs,), device=self.device)
817
+ t = (1 - t ** 3) * self.num_timesteps
818
+ t = t.long()
819
+ else:
820
+ raise NotImplementedError
821
+ t = torch.clamp(t, min=0, max=self.num_timesteps-1)
822
+ return t
823
+
824
+ def forward(self, x, c, *args, **kwargs):
825
+ if 't' not in kwargs:
826
+ t = torch.randint(0, self.num_timesteps, (x.shape[0], ), device=self.device).long()
827
+ else:
828
+ t = kwargs.pop('t')
829
+
830
+ return self.p_losses(x, c, t, *args, **kwargs)
831
+
832
+ def apply_model(self, x_noisy, t, cond, return_ids=False, **kwargs):
833
+ if isinstance(cond, dict):
834
+ # hybrid case, cond is expected to be a dict
835
+ pass
836
+ else:
837
+ if not isinstance(cond, list):
838
+ cond = [cond]
839
+ key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
840
+ cond = {key: cond}
841
+
842
+ x_recon = self.model(x_noisy, t, **cond, **kwargs)
843
+
844
+ if isinstance(x_recon, tuple) and not return_ids:
845
+ return x_recon[0]
846
+ else:
847
+ return x_recon
848
+
849
+ def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
850
+ return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
851
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
852
+
853
+ def _prior_bpd(self, x_start):
854
+ """
855
+ Get the prior KL term for the variational lower-bound, measured in
856
+ bits-per-dim.
857
+ This term can't be optimized, as it only depends on the encoder.
858
+ :param x_start: the [N x C x ...] tensor of inputs.
859
+ :return: a batch of [N] KL values (in bits), one per batch element.
860
+ """
861
+ batch_size = x_start.shape[0]
862
+ t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
863
+ qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
864
+ kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
865
+ return mean_flat(kl_prior) / np.log(2.0)
866
+
867
+ def p_losses(self, x_start, cond, t, noise=None, **kwargs):
868
+ noise = default(noise, lambda: torch.randn_like(x_start))
869
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
870
+ model_output = self.apply_model(x_noisy, t, cond, **kwargs)
871
+
872
+ loss_dict = {}
873
+ prefix = 'train' if self.training else 'val'
874
+
875
+ if self.parameterization == "x0":
876
+ target = x_start
877
+ elif self.parameterization == "eps":
878
+ target = noise
879
+ elif self.parameterization == "v":
880
+ target = self.get_v(x_start, noise, t)
881
+ else:
882
+ raise NotImplementedError()
883
+
884
+ loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
885
+ loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
886
+
887
+ logvar_t = self.logvar[t].to(self.device)
888
+ loss = loss_simple / torch.exp(logvar_t) + logvar_t
889
+ # loss = loss_simple / torch.exp(self.logvar) + self.logvar
890
+ if self.learn_logvar:
891
+ loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
892
+ loss_dict.update({'logvar': self.logvar.data.mean()})
893
+
894
+ loss = self.l_simple_weight * loss.mean()
895
+
896
+ loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3))
897
+ loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
898
+ loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
899
+ loss += (self.original_elbo_weight * loss_vlb)
900
+ loss_dict.update({f'{prefix}/loss': loss})
901
+
902
+ return loss, loss_dict
903
+
904
+ def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False,
905
+ return_x0=False, score_corrector=None, corrector_kwargs=None):
906
+ t_in = t
907
+ model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
908
+
909
+ if score_corrector is not None:
910
+ assert self.parameterization == "eps"
911
+ model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs)
912
+
913
+ if return_codebook_ids:
914
+ model_out, logits = model_out
915
+
916
+ if self.parameterization == "eps":
917
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
918
+ elif self.parameterization == "x0":
919
+ x_recon = model_out
920
+ else:
921
+ raise NotImplementedError()
922
+
923
+ if clip_denoised:
924
+ x_recon.clamp_(-1., 1.)
925
+ if quantize_denoised:
926
+ x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
927
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
928
+ if return_codebook_ids:
929
+ return model_mean, posterior_variance, posterior_log_variance, logits
930
+ elif return_x0:
931
+ return model_mean, posterior_variance, posterior_log_variance, x_recon
932
+ else:
933
+ return model_mean, posterior_variance, posterior_log_variance
934
+
935
+ @torch.no_grad()
936
+ def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
937
+ return_codebook_ids=False, quantize_denoised=False, return_x0=False,
938
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None):
939
+ b, *_, device = *x.shape, x.device
940
+ outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
941
+ return_codebook_ids=return_codebook_ids,
942
+ quantize_denoised=quantize_denoised,
943
+ return_x0=return_x0,
944
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
945
+ if return_codebook_ids:
946
+ raise DeprecationWarning("Support dropped.")
947
+ model_mean, _, model_log_variance, logits = outputs
948
+ elif return_x0:
949
+ model_mean, _, model_log_variance, x0 = outputs
950
+ else:
951
+ model_mean, _, model_log_variance = outputs
952
+
953
+ noise = noise_like(x.shape, device, repeat_noise) * temperature
954
+ if noise_dropout > 0.:
955
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
956
+ # no noise when t == 0
957
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
958
+
959
+ if return_codebook_ids:
960
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
961
+ if return_x0:
962
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
963
+ else:
964
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
965
+
966
+ @torch.no_grad()
967
+ def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False,
968
+ img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
969
+ score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
970
+ log_every_t=None):
971
+ if not log_every_t:
972
+ log_every_t = self.log_every_t
973
+ timesteps = self.num_timesteps
974
+ if batch_size is not None:
975
+ b = batch_size if batch_size is not None else shape[0]
976
+ shape = [batch_size] + list(shape)
977
+ else:
978
+ b = batch_size = shape[0]
979
+ if x_T is None:
980
+ img = torch.randn(shape, device=self.device)
981
+ else:
982
+ img = x_T
983
+ intermediates = []
984
+ if cond is not None:
985
+ if isinstance(cond, dict):
986
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
987
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
988
+ else:
989
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
990
+
991
+ if start_T is not None:
992
+ timesteps = min(timesteps, start_T)
993
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation',
994
+ total=timesteps) if verbose else reversed(
995
+ range(0, timesteps))
996
+ if type(temperature) == float:
997
+ temperature = [temperature] * timesteps
998
+
999
+ for i in iterator:
1000
+ ts = torch.full((b,), i, device=self.device, dtype=torch.long)
1001
+ if self.shorten_cond_schedule:
1002
+ assert self.model.conditioning_key != 'hybrid'
1003
+ tc = self.cond_ids[ts].to(cond.device)
1004
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1005
+
1006
+ img, x0_partial = self.p_sample(img, cond, ts,
1007
+ clip_denoised=self.clip_denoised,
1008
+ quantize_denoised=quantize_denoised, return_x0=True,
1009
+ temperature=temperature[i], noise_dropout=noise_dropout,
1010
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
1011
+ if mask is not None:
1012
+ assert x0 is not None
1013
+ img_orig = self.q_sample(x0, ts)
1014
+ img = img_orig * mask + (1. - mask) * img
1015
+
1016
+ if i % log_every_t == 0 or i == timesteps - 1:
1017
+ intermediates.append(x0_partial)
1018
+ if callback: callback(i)
1019
+ if img_callback: img_callback(img, i)
1020
+ return img, intermediates
1021
+
1022
+ @torch.no_grad()
1023
+ def p_sample_loop(self, cond, shape, return_intermediates=False,
1024
+ x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False,
1025
+ mask=None, x0=None, img_callback=None, start_T=None,
1026
+ log_every_t=None):
1027
+
1028
+ if not log_every_t:
1029
+ log_every_t = self.log_every_t
1030
+ device = self.betas.device
1031
+ b = shape[0]
1032
+ if x_T is None:
1033
+ img = torch.randn(shape, device=device)
1034
+ else:
1035
+ img = x_T
1036
+
1037
+ intermediates = [img]
1038
+ if timesteps is None:
1039
+ timesteps = self.num_timesteps
1040
+
1041
+ if start_T is not None:
1042
+ timesteps = min(timesteps, start_T)
1043
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(
1044
+ range(0, timesteps))
1045
+
1046
+ if mask is not None:
1047
+ assert x0 is not None
1048
+ assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match
1049
+
1050
+ for i in iterator:
1051
+ ts = torch.full((b,), i, device=device, dtype=torch.long)
1052
+ if self.shorten_cond_schedule:
1053
+ assert self.model.conditioning_key != 'hybrid'
1054
+ tc = self.cond_ids[ts].to(cond.device)
1055
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1056
+
1057
+ img = self.p_sample(img, cond, ts,
1058
+ clip_denoised=self.clip_denoised,
1059
+ quantize_denoised=quantize_denoised)
1060
+ if mask is not None:
1061
+ img_orig = self.q_sample(x0, ts)
1062
+ img = img_orig * mask + (1. - mask) * img
1063
+
1064
+ if i % log_every_t == 0 or i == timesteps - 1:
1065
+ intermediates.append(img)
1066
+ if callback: callback(i)
1067
+ if img_callback: img_callback(img, i)
1068
+
1069
+ if return_intermediates:
1070
+ return img, intermediates
1071
+ return img
1072
+
1073
+ @torch.no_grad()
1074
+ def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
1075
+ verbose=True, timesteps=None, quantize_denoised=False,
1076
+ mask=None, x0=None, shape=None, **kwargs):
1077
+ if shape is None:
1078
+ shape = (batch_size, self.channels, self.image_size, self.image_size)
1079
+ if cond is not None:
1080
+ if isinstance(cond, dict):
1081
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
1082
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
1083
+ else:
1084
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
1085
+ return self.p_sample_loop(cond,
1086
+ shape,
1087
+ return_intermediates=return_intermediates, x_T=x_T,
1088
+ verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
1089
+ mask=mask, x0=x0)
1090
+
1091
+ @torch.no_grad()
1092
+ def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
1093
+ if ddim:
1094
+ ddim_sampler = DDIMSampler(self)
1095
+ shape = (self.channels, self.image_size, self.image_size)
1096
+ samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size,
1097
+ shape, cond, verbose=False, **kwargs)
1098
+
1099
+ else:
1100
+ samples, intermediates = self.sample(cond=cond, batch_size=batch_size,
1101
+ return_intermediates=True, **kwargs)
1102
+
1103
+ return samples, intermediates
1104
+
1105
+ @torch.no_grad()
1106
+ def get_unconditional_conditioning(self, batch_size, null_label=None):
1107
+ if null_label is not None:
1108
+ xc = null_label
1109
+ if isinstance(xc, ListConfig):
1110
+ xc = list(xc)
1111
+ if isinstance(xc, dict) or isinstance(xc, list):
1112
+ c = self.get_learned_conditioning(xc)
1113
+ else:
1114
+ if hasattr(xc, "to"):
1115
+ xc = xc.to(self.device)
1116
+ c = self.get_learned_conditioning(xc)
1117
+ else:
1118
+ if self.cond_stage_key in ["class_label", "cls"]:
1119
+ xc = self.cond_stage_model.get_unconditional_conditioning(batch_size, device=self.device)
1120
+ return self.get_learned_conditioning(xc)
1121
+ else:
1122
+ raise NotImplementedError("todo")
1123
+ if isinstance(c, list): # in case the encoder gives us a list
1124
+ for i in range(len(c)):
1125
+ c[i] = repeat(c[i], '1 ... -> b ...', b=batch_size).to(self.device)
1126
+ else:
1127
+ c = repeat(c, '1 ... -> b ...', b=batch_size).to(self.device)
1128
+ return c
1129
+
1130
+ @torch.no_grad()
1131
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=50, ddim_eta=0., return_keys=None,
1132
+ quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
1133
+ plot_diffusion_rows=True, unconditional_guidance_scale=1., unconditional_guidance_label=None,
1134
+ use_ema_scope=True,
1135
+ **kwargs):
1136
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1137
+ use_ddim = ddim_steps is not None
1138
+
1139
+ log = dict()
1140
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key,
1141
+ return_first_stage_outputs=True,
1142
+ force_c_encode=True,
1143
+ return_original_cond=True,
1144
+ bs=N)
1145
+ N = min(x.shape[0], N)
1146
+ n_row = min(x.shape[0], n_row)
1147
+ log["inputs"] = x
1148
+ log["reconstruction"] = xrec
1149
+ if self.model.conditioning_key is not None:
1150
+ if hasattr(self.cond_stage_model, "decode"):
1151
+ xc = self.cond_stage_model.decode(c)
1152
+ log["conditioning"] = xc
1153
+ elif self.cond_stage_key in ["caption", "txt"]:
1154
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1155
+ log["conditioning"] = xc
1156
+ elif self.cond_stage_key in ['class_label', "cls"]:
1157
+ try:
1158
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25)
1159
+ log['conditioning'] = xc
1160
+ except KeyError:
1161
+ # probably no "human_label" in batch
1162
+ pass
1163
+ elif isimage(xc):
1164
+ log["conditioning"] = xc
1165
+ if ismap(xc):
1166
+ log["original_conditioning"] = self.to_rgb(xc)
1167
+
1168
+ if plot_diffusion_rows:
1169
+ # get diffusion row
1170
+ diffusion_row = list()
1171
+ z_start = z[:n_row]
1172
+ for t in range(self.num_timesteps):
1173
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1174
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1175
+ t = t.to(self.device).long()
1176
+ noise = torch.randn_like(z_start)
1177
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1178
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1179
+
1180
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1181
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1182
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1183
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1184
+ log["diffusion_row"] = diffusion_grid
1185
+
1186
+ if sample:
1187
+ # get denoise row
1188
+ with ema_scope("Sampling"):
1189
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1190
+ ddim_steps=ddim_steps, eta=ddim_eta)
1191
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1192
+ x_samples = self.decode_first_stage(samples)
1193
+ log["samples"] = x_samples
1194
+ if plot_denoise_rows:
1195
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1196
+ log["denoise_row"] = denoise_grid
1197
+
1198
+ if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance(
1199
+ self.first_stage_model, IdentityFirstStage):
1200
+ # also display when quantizing x0 while sampling
1201
+ with ema_scope("Plotting Quantized Denoised"):
1202
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1203
+ ddim_steps=ddim_steps, eta=ddim_eta,
1204
+ quantize_denoised=True)
1205
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True,
1206
+ # quantize_denoised=True)
1207
+ x_samples = self.decode_first_stage(samples.to(self.device))
1208
+ log["samples_x0_quantized"] = x_samples
1209
+
1210
+ if unconditional_guidance_scale > 1.0:
1211
+ uc = self.get_unconditional_conditioning(N, unconditional_guidance_label)
1212
+ if self.model.conditioning_key == "crossattn-adm":
1213
+ uc = {"c_crossattn": [uc], "c_adm": c["c_adm"]}
1214
+ with ema_scope("Sampling with classifier-free guidance"):
1215
+ samples_cfg, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1216
+ ddim_steps=ddim_steps, eta=ddim_eta,
1217
+ unconditional_guidance_scale=unconditional_guidance_scale,
1218
+ unconditional_conditioning=uc,
1219
+ )
1220
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1221
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1222
+
1223
+ if inpaint:
1224
+ # make a simple center square
1225
+ b, h, w = z.shape[0], z.shape[2], z.shape[3]
1226
+ mask = torch.ones(N, h, w).to(self.device)
1227
+ # zeros will be filled in
1228
+ mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
1229
+ mask = mask[:, None, ...]
1230
+ with ema_scope("Plotting Inpaint"):
1231
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
1232
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1233
+ x_samples = self.decode_first_stage(samples.to(self.device))
1234
+ log["samples_inpainting"] = x_samples
1235
+ log["mask"] = mask
1236
+
1237
+ # outpaint
1238
+ mask = 1. - mask
1239
+ with ema_scope("Plotting Outpaint"):
1240
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
1241
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1242
+ x_samples = self.decode_first_stage(samples.to(self.device))
1243
+ log["samples_outpainting"] = x_samples
1244
+
1245
+ if plot_progressive_rows:
1246
+ with ema_scope("Plotting Progressives"):
1247
+ img, progressives = self.progressive_denoising(c,
1248
+ shape=(self.channels, self.image_size, self.image_size),
1249
+ batch_size=N)
1250
+ prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
1251
+ log["progressive_row"] = prog_row
1252
+
1253
+ if return_keys:
1254
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
1255
+ return log
1256
+ else:
1257
+ return {key: log[key] for key in return_keys}
1258
+ return log
1259
+
1260
+ def configure_optimizers(self):
1261
+ lr = self.learning_rate
1262
+ params = list(self.model.parameters())
1263
+ if self.cond_stage_trainable:
1264
+ print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
1265
+ params = params + list(self.cond_stage_model.parameters())
1266
+ if self.learn_logvar:
1267
+ print('Diffusion model optimizing logvar')
1268
+ params.append(self.logvar)
1269
+ opt = torch.optim.AdamW(params, lr=lr)
1270
+ if self.use_scheduler:
1271
+ assert 'target' in self.scheduler_config
1272
+ scheduler = instantiate_from_config(self.scheduler_config)
1273
+
1274
+ print("Setting up LambdaLR scheduler...")
1275
+ scheduler = [
1276
+ {
1277
+ 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
1278
+ 'interval': 'step',
1279
+ 'frequency': 1
1280
+ }]
1281
+ return [opt], scheduler
1282
+ return opt
1283
+
1284
+ @torch.no_grad()
1285
+ def to_rgb(self, x):
1286
+ x = x.float()
1287
+ if not hasattr(self, "colorize"):
1288
+ self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
1289
+ x = nn.functional.conv2d(x, weight=self.colorize)
1290
+ x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
1291
+ return x
1292
+
1293
+
1294
+ class DiffusionWrapper(pl.LightningModule):
1295
+ def __init__(self, diff_model_config, conditioning_key):
1296
+ super().__init__()
1297
+ self.diffusion_model = instantiate_from_config(diff_model_config)
1298
+ self.conditioning_key = conditioning_key
1299
+ assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'hybrid-adm', 'crossattn-adm']
1300
+
1301
+ def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None, **kwargs):
1302
+ if self.conditioning_key is None:
1303
+ out = self.diffusion_model(x, t, **kwargs)
1304
+ elif self.conditioning_key == 'concat':
1305
+ xc = torch.cat([x] + c_concat, dim=1)
1306
+ out = self.diffusion_model(xc, t, **kwargs)
1307
+ elif self.conditioning_key == 'crossattn':
1308
+ cc = torch.cat(c_crossattn, 1)
1309
+ out = self.diffusion_model(x, t, context=cc, **kwargs)
1310
+ elif self.conditioning_key == 'hybrid':
1311
+ xc = torch.cat([x] + c_concat, dim=1)
1312
+ cc = torch.cat(c_crossattn, 1)
1313
+ out = self.diffusion_model(xc, t, context=cc, **kwargs)
1314
+ elif self.conditioning_key == 'hybrid-adm':
1315
+ assert c_adm is not None
1316
+ xc = torch.cat([x] + c_concat, dim=1)
1317
+ cc = torch.cat(c_crossattn, 1)
1318
+ out = self.diffusion_model(xc, t, context=cc, y=c_adm, **kwargs)
1319
+ elif self.conditioning_key == 'crossattn-adm':
1320
+ assert c_adm is not None
1321
+ cc = torch.cat(c_crossattn, 1)
1322
+ out = self.diffusion_model(x, t, context=cc, y=c_adm, **kwargs)
1323
+ elif self.conditioning_key == 'adm':
1324
+ cc = c_crossattn[0]
1325
+ out = self.diffusion_model(x, t, y=cc, **kwargs)
1326
+ else:
1327
+ raise NotImplementedError()
1328
+
1329
+ return out
ldm/models/diffusion/dpm_solver/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .sampler import DPMSolverSampler
ldm/models/diffusion/dpm_solver/dpm_solver.py ADDED
@@ -0,0 +1,1217 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn.functional as F
3
+ import math
4
+ from tqdm import tqdm
5
+
6
+
7
+ class NoiseScheduleVP:
8
+ def __init__(
9
+ self,
10
+ schedule='discrete',
11
+ betas=None,
12
+ alphas_cumprod=None,
13
+ continuous_beta_0=0.1,
14
+ continuous_beta_1=20.,
15
+ ):
16
+ """Create a wrapper class for the forward SDE (VP type).
17
+
18
+ ***
19
+ Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
20
+ We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
21
+ ***
22
+
23
+ The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
24
+ We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
25
+ Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
26
+
27
+ log_alpha_t = self.marginal_log_mean_coeff(t)
28
+ sigma_t = self.marginal_std(t)
29
+ lambda_t = self.marginal_lambda(t)
30
+
31
+ Moreover, as lambda(t) is an invertible function, we also support its inverse function:
32
+
33
+ t = self.inverse_lambda(lambda_t)
34
+
35
+ ===============================================================
36
+
37
+ We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
38
+
39
+ 1. For discrete-time DPMs:
40
+
41
+ For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
42
+ t_i = (i + 1) / N
43
+ e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
44
+ We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
45
+
46
+ Args:
47
+ betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
48
+ alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
49
+
50
+ Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
51
+
52
+ **Important**: Please pay special attention for the args for `alphas_cumprod`:
53
+ The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
54
+ q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
55
+ Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
56
+ alpha_{t_n} = \sqrt{\hat{alpha_n}},
57
+ and
58
+ log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
59
+
60
+
61
+ 2. For continuous-time DPMs:
62
+
63
+ We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
64
+ schedule are the default settings in DDPM and improved-DDPM:
65
+
66
+ Args:
67
+ beta_min: A `float` number. The smallest beta for the linear schedule.
68
+ beta_max: A `float` number. The largest beta for the linear schedule.
69
+ cosine_s: A `float` number. The hyperparameter in the cosine schedule.
70
+ cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
71
+ T: A `float` number. The ending time of the forward process.
72
+
73
+ ===============================================================
74
+
75
+ Args:
76
+ schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
77
+ 'linear' or 'cosine' for continuous-time DPMs.
78
+ Returns:
79
+ A wrapper object of the forward SDE (VP type).
80
+
81
+ ===============================================================
82
+
83
+ Example:
84
+
85
+ # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
86
+ >>> ns = NoiseScheduleVP('discrete', betas=betas)
87
+
88
+ # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
89
+ >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
90
+
91
+ # For continuous-time DPMs (VPSDE), linear schedule:
92
+ >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
93
+
94
+ """
95
+
96
+ if schedule not in ['discrete', 'linear', 'cosine']:
97
+ raise ValueError(
98
+ "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
99
+ schedule))
100
+
101
+ self.schedule = schedule
102
+ if schedule == 'discrete':
103
+ if betas is not None:
104
+ log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
105
+ else:
106
+ assert alphas_cumprod is not None
107
+ log_alphas = 0.5 * torch.log(alphas_cumprod)
108
+ self.total_N = len(log_alphas)
109
+ self.T = 1.
110
+ self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
111
+ self.log_alpha_array = log_alphas.reshape((1, -1,))
112
+ else:
113
+ self.total_N = 1000
114
+ self.beta_0 = continuous_beta_0
115
+ self.beta_1 = continuous_beta_1
116
+ self.cosine_s = 0.008
117
+ self.cosine_beta_max = 999.
118
+ self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (
119
+ 1. + self.cosine_s) / math.pi - self.cosine_s
120
+ self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
121
+ self.schedule = schedule
122
+ if schedule == 'cosine':
123
+ # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
124
+ # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
125
+ self.T = 0.9946
126
+ else:
127
+ self.T = 1.
128
+
129
+ def marginal_log_mean_coeff(self, t):
130
+ """
131
+ Compute log(alpha_t) of a given continuous-time label t in [0, T].
132
+ """
133
+ if self.schedule == 'discrete':
134
+ return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device),
135
+ self.log_alpha_array.to(t.device)).reshape((-1))
136
+ elif self.schedule == 'linear':
137
+ return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
138
+ elif self.schedule == 'cosine':
139
+ log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
140
+ log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
141
+ return log_alpha_t
142
+
143
+ def marginal_alpha(self, t):
144
+ """
145
+ Compute alpha_t of a given continuous-time label t in [0, T].
146
+ """
147
+ return torch.exp(self.marginal_log_mean_coeff(t))
148
+
149
+ def marginal_std(self, t):
150
+ """
151
+ Compute sigma_t of a given continuous-time label t in [0, T].
152
+ """
153
+ return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
154
+
155
+ def marginal_lambda(self, t):
156
+ """
157
+ Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
158
+ """
159
+ log_mean_coeff = self.marginal_log_mean_coeff(t)
160
+ log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
161
+ return log_mean_coeff - log_std
162
+
163
+ def inverse_lambda(self, lamb):
164
+ """
165
+ Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
166
+ """
167
+ if self.schedule == 'linear':
168
+ tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
169
+ Delta = self.beta_0 ** 2 + tmp
170
+ return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
171
+ elif self.schedule == 'discrete':
172
+ log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
173
+ t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]),
174
+ torch.flip(self.t_array.to(lamb.device), [1]))
175
+ return t.reshape((-1,))
176
+ else:
177
+ log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
178
+ t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (
179
+ 1. + self.cosine_s) / math.pi - self.cosine_s
180
+ t = t_fn(log_alpha)
181
+ return t
182
+
183
+
184
+ def model_wrapper(
185
+ model,
186
+ noise_schedule,
187
+ model_type="noise",
188
+ model_kwargs={},
189
+ guidance_type="uncond",
190
+ condition=None,
191
+ unconditional_condition=None,
192
+ guidance_scale=1.,
193
+ classifier_fn=None,
194
+ classifier_kwargs={},
195
+ ):
196
+ """Create a wrapper function for the noise prediction model.
197
+
198
+ DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
199
+ firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
200
+
201
+ We support four types of the diffusion model by setting `model_type`:
202
+
203
+ 1. "noise": noise prediction model. (Trained by predicting noise).
204
+
205
+ 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
206
+
207
+ 3. "v": velocity prediction model. (Trained by predicting the velocity).
208
+ The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
209
+
210
+ [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
211
+ arXiv preprint arXiv:2202.00512 (2022).
212
+ [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
213
+ arXiv preprint arXiv:2210.02303 (2022).
214
+
215
+ 4. "score": marginal score function. (Trained by denoising score matching).
216
+ Note that the score function and the noise prediction model follows a simple relationship:
217
+ ```
218
+ noise(x_t, t) = -sigma_t * score(x_t, t)
219
+ ```
220
+
221
+ We support three types of guided sampling by DPMs by setting `guidance_type`:
222
+ 1. "uncond": unconditional sampling by DPMs.
223
+ The input `model` has the following format:
224
+ ``
225
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
226
+ ``
227
+
228
+ 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
229
+ The input `model` has the following format:
230
+ ``
231
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
232
+ ``
233
+
234
+ The input `classifier_fn` has the following format:
235
+ ``
236
+ classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
237
+ ``
238
+
239
+ [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
240
+ in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
241
+
242
+ 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
243
+ The input `model` has the following format:
244
+ ``
245
+ model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
246
+ ``
247
+ And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
248
+
249
+ [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
250
+ arXiv preprint arXiv:2207.12598 (2022).
251
+
252
+
253
+ The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
254
+ or continuous-time labels (i.e. epsilon to T).
255
+
256
+ We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
257
+ ``
258
+ def model_fn(x, t_continuous) -> noise:
259
+ t_input = get_model_input_time(t_continuous)
260
+ return noise_pred(model, x, t_input, **model_kwargs)
261
+ ``
262
+ where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
263
+
264
+ ===============================================================
265
+
266
+ Args:
267
+ model: A diffusion model with the corresponding format described above.
268
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
269
+ model_type: A `str`. The parameterization type of the diffusion model.
270
+ "noise" or "x_start" or "v" or "score".
271
+ model_kwargs: A `dict`. A dict for the other inputs of the model function.
272
+ guidance_type: A `str`. The type of the guidance for sampling.
273
+ "uncond" or "classifier" or "classifier-free".
274
+ condition: A pytorch tensor. The condition for the guided sampling.
275
+ Only used for "classifier" or "classifier-free" guidance type.
276
+ unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
277
+ Only used for "classifier-free" guidance type.
278
+ guidance_scale: A `float`. The scale for the guided sampling.
279
+ classifier_fn: A classifier function. Only used for the classifier guidance.
280
+ classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
281
+ Returns:
282
+ A noise prediction model that accepts the noised data and the continuous time as the inputs.
283
+ """
284
+
285
+ def get_model_input_time(t_continuous):
286
+ """
287
+ Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
288
+ For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
289
+ For continuous-time DPMs, we just use `t_continuous`.
290
+ """
291
+ if noise_schedule.schedule == 'discrete':
292
+ return (t_continuous - 1. / noise_schedule.total_N) * 1000.
293
+ else:
294
+ return t_continuous
295
+
296
+ def noise_pred_fn(x, t_continuous, cond=None):
297
+ if t_continuous.reshape((-1,)).shape[0] == 1:
298
+ t_continuous = t_continuous.expand((x.shape[0]))
299
+ t_input = get_model_input_time(t_continuous)
300
+ if cond is None:
301
+ output = model(x, t_input, **model_kwargs)
302
+ else:
303
+ output = model(x, t_input, cond, **model_kwargs)
304
+ if model_type == "noise":
305
+ return output
306
+ elif model_type == "x_start":
307
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
308
+ dims = x.dim()
309
+ return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
310
+ elif model_type == "v":
311
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
312
+ dims = x.dim()
313
+ return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
314
+ elif model_type == "score":
315
+ sigma_t = noise_schedule.marginal_std(t_continuous)
316
+ dims = x.dim()
317
+ return -expand_dims(sigma_t, dims) * output
318
+
319
+ def cond_grad_fn(x, t_input):
320
+ """
321
+ Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
322
+ """
323
+ with torch.enable_grad():
324
+ x_in = x.detach().requires_grad_(True)
325
+ log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
326
+ return torch.autograd.grad(log_prob.sum(), x_in)[0]
327
+
328
+ def model_fn(x, t_continuous):
329
+ """
330
+ The noise predicition model function that is used for DPM-Solver.
331
+ """
332
+ if t_continuous.reshape((-1,)).shape[0] == 1:
333
+ t_continuous = t_continuous.expand((x.shape[0]))
334
+ if guidance_type == "uncond":
335
+ return noise_pred_fn(x, t_continuous)
336
+ elif guidance_type == "classifier":
337
+ assert classifier_fn is not None
338
+ t_input = get_model_input_time(t_continuous)
339
+ cond_grad = cond_grad_fn(x, t_input)
340
+ sigma_t = noise_schedule.marginal_std(t_continuous)
341
+ noise = noise_pred_fn(x, t_continuous)
342
+ return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
343
+ elif guidance_type == "classifier-free":
344
+ if guidance_scale == 1. or unconditional_condition is None:
345
+ return noise_pred_fn(x, t_continuous, cond=condition)
346
+ else:
347
+ x_in = torch.cat([x] * 2)
348
+ t_in = torch.cat([t_continuous] * 2)
349
+ c_in = torch.cat([unconditional_condition, condition])
350
+ noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
351
+ return noise_uncond + guidance_scale * (noise - noise_uncond)
352
+
353
+ assert model_type in ["noise", "x_start", "v"]
354
+ assert guidance_type in ["uncond", "classifier", "classifier-free"]
355
+ return model_fn
356
+
357
+
358
+ class DPM_Solver:
359
+ def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):
360
+ """Construct a DPM-Solver.
361
+
362
+ We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
363
+ If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
364
+ If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
365
+ In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
366
+ The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.
367
+
368
+ Args:
369
+ model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
370
+ ``
371
+ def model_fn(x, t_continuous):
372
+ return noise
373
+ ``
374
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
375
+ predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
376
+ thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
377
+ max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.
378
+
379
+ [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
380
+ """
381
+ self.model = model_fn
382
+ self.noise_schedule = noise_schedule
383
+ self.predict_x0 = predict_x0
384
+ self.thresholding = thresholding
385
+ self.max_val = max_val
386
+
387
+ def noise_prediction_fn(self, x, t):
388
+ """
389
+ Return the noise prediction model.
390
+ """
391
+ return self.model(x, t)
392
+
393
+ def data_prediction_fn(self, x, t):
394
+ """
395
+ Return the data prediction model (with thresholding).
396
+ """
397
+ noise = self.noise_prediction_fn(x, t)
398
+ dims = x.dim()
399
+ alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
400
+ x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
401
+ if self.thresholding:
402
+ p = 0.995 # A hyperparameter in the paper of "Imagen" [1].
403
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
404
+ s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
405
+ x0 = torch.clamp(x0, -s, s) / s
406
+ return x0
407
+
408
+ def model_fn(self, x, t):
409
+ """
410
+ Convert the model to the noise prediction model or the data prediction model.
411
+ """
412
+ if self.predict_x0:
413
+ return self.data_prediction_fn(x, t)
414
+ else:
415
+ return self.noise_prediction_fn(x, t)
416
+
417
+ def get_time_steps(self, skip_type, t_T, t_0, N, device):
418
+ """Compute the intermediate time steps for sampling.
419
+
420
+ Args:
421
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
422
+ - 'logSNR': uniform logSNR for the time steps.
423
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
424
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
425
+ t_T: A `float`. The starting time of the sampling (default is T).
426
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
427
+ N: A `int`. The total number of the spacing of the time steps.
428
+ device: A torch device.
429
+ Returns:
430
+ A pytorch tensor of the time steps, with the shape (N + 1,).
431
+ """
432
+ if skip_type == 'logSNR':
433
+ lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
434
+ lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
435
+ logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
436
+ return self.noise_schedule.inverse_lambda(logSNR_steps)
437
+ elif skip_type == 'time_uniform':
438
+ return torch.linspace(t_T, t_0, N + 1).to(device)
439
+ elif skip_type == 'time_quadratic':
440
+ t_order = 2
441
+ t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)
442
+ return t
443
+ else:
444
+ raise ValueError(
445
+ "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
446
+
447
+ def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
448
+ """
449
+ Get the order of each step for sampling by the singlestep DPM-Solver.
450
+
451
+ We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
452
+ Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
453
+ - If order == 1:
454
+ We take `steps` of DPM-Solver-1 (i.e. DDIM).
455
+ - If order == 2:
456
+ - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
457
+ - If steps % 2 == 0, we use K steps of DPM-Solver-2.
458
+ - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
459
+ - If order == 3:
460
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
461
+ - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1.
462
+ - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
463
+ - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
464
+
465
+ ============================================
466
+ Args:
467
+ order: A `int`. The max order for the solver (2 or 3).
468
+ steps: A `int`. The total number of function evaluations (NFE).
469
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
470
+ - 'logSNR': uniform logSNR for the time steps.
471
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
472
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
473
+ t_T: A `float`. The starting time of the sampling (default is T).
474
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
475
+ device: A torch device.
476
+ Returns:
477
+ orders: A list of the solver order of each step.
478
+ """
479
+ if order == 3:
480
+ K = steps // 3 + 1
481
+ if steps % 3 == 0:
482
+ orders = [3, ] * (K - 2) + [2, 1]
483
+ elif steps % 3 == 1:
484
+ orders = [3, ] * (K - 1) + [1]
485
+ else:
486
+ orders = [3, ] * (K - 1) + [2]
487
+ elif order == 2:
488
+ if steps % 2 == 0:
489
+ K = steps // 2
490
+ orders = [2, ] * K
491
+ else:
492
+ K = steps // 2 + 1
493
+ orders = [2, ] * (K - 1) + [1]
494
+ elif order == 1:
495
+ K = 1
496
+ orders = [1, ] * steps
497
+ else:
498
+ raise ValueError("'order' must be '1' or '2' or '3'.")
499
+ if skip_type == 'logSNR':
500
+ # To reproduce the results in DPM-Solver paper
501
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
502
+ else:
503
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[
504
+ torch.cumsum(torch.tensor([0, ] + orders)).to(device)]
505
+ return timesteps_outer, orders
506
+
507
+ def denoise_to_zero_fn(self, x, s):
508
+ """
509
+ Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
510
+ """
511
+ return self.data_prediction_fn(x, s)
512
+
513
+ def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
514
+ """
515
+ DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
516
+
517
+ Args:
518
+ x: A pytorch tensor. The initial value at time `s`.
519
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
520
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
521
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
522
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
523
+ return_intermediate: A `bool`. If true, also return the model value at time `s`.
524
+ Returns:
525
+ x_t: A pytorch tensor. The approximated solution at time `t`.
526
+ """
527
+ ns = self.noise_schedule
528
+ dims = x.dim()
529
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
530
+ h = lambda_t - lambda_s
531
+ log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
532
+ sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
533
+ alpha_t = torch.exp(log_alpha_t)
534
+
535
+ if self.predict_x0:
536
+ phi_1 = torch.expm1(-h)
537
+ if model_s is None:
538
+ model_s = self.model_fn(x, s)
539
+ x_t = (
540
+ expand_dims(sigma_t / sigma_s, dims) * x
541
+ - expand_dims(alpha_t * phi_1, dims) * model_s
542
+ )
543
+ if return_intermediate:
544
+ return x_t, {'model_s': model_s}
545
+ else:
546
+ return x_t
547
+ else:
548
+ phi_1 = torch.expm1(h)
549
+ if model_s is None:
550
+ model_s = self.model_fn(x, s)
551
+ x_t = (
552
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
553
+ - expand_dims(sigma_t * phi_1, dims) * model_s
554
+ )
555
+ if return_intermediate:
556
+ return x_t, {'model_s': model_s}
557
+ else:
558
+ return x_t
559
+
560
+ def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False,
561
+ solver_type='dpm_solver'):
562
+ """
563
+ Singlestep solver DPM-Solver-2 from time `s` to time `t`.
564
+
565
+ Args:
566
+ x: A pytorch tensor. The initial value at time `s`.
567
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
568
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
569
+ r1: A `float`. The hyperparameter of the second-order solver.
570
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
571
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
572
+ return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
573
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
574
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
575
+ Returns:
576
+ x_t: A pytorch tensor. The approximated solution at time `t`.
577
+ """
578
+ if solver_type not in ['dpm_solver', 'taylor']:
579
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
580
+ if r1 is None:
581
+ r1 = 0.5
582
+ ns = self.noise_schedule
583
+ dims = x.dim()
584
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
585
+ h = lambda_t - lambda_s
586
+ lambda_s1 = lambda_s + r1 * h
587
+ s1 = ns.inverse_lambda(lambda_s1)
588
+ log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(
589
+ s1), ns.marginal_log_mean_coeff(t)
590
+ sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
591
+ alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
592
+
593
+ if self.predict_x0:
594
+ phi_11 = torch.expm1(-r1 * h)
595
+ phi_1 = torch.expm1(-h)
596
+
597
+ if model_s is None:
598
+ model_s = self.model_fn(x, s)
599
+ x_s1 = (
600
+ expand_dims(sigma_s1 / sigma_s, dims) * x
601
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
602
+ )
603
+ model_s1 = self.model_fn(x_s1, s1)
604
+ if solver_type == 'dpm_solver':
605
+ x_t = (
606
+ expand_dims(sigma_t / sigma_s, dims) * x
607
+ - expand_dims(alpha_t * phi_1, dims) * model_s
608
+ - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
609
+ )
610
+ elif solver_type == 'taylor':
611
+ x_t = (
612
+ expand_dims(sigma_t / sigma_s, dims) * x
613
+ - expand_dims(alpha_t * phi_1, dims) * model_s
614
+ + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (
615
+ model_s1 - model_s)
616
+ )
617
+ else:
618
+ phi_11 = torch.expm1(r1 * h)
619
+ phi_1 = torch.expm1(h)
620
+
621
+ if model_s is None:
622
+ model_s = self.model_fn(x, s)
623
+ x_s1 = (
624
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
625
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
626
+ )
627
+ model_s1 = self.model_fn(x_s1, s1)
628
+ if solver_type == 'dpm_solver':
629
+ x_t = (
630
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
631
+ - expand_dims(sigma_t * phi_1, dims) * model_s
632
+ - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
633
+ )
634
+ elif solver_type == 'taylor':
635
+ x_t = (
636
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
637
+ - expand_dims(sigma_t * phi_1, dims) * model_s
638
+ - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
639
+ )
640
+ if return_intermediate:
641
+ return x_t, {'model_s': model_s, 'model_s1': model_s1}
642
+ else:
643
+ return x_t
644
+
645
+ def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None,
646
+ return_intermediate=False, solver_type='dpm_solver'):
647
+ """
648
+ Singlestep solver DPM-Solver-3 from time `s` to time `t`.
649
+
650
+ Args:
651
+ x: A pytorch tensor. The initial value at time `s`.
652
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
653
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
654
+ r1: A `float`. The hyperparameter of the third-order solver.
655
+ r2: A `float`. The hyperparameter of the third-order solver.
656
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
657
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
658
+ model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
659
+ If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
660
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
661
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
662
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
663
+ Returns:
664
+ x_t: A pytorch tensor. The approximated solution at time `t`.
665
+ """
666
+ if solver_type not in ['dpm_solver', 'taylor']:
667
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
668
+ if r1 is None:
669
+ r1 = 1. / 3.
670
+ if r2 is None:
671
+ r2 = 2. / 3.
672
+ ns = self.noise_schedule
673
+ dims = x.dim()
674
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
675
+ h = lambda_t - lambda_s
676
+ lambda_s1 = lambda_s + r1 * h
677
+ lambda_s2 = lambda_s + r2 * h
678
+ s1 = ns.inverse_lambda(lambda_s1)
679
+ s2 = ns.inverse_lambda(lambda_s2)
680
+ log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(
681
+ s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
682
+ sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(
683
+ s2), ns.marginal_std(t)
684
+ alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)
685
+
686
+ if self.predict_x0:
687
+ phi_11 = torch.expm1(-r1 * h)
688
+ phi_12 = torch.expm1(-r2 * h)
689
+ phi_1 = torch.expm1(-h)
690
+ phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
691
+ phi_2 = phi_1 / h + 1.
692
+ phi_3 = phi_2 / h - 0.5
693
+
694
+ if model_s is None:
695
+ model_s = self.model_fn(x, s)
696
+ if model_s1 is None:
697
+ x_s1 = (
698
+ expand_dims(sigma_s1 / sigma_s, dims) * x
699
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
700
+ )
701
+ model_s1 = self.model_fn(x_s1, s1)
702
+ x_s2 = (
703
+ expand_dims(sigma_s2 / sigma_s, dims) * x
704
+ - expand_dims(alpha_s2 * phi_12, dims) * model_s
705
+ + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
706
+ )
707
+ model_s2 = self.model_fn(x_s2, s2)
708
+ if solver_type == 'dpm_solver':
709
+ x_t = (
710
+ expand_dims(sigma_t / sigma_s, dims) * x
711
+ - expand_dims(alpha_t * phi_1, dims) * model_s
712
+ + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
713
+ )
714
+ elif solver_type == 'taylor':
715
+ D1_0 = (1. / r1) * (model_s1 - model_s)
716
+ D1_1 = (1. / r2) * (model_s2 - model_s)
717
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
718
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
719
+ x_t = (
720
+ expand_dims(sigma_t / sigma_s, dims) * x
721
+ - expand_dims(alpha_t * phi_1, dims) * model_s
722
+ + expand_dims(alpha_t * phi_2, dims) * D1
723
+ - expand_dims(alpha_t * phi_3, dims) * D2
724
+ )
725
+ else:
726
+ phi_11 = torch.expm1(r1 * h)
727
+ phi_12 = torch.expm1(r2 * h)
728
+ phi_1 = torch.expm1(h)
729
+ phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
730
+ phi_2 = phi_1 / h - 1.
731
+ phi_3 = phi_2 / h - 0.5
732
+
733
+ if model_s is None:
734
+ model_s = self.model_fn(x, s)
735
+ if model_s1 is None:
736
+ x_s1 = (
737
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
738
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
739
+ )
740
+ model_s1 = self.model_fn(x_s1, s1)
741
+ x_s2 = (
742
+ expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
743
+ - expand_dims(sigma_s2 * phi_12, dims) * model_s
744
+ - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
745
+ )
746
+ model_s2 = self.model_fn(x_s2, s2)
747
+ if solver_type == 'dpm_solver':
748
+ x_t = (
749
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
750
+ - expand_dims(sigma_t * phi_1, dims) * model_s
751
+ - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
752
+ )
753
+ elif solver_type == 'taylor':
754
+ D1_0 = (1. / r1) * (model_s1 - model_s)
755
+ D1_1 = (1. / r2) * (model_s2 - model_s)
756
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
757
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
758
+ x_t = (
759
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
760
+ - expand_dims(sigma_t * phi_1, dims) * model_s
761
+ - expand_dims(sigma_t * phi_2, dims) * D1
762
+ - expand_dims(sigma_t * phi_3, dims) * D2
763
+ )
764
+
765
+ if return_intermediate:
766
+ return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
767
+ else:
768
+ return x_t
769
+
770
+ def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
771
+ """
772
+ Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
773
+
774
+ Args:
775
+ x: A pytorch tensor. The initial value at time `s`.
776
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
777
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
778
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
779
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
780
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
781
+ Returns:
782
+ x_t: A pytorch tensor. The approximated solution at time `t`.
783
+ """
784
+ if solver_type not in ['dpm_solver', 'taylor']:
785
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
786
+ ns = self.noise_schedule
787
+ dims = x.dim()
788
+ model_prev_1, model_prev_0 = model_prev_list
789
+ t_prev_1, t_prev_0 = t_prev_list
790
+ lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(
791
+ t_prev_0), ns.marginal_lambda(t)
792
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
793
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
794
+ alpha_t = torch.exp(log_alpha_t)
795
+
796
+ h_0 = lambda_prev_0 - lambda_prev_1
797
+ h = lambda_t - lambda_prev_0
798
+ r0 = h_0 / h
799
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
800
+ if self.predict_x0:
801
+ if solver_type == 'dpm_solver':
802
+ x_t = (
803
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
804
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
805
+ - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
806
+ )
807
+ elif solver_type == 'taylor':
808
+ x_t = (
809
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
810
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
811
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
812
+ )
813
+ else:
814
+ if solver_type == 'dpm_solver':
815
+ x_t = (
816
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
817
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
818
+ - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
819
+ )
820
+ elif solver_type == 'taylor':
821
+ x_t = (
822
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
823
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
824
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
825
+ )
826
+ return x_t
827
+
828
+ def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
829
+ """
830
+ Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
831
+
832
+ Args:
833
+ x: A pytorch tensor. The initial value at time `s`.
834
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
835
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
836
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
837
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
838
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
839
+ Returns:
840
+ x_t: A pytorch tensor. The approximated solution at time `t`.
841
+ """
842
+ ns = self.noise_schedule
843
+ dims = x.dim()
844
+ model_prev_2, model_prev_1, model_prev_0 = model_prev_list
845
+ t_prev_2, t_prev_1, t_prev_0 = t_prev_list
846
+ lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(
847
+ t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
848
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
849
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
850
+ alpha_t = torch.exp(log_alpha_t)
851
+
852
+ h_1 = lambda_prev_1 - lambda_prev_2
853
+ h_0 = lambda_prev_0 - lambda_prev_1
854
+ h = lambda_t - lambda_prev_0
855
+ r0, r1 = h_0 / h, h_1 / h
856
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
857
+ D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
858
+ D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
859
+ D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
860
+ if self.predict_x0:
861
+ x_t = (
862
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
863
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
864
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
865
+ - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2
866
+ )
867
+ else:
868
+ x_t = (
869
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
870
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
871
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
872
+ - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2
873
+ )
874
+ return x_t
875
+
876
+ def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None,
877
+ r2=None):
878
+ """
879
+ Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
880
+
881
+ Args:
882
+ x: A pytorch tensor. The initial value at time `s`.
883
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
884
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
885
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
886
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
887
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
888
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
889
+ r1: A `float`. The hyperparameter of the second-order or third-order solver.
890
+ r2: A `float`. The hyperparameter of the third-order solver.
891
+ Returns:
892
+ x_t: A pytorch tensor. The approximated solution at time `t`.
893
+ """
894
+ if order == 1:
895
+ return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
896
+ elif order == 2:
897
+ return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate,
898
+ solver_type=solver_type, r1=r1)
899
+ elif order == 3:
900
+ return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate,
901
+ solver_type=solver_type, r1=r1, r2=r2)
902
+ else:
903
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
904
+
905
+ def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):
906
+ """
907
+ Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
908
+
909
+ Args:
910
+ x: A pytorch tensor. The initial value at time `s`.
911
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
912
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
913
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
914
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
915
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
916
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
917
+ Returns:
918
+ x_t: A pytorch tensor. The approximated solution at time `t`.
919
+ """
920
+ if order == 1:
921
+ return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
922
+ elif order == 2:
923
+ return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
924
+ elif order == 3:
925
+ return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
926
+ else:
927
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
928
+
929
+ def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5,
930
+ solver_type='dpm_solver'):
931
+ """
932
+ The adaptive step size solver based on singlestep DPM-Solver.
933
+
934
+ Args:
935
+ x: A pytorch tensor. The initial value at time `t_T`.
936
+ order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
937
+ t_T: A `float`. The starting time of the sampling (default is T).
938
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
939
+ h_init: A `float`. The initial step size (for logSNR).
940
+ atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
941
+ rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
942
+ theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
943
+ t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the
944
+ current time and `t_0` is less than `t_err`. The default setting is 1e-5.
945
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
946
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
947
+ Returns:
948
+ x_0: A pytorch tensor. The approximated solution at time `t_0`.
949
+
950
+ [1] A. Jolicoeur-Martineau, K. Li, R. PichΓ©-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.
951
+ """
952
+ ns = self.noise_schedule
953
+ s = t_T * torch.ones((x.shape[0],)).to(x)
954
+ lambda_s = ns.marginal_lambda(s)
955
+ lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
956
+ h = h_init * torch.ones_like(s).to(x)
957
+ x_prev = x
958
+ nfe = 0
959
+ if order == 2:
960
+ r1 = 0.5
961
+ lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
962
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
963
+ solver_type=solver_type,
964
+ **kwargs)
965
+ elif order == 3:
966
+ r1, r2 = 1. / 3., 2. / 3.
967
+ lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
968
+ return_intermediate=True,
969
+ solver_type=solver_type)
970
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,
971
+ solver_type=solver_type,
972
+ **kwargs)
973
+ else:
974
+ raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
975
+ while torch.abs((s - t_0)).mean() > t_err:
976
+ t = ns.inverse_lambda(lambda_s + h)
977
+ x_lower, lower_noise_kwargs = lower_update(x, s, t)
978
+ x_higher = higher_update(x, s, t, **lower_noise_kwargs)
979
+ delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
980
+ norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
981
+ E = norm_fn((x_higher - x_lower) / delta).max()
982
+ if torch.all(E <= 1.):
983
+ x = x_higher
984
+ s = t
985
+ x_prev = x_lower
986
+ lambda_s = ns.marginal_lambda(s)
987
+ h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
988
+ nfe += order
989
+ print('adaptive solver nfe', nfe)
990
+ return x
991
+
992
+ def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
993
+ method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
994
+ atol=0.0078, rtol=0.05,
995
+ ):
996
+ """
997
+ Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
998
+
999
+ =====================================================
1000
+
1001
+ We support the following algorithms for both noise prediction model and data prediction model:
1002
+ - 'singlestep':
1003
+ Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver.
1004
+ We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
1005
+ The total number of function evaluations (NFE) == `steps`.
1006
+ Given a fixed NFE == `steps`, the sampling procedure is:
1007
+ - If `order` == 1:
1008
+ - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
1009
+ - If `order` == 2:
1010
+ - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
1011
+ - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
1012
+ - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
1013
+ - If `order` == 3:
1014
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
1015
+ - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
1016
+ - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
1017
+ - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
1018
+ - 'multistep':
1019
+ Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
1020
+ We initialize the first `order` values by lower order multistep solvers.
1021
+ Given a fixed NFE == `steps`, the sampling procedure is:
1022
+ Denote K = steps.
1023
+ - If `order` == 1:
1024
+ - We use K steps of DPM-Solver-1 (i.e. DDIM).
1025
+ - If `order` == 2:
1026
+ - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
1027
+ - If `order` == 3:
1028
+ - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
1029
+ - 'singlestep_fixed':
1030
+ Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
1031
+ We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
1032
+ - 'adaptive':
1033
+ Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
1034
+ We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
1035
+ You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
1036
+ (NFE) and the sample quality.
1037
+ - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
1038
+ - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
1039
+
1040
+ =====================================================
1041
+
1042
+ Some advices for choosing the algorithm:
1043
+ - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
1044
+ Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
1045
+ e.g.
1046
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
1047
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
1048
+ skip_type='time_uniform', method='singlestep')
1049
+ - For **guided sampling with large guidance scale** by DPMs:
1050
+ Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.
1051
+ e.g.
1052
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
1053
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
1054
+ skip_type='time_uniform', method='multistep')
1055
+
1056
+ We support three types of `skip_type`:
1057
+ - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
1058
+ - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
1059
+ - 'time_quadratic': quadratic time for the time steps.
1060
+
1061
+ =====================================================
1062
+ Args:
1063
+ x: A pytorch tensor. The initial value at time `t_start`
1064
+ e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
1065
+ steps: A `int`. The total number of function evaluations (NFE).
1066
+ t_start: A `float`. The starting time of the sampling.
1067
+ If `T` is None, we use self.noise_schedule.T (default is 1.0).
1068
+ t_end: A `float`. The ending time of the sampling.
1069
+ If `t_end` is None, we use 1. / self.noise_schedule.total_N.
1070
+ e.g. if total_N == 1000, we have `t_end` == 1e-3.
1071
+ For discrete-time DPMs:
1072
+ - We recommend `t_end` == 1. / self.noise_schedule.total_N.
1073
+ For continuous-time DPMs:
1074
+ - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
1075
+ order: A `int`. The order of DPM-Solver.
1076
+ skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
1077
+ method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
1078
+ denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.
1079
+ Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).
1080
+
1081
+ This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and
1082
+ score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID
1083
+ for diffusion models sampling by diffusion SDEs for low-resolutional images
1084
+ (such as CIFAR-10). However, we observed that such trick does not matter for
1085
+ high-resolutional images. As it needs an additional NFE, we do not recommend
1086
+ it for high-resolutional images.
1087
+ lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.
1088
+ Only valid for `method=multistep` and `steps < 15`. We empirically find that
1089
+ this trick is a key to stabilizing the sampling by DPM-Solver with very few steps
1090
+ (especially for steps <= 10). So we recommend to set it to be `True`.
1091
+ solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
1092
+ atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1093
+ rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1094
+ Returns:
1095
+ x_end: A pytorch tensor. The approximated solution at time `t_end`.
1096
+
1097
+ """
1098
+ t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
1099
+ t_T = self.noise_schedule.T if t_start is None else t_start
1100
+ device = x.device
1101
+ if method == 'adaptive':
1102
+ with torch.no_grad():
1103
+ x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol,
1104
+ solver_type=solver_type)
1105
+ elif method == 'multistep':
1106
+ assert steps >= order
1107
+ timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
1108
+ assert timesteps.shape[0] - 1 == steps
1109
+ with torch.no_grad():
1110
+ vec_t = timesteps[0].expand((x.shape[0]))
1111
+ model_prev_list = [self.model_fn(x, vec_t)]
1112
+ t_prev_list = [vec_t]
1113
+ # Init the first `order` values by lower order multistep DPM-Solver.
1114
+ for init_order in tqdm(range(1, order), desc="DPM init order"):
1115
+ vec_t = timesteps[init_order].expand(x.shape[0])
1116
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,
1117
+ solver_type=solver_type)
1118
+ model_prev_list.append(self.model_fn(x, vec_t))
1119
+ t_prev_list.append(vec_t)
1120
+ # Compute the remaining values by `order`-th order multistep DPM-Solver.
1121
+ for step in tqdm(range(order, steps + 1), desc="DPM multistep"):
1122
+ vec_t = timesteps[step].expand(x.shape[0])
1123
+ if lower_order_final and steps < 15:
1124
+ step_order = min(order, steps + 1 - step)
1125
+ else:
1126
+ step_order = order
1127
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, step_order,
1128
+ solver_type=solver_type)
1129
+ for i in range(order - 1):
1130
+ t_prev_list[i] = t_prev_list[i + 1]
1131
+ model_prev_list[i] = model_prev_list[i + 1]
1132
+ t_prev_list[-1] = vec_t
1133
+ # We do not need to evaluate the final model value.
1134
+ if step < steps:
1135
+ model_prev_list[-1] = self.model_fn(x, vec_t)
1136
+ elif method in ['singlestep', 'singlestep_fixed']:
1137
+ if method == 'singlestep':
1138
+ timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order,
1139
+ skip_type=skip_type,
1140
+ t_T=t_T, t_0=t_0,
1141
+ device=device)
1142
+ elif method == 'singlestep_fixed':
1143
+ K = steps // order
1144
+ orders = [order, ] * K
1145
+ timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
1146
+ for i, order in enumerate(orders):
1147
+ t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
1148
+ timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),
1149
+ N=order, device=device)
1150
+ lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
1151
+ vec_s, vec_t = t_T_inner.tile(x.shape[0]), t_0_inner.tile(x.shape[0])
1152
+ h = lambda_inner[-1] - lambda_inner[0]
1153
+ r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
1154
+ r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
1155
+ x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
1156
+ if denoise_to_zero:
1157
+ x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
1158
+ return x
1159
+
1160
+
1161
+ #############################################################
1162
+ # other utility functions
1163
+ #############################################################
1164
+
1165
+ def interpolate_fn(x, xp, yp):
1166
+ """
1167
+ A piecewise linear function y = f(x), using xp and yp as keypoints.
1168
+ We implement f(x) in a differentiable way (i.e. applicable for autograd).
1169
+ The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
1170
+
1171
+ Args:
1172
+ x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
1173
+ xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
1174
+ yp: PyTorch tensor with shape [C, K].
1175
+ Returns:
1176
+ The function values f(x), with shape [N, C].
1177
+ """
1178
+ N, K = x.shape[0], xp.shape[1]
1179
+ all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
1180
+ sorted_all_x, x_indices = torch.sort(all_x, dim=2)
1181
+ x_idx = torch.argmin(x_indices, dim=2)
1182
+ cand_start_idx = x_idx - 1
1183
+ start_idx = torch.where(
1184
+ torch.eq(x_idx, 0),
1185
+ torch.tensor(1, device=x.device),
1186
+ torch.where(
1187
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1188
+ ),
1189
+ )
1190
+ end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
1191
+ start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
1192
+ end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
1193
+ start_idx2 = torch.where(
1194
+ torch.eq(x_idx, 0),
1195
+ torch.tensor(0, device=x.device),
1196
+ torch.where(
1197
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1198
+ ),
1199
+ )
1200
+ y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
1201
+ start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
1202
+ end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
1203
+ cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
1204
+ return cand
1205
+
1206
+
1207
+ def expand_dims(v, dims):
1208
+ """
1209
+ Expand the tensor `v` to the dim `dims`.
1210
+
1211
+ Args:
1212
+ `v`: a PyTorch tensor with shape [N].
1213
+ `dim`: a `int`.
1214
+ Returns:
1215
+ a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
1216
+ """
1217
+ return v[(...,) + (None,) * (dims - 1)]
ldm/models/diffusion/dpm_solver/sampler.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+ import torch
3
+
4
+ from .dpm_solver import NoiseScheduleVP, model_wrapper, DPM_Solver
5
+
6
+
7
+ MODEL_TYPES = {
8
+ "eps": "noise",
9
+ "v": "v"
10
+ }
11
+
12
+
13
+ class DPMSolverSampler(object):
14
+ def __init__(self, model, **kwargs):
15
+ super().__init__()
16
+ self.model = model
17
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
18
+ self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))
19
+
20
+ def register_buffer(self, name, attr):
21
+ if type(attr) == torch.Tensor:
22
+ if attr.device != torch.device("cuda"):
23
+ attr = attr.to(torch.device("cuda"))
24
+ setattr(self, name, attr)
25
+
26
+ @torch.no_grad()
27
+ def sample(self,
28
+ S,
29
+ batch_size,
30
+ shape,
31
+ conditioning=None,
32
+ callback=None,
33
+ normals_sequence=None,
34
+ img_callback=None,
35
+ quantize_x0=False,
36
+ eta=0.,
37
+ mask=None,
38
+ x0=None,
39
+ temperature=1.,
40
+ noise_dropout=0.,
41
+ score_corrector=None,
42
+ corrector_kwargs=None,
43
+ verbose=True,
44
+ x_T=None,
45
+ log_every_t=100,
46
+ unconditional_guidance_scale=1.,
47
+ unconditional_conditioning=None,
48
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
49
+ **kwargs
50
+ ):
51
+ if conditioning is not None:
52
+ if isinstance(conditioning, dict):
53
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
54
+ if cbs != batch_size:
55
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
56
+ else:
57
+ if conditioning.shape[0] != batch_size:
58
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
59
+
60
+ # sampling
61
+ C, H, W = shape
62
+ size = (batch_size, C, H, W)
63
+
64
+ print(f'Data shape for DPM-Solver sampling is {size}, sampling steps {S}')
65
+
66
+ device = self.model.betas.device
67
+ if x_T is None:
68
+ img = torch.randn(size, device=device)
69
+ else:
70
+ img = x_T
71
+
72
+ ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
73
+
74
+ model_fn = model_wrapper(
75
+ lambda x, t, c: self.model.apply_model(x, t, c),
76
+ ns,
77
+ model_type=MODEL_TYPES[self.model.parameterization],
78
+ guidance_type="classifier-free",
79
+ condition=conditioning,
80
+ unconditional_condition=unconditional_conditioning,
81
+ guidance_scale=unconditional_guidance_scale,
82
+ )
83
+
84
+ dpm_solver = DPM_Solver(model_fn, ns, predict_x0=True, thresholding=False)
85
+ x = dpm_solver.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=2, lower_order_final=True)
86
+
87
+ return x.to(device), None
ldm/models/diffusion/plms.py ADDED
@@ -0,0 +1,243 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
7
+
8
+
9
+ class PLMSSampler(object):
10
+ def __init__(self, model, schedule="linear", **kwargs):
11
+ super().__init__()
12
+ self.model = model
13
+ self.ddpm_num_timesteps = model.num_timesteps
14
+ self.schedule = schedule
15
+
16
+ def register_buffer(self, name, attr):
17
+ if type(attr) == torch.Tensor:
18
+ if attr.device != torch.device("cuda"):
19
+ attr = attr.to(torch.device("cuda"))
20
+ setattr(self, name, attr)
21
+
22
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
23
+ if ddim_eta != 0:
24
+ raise ValueError('ddim_eta must be 0 for PLMS')
25
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
26
+ num_ddpm_timesteps=self.ddpm_num_timesteps, verbose=verbose)
27
+ alphas_cumprod = self.model.alphas_cumprod
28
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
29
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
30
+
31
+ self.register_buffer('betas', to_torch(self.model.betas))
32
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
33
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
34
+
35
+ # calculations for diffusion q(x_t | x_{t-1}) and others
36
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
37
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
38
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
39
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
40
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
41
+
42
+ # ddim sampling parameters
43
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
44
+ ddim_timesteps=self.ddim_timesteps,
45
+ eta=ddim_eta, verbose=verbose)
46
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
47
+ self.register_buffer('ddim_alphas', ddim_alphas)
48
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
49
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
50
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
51
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
52
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
53
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
54
+
55
+ @torch.no_grad()
56
+ def sample(self,
57
+ S,
58
+ batch_size,
59
+ shape,
60
+ conditioning=None,
61
+ callback=None,
62
+ normals_sequence=None,
63
+ img_callback=None,
64
+ quantize_x0=False,
65
+ eta=0.,
66
+ mask=None,
67
+ x0=None,
68
+ temperature=1.,
69
+ noise_dropout=0.,
70
+ score_corrector=None,
71
+ corrector_kwargs=None,
72
+ verbose=True,
73
+ x_T=None,
74
+ log_every_t=100,
75
+ unconditional_guidance_scale=1.,
76
+ unconditional_conditioning=None,
77
+ features_adapter=None,
78
+ cond_tau=0.4,
79
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
80
+ **kwargs
81
+ ):
82
+ # print('*'*20,x_T)
83
+ # exit(0)
84
+ if conditioning is not None:
85
+ if isinstance(conditioning, dict):
86
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
87
+ if cbs != batch_size:
88
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
89
+ else:
90
+ if conditioning.shape[0] != batch_size:
91
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
92
+
93
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
94
+ C, H, W = shape
95
+ size = (batch_size, C, H, W)
96
+ print(f'Data shape for PLMS sampling is {size}')
97
+
98
+ samples, intermediates = self.plms_sampling(conditioning, size,
99
+ callback=callback,
100
+ img_callback=img_callback,
101
+ quantize_denoised=quantize_x0,
102
+ mask=mask, x0=x0,
103
+ ddim_use_original_steps=False,
104
+ noise_dropout=noise_dropout,
105
+ temperature=temperature,
106
+ score_corrector=score_corrector,
107
+ corrector_kwargs=corrector_kwargs,
108
+ x_T=x_T,
109
+ log_every_t=log_every_t,
110
+ unconditional_guidance_scale=unconditional_guidance_scale,
111
+ unconditional_conditioning=unconditional_conditioning,
112
+ features_adapter=features_adapter,
113
+ cond_tau=cond_tau
114
+ )
115
+ return samples, intermediates
116
+
117
+ @torch.no_grad()
118
+ def plms_sampling(self, cond, shape,
119
+ x_T=None, ddim_use_original_steps=False,
120
+ callback=None, timesteps=None, quantize_denoised=False,
121
+ mask=None, x0=None, img_callback=None, log_every_t=100,
122
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
123
+ unconditional_guidance_scale=1., unconditional_conditioning=None, features_adapter=None,
124
+ cond_tau=0.4):
125
+ device = self.model.betas.device
126
+ b = shape[0]
127
+ if x_T is None:
128
+ img = torch.randn(shape, device=device)
129
+ else:
130
+ img = x_T
131
+ if timesteps is None:
132
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
133
+ elif timesteps is not None and not ddim_use_original_steps:
134
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
135
+ timesteps = self.ddim_timesteps[:subset_end]
136
+
137
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
138
+ time_range = list(reversed(range(0, timesteps))) if ddim_use_original_steps else np.flip(timesteps)
139
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
140
+ print(f"Running PLMS Sampling with {total_steps} timesteps")
141
+
142
+ iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
143
+ old_eps = []
144
+
145
+ for i, step in enumerate(iterator):
146
+ index = total_steps - i - 1
147
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
148
+ ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long)
149
+
150
+ if mask is not None: # and index>=10:
151
+ assert x0 is not None
152
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
153
+ img = img_orig * mask + (1. - mask) * img
154
+
155
+ outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
156
+ quantize_denoised=quantize_denoised, temperature=temperature,
157
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
158
+ corrector_kwargs=corrector_kwargs,
159
+ unconditional_guidance_scale=unconditional_guidance_scale,
160
+ unconditional_conditioning=unconditional_conditioning,
161
+ old_eps=old_eps, t_next=ts_next,
162
+ features_adapter=None if index < int(
163
+ (1 - cond_tau) * total_steps) else features_adapter)
164
+
165
+ img, pred_x0, e_t = outs
166
+ old_eps.append(e_t)
167
+ if len(old_eps) >= 4:
168
+ old_eps.pop(0)
169
+ if callback: callback(i)
170
+ if img_callback: img_callback(pred_x0, i)
171
+
172
+ if index % log_every_t == 0 or index == total_steps - 1:
173
+ intermediates['x_inter'].append(img)
174
+ intermediates['pred_x0'].append(pred_x0)
175
+
176
+ return img, intermediates
177
+
178
+ @torch.no_grad()
179
+ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
180
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
181
+ unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None,
182
+ features_adapter=None):
183
+ b, *_, device = *x.shape, x.device
184
+
185
+ def get_model_output(x, t):
186
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
187
+ e_t = self.model.apply_model(x, t, c, features_adapter=features_adapter)
188
+ else:
189
+ x_in = torch.cat([x] * 2)
190
+ t_in = torch.cat([t] * 2)
191
+ c_in = torch.cat([unconditional_conditioning, c])
192
+ e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in, features_adapter=features_adapter).chunk(2)
193
+ e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
194
+
195
+ if score_corrector is not None:
196
+ assert self.model.parameterization == "eps"
197
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
198
+
199
+ return e_t
200
+
201
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
202
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
203
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
204
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
205
+
206
+ def get_x_prev_and_pred_x0(e_t, index):
207
+ # select parameters corresponding to the currently considered timestep
208
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
209
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
210
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
211
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index], device=device)
212
+
213
+ # current prediction for x_0
214
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
215
+ if quantize_denoised:
216
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
217
+ # direction pointing to x_t
218
+ dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
219
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
220
+ if noise_dropout > 0.:
221
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
222
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
223
+ return x_prev, pred_x0
224
+
225
+ e_t = get_model_output(x, t)
226
+ if len(old_eps) == 0:
227
+ # Pseudo Improved Euler (2nd order)
228
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
229
+ e_t_next = get_model_output(x_prev, t_next)
230
+ e_t_prime = (e_t + e_t_next) / 2
231
+ elif len(old_eps) == 1:
232
+ # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
233
+ e_t_prime = (3 * e_t - old_eps[-1]) / 2
234
+ elif len(old_eps) == 2:
235
+ # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
236
+ e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
237
+ elif len(old_eps) >= 3:
238
+ # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
239
+ e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
240
+
241
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
242
+
243
+ return x_prev, pred_x0, e_t
ldm/modules/__pycache__/attention.cpython-38.pyc ADDED
Binary file (10.6 kB). View file
 
ldm/modules/__pycache__/ema.cpython-38.pyc ADDED
Binary file (3.23 kB). View file
 
ldm/modules/attention.py ADDED
@@ -0,0 +1,344 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ from typing import Optional, Any
8
+
9
+ from ldm.modules.diffusionmodules.util import checkpoint
10
+
11
+
12
+ try:
13
+ import xformers
14
+ import xformers.ops
15
+ XFORMERS_IS_AVAILBLE = True
16
+ except:
17
+ XFORMERS_IS_AVAILBLE = False
18
+
19
+ # CrossAttn precision handling
20
+ import os
21
+ _ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
22
+
23
+ if os.environ.get("DISABLE_XFORMERS", "false").lower() == 'true':
24
+ XFORMERS_IS_AVAILBLE = False
25
+
26
+
27
+ def exists(val):
28
+ return val is not None
29
+
30
+
31
+ def uniq(arr):
32
+ return{el: True for el in arr}.keys()
33
+
34
+
35
+ def default(val, d):
36
+ if exists(val):
37
+ return val
38
+ return d() if isfunction(d) else d
39
+
40
+
41
+ def max_neg_value(t):
42
+ return -torch.finfo(t.dtype).max
43
+
44
+
45
+ def init_(tensor):
46
+ dim = tensor.shape[-1]
47
+ std = 1 / math.sqrt(dim)
48
+ tensor.uniform_(-std, std)
49
+ return tensor
50
+
51
+
52
+ # feedforward
53
+ class GEGLU(nn.Module):
54
+ def __init__(self, dim_in, dim_out):
55
+ super().__init__()
56
+ self.proj = nn.Linear(dim_in, dim_out * 2)
57
+
58
+ def forward(self, x):
59
+ x, gate = self.proj(x).chunk(2, dim=-1)
60
+ return x * F.gelu(gate)
61
+
62
+
63
+ class FeedForward(nn.Module):
64
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
65
+ super().__init__()
66
+ inner_dim = int(dim * mult)
67
+ dim_out = default(dim_out, dim)
68
+ project_in = nn.Sequential(
69
+ nn.Linear(dim, inner_dim),
70
+ nn.GELU()
71
+ ) if not glu else GEGLU(dim, inner_dim)
72
+
73
+ self.net = nn.Sequential(
74
+ project_in,
75
+ nn.Dropout(dropout),
76
+ nn.Linear(inner_dim, dim_out)
77
+ )
78
+
79
+ def forward(self, x):
80
+ return self.net(x)
81
+
82
+
83
+ def zero_module(module):
84
+ """
85
+ Zero out the parameters of a module and return it.
86
+ """
87
+ for p in module.parameters():
88
+ p.detach().zero_()
89
+ return module
90
+
91
+
92
+ def Normalize(in_channels):
93
+ return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
94
+
95
+
96
+ class SpatialSelfAttention(nn.Module):
97
+ def __init__(self, in_channels):
98
+ super().__init__()
99
+ self.in_channels = in_channels
100
+
101
+ self.norm = Normalize(in_channels)
102
+ self.q = torch.nn.Conv2d(in_channels,
103
+ in_channels,
104
+ kernel_size=1,
105
+ stride=1,
106
+ padding=0)
107
+ self.k = torch.nn.Conv2d(in_channels,
108
+ in_channels,
109
+ kernel_size=1,
110
+ stride=1,
111
+ padding=0)
112
+ self.v = torch.nn.Conv2d(in_channels,
113
+ in_channels,
114
+ kernel_size=1,
115
+ stride=1,
116
+ padding=0)
117
+ self.proj_out = torch.nn.Conv2d(in_channels,
118
+ in_channels,
119
+ kernel_size=1,
120
+ stride=1,
121
+ padding=0)
122
+
123
+ def forward(self, x):
124
+ h_ = x
125
+ h_ = self.norm(h_)
126
+ q = self.q(h_)
127
+ k = self.k(h_)
128
+ v = self.v(h_)
129
+
130
+ # compute attention
131
+ b,c,h,w = q.shape
132
+ q = rearrange(q, 'b c h w -> b (h w) c')
133
+ k = rearrange(k, 'b c h w -> b c (h w)')
134
+ w_ = torch.einsum('bij,bjk->bik', q, k)
135
+
136
+ w_ = w_ * (int(c)**(-0.5))
137
+ w_ = torch.nn.functional.softmax(w_, dim=2)
138
+
139
+ # attend to values
140
+ v = rearrange(v, 'b c h w -> b c (h w)')
141
+ w_ = rearrange(w_, 'b i j -> b j i')
142
+ h_ = torch.einsum('bij,bjk->bik', v, w_)
143
+ h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
144
+ h_ = self.proj_out(h_)
145
+
146
+ return x+h_
147
+
148
+
149
+ class CrossAttention(nn.Module):
150
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
151
+ super().__init__()
152
+ inner_dim = dim_head * heads
153
+ context_dim = default(context_dim, query_dim)
154
+
155
+ self.scale = dim_head ** -0.5
156
+ self.heads = heads
157
+
158
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
159
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
160
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
161
+
162
+ self.to_out = nn.Sequential(
163
+ nn.Linear(inner_dim, query_dim),
164
+ nn.Dropout(dropout)
165
+ )
166
+
167
+ def forward(self, x, context=None, mask=None):
168
+ h = self.heads
169
+
170
+ q = self.to_q(x)
171
+ context = default(context, x)
172
+ k = self.to_k(context)
173
+ v = self.to_v(context)
174
+
175
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
176
+
177
+ # force cast to fp32 to avoid overflowing
178
+ if _ATTN_PRECISION =="fp32":
179
+ with torch.autocast(enabled=False, device_type = 'cuda'):
180
+ q, k = q.float(), k.float()
181
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
182
+ else:
183
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
184
+
185
+ del q, k
186
+
187
+ if exists(mask):
188
+ mask = rearrange(mask, 'b ... -> b (...)')
189
+ max_neg_value = -torch.finfo(sim.dtype).max
190
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
191
+ sim.masked_fill_(~mask, max_neg_value)
192
+
193
+ # attention, what we cannot get enough of
194
+ sim = sim.softmax(dim=-1)
195
+
196
+ out = einsum('b i j, b j d -> b i d', sim, v)
197
+ out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
198
+ return self.to_out(out)
199
+
200
+
201
+ class MemoryEfficientCrossAttention(nn.Module):
202
+ # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
203
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
204
+ super().__init__()
205
+ print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
206
+ f"{heads} heads.")
207
+ inner_dim = dim_head * heads
208
+ context_dim = default(context_dim, query_dim)
209
+
210
+ self.heads = heads
211
+ self.dim_head = dim_head
212
+
213
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
214
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
215
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
216
+
217
+ self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
218
+ self.attention_op: Optional[Any] = None
219
+
220
+ def forward(self, x, context=None, mask=None):
221
+ q = self.to_q(x)
222
+ context = default(context, x)
223
+ k = self.to_k(context)
224
+ v = self.to_v(context)
225
+
226
+ b, _, _ = q.shape
227
+ q, k, v = map(
228
+ lambda t: t.unsqueeze(3)
229
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
230
+ .permute(0, 2, 1, 3)
231
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
232
+ .contiguous(),
233
+ (q, k, v),
234
+ )
235
+
236
+ # actually compute the attention, what we cannot get enough of
237
+ out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
238
+
239
+ if exists(mask):
240
+ raise NotImplementedError
241
+ out = (
242
+ out.unsqueeze(0)
243
+ .reshape(b, self.heads, out.shape[1], self.dim_head)
244
+ .permute(0, 2, 1, 3)
245
+ .reshape(b, out.shape[1], self.heads * self.dim_head)
246
+ )
247
+ return self.to_out(out)
248
+
249
+
250
+ class BasicTransformerBlock(nn.Module):
251
+ ATTENTION_MODES = {
252
+ "softmax": CrossAttention, # vanilla attention
253
+ "softmax-xformers": MemoryEfficientCrossAttention
254
+ }
255
+ def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
256
+ disable_self_attn=False):
257
+ super().__init__()
258
+ attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
259
+ assert attn_mode in self.ATTENTION_MODES
260
+ attn_cls = self.ATTENTION_MODES[attn_mode]
261
+ self.disable_self_attn = disable_self_attn
262
+ self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
263
+ context_dim=context_dim if self.disable_self_attn else None) # is a self-attention if not self.disable_self_attn
264
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
265
+ self.attn2 = attn_cls(query_dim=dim, context_dim=context_dim,
266
+ heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none
267
+ self.norm1 = nn.LayerNorm(dim)
268
+ self.norm2 = nn.LayerNorm(dim)
269
+ self.norm3 = nn.LayerNorm(dim)
270
+ self.checkpoint = checkpoint
271
+
272
+ def forward(self, x, context=None):
273
+ return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
274
+
275
+ def _forward(self, x, context=None):
276
+ x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
277
+ x = self.attn2(self.norm2(x), context=context) + x
278
+ x = self.ff(self.norm3(x)) + x
279
+ return x
280
+
281
+
282
+ class SpatialTransformer(nn.Module):
283
+ """
284
+ Transformer block for image-like data.
285
+ First, project the input (aka embedding)
286
+ and reshape to b, t, d.
287
+ Then apply standard transformer action.
288
+ Finally, reshape to image
289
+ NEW: use_linear for more efficiency instead of the 1x1 convs
290
+ """
291
+ def __init__(self, in_channels, n_heads, d_head,
292
+ depth=1, dropout=0., context_dim=None,
293
+ disable_self_attn=False, use_linear=False,
294
+ use_checkpoint=True):
295
+ super().__init__()
296
+ if exists(context_dim) and not isinstance(context_dim, list):
297
+ context_dim = [context_dim]
298
+ self.in_channels = in_channels
299
+ inner_dim = n_heads * d_head
300
+ self.norm = Normalize(in_channels)
301
+ if not use_linear:
302
+ self.proj_in = nn.Conv2d(in_channels,
303
+ inner_dim,
304
+ kernel_size=1,
305
+ stride=1,
306
+ padding=0)
307
+ else:
308
+ self.proj_in = nn.Linear(in_channels, inner_dim)
309
+
310
+ self.transformer_blocks = nn.ModuleList(
311
+ [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d],
312
+ disable_self_attn=disable_self_attn, checkpoint=use_checkpoint)
313
+ for d in range(depth)]
314
+ )
315
+ if not use_linear:
316
+ self.proj_out = zero_module(nn.Conv2d(inner_dim,
317
+ in_channels,
318
+ kernel_size=1,
319
+ stride=1,
320
+ padding=0))
321
+ else:
322
+ self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
323
+ self.use_linear = use_linear
324
+
325
+ def forward(self, x, context=None):
326
+ # note: if no context is given, cross-attention defaults to self-attention
327
+ if not isinstance(context, list):
328
+ context = [context]
329
+ b, c, h, w = x.shape
330
+ x_in = x
331
+ x = self.norm(x)
332
+ if not self.use_linear:
333
+ x = self.proj_in(x)
334
+ x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
335
+ if self.use_linear:
336
+ x = self.proj_in(x)
337
+ for i, block in enumerate(self.transformer_blocks):
338
+ x = block(x, context=context[i])
339
+ if self.use_linear:
340
+ x = self.proj_out(x)
341
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
342
+ if not self.use_linear:
343
+ x = self.proj_out(x)
344
+ return x + x_in
ldm/modules/diffusionmodules/__init__.py ADDED
File without changes
ldm/modules/diffusionmodules/__pycache__/__init__.cpython-38.pyc ADDED
Binary file (214 Bytes). View file
 
ldm/modules/diffusionmodules/__pycache__/model.cpython-38.pyc ADDED
Binary file (21.3 kB). View file
 
ldm/modules/diffusionmodules/__pycache__/openaimodel.cpython-38.pyc ADDED
Binary file (21.1 kB). View file
 
ldm/modules/diffusionmodules/__pycache__/util.cpython-38.pyc ADDED
Binary file (9.71 kB). View file
 
ldm/modules/diffusionmodules/model.py ADDED
@@ -0,0 +1,852 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ from typing import Optional, Any
8
+
9
+ from ldm.modules.attention import MemoryEfficientCrossAttention
10
+
11
+ try:
12
+ import xformers
13
+ import xformers.ops
14
+ XFORMERS_IS_AVAILBLE = True
15
+ except:
16
+ XFORMERS_IS_AVAILBLE = False
17
+ print("No module 'xformers'. Proceeding without it.")
18
+
19
+
20
+ def get_timestep_embedding(timesteps, embedding_dim):
21
+ """
22
+ This matches the implementation in Denoising Diffusion Probabilistic Models:
23
+ From Fairseq.
24
+ Build sinusoidal embeddings.
25
+ This matches the implementation in tensor2tensor, but differs slightly
26
+ from the description in Section 3.5 of "Attention Is All You Need".
27
+ """
28
+ assert len(timesteps.shape) == 1
29
+
30
+ half_dim = embedding_dim // 2
31
+ emb = math.log(10000) / (half_dim - 1)
32
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
33
+ emb = emb.to(device=timesteps.device)
34
+ emb = timesteps.float()[:, None] * emb[None, :]
35
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
36
+ if embedding_dim % 2 == 1: # zero pad
37
+ emb = torch.nn.functional.pad(emb, (0,1,0,0))
38
+ return emb
39
+
40
+
41
+ def nonlinearity(x):
42
+ # swish
43
+ return x*torch.sigmoid(x)
44
+
45
+
46
+ def Normalize(in_channels, num_groups=32):
47
+ return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
48
+
49
+
50
+ class Upsample(nn.Module):
51
+ def __init__(self, in_channels, with_conv):
52
+ super().__init__()
53
+ self.with_conv = with_conv
54
+ if self.with_conv:
55
+ self.conv = torch.nn.Conv2d(in_channels,
56
+ in_channels,
57
+ kernel_size=3,
58
+ stride=1,
59
+ padding=1)
60
+
61
+ def forward(self, x):
62
+ x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
63
+ if self.with_conv:
64
+ x = self.conv(x)
65
+ return x
66
+
67
+
68
+ class Downsample(nn.Module):
69
+ def __init__(self, in_channels, with_conv):
70
+ super().__init__()
71
+ self.with_conv = with_conv
72
+ if self.with_conv:
73
+ # no asymmetric padding in torch conv, must do it ourselves
74
+ self.conv = torch.nn.Conv2d(in_channels,
75
+ in_channels,
76
+ kernel_size=3,
77
+ stride=2,
78
+ padding=0)
79
+
80
+ def forward(self, x):
81
+ if self.with_conv:
82
+ pad = (0,1,0,1)
83
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
84
+ x = self.conv(x)
85
+ else:
86
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
87
+ return x
88
+
89
+
90
+ class ResnetBlock(nn.Module):
91
+ def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
92
+ dropout, temb_channels=512):
93
+ super().__init__()
94
+ self.in_channels = in_channels
95
+ out_channels = in_channels if out_channels is None else out_channels
96
+ self.out_channels = out_channels
97
+ self.use_conv_shortcut = conv_shortcut
98
+
99
+ self.norm1 = Normalize(in_channels)
100
+ self.conv1 = torch.nn.Conv2d(in_channels,
101
+ out_channels,
102
+ kernel_size=3,
103
+ stride=1,
104
+ padding=1)
105
+ if temb_channels > 0:
106
+ self.temb_proj = torch.nn.Linear(temb_channels,
107
+ out_channels)
108
+ self.norm2 = Normalize(out_channels)
109
+ self.dropout = torch.nn.Dropout(dropout)
110
+ self.conv2 = torch.nn.Conv2d(out_channels,
111
+ out_channels,
112
+ kernel_size=3,
113
+ stride=1,
114
+ padding=1)
115
+ if self.in_channels != self.out_channels:
116
+ if self.use_conv_shortcut:
117
+ self.conv_shortcut = torch.nn.Conv2d(in_channels,
118
+ out_channels,
119
+ kernel_size=3,
120
+ stride=1,
121
+ padding=1)
122
+ else:
123
+ self.nin_shortcut = torch.nn.Conv2d(in_channels,
124
+ out_channels,
125
+ kernel_size=1,
126
+ stride=1,
127
+ padding=0)
128
+
129
+ def forward(self, x, temb):
130
+ h = x
131
+ h = self.norm1(h)
132
+ h = nonlinearity(h)
133
+ h = self.conv1(h)
134
+
135
+ if temb is not None:
136
+ h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
137
+
138
+ h = self.norm2(h)
139
+ h = nonlinearity(h)
140
+ h = self.dropout(h)
141
+ h = self.conv2(h)
142
+
143
+ if self.in_channels != self.out_channels:
144
+ if self.use_conv_shortcut:
145
+ x = self.conv_shortcut(x)
146
+ else:
147
+ x = self.nin_shortcut(x)
148
+
149
+ return x+h
150
+
151
+
152
+ class AttnBlock(nn.Module):
153
+ def __init__(self, in_channels):
154
+ super().__init__()
155
+ self.in_channels = in_channels
156
+
157
+ self.norm = Normalize(in_channels)
158
+ self.q = torch.nn.Conv2d(in_channels,
159
+ in_channels,
160
+ kernel_size=1,
161
+ stride=1,
162
+ padding=0)
163
+ self.k = torch.nn.Conv2d(in_channels,
164
+ in_channels,
165
+ kernel_size=1,
166
+ stride=1,
167
+ padding=0)
168
+ self.v = torch.nn.Conv2d(in_channels,
169
+ in_channels,
170
+ kernel_size=1,
171
+ stride=1,
172
+ padding=0)
173
+ self.proj_out = torch.nn.Conv2d(in_channels,
174
+ in_channels,
175
+ kernel_size=1,
176
+ stride=1,
177
+ padding=0)
178
+
179
+ def forward(self, x):
180
+ h_ = x
181
+ h_ = self.norm(h_)
182
+ q = self.q(h_)
183
+ k = self.k(h_)
184
+ v = self.v(h_)
185
+
186
+ # compute attention
187
+ b,c,h,w = q.shape
188
+ q = q.reshape(b,c,h*w)
189
+ q = q.permute(0,2,1) # b,hw,c
190
+ k = k.reshape(b,c,h*w) # b,c,hw
191
+ w_ = torch.bmm(q,k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
192
+ w_ = w_ * (int(c)**(-0.5))
193
+ w_ = torch.nn.functional.softmax(w_, dim=2)
194
+
195
+ # attend to values
196
+ v = v.reshape(b,c,h*w)
197
+ w_ = w_.permute(0,2,1) # b,hw,hw (first hw of k, second of q)
198
+ 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]
199
+ h_ = h_.reshape(b,c,h,w)
200
+
201
+ h_ = self.proj_out(h_)
202
+
203
+ return x+h_
204
+
205
+ class MemoryEfficientAttnBlock(nn.Module):
206
+ """
207
+ Uses xformers efficient implementation,
208
+ see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
209
+ Note: this is a single-head self-attention operation
210
+ """
211
+ #
212
+ def __init__(self, in_channels):
213
+ super().__init__()
214
+ self.in_channels = in_channels
215
+
216
+ self.norm = Normalize(in_channels)
217
+ self.q = torch.nn.Conv2d(in_channels,
218
+ in_channels,
219
+ kernel_size=1,
220
+ stride=1,
221
+ padding=0)
222
+ self.k = torch.nn.Conv2d(in_channels,
223
+ in_channels,
224
+ kernel_size=1,
225
+ stride=1,
226
+ padding=0)
227
+ self.v = torch.nn.Conv2d(in_channels,
228
+ in_channels,
229
+ kernel_size=1,
230
+ stride=1,
231
+ padding=0)
232
+ self.proj_out = torch.nn.Conv2d(in_channels,
233
+ in_channels,
234
+ kernel_size=1,
235
+ stride=1,
236
+ padding=0)
237
+ self.attention_op: Optional[Any] = None
238
+
239
+ def forward(self, x):
240
+ h_ = x
241
+ h_ = self.norm(h_)
242
+ q = self.q(h_)
243
+ k = self.k(h_)
244
+ v = self.v(h_)
245
+
246
+ # compute attention
247
+ B, C, H, W = q.shape
248
+ q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v))
249
+
250
+ q, k, v = map(
251
+ lambda t: t.unsqueeze(3)
252
+ .reshape(B, t.shape[1], 1, C)
253
+ .permute(0, 2, 1, 3)
254
+ .reshape(B * 1, t.shape[1], C)
255
+ .contiguous(),
256
+ (q, k, v),
257
+ )
258
+ out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
259
+
260
+ out = (
261
+ out.unsqueeze(0)
262
+ .reshape(B, 1, out.shape[1], C)
263
+ .permute(0, 2, 1, 3)
264
+ .reshape(B, out.shape[1], C)
265
+ )
266
+ out = rearrange(out, 'b (h w) c -> b c h w', b=B, h=H, w=W, c=C)
267
+ out = self.proj_out(out)
268
+ return x+out
269
+
270
+
271
+ class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention):
272
+ def forward(self, x, context=None, mask=None):
273
+ b, c, h, w = x.shape
274
+ x = rearrange(x, 'b c h w -> b (h w) c')
275
+ out = super().forward(x, context=context, mask=mask)
276
+ out = rearrange(out, 'b (h w) c -> b c h w', h=h, w=w, c=c)
277
+ return x + out
278
+
279
+
280
+ def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
281
+ assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown'
282
+ if XFORMERS_IS_AVAILBLE and attn_type == "vanilla":
283
+ attn_type = "vanilla-xformers"
284
+ print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
285
+ if attn_type == "vanilla":
286
+ assert attn_kwargs is None
287
+ return AttnBlock(in_channels)
288
+ elif attn_type == "vanilla-xformers":
289
+ print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...")
290
+ return MemoryEfficientAttnBlock(in_channels)
291
+ elif type == "memory-efficient-cross-attn":
292
+ attn_kwargs["query_dim"] = in_channels
293
+ return MemoryEfficientCrossAttentionWrapper(**attn_kwargs)
294
+ elif attn_type == "none":
295
+ return nn.Identity(in_channels)
296
+ else:
297
+ raise NotImplementedError()
298
+
299
+
300
+ class Model(nn.Module):
301
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
302
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
303
+ resolution, use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
304
+ super().__init__()
305
+ if use_linear_attn: attn_type = "linear"
306
+ self.ch = ch
307
+ self.temb_ch = self.ch*4
308
+ self.num_resolutions = len(ch_mult)
309
+ self.num_res_blocks = num_res_blocks
310
+ self.resolution = resolution
311
+ self.in_channels = in_channels
312
+
313
+ self.use_timestep = use_timestep
314
+ if self.use_timestep:
315
+ # timestep embedding
316
+ self.temb = nn.Module()
317
+ self.temb.dense = nn.ModuleList([
318
+ torch.nn.Linear(self.ch,
319
+ self.temb_ch),
320
+ torch.nn.Linear(self.temb_ch,
321
+ self.temb_ch),
322
+ ])
323
+
324
+ # downsampling
325
+ self.conv_in = torch.nn.Conv2d(in_channels,
326
+ self.ch,
327
+ kernel_size=3,
328
+ stride=1,
329
+ padding=1)
330
+
331
+ curr_res = resolution
332
+ in_ch_mult = (1,)+tuple(ch_mult)
333
+ self.down = nn.ModuleList()
334
+ for i_level in range(self.num_resolutions):
335
+ block = nn.ModuleList()
336
+ attn = nn.ModuleList()
337
+ block_in = ch*in_ch_mult[i_level]
338
+ block_out = ch*ch_mult[i_level]
339
+ for i_block in range(self.num_res_blocks):
340
+ block.append(ResnetBlock(in_channels=block_in,
341
+ out_channels=block_out,
342
+ temb_channels=self.temb_ch,
343
+ dropout=dropout))
344
+ block_in = block_out
345
+ if curr_res in attn_resolutions:
346
+ attn.append(make_attn(block_in, attn_type=attn_type))
347
+ down = nn.Module()
348
+ down.block = block
349
+ down.attn = attn
350
+ if i_level != self.num_resolutions-1:
351
+ down.downsample = Downsample(block_in, resamp_with_conv)
352
+ curr_res = curr_res // 2
353
+ self.down.append(down)
354
+
355
+ # middle
356
+ self.mid = nn.Module()
357
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
358
+ out_channels=block_in,
359
+ temb_channels=self.temb_ch,
360
+ dropout=dropout)
361
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
362
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
363
+ out_channels=block_in,
364
+ temb_channels=self.temb_ch,
365
+ dropout=dropout)
366
+
367
+ # upsampling
368
+ self.up = nn.ModuleList()
369
+ for i_level in reversed(range(self.num_resolutions)):
370
+ block = nn.ModuleList()
371
+ attn = nn.ModuleList()
372
+ block_out = ch*ch_mult[i_level]
373
+ skip_in = ch*ch_mult[i_level]
374
+ for i_block in range(self.num_res_blocks+1):
375
+ if i_block == self.num_res_blocks:
376
+ skip_in = ch*in_ch_mult[i_level]
377
+ block.append(ResnetBlock(in_channels=block_in+skip_in,
378
+ out_channels=block_out,
379
+ temb_channels=self.temb_ch,
380
+ dropout=dropout))
381
+ block_in = block_out
382
+ if curr_res in attn_resolutions:
383
+ attn.append(make_attn(block_in, attn_type=attn_type))
384
+ up = nn.Module()
385
+ up.block = block
386
+ up.attn = attn
387
+ if i_level != 0:
388
+ up.upsample = Upsample(block_in, resamp_with_conv)
389
+ curr_res = curr_res * 2
390
+ self.up.insert(0, up) # prepend to get consistent order
391
+
392
+ # end
393
+ self.norm_out = Normalize(block_in)
394
+ self.conv_out = torch.nn.Conv2d(block_in,
395
+ out_ch,
396
+ kernel_size=3,
397
+ stride=1,
398
+ padding=1)
399
+
400
+ def forward(self, x, t=None, context=None):
401
+ #assert x.shape[2] == x.shape[3] == self.resolution
402
+ if context is not None:
403
+ # assume aligned context, cat along channel axis
404
+ x = torch.cat((x, context), dim=1)
405
+ if self.use_timestep:
406
+ # timestep embedding
407
+ assert t is not None
408
+ temb = get_timestep_embedding(t, self.ch)
409
+ temb = self.temb.dense[0](temb)
410
+ temb = nonlinearity(temb)
411
+ temb = self.temb.dense[1](temb)
412
+ else:
413
+ temb = None
414
+
415
+ # downsampling
416
+ hs = [self.conv_in(x)]
417
+ for i_level in range(self.num_resolutions):
418
+ for i_block in range(self.num_res_blocks):
419
+ h = self.down[i_level].block[i_block](hs[-1], temb)
420
+ if len(self.down[i_level].attn) > 0:
421
+ h = self.down[i_level].attn[i_block](h)
422
+ hs.append(h)
423
+ if i_level != self.num_resolutions-1:
424
+ hs.append(self.down[i_level].downsample(hs[-1]))
425
+
426
+ # middle
427
+ h = hs[-1]
428
+ h = self.mid.block_1(h, temb)
429
+ h = self.mid.attn_1(h)
430
+ h = self.mid.block_2(h, temb)
431
+
432
+ # upsampling
433
+ for i_level in reversed(range(self.num_resolutions)):
434
+ for i_block in range(self.num_res_blocks+1):
435
+ h = self.up[i_level].block[i_block](
436
+ torch.cat([h, hs.pop()], dim=1), temb)
437
+ if len(self.up[i_level].attn) > 0:
438
+ h = self.up[i_level].attn[i_block](h)
439
+ if i_level != 0:
440
+ h = self.up[i_level].upsample(h)
441
+
442
+ # end
443
+ h = self.norm_out(h)
444
+ h = nonlinearity(h)
445
+ h = self.conv_out(h)
446
+ return h
447
+
448
+ def get_last_layer(self):
449
+ return self.conv_out.weight
450
+
451
+
452
+ class Encoder(nn.Module):
453
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
454
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
455
+ resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
456
+ **ignore_kwargs):
457
+ super().__init__()
458
+ if use_linear_attn: attn_type = "linear"
459
+ self.ch = ch
460
+ self.temb_ch = 0
461
+ self.num_resolutions = len(ch_mult)
462
+ self.num_res_blocks = num_res_blocks
463
+ self.resolution = resolution
464
+ self.in_channels = in_channels
465
+
466
+ # downsampling
467
+ self.conv_in = torch.nn.Conv2d(in_channels,
468
+ self.ch,
469
+ kernel_size=3,
470
+ stride=1,
471
+ padding=1)
472
+
473
+ curr_res = resolution
474
+ in_ch_mult = (1,)+tuple(ch_mult)
475
+ self.in_ch_mult = in_ch_mult
476
+ self.down = nn.ModuleList()
477
+ for i_level in range(self.num_resolutions):
478
+ block = nn.ModuleList()
479
+ attn = nn.ModuleList()
480
+ block_in = ch*in_ch_mult[i_level]
481
+ block_out = ch*ch_mult[i_level]
482
+ for i_block in range(self.num_res_blocks):
483
+ block.append(ResnetBlock(in_channels=block_in,
484
+ out_channels=block_out,
485
+ temb_channels=self.temb_ch,
486
+ dropout=dropout))
487
+ block_in = block_out
488
+ if curr_res in attn_resolutions:
489
+ attn.append(make_attn(block_in, attn_type=attn_type))
490
+ down = nn.Module()
491
+ down.block = block
492
+ down.attn = attn
493
+ if i_level != self.num_resolutions-1:
494
+ down.downsample = Downsample(block_in, resamp_with_conv)
495
+ curr_res = curr_res // 2
496
+ self.down.append(down)
497
+
498
+ # middle
499
+ self.mid = nn.Module()
500
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
501
+ out_channels=block_in,
502
+ temb_channels=self.temb_ch,
503
+ dropout=dropout)
504
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
505
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
506
+ out_channels=block_in,
507
+ temb_channels=self.temb_ch,
508
+ dropout=dropout)
509
+
510
+ # end
511
+ self.norm_out = Normalize(block_in)
512
+ self.conv_out = torch.nn.Conv2d(block_in,
513
+ 2*z_channels if double_z else z_channels,
514
+ kernel_size=3,
515
+ stride=1,
516
+ padding=1)
517
+
518
+ def forward(self, x):
519
+ # timestep embedding
520
+ temb = None
521
+
522
+ # downsampling
523
+ hs = [self.conv_in(x)]
524
+ for i_level in range(self.num_resolutions):
525
+ for i_block in range(self.num_res_blocks):
526
+ h = self.down[i_level].block[i_block](hs[-1], temb)
527
+ if len(self.down[i_level].attn) > 0:
528
+ h = self.down[i_level].attn[i_block](h)
529
+ hs.append(h)
530
+ if i_level != self.num_resolutions-1:
531
+ hs.append(self.down[i_level].downsample(hs[-1]))
532
+
533
+ # middle
534
+ h = hs[-1]
535
+ h = self.mid.block_1(h, temb)
536
+ h = self.mid.attn_1(h)
537
+ h = self.mid.block_2(h, temb)
538
+
539
+ # end
540
+ h = self.norm_out(h)
541
+ h = nonlinearity(h)
542
+ h = self.conv_out(h)
543
+ return h
544
+
545
+
546
+ class Decoder(nn.Module):
547
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
548
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
549
+ resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
550
+ attn_type="vanilla", **ignorekwargs):
551
+ super().__init__()
552
+ if use_linear_attn: attn_type = "linear"
553
+ self.ch = ch
554
+ self.temb_ch = 0
555
+ self.num_resolutions = len(ch_mult)
556
+ self.num_res_blocks = num_res_blocks
557
+ self.resolution = resolution
558
+ self.in_channels = in_channels
559
+ self.give_pre_end = give_pre_end
560
+ self.tanh_out = tanh_out
561
+
562
+ # compute in_ch_mult, block_in and curr_res at lowest res
563
+ in_ch_mult = (1,)+tuple(ch_mult)
564
+ block_in = ch*ch_mult[self.num_resolutions-1]
565
+ curr_res = resolution // 2**(self.num_resolutions-1)
566
+ self.z_shape = (1,z_channels,curr_res,curr_res)
567
+ print("Working with z of shape {} = {} dimensions.".format(
568
+ self.z_shape, np.prod(self.z_shape)))
569
+
570
+ # z to block_in
571
+ self.conv_in = torch.nn.Conv2d(z_channels,
572
+ block_in,
573
+ kernel_size=3,
574
+ stride=1,
575
+ padding=1)
576
+
577
+ # middle
578
+ self.mid = nn.Module()
579
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
580
+ out_channels=block_in,
581
+ temb_channels=self.temb_ch,
582
+ dropout=dropout)
583
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
584
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
585
+ out_channels=block_in,
586
+ temb_channels=self.temb_ch,
587
+ dropout=dropout)
588
+
589
+ # upsampling
590
+ self.up = nn.ModuleList()
591
+ for i_level in reversed(range(self.num_resolutions)):
592
+ block = nn.ModuleList()
593
+ attn = nn.ModuleList()
594
+ block_out = ch*ch_mult[i_level]
595
+ for i_block in range(self.num_res_blocks+1):
596
+ block.append(ResnetBlock(in_channels=block_in,
597
+ out_channels=block_out,
598
+ temb_channels=self.temb_ch,
599
+ dropout=dropout))
600
+ block_in = block_out
601
+ if curr_res in attn_resolutions:
602
+ attn.append(make_attn(block_in, attn_type=attn_type))
603
+ up = nn.Module()
604
+ up.block = block
605
+ up.attn = attn
606
+ if i_level != 0:
607
+ up.upsample = Upsample(block_in, resamp_with_conv)
608
+ curr_res = curr_res * 2
609
+ self.up.insert(0, up) # prepend to get consistent order
610
+
611
+ # end
612
+ self.norm_out = Normalize(block_in)
613
+ self.conv_out = torch.nn.Conv2d(block_in,
614
+ out_ch,
615
+ kernel_size=3,
616
+ stride=1,
617
+ padding=1)
618
+
619
+ def forward(self, z):
620
+ #assert z.shape[1:] == self.z_shape[1:]
621
+ self.last_z_shape = z.shape
622
+
623
+ # timestep embedding
624
+ temb = None
625
+
626
+ # z to block_in
627
+ h = self.conv_in(z)
628
+
629
+ # middle
630
+ h = self.mid.block_1(h, temb)
631
+ h = self.mid.attn_1(h)
632
+ h = self.mid.block_2(h, temb)
633
+
634
+ # upsampling
635
+ for i_level in reversed(range(self.num_resolutions)):
636
+ for i_block in range(self.num_res_blocks+1):
637
+ h = self.up[i_level].block[i_block](h, temb)
638
+ if len(self.up[i_level].attn) > 0:
639
+ h = self.up[i_level].attn[i_block](h)
640
+ if i_level != 0:
641
+ h = self.up[i_level].upsample(h)
642
+
643
+ # end
644
+ if self.give_pre_end:
645
+ return h
646
+
647
+ h = self.norm_out(h)
648
+ h = nonlinearity(h)
649
+ h = self.conv_out(h)
650
+ if self.tanh_out:
651
+ h = torch.tanh(h)
652
+ return h
653
+
654
+
655
+ class SimpleDecoder(nn.Module):
656
+ def __init__(self, in_channels, out_channels, *args, **kwargs):
657
+ super().__init__()
658
+ self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
659
+ ResnetBlock(in_channels=in_channels,
660
+ out_channels=2 * in_channels,
661
+ temb_channels=0, dropout=0.0),
662
+ ResnetBlock(in_channels=2 * in_channels,
663
+ out_channels=4 * in_channels,
664
+ temb_channels=0, dropout=0.0),
665
+ ResnetBlock(in_channels=4 * in_channels,
666
+ out_channels=2 * in_channels,
667
+ temb_channels=0, dropout=0.0),
668
+ nn.Conv2d(2*in_channels, in_channels, 1),
669
+ Upsample(in_channels, with_conv=True)])
670
+ # end
671
+ self.norm_out = Normalize(in_channels)
672
+ self.conv_out = torch.nn.Conv2d(in_channels,
673
+ out_channels,
674
+ kernel_size=3,
675
+ stride=1,
676
+ padding=1)
677
+
678
+ def forward(self, x):
679
+ for i, layer in enumerate(self.model):
680
+ if i in [1,2,3]:
681
+ x = layer(x, None)
682
+ else:
683
+ x = layer(x)
684
+
685
+ h = self.norm_out(x)
686
+ h = nonlinearity(h)
687
+ x = self.conv_out(h)
688
+ return x
689
+
690
+
691
+ class UpsampleDecoder(nn.Module):
692
+ def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
693
+ ch_mult=(2,2), dropout=0.0):
694
+ super().__init__()
695
+ # upsampling
696
+ self.temb_ch = 0
697
+ self.num_resolutions = len(ch_mult)
698
+ self.num_res_blocks = num_res_blocks
699
+ block_in = in_channels
700
+ curr_res = resolution // 2 ** (self.num_resolutions - 1)
701
+ self.res_blocks = nn.ModuleList()
702
+ self.upsample_blocks = nn.ModuleList()
703
+ for i_level in range(self.num_resolutions):
704
+ res_block = []
705
+ block_out = ch * ch_mult[i_level]
706
+ for i_block in range(self.num_res_blocks + 1):
707
+ res_block.append(ResnetBlock(in_channels=block_in,
708
+ out_channels=block_out,
709
+ temb_channels=self.temb_ch,
710
+ dropout=dropout))
711
+ block_in = block_out
712
+ self.res_blocks.append(nn.ModuleList(res_block))
713
+ if i_level != self.num_resolutions - 1:
714
+ self.upsample_blocks.append(Upsample(block_in, True))
715
+ curr_res = curr_res * 2
716
+
717
+ # end
718
+ self.norm_out = Normalize(block_in)
719
+ self.conv_out = torch.nn.Conv2d(block_in,
720
+ out_channels,
721
+ kernel_size=3,
722
+ stride=1,
723
+ padding=1)
724
+
725
+ def forward(self, x):
726
+ # upsampling
727
+ h = x
728
+ for k, i_level in enumerate(range(self.num_resolutions)):
729
+ for i_block in range(self.num_res_blocks + 1):
730
+ h = self.res_blocks[i_level][i_block](h, None)
731
+ if i_level != self.num_resolutions - 1:
732
+ h = self.upsample_blocks[k](h)
733
+ h = self.norm_out(h)
734
+ h = nonlinearity(h)
735
+ h = self.conv_out(h)
736
+ return h
737
+
738
+
739
+ class LatentRescaler(nn.Module):
740
+ def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
741
+ super().__init__()
742
+ # residual block, interpolate, residual block
743
+ self.factor = factor
744
+ self.conv_in = nn.Conv2d(in_channels,
745
+ mid_channels,
746
+ kernel_size=3,
747
+ stride=1,
748
+ padding=1)
749
+ self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
750
+ out_channels=mid_channels,
751
+ temb_channels=0,
752
+ dropout=0.0) for _ in range(depth)])
753
+ self.attn = AttnBlock(mid_channels)
754
+ self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
755
+ out_channels=mid_channels,
756
+ temb_channels=0,
757
+ dropout=0.0) for _ in range(depth)])
758
+
759
+ self.conv_out = nn.Conv2d(mid_channels,
760
+ out_channels,
761
+ kernel_size=1,
762
+ )
763
+
764
+ def forward(self, x):
765
+ x = self.conv_in(x)
766
+ for block in self.res_block1:
767
+ x = block(x, None)
768
+ x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor))))
769
+ x = self.attn(x)
770
+ for block in self.res_block2:
771
+ x = block(x, None)
772
+ x = self.conv_out(x)
773
+ return x
774
+
775
+
776
+ class MergedRescaleEncoder(nn.Module):
777
+ def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks,
778
+ attn_resolutions, dropout=0.0, resamp_with_conv=True,
779
+ ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1):
780
+ super().__init__()
781
+ intermediate_chn = ch * ch_mult[-1]
782
+ self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
783
+ z_channels=intermediate_chn, double_z=False, resolution=resolution,
784
+ attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv,
785
+ out_ch=None)
786
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
787
+ mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
788
+
789
+ def forward(self, x):
790
+ x = self.encoder(x)
791
+ x = self.rescaler(x)
792
+ return x
793
+
794
+
795
+ class MergedRescaleDecoder(nn.Module):
796
+ def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8),
797
+ dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
798
+ super().__init__()
799
+ tmp_chn = z_channels*ch_mult[-1]
800
+ self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout,
801
+ resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
802
+ ch_mult=ch_mult, resolution=resolution, ch=ch)
803
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
804
+ out_channels=tmp_chn, depth=rescale_module_depth)
805
+
806
+ def forward(self, x):
807
+ x = self.rescaler(x)
808
+ x = self.decoder(x)
809
+ return x
810
+
811
+
812
+ class Upsampler(nn.Module):
813
+ def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
814
+ super().__init__()
815
+ assert out_size >= in_size
816
+ num_blocks = int(np.log2(out_size//in_size))+1
817
+ factor_up = 1.+ (out_size % in_size)
818
+ print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
819
+ self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels,
820
+ out_channels=in_channels)
821
+ self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2,
822
+ attn_resolutions=[], in_channels=None, ch=in_channels,
823
+ ch_mult=[ch_mult for _ in range(num_blocks)])
824
+
825
+ def forward(self, x):
826
+ x = self.rescaler(x)
827
+ x = self.decoder(x)
828
+ return x
829
+
830
+
831
+ class Resize(nn.Module):
832
+ def __init__(self, in_channels=None, learned=False, mode="bilinear"):
833
+ super().__init__()
834
+ self.with_conv = learned
835
+ self.mode = mode
836
+ if self.with_conv:
837
+ print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
838
+ raise NotImplementedError()
839
+ assert in_channels is not None
840
+ # no asymmetric padding in torch conv, must do it ourselves
841
+ self.conv = torch.nn.Conv2d(in_channels,
842
+ in_channels,
843
+ kernel_size=4,
844
+ stride=2,
845
+ padding=1)
846
+
847
+ def forward(self, x, scale_factor=1.0):
848
+ if scale_factor==1.0:
849
+ return x
850
+ else:
851
+ x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
852
+ return x
ldm/modules/diffusionmodules/openaimodel.py ADDED
@@ -0,0 +1,798 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from abc import abstractmethod
2
+ import math
3
+ import torch
4
+
5
+ import numpy as np
6
+ import torch as th
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+
10
+ from ldm.modules.diffusionmodules.util import (
11
+ checkpoint,
12
+ conv_nd,
13
+ linear,
14
+ avg_pool_nd,
15
+ zero_module,
16
+ normalization,
17
+ timestep_embedding,
18
+ )
19
+ from ldm.modules.attention import SpatialTransformer
20
+ from ldm.util import exists
21
+
22
+
23
+ # dummy replace
24
+ def convert_module_to_f16(x):
25
+ pass
26
+
27
+ def convert_module_to_f32(x):
28
+ pass
29
+
30
+
31
+ ## go
32
+ class AttentionPool2d(nn.Module):
33
+ """
34
+ Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
35
+ """
36
+
37
+ def __init__(
38
+ self,
39
+ spacial_dim: int,
40
+ embed_dim: int,
41
+ num_heads_channels: int,
42
+ output_dim: int = None,
43
+ ):
44
+ super().__init__()
45
+ self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
46
+ self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
47
+ self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
48
+ self.num_heads = embed_dim // num_heads_channels
49
+ self.attention = QKVAttention(self.num_heads)
50
+
51
+ def forward(self, x):
52
+ b, c, *_spatial = x.shape
53
+ x = x.reshape(b, c, -1) # NC(HW)
54
+ x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)
55
+ x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)
56
+ x = self.qkv_proj(x)
57
+ x = self.attention(x)
58
+ x = self.c_proj(x)
59
+ return x[:, :, 0]
60
+
61
+
62
+ class TimestepBlock(nn.Module):
63
+ """
64
+ Any module where forward() takes timestep embeddings as a second argument.
65
+ """
66
+
67
+ @abstractmethod
68
+ def forward(self, x, emb):
69
+ """
70
+ Apply the module to `x` given `emb` timestep embeddings.
71
+ """
72
+
73
+
74
+ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
75
+ """
76
+ A sequential module that passes timestep embeddings to the children that
77
+ support it as an extra input.
78
+ """
79
+
80
+ def forward(self, x, emb, context=None):
81
+ for layer in self:
82
+ if isinstance(layer, TimestepBlock):
83
+ x = layer(x, emb)
84
+ elif isinstance(layer, SpatialTransformer):
85
+ x = layer(x, context)
86
+ else:
87
+ x = layer(x)
88
+ return x
89
+
90
+
91
+ class Upsample(nn.Module):
92
+ """
93
+ An upsampling layer with an optional convolution.
94
+ :param channels: channels in the inputs and outputs.
95
+ :param use_conv: a bool determining if a convolution is applied.
96
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
97
+ upsampling occurs in the inner-two dimensions.
98
+ """
99
+
100
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
101
+ super().__init__()
102
+ self.channels = channels
103
+ self.out_channels = out_channels or channels
104
+ self.use_conv = use_conv
105
+ self.dims = dims
106
+ if use_conv:
107
+ self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
108
+
109
+ def forward(self, x):
110
+ assert x.shape[1] == self.channels
111
+ if self.dims == 3:
112
+ x = F.interpolate(
113
+ x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
114
+ )
115
+ else:
116
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
117
+ if self.use_conv:
118
+ x = self.conv(x)
119
+ return x
120
+
121
+ class TransposedUpsample(nn.Module):
122
+ 'Learned 2x upsampling without padding'
123
+ def __init__(self, channels, out_channels=None, ks=5):
124
+ super().__init__()
125
+ self.channels = channels
126
+ self.out_channels = out_channels or channels
127
+
128
+ self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
129
+
130
+ def forward(self,x):
131
+ return self.up(x)
132
+
133
+
134
+ class Downsample(nn.Module):
135
+ """
136
+ A downsampling layer with an optional convolution.
137
+ :param channels: channels in the inputs and outputs.
138
+ :param use_conv: a bool determining if a convolution is applied.
139
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
140
+ downsampling occurs in the inner-two dimensions.
141
+ """
142
+
143
+ def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1):
144
+ super().__init__()
145
+ self.channels = channels
146
+ self.out_channels = out_channels or channels
147
+ self.use_conv = use_conv
148
+ self.dims = dims
149
+ stride = 2 if dims != 3 else (1, 2, 2)
150
+ if use_conv:
151
+ self.op = conv_nd(
152
+ dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
153
+ )
154
+ else:
155
+ assert self.channels == self.out_channels
156
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
157
+
158
+ def forward(self, x):
159
+ assert x.shape[1] == self.channels
160
+ return self.op(x)
161
+
162
+
163
+ class ResBlock(TimestepBlock):
164
+ """
165
+ A residual block that can optionally change the number of channels.
166
+ :param channels: the number of input channels.
167
+ :param emb_channels: the number of timestep embedding channels.
168
+ :param dropout: the rate of dropout.
169
+ :param out_channels: if specified, the number of out channels.
170
+ :param use_conv: if True and out_channels is specified, use a spatial
171
+ convolution instead of a smaller 1x1 convolution to change the
172
+ channels in the skip connection.
173
+ :param dims: determines if the signal is 1D, 2D, or 3D.
174
+ :param use_checkpoint: if True, use gradient checkpointing on this module.
175
+ :param up: if True, use this block for upsampling.
176
+ :param down: if True, use this block for downsampling.
177
+ """
178
+
179
+ def __init__(
180
+ self,
181
+ channels,
182
+ emb_channels,
183
+ dropout,
184
+ out_channels=None,
185
+ use_conv=False,
186
+ use_scale_shift_norm=False,
187
+ dims=2,
188
+ use_checkpoint=False,
189
+ up=False,
190
+ down=False,
191
+ ):
192
+ super().__init__()
193
+ self.channels = channels
194
+ self.emb_channels = emb_channels
195
+ self.dropout = dropout
196
+ self.out_channels = out_channels or channels
197
+ self.use_conv = use_conv
198
+ self.use_checkpoint = use_checkpoint
199
+ self.use_scale_shift_norm = use_scale_shift_norm
200
+
201
+ self.in_layers = nn.Sequential(
202
+ normalization(channels),
203
+ nn.SiLU(),
204
+ conv_nd(dims, channels, self.out_channels, 3, padding=1),
205
+ )
206
+
207
+ self.updown = up or down
208
+
209
+ if up:
210
+ self.h_upd = Upsample(channels, False, dims)
211
+ self.x_upd = Upsample(channels, False, dims)
212
+ elif down:
213
+ self.h_upd = Downsample(channels, False, dims)
214
+ self.x_upd = Downsample(channels, False, dims)
215
+ else:
216
+ self.h_upd = self.x_upd = nn.Identity()
217
+
218
+ self.emb_layers = nn.Sequential(
219
+ nn.SiLU(),
220
+ linear(
221
+ emb_channels,
222
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
223
+ ),
224
+ )
225
+ self.out_layers = nn.Sequential(
226
+ normalization(self.out_channels),
227
+ nn.SiLU(),
228
+ nn.Dropout(p=dropout),
229
+ zero_module(
230
+ conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
231
+ ),
232
+ )
233
+
234
+ if self.out_channels == channels:
235
+ self.skip_connection = nn.Identity()
236
+ elif use_conv:
237
+ self.skip_connection = conv_nd(
238
+ dims, channels, self.out_channels, 3, padding=1
239
+ )
240
+ else:
241
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
242
+
243
+ def forward(self, x, emb):
244
+ """
245
+ Apply the block to a Tensor, conditioned on a timestep embedding.
246
+ :param x: an [N x C x ...] Tensor of features.
247
+ :param emb: an [N x emb_channels] Tensor of timestep embeddings.
248
+ :return: an [N x C x ...] Tensor of outputs.
249
+ """
250
+ return checkpoint(
251
+ self._forward, (x, emb), self.parameters(), self.use_checkpoint
252
+ )
253
+
254
+
255
+ def _forward(self, x, emb):
256
+ if self.updown:
257
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
258
+ h = in_rest(x)
259
+ h = self.h_upd(h)
260
+ x = self.x_upd(x)
261
+ h = in_conv(h)
262
+ else:
263
+ h = self.in_layers(x)
264
+ emb_out = self.emb_layers(emb).type(h.dtype)
265
+ while len(emb_out.shape) < len(h.shape):
266
+ emb_out = emb_out[..., None]
267
+ if self.use_scale_shift_norm:
268
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
269
+ scale, shift = th.chunk(emb_out, 2, dim=1)
270
+ h = out_norm(h) * (1 + scale) + shift
271
+ h = out_rest(h)
272
+ else:
273
+ h = h + emb_out
274
+ h = self.out_layers(h)
275
+ return self.skip_connection(x) + h
276
+
277
+
278
+ class AttentionBlock(nn.Module):
279
+ """
280
+ An attention block that allows spatial positions to attend to each other.
281
+ Originally ported from here, but adapted to the N-d case.
282
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
283
+ """
284
+
285
+ def __init__(
286
+ self,
287
+ channels,
288
+ num_heads=1,
289
+ num_head_channels=-1,
290
+ use_checkpoint=False,
291
+ use_new_attention_order=False,
292
+ ):
293
+ super().__init__()
294
+ self.channels = channels
295
+ if num_head_channels == -1:
296
+ self.num_heads = num_heads
297
+ else:
298
+ assert (
299
+ channels % num_head_channels == 0
300
+ ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
301
+ self.num_heads = channels // num_head_channels
302
+ self.use_checkpoint = use_checkpoint
303
+ self.norm = normalization(channels)
304
+ self.qkv = conv_nd(1, channels, channels * 3, 1)
305
+ if use_new_attention_order:
306
+ # split qkv before split heads
307
+ self.attention = QKVAttention(self.num_heads)
308
+ else:
309
+ # split heads before split qkv
310
+ self.attention = QKVAttentionLegacy(self.num_heads)
311
+
312
+ self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
313
+
314
+ def forward(self, x):
315
+ return checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
316
+ #return pt_checkpoint(self._forward, x) # pytorch
317
+
318
+ def _forward(self, x):
319
+ b, c, *spatial = x.shape
320
+ x = x.reshape(b, c, -1)
321
+ qkv = self.qkv(self.norm(x))
322
+ h = self.attention(qkv)
323
+ h = self.proj_out(h)
324
+ return (x + h).reshape(b, c, *spatial)
325
+
326
+
327
+ def count_flops_attn(model, _x, y):
328
+ """
329
+ A counter for the `thop` package to count the operations in an
330
+ attention operation.
331
+ Meant to be used like:
332
+ macs, params = thop.profile(
333
+ model,
334
+ inputs=(inputs, timestamps),
335
+ custom_ops={QKVAttention: QKVAttention.count_flops},
336
+ )
337
+ """
338
+ b, c, *spatial = y[0].shape
339
+ num_spatial = int(np.prod(spatial))
340
+ # We perform two matmuls with the same number of ops.
341
+ # The first computes the weight matrix, the second computes
342
+ # the combination of the value vectors.
343
+ matmul_ops = 2 * b * (num_spatial ** 2) * c
344
+ model.total_ops += th.DoubleTensor([matmul_ops])
345
+
346
+
347
+ class QKVAttentionLegacy(nn.Module):
348
+ """
349
+ A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
350
+ """
351
+
352
+ def __init__(self, n_heads):
353
+ super().__init__()
354
+ self.n_heads = n_heads
355
+
356
+ def forward(self, qkv):
357
+ """
358
+ Apply QKV attention.
359
+ :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
360
+ :return: an [N x (H * C) x T] tensor after attention.
361
+ """
362
+ bs, width, length = qkv.shape
363
+ assert width % (3 * self.n_heads) == 0
364
+ ch = width // (3 * self.n_heads)
365
+ q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
366
+ scale = 1 / math.sqrt(math.sqrt(ch))
367
+ weight = th.einsum(
368
+ "bct,bcs->bts", q * scale, k * scale
369
+ ) # More stable with f16 than dividing afterwards
370
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
371
+ a = th.einsum("bts,bcs->bct", weight, v)
372
+ return a.reshape(bs, -1, length)
373
+
374
+ @staticmethod
375
+ def count_flops(model, _x, y):
376
+ return count_flops_attn(model, _x, y)
377
+
378
+
379
+ class QKVAttention(nn.Module):
380
+ """
381
+ A module which performs QKV attention and splits in a different order.
382
+ """
383
+
384
+ def __init__(self, n_heads):
385
+ super().__init__()
386
+ self.n_heads = n_heads
387
+
388
+ def forward(self, qkv):
389
+ """
390
+ Apply QKV attention.
391
+ :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
392
+ :return: an [N x (H * C) x T] tensor after attention.
393
+ """
394
+ bs, width, length = qkv.shape
395
+ assert width % (3 * self.n_heads) == 0
396
+ ch = width // (3 * self.n_heads)
397
+ q, k, v = qkv.chunk(3, dim=1)
398
+ scale = 1 / math.sqrt(math.sqrt(ch))
399
+ weight = th.einsum(
400
+ "bct,bcs->bts",
401
+ (q * scale).view(bs * self.n_heads, ch, length),
402
+ (k * scale).view(bs * self.n_heads, ch, length),
403
+ ) # More stable with f16 than dividing afterwards
404
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
405
+ a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
406
+ return a.reshape(bs, -1, length)
407
+
408
+ @staticmethod
409
+ def count_flops(model, _x, y):
410
+ return count_flops_attn(model, _x, y)
411
+
412
+
413
+ class UNetModel(nn.Module):
414
+ """
415
+ The full UNet model with attention and timestep embedding.
416
+ :param in_channels: channels in the input Tensor.
417
+ :param model_channels: base channel count for the model.
418
+ :param out_channels: channels in the output Tensor.
419
+ :param num_res_blocks: number of residual blocks per downsample.
420
+ :param attention_resolutions: a collection of downsample rates at which
421
+ attention will take place. May be a set, list, or tuple.
422
+ For example, if this contains 4, then at 4x downsampling, attention
423
+ will be used.
424
+ :param dropout: the dropout probability.
425
+ :param channel_mult: channel multiplier for each level of the UNet.
426
+ :param conv_resample: if True, use learned convolutions for upsampling and
427
+ downsampling.
428
+ :param dims: determines if the signal is 1D, 2D, or 3D.
429
+ :param num_classes: if specified (as an int), then this model will be
430
+ class-conditional with `num_classes` classes.
431
+ :param use_checkpoint: use gradient checkpointing to reduce memory usage.
432
+ :param num_heads: the number of attention heads in each attention layer.
433
+ :param num_heads_channels: if specified, ignore num_heads and instead use
434
+ a fixed channel width per attention head.
435
+ :param num_heads_upsample: works with num_heads to set a different number
436
+ of heads for upsampling. Deprecated.
437
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
438
+ :param resblock_updown: use residual blocks for up/downsampling.
439
+ :param use_new_attention_order: use a different attention pattern for potentially
440
+ increased efficiency.
441
+ """
442
+
443
+ def __init__(
444
+ self,
445
+ image_size,
446
+ in_channels,
447
+ model_channels,
448
+ out_channels,
449
+ num_res_blocks,
450
+ attention_resolutions,
451
+ dropout=0,
452
+ channel_mult=(1, 2, 4, 8),
453
+ conv_resample=True,
454
+ dims=2,
455
+ num_classes=None,
456
+ use_checkpoint=False,
457
+ use_fp16=False,
458
+ num_heads=-1,
459
+ num_head_channels=-1,
460
+ num_heads_upsample=-1,
461
+ use_scale_shift_norm=False,
462
+ resblock_updown=False,
463
+ use_new_attention_order=False,
464
+ use_spatial_transformer=False, # custom transformer support
465
+ transformer_depth=1, # custom transformer support
466
+ context_dim=None, # custom transformer support
467
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
468
+ legacy=True,
469
+ disable_self_attentions=None,
470
+ num_attention_blocks=None,
471
+ disable_middle_self_attn=False,
472
+ use_linear_in_transformer=False,
473
+ ):
474
+ super().__init__()
475
+ if use_spatial_transformer:
476
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
477
+
478
+ if context_dim is not None:
479
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
480
+ from omegaconf.listconfig import ListConfig
481
+ if type(context_dim) == ListConfig:
482
+ context_dim = list(context_dim)
483
+
484
+ if num_heads_upsample == -1:
485
+ num_heads_upsample = num_heads
486
+
487
+ if num_heads == -1:
488
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
489
+
490
+ if num_head_channels == -1:
491
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
492
+
493
+ self.image_size = image_size
494
+ self.in_channels = in_channels
495
+ self.model_channels = model_channels
496
+ self.out_channels = out_channels
497
+ if isinstance(num_res_blocks, int):
498
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
499
+ else:
500
+ if len(num_res_blocks) != len(channel_mult):
501
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
502
+ "as a list/tuple (per-level) with the same length as channel_mult")
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.")
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
+ if isinstance(self.num_classes, int):
536
+ self.label_emb = nn.Embedding(num_classes, time_embed_dim)
537
+ elif self.num_classes == "continuous":
538
+ print("setting up linear c_adm embedding layer")
539
+ self.label_emb = nn.Linear(1, time_embed_dim)
540
+ else:
541
+ raise ValueError()
542
+
543
+ self.input_blocks = nn.ModuleList(
544
+ [
545
+ TimestepEmbedSequential(
546
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
547
+ )
548
+ ]
549
+ )
550
+ self._feature_size = model_channels
551
+ input_block_chans = [model_channels]
552
+ ch = model_channels
553
+ ds = 1
554
+ for level, mult in enumerate(channel_mult):
555
+ for nr in range(self.num_res_blocks[level]):
556
+ layers = [
557
+ ResBlock(
558
+ ch,
559
+ time_embed_dim,
560
+ dropout,
561
+ out_channels=mult * model_channels,
562
+ dims=dims,
563
+ use_checkpoint=use_checkpoint,
564
+ use_scale_shift_norm=use_scale_shift_norm,
565
+ )
566
+ ]
567
+ ch = mult * model_channels
568
+ if ds in attention_resolutions:
569
+ if num_head_channels == -1:
570
+ dim_head = ch // num_heads
571
+ else:
572
+ num_heads = ch // num_head_channels
573
+ dim_head = num_head_channels
574
+ if legacy:
575
+ #num_heads = 1
576
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
577
+ if exists(disable_self_attentions):
578
+ disabled_sa = disable_self_attentions[level]
579
+ else:
580
+ disabled_sa = False
581
+
582
+ if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
583
+ layers.append(
584
+ AttentionBlock(
585
+ ch,
586
+ use_checkpoint=use_checkpoint,
587
+ num_heads=num_heads,
588
+ num_head_channels=dim_head,
589
+ use_new_attention_order=use_new_attention_order,
590
+ ) if not use_spatial_transformer else SpatialTransformer(
591
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
592
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
593
+ use_checkpoint=use_checkpoint
594
+ )
595
+ )
596
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
597
+ self._feature_size += ch
598
+ input_block_chans.append(ch)
599
+ if level != len(channel_mult) - 1:
600
+ out_ch = ch
601
+ self.input_blocks.append(
602
+ TimestepEmbedSequential(
603
+ ResBlock(
604
+ ch,
605
+ time_embed_dim,
606
+ dropout,
607
+ out_channels=out_ch,
608
+ dims=dims,
609
+ use_checkpoint=use_checkpoint,
610
+ use_scale_shift_norm=use_scale_shift_norm,
611
+ down=True,
612
+ )
613
+ if resblock_updown
614
+ else Downsample(
615
+ ch, conv_resample, dims=dims, out_channels=out_ch
616
+ )
617
+ )
618
+ )
619
+ ch = out_ch
620
+ input_block_chans.append(ch)
621
+ ds *= 2
622
+ self._feature_size += ch
623
+
624
+ if num_head_channels == -1:
625
+ dim_head = ch // num_heads
626
+ else:
627
+ num_heads = ch // num_head_channels
628
+ dim_head = num_head_channels
629
+ if legacy:
630
+ #num_heads = 1
631
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
632
+ self.middle_block = TimestepEmbedSequential(
633
+ ResBlock(
634
+ ch,
635
+ time_embed_dim,
636
+ dropout,
637
+ dims=dims,
638
+ use_checkpoint=use_checkpoint,
639
+ use_scale_shift_norm=use_scale_shift_norm,
640
+ ),
641
+ AttentionBlock(
642
+ ch,
643
+ use_checkpoint=use_checkpoint,
644
+ num_heads=num_heads,
645
+ num_head_channels=dim_head,
646
+ use_new_attention_order=use_new_attention_order,
647
+ ) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn
648
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
649
+ disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
650
+ use_checkpoint=use_checkpoint
651
+ ),
652
+ ResBlock(
653
+ ch,
654
+ time_embed_dim,
655
+ dropout,
656
+ dims=dims,
657
+ use_checkpoint=use_checkpoint,
658
+ use_scale_shift_norm=use_scale_shift_norm,
659
+ ),
660
+ )
661
+ self._feature_size += ch
662
+
663
+ self.output_blocks = nn.ModuleList([])
664
+ for level, mult in list(enumerate(channel_mult))[::-1]:
665
+ for i in range(self.num_res_blocks[level] + 1):
666
+ ich = input_block_chans.pop()
667
+ layers = [
668
+ ResBlock(
669
+ ch + ich,
670
+ time_embed_dim,
671
+ dropout,
672
+ out_channels=model_channels * mult,
673
+ dims=dims,
674
+ use_checkpoint=use_checkpoint,
675
+ use_scale_shift_norm=use_scale_shift_norm,
676
+ )
677
+ ]
678
+ ch = model_channels * mult
679
+ if ds in attention_resolutions:
680
+ if num_head_channels == -1:
681
+ dim_head = ch // num_heads
682
+ else:
683
+ num_heads = ch // num_head_channels
684
+ dim_head = num_head_channels
685
+ if legacy:
686
+ #num_heads = 1
687
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
688
+ if exists(disable_self_attentions):
689
+ disabled_sa = disable_self_attentions[level]
690
+ else:
691
+ disabled_sa = False
692
+
693
+ if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
694
+ layers.append(
695
+ AttentionBlock(
696
+ ch,
697
+ use_checkpoint=use_checkpoint,
698
+ num_heads=num_heads_upsample,
699
+ num_head_channels=dim_head,
700
+ use_new_attention_order=use_new_attention_order,
701
+ ) if not use_spatial_transformer else SpatialTransformer(
702
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
703
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
704
+ use_checkpoint=use_checkpoint
705
+ )
706
+ )
707
+ if level and i == self.num_res_blocks[level]:
708
+ out_ch = ch
709
+ layers.append(
710
+ ResBlock(
711
+ ch,
712
+ time_embed_dim,
713
+ dropout,
714
+ out_channels=out_ch,
715
+ dims=dims,
716
+ use_checkpoint=use_checkpoint,
717
+ use_scale_shift_norm=use_scale_shift_norm,
718
+ up=True,
719
+ )
720
+ if resblock_updown
721
+ else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
722
+ )
723
+ ds //= 2
724
+ self.output_blocks.append(TimestepEmbedSequential(*layers))
725
+ self._feature_size += ch
726
+
727
+ self.out = nn.Sequential(
728
+ normalization(ch),
729
+ nn.SiLU(),
730
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
731
+ )
732
+ if self.predict_codebook_ids:
733
+ self.id_predictor = nn.Sequential(
734
+ normalization(ch),
735
+ conv_nd(dims, model_channels, n_embed, 1),
736
+ #nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
737
+ )
738
+
739
+ def convert_to_fp16(self):
740
+ """
741
+ Convert the torso of the model to float16.
742
+ """
743
+ self.input_blocks.apply(convert_module_to_f16)
744
+ self.middle_block.apply(convert_module_to_f16)
745
+ self.output_blocks.apply(convert_module_to_f16)
746
+
747
+ def convert_to_fp32(self):
748
+ """
749
+ Convert the torso of the model to float32.
750
+ """
751
+ self.input_blocks.apply(convert_module_to_f32)
752
+ self.middle_block.apply(convert_module_to_f32)
753
+ self.output_blocks.apply(convert_module_to_f32)
754
+
755
+ def forward(self, x, timesteps=None, context=None, y=None, features_adapter=None, append_to_context=None, **kwargs):
756
+ """
757
+ Apply the model to an input batch.
758
+ :param x: an [N x C x ...] Tensor of inputs.
759
+ :param timesteps: a 1-D batch of timesteps.
760
+ :param context: conditioning plugged in via crossattn
761
+ :param y: an [N] Tensor of labels, if class-conditional.
762
+ :return: an [N x C x ...] Tensor of outputs.
763
+ """
764
+ assert (y is not None) == (
765
+ self.num_classes is not None
766
+ ), "must specify y if and only if the model is class-conditional"
767
+ hs = []
768
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
769
+ emb = self.time_embed(t_emb)
770
+
771
+ if self.num_classes is not None:
772
+ assert y.shape[0] == x.shape[0]
773
+ emb = emb + self.label_emb(y)
774
+
775
+ h = x.type(self.dtype)
776
+
777
+ if append_to_context is not None:
778
+ context = torch.cat([context, append_to_context], dim=1)
779
+
780
+ adapter_idx = 0
781
+ for id, module in enumerate(self.input_blocks):
782
+ h = module(h, emb, context)
783
+ if ((id+1)%3 == 0) and features_adapter is not None:
784
+ h = h + features_adapter[adapter_idx]
785
+ adapter_idx += 1
786
+ hs.append(h)
787
+ if features_adapter is not None:
788
+ assert len(features_adapter)==adapter_idx, 'Wrong features_adapter'
789
+
790
+ h = self.middle_block(h, emb, context)
791
+ for module in self.output_blocks:
792
+ h = th.cat([h, hs.pop()], dim=1)
793
+ h = module(h, emb, context)
794
+ h = h.type(x.dtype)
795
+ if self.predict_codebook_ids:
796
+ return self.id_predictor(h)
797
+ else:
798
+ return self.out(h)
ldm/modules/diffusionmodules/util.py ADDED
@@ -0,0 +1,270 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ ctx.gpu_autocast_kwargs = {"enabled": torch.is_autocast_enabled(),
126
+ "dtype": torch.get_autocast_gpu_dtype(),
127
+ "cache_enabled": torch.is_autocast_cache_enabled()}
128
+ with torch.no_grad():
129
+ output_tensors = ctx.run_function(*ctx.input_tensors)
130
+ return output_tensors
131
+
132
+ @staticmethod
133
+ def backward(ctx, *output_grads):
134
+ ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
135
+ with torch.enable_grad(), \
136
+ torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
137
+ # Fixes a bug where the first op in run_function modifies the
138
+ # Tensor storage in place, which is not allowed for detach()'d
139
+ # Tensors.
140
+ shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
141
+ output_tensors = ctx.run_function(*shallow_copies)
142
+ input_grads = torch.autograd.grad(
143
+ output_tensors,
144
+ ctx.input_tensors + ctx.input_params,
145
+ output_grads,
146
+ allow_unused=True,
147
+ )
148
+ del ctx.input_tensors
149
+ del ctx.input_params
150
+ del output_tensors
151
+ return (None, None) + input_grads
152
+
153
+
154
+ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
155
+ """
156
+ Create sinusoidal timestep embeddings.
157
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
158
+ These may be fractional.
159
+ :param dim: the dimension of the output.
160
+ :param max_period: controls the minimum frequency of the embeddings.
161
+ :return: an [N x dim] Tensor of positional embeddings.
162
+ """
163
+ if not repeat_only:
164
+ half = dim // 2
165
+ freqs = torch.exp(
166
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
167
+ ).to(device=timesteps.device)
168
+ args = timesteps[:, None].float() * freqs[None]
169
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
170
+ if dim % 2:
171
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
172
+ else:
173
+ embedding = repeat(timesteps, 'b -> b d', d=dim)
174
+ return embedding
175
+
176
+
177
+ def zero_module(module):
178
+ """
179
+ Zero out the parameters of a module and return it.
180
+ """
181
+ for p in module.parameters():
182
+ p.detach().zero_()
183
+ return module
184
+
185
+
186
+ def scale_module(module, scale):
187
+ """
188
+ Scale the parameters of a module and return it.
189
+ """
190
+ for p in module.parameters():
191
+ p.detach().mul_(scale)
192
+ return module
193
+
194
+
195
+ def mean_flat(tensor):
196
+ """
197
+ Take the mean over all non-batch dimensions.
198
+ """
199
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
200
+
201
+
202
+ def normalization(channels):
203
+ """
204
+ Make a standard normalization layer.
205
+ :param channels: number of input channels.
206
+ :return: an nn.Module for normalization.
207
+ """
208
+ return GroupNorm32(32, channels)
209
+
210
+
211
+ # PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
212
+ class SiLU(nn.Module):
213
+ def forward(self, x):
214
+ return x * torch.sigmoid(x)
215
+
216
+
217
+ class GroupNorm32(nn.GroupNorm):
218
+ def forward(self, x):
219
+ return super().forward(x.float()).type(x.dtype)
220
+
221
+ def conv_nd(dims, *args, **kwargs):
222
+ """
223
+ Create a 1D, 2D, or 3D convolution module.
224
+ """
225
+ if dims == 1:
226
+ return nn.Conv1d(*args, **kwargs)
227
+ elif dims == 2:
228
+ return nn.Conv2d(*args, **kwargs)
229
+ elif dims == 3:
230
+ return nn.Conv3d(*args, **kwargs)
231
+ raise ValueError(f"unsupported dimensions: {dims}")
232
+
233
+
234
+ def linear(*args, **kwargs):
235
+ """
236
+ Create a linear module.
237
+ """
238
+ return nn.Linear(*args, **kwargs)
239
+
240
+
241
+ def avg_pool_nd(dims, *args, **kwargs):
242
+ """
243
+ Create a 1D, 2D, or 3D average pooling module.
244
+ """
245
+ if dims == 1:
246
+ return nn.AvgPool1d(*args, **kwargs)
247
+ elif dims == 2:
248
+ return nn.AvgPool2d(*args, **kwargs)
249
+ elif dims == 3:
250
+ return nn.AvgPool3d(*args, **kwargs)
251
+ raise ValueError(f"unsupported dimensions: {dims}")
252
+
253
+
254
+ class HybridConditioner(nn.Module):
255
+
256
+ def __init__(self, c_concat_config, c_crossattn_config):
257
+ super().__init__()
258
+ self.concat_conditioner = instantiate_from_config(c_concat_config)
259
+ self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
260
+
261
+ def forward(self, c_concat, c_crossattn):
262
+ c_concat = self.concat_conditioner(c_concat)
263
+ c_crossattn = self.crossattn_conditioner(c_crossattn)
264
+ return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
265
+
266
+
267
+ def noise_like(shape, device, repeat=False):
268
+ repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
269
+ noise = lambda: torch.randn(shape, device=device)
270
+ return repeat_noise() if repeat else noise()
ldm/modules/distributions/__init__.py ADDED
File without changes
ldm/modules/distributions/__pycache__/__init__.cpython-38.pyc ADDED
Binary file (211 Bytes). View file
 
ldm/modules/distributions/__pycache__/distributions.cpython-38.pyc ADDED
Binary file (3.85 kB). View file
 
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,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 reset_num_updates(self):
26
+ del self.num_updates
27
+ self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int))
28
+
29
+ def forward(self, model):
30
+ decay = self.decay
31
+
32
+ if self.num_updates >= 0:
33
+ self.num_updates += 1
34
+ decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))
35
+
36
+ one_minus_decay = 1.0 - decay
37
+
38
+ with torch.no_grad():
39
+ m_param = dict(model.named_parameters())
40
+ shadow_params = dict(self.named_buffers())
41
+
42
+ for key in m_param:
43
+ if m_param[key].requires_grad:
44
+ sname = self.m_name2s_name[key]
45
+ shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
46
+ shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
47
+ else:
48
+ assert not key in self.m_name2s_name
49
+
50
+ def copy_to(self, model):
51
+ m_param = dict(model.named_parameters())
52
+ shadow_params = dict(self.named_buffers())
53
+ for key in m_param:
54
+ if m_param[key].requires_grad:
55
+ m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
56
+ else:
57
+ assert not key in self.m_name2s_name
58
+
59
+ def store(self, parameters):
60
+ """
61
+ Save the current parameters for restoring later.
62
+ Args:
63
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
64
+ temporarily stored.
65
+ """
66
+ self.collected_params = [param.clone() for param in parameters]
67
+
68
+ def restore(self, parameters):
69
+ """
70
+ Restore the parameters stored with the `store` method.
71
+ Useful to validate the model with EMA parameters without affecting the
72
+ original optimization process. Store the parameters before the
73
+ `copy_to` method. After validation (or model saving), use this to
74
+ restore the former parameters.
75
+ Args:
76
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
77
+ updated with the stored parameters.
78
+ """
79
+ for c_param, param in zip(self.collected_params, parameters):
80
+ param.data.copy_(c_param.data)
ldm/modules/encoders/__init__.py ADDED
File without changes
ldm/modules/encoders/__pycache__/__init__.cpython-38.pyc ADDED
Binary file (206 Bytes). View file
 
ldm/modules/encoders/__pycache__/adapter.cpython-38.pyc ADDED
Binary file (11.6 kB). View file
 
ldm/modules/encoders/__pycache__/modules.cpython-38.pyc ADDED
Binary file (14.9 kB). View file